불확실한 영상 분할을 위한 확률적 U‑Net

초록

본 논문은 의료 영상·자율주행 등에서 발생하는 본질적 불확실성을 고려해, 입력 이미지 하나당 다수의 가능한 세그멘테이션을 확률적으로 생성할 수 있는 “Probabilistic U‑Net”을 제안한다. Conditional VAE와 U‑Net을 결합해 저차원 라틴트 공간을 학습하고, 이 공간에서 샘플을 추출해 일관된 전체 이미지 마스크를 생성한다. LIDC‑IDRI 폐 CT와 Cityscapes 데이터에서 기존 드롭아웃, 앙상블, M‑head 등과 비교해 에너지 거리 지표와 IoU 기반 평가에서 우수한 성능을 보이며, 희귀 모드까지 정확히 모델링한다.

상세 분석

본 연구는 “세그멘테이션은 픽셀 단위 확률이 아니라 전체 이미지 수준에서의 다중 모드 분포”라는 전제를 명확히 하고, 이를 구현하기 위해 Conditional Variational Auto‑Encoder(CVAE)와 최신 U‑Net 구조를 결합한 새로운 아키텍처를 설계하였다. 핵심은 저차원 라틴트 공간 ( \mathbb{R}^N ) (N≈6)이다. Prior net은 입력 이미지 (X) 에 조건화된 가우시안 (P(z|X)=\mathcal{N}(\mu_{\text{prior}}(X),\operatorname{diag}(\sigma_{\text{prior}}(X)))) 를 출력하고, Posterior net은 (X) 와 해당 라벨 (Y) 를 함께 입력받아 (Q(z|X,Y)=\mathcal{N}(\mu_{\text{post}}(X,Y),\operatorname{diag}(\sigma_{\text{post}}(X,Y)))) 를 학습한다. 학습 단계에서는 ELBO를 최적화하는데, 여기서 교차 엔트로피 손실은 샘플 (z\sim Q) 가 생성한 마스크 (S) 와 실제 라벨 (Y) 간의 픽셀‑와이드 일치성을 강제하고, KL‑다이버전스는 Posterior와 Prior를 정규화한다. 이 과정은 라벨 간 변이성을 라틴트 공간에 압축시키면서, Prior가 전체 라벨 변이 분포를 포괄하도록 만든다.

샘플링 단계에서는 Prior에서 (z) 를 한 번 추출하고, 이를 1×1 컨볼루션을 통해 (N) 채널 특성 맵으로 브로드캐스트한 뒤 U‑Net의 마지막 피처와 결합한다. 결합된 피처는 또 다른 1×1 컨볼루션 스택 (f_{\text{comb}}) 을 거쳐 클래스별 로짓을 출력한다. 이 구조는 동일 이미지에 대해 (m) 번 샘플링할 경우 U‑Net의 대부분 연산을 재사용하므로, 다중 가설을 빠르게 생성할 수 있다.

성능 평가에서는 기존 픽셀‑와이드 불확실성 추정 방법(드롭아웃, 앙상블)과 달리, 전체 마스크 간의 분포 차이를 측정하는 Generalized Energy Distance(GED)를 채택했다. GED는 두 분포 (P_{\text{gt}},P_{\text{out}}) 에서 독립 샘플을 뽑아 (d=1-\text{IoU}) 거리로 계산한다. 실험 결과, LIDC‑IDRI 데이터에서 4명의 방사선과 전문가가 제공한 라벨을 정확히 재현하고, 희귀 라벨(예: 1% 빈도)까지도 Prior가 학습하도록 성공했다. Cityscapes에서는 라벨을 인위적으로 교체해 다중 모드 상황을 만들었으며, Probabilistic U‑Net은 교체 비율을 정확히 추정해 기존 베이스라인보다 GED가 현저히 낮았다.

또한, 모델의 캘리브레이션 능력을 검증하기 위해 라벨 빈도와 샘플링 빈도 간 상관관계를 분석했으며, 선형 회귀 결과 높은 (R^2) 값을 보였다. 이는 라틴트 공간이 실제 라벨 모드의 확률 질량을 잘 반영한다는 증거다. 마지막으로, 샘플당 연산 비용이 U‑Net 한 번 실행에 비해 거의 동일하므로, 실시간 의료·자율주행 시스템에 적용 가능성이 높다.

이 논문의 주요 기여는 다음과 같다. (1) 픽셀‑와이드 확률이 아니라 전체 마스크 수준의 일관된 다중 가설을 제공한다. (2) 저차원 라틴트 공간을 통해 복잡하고 희귀한 모드까지 학습·샘플링한다. (3) 샘플링 비용이 낮아 실시간 응용에 적합하다. (4) 에너지 거리 기반 정량적 평가를 통해 모드 커버리지와 캘리브레이션을 객관적으로 검증한다.