[논문 리뷰] SegDiff: Image Segmentation with Diffusion Probabilistic Models

Tomer Amit, Tal Shaharbany외 2명|arXiv (Cornell University)|2021. 12. 01.

Advanced Neural Network Applications참고 문헌 47인용 수 13

한 줄 요약

이 논문은 사전 학습된 백본 없이 엔드 투 엔드 이미지 세그멘테이션을 위한 새로운 확산 확률 모델인 SegDiff를 소개한다. 이 모델은 입력 이미지의 인코딩된 특징과 현재 세그멘테이션 예측값의 특징을 합산하여 노이즈 제거 과정을 조건화하며, 반복적으로 출력을 정교화한다; 이는 다수의 확률적 생성과 모델 평균화를 활용하여 Cityscapes, Vaihingen, MoNuSeg 벤치마크에서 최신 기술 수준(SOTA) 성능을 달성한다.

ABSTRACT

Diffusion Probabilistic Methods are employed for state-of-the-art image generation. In this work, we present a method for extending such models for performing image segmentation. The method learns end-to-end, without relying on a pre-trained backbone. The information in the input image and in the current estimation of the segmentation map is merged by summing the output of two encoders. Additional encoding layers and a decoder are then used to iteratively refine the segmentation map, using a diffusion model. Since the diffusion model is probabilistic, it is applied multiple times, and the results are merged into a final segmentation map. The new method produces state-of-the-art results on the Cityscapes validation set, the Vaihingen building segmentation benchmark, and the MoNuSeg dataset.

연구 동기 및 목표

• 이미지 생성에 주로 사용된 확산 확률 모델을, 출력이 유일하고 결정적인 이미지 세그멘테이션 작업으로 확장하고자 한다.
• 모든 모델을 처음부터 엔드 투 엔드로 훈련함으로써 사전 학습된 백본에 의존하지 않도록 하고자 한다.
• 입력 이미지와 현재 세그멘테이션 추정치를 특징 합산을 통해 융합하는 새로운 조건화 메커니즘을 개발하고자 한다.
• 다수의 확률적 출력을 생성하고 평균화하여 정확도와 모델 캘리브레이션을 향상시키고자 한다.
• 특히 소규모 데이터셋에서의 성능을 포함하여 다양한 세그멘테이션 벤치마크에서 최신 기술 수준의 성능을 입증하고자 한다.

제안 방법

• 모델은 노이즈 제거 네트워크로 U-Net 아키텍처를 사용하며, 입력 이미지 I와 현재 세그멘테이션 추정치 xₜ에 대한 별도의 인코더 두 개를 갖는다.
• 두 인코더의 특징은 U-Net 인코더 경로의 초기 단계에서 합산되어 이미지와 세그멘테이션 정보의 공동 모델링을 가능하게 한다.
• 노이즈 제거 네트워크는 합산된 특징을 사용하여 다음 상태 xₜ₋₁를 예측하며, T개의 확산 단계 동안 반복적으로 세그멘테이션 맵을 정교화한다.
• 역과정 손실을 최소화하기 위해 변분 추론 프레임워크를 사용하여 모델을 훈련하며, 노이즈 제거 목표를 최적화한다.
• 입력당 다수의 독립적인 생성을 수행하고, 결과를 평균화하여 최종 세그멘테이션을 도출함으로써 정확도와 견고성을 향상시킨다.
• 잔차형 조건화 메커니즘을 적용하여, 입력 이미지를 인코딩하고 현재 예측의 표현에 더하는 방식을 사용하며, 연결(concatenation)이 아닌 합산(summation)을 사용한다.

실험 결과

연구 질문

• RQ1 유일한 정답이 존재하는 이미지 세그멘테이션 작업에 대해, 이전에는 주로 생성 작업에 사용된 확산 확률 모델이 효과적으로 적용될 수 있는가?
• RQ2 특히 사전 학습된 백본을 피하기 위해, 입력 이미지를 확산 모델에서 효과적으로 조건화할 수 있는가?
• RQ3 다수의 확률적 샘플을 생성하고 평균화하는 것이 단일 결정적 예측에 비해 정확도와 모델 캘리브레이션을 향상시키는가?
• RQ4 특징 융합 메커니즘의 선택(예: 합산 대비 연결)이 세그멘테이션 확산 설정에서 성능에 어떤 영향을 미치는가?
• RQ5 이 새로운 프레임워크에서 추론 속도(확산 단계 수)와 세그멘테이션 정확도 사이의 상충 관계는 어떠한가?

주요 결과

• SegDiff는 Cityscapes 검증 세트에서 최신 기술 수준의 성능을 달성하였으며, 특히 '버스'와 같은 도전적인 카테고리에서 이전 방법에 비해 mIoU가 향상되었다.
• Vaihingen 빌딩 세그멘테이션 벤치마크에서, 조건화에 특징 합산을 사용할 경우 기존 접근법에 비해 뚜렷한 성능 향상을 보였다.
• MoNuSeg 데이터셋에서 세포핵 세그멘테이션 작업에 대해 강력한 성능을 기록하여 다양한 의료 영상 도메인으로의 일반화 능력을 입증하였다.
• 확산 단계 수를 100에서 25로 줄였을 때 성능 저하가 1–2 mIoU 이내로 발생하여 최대 4배의 속도 향상과 함께 정확도 손실 최소화를 달성하였다.
• RRDB 블록의 수는 성능에 미치는 영향이 제한적이었으며, 최적 설정 간 mIoU 차이가 1 미만이었다.
• 9개의 생성 샘플을 평균화하면 정확도와 모델 캘리브레이션 모두가 크게 향상되어, 결정적 작업에서의 확률적 샘플링의 유용성을 확인하였다.