[논문 리뷰] A Hierarchical Probabilistic U-Net for Modeling Multi-Scale Ambiguities

Simon Köhl, Bernardino Romera‐Paredes외 6명|arXiv (Cornell University)|2019. 05. 30.

AI in cancer detection참고 문헌 42인용 수 27

한 줄 요약

이 논문은 의료 영상 및 자연 영상 분할에서 다중 척도의 모호성을 모델링하기 위해 공간적으로 구조화된 잠재 변수 맵의 군집적 계층 구조를 사용하는 조건부 변동형 자동에코더(conditional variational autoencoder)를 U-Net 아키텍처와 통합한 계층적 확률 U-Net(HPU-Net)을 제안한다. 이 방법은 높은 정밀도의 재구성과 세밀한 분할의 샘플링을 가능하게 하며, 다양한 척도에서 복잡하고 구조화된 분포를 학습하여, 인스턴스 분할 및 불확실성 인식 예측 과제에서 이전 모델들을 능가한다.

ABSTRACT

Medical imaging only indirectly measures the molecular identity of the tissue within each voxel, which often produces only ambiguous image evidence for target measures of interest, like semantic segmentation. This diversity and the variations of plausible interpretations are often specific to given image regions and may thus manifest on various scales, spanning all the way from the pixel to the image level. In order to learn a flexible distribution that can account for multiple scales of variations, we propose the Hierarchical Probabilistic U-Net, a segmentation network with a conditional variational auto-encoder (cVAE) that uses a hierarchical latent space decomposition. We show that this model formulation enables sampling and reconstruction of segmenations with high fidelity, i.e. with finely resolved detail, while providing the flexibility to learn complex structured distributions across scales. We demonstrate these abilities on the task of segmenting ambiguous medical scans as well as on instance segmentation of neurobiological and natural images. Our model automatically separates independent factors across scales, an inductive bias that we deem beneficial in structured output prediction tasks beyond segmentation.

연구 동기 및 목표

• 이미지 증거가 단일 정확한 해석을 도출하기에 부족한 경우, 의료 영상 분할에서 복잡하고 다중 척도의 모호성을 모델링하는 데 도전하는 것.
• 계층적 잠재 공간 분해를 통해 세밀한 구조적 세부 사항을 유지함으로써 생성 및 재구성된 분할 맵의 정밀도를 향상시키는 것.
• 공간적 척도와 위치에 따라 독립적으로 변하는, 구조화되고 분리된 분할에 대한 분포를 학습하는 것.
• 모호하거나 가림을 입은 영역에서 불확실성 인식 분할 및 인스턴스 수준의 레이블링을 가능하게 하는 것.
• 다양한 가능성이 있는 분할 가설을 생성함으로써 진단 예측과 같은 후행 작업을 지원하는 것.

제안 방법

• 모델은 다중 해상도에서 공간적으로 구조화된 잠재 맵을 사용하는 계층적 조건부 변동형 자동에코더(cVAE)와 U-Net 인코더-디코더를 통합한다.
• 각 척도에서의 잠재 변수는 U-Net 디코더 경로에 삽입되어, 불확실성과 구조에 대한 척도별 모델링을 가능하게 한다.
• 계층적 사전 분포와 사후 분포는 공간적으로 배열된 잠재 맵 위에 정의되어, 분할 맵의 군집적 계층적 생성을 가능하게 한다.
• 분할 샘플에 대해 히브리드 거리 기반 군집 알고리즘을 적용하여, 샘플 간 일관된 클래스 레이블을 가진 픽셀들을 그룹화함으로써 인스턴스 수준의 레이블링을 수행한다.
• 후처리로는 잡음 제거를 위한 침식 기반 필터링과 경계 아티팩트 제거를 위한 다수결 레이블 대체가 포함된다.
• 모델은 재구성 손실과 KL 발산 손실을 사용하여 엔드 투 엔드로 훈련되며, 계층적 잠재 변수들이 공간적 및 척도별 변동성의 분리된 모델링을 가능하게 한다.

실험 결과

연구 질문

• RQ1 계층적 잠재 공간은 모호한 의료 영상에서 생성 및 재구성된 분할 맵의 정밀도를 향상시키는가?
• RQ2 모델은 다중 공간 척도에서 독립적인 변동성을 체계적인 출력 예측 과제에서 학습하고 분리할 수 있는가?
• RQ3 계층적 cVAE 아키텍처는 복잡한 시나리오에서 가림과 모호성이 존재할 경우, 전반적인 잠재 변수 모델보다 더 나은 인스턴스 분할 성능을 보이는가?
• RQ4 모델은 불확실성을 반영하는 일관되고 다양한 분할 가설을 생성할 수 있는가, 특히 부분적으로 가려진 영역에서 유사한가?
• RQ5 모델이 척도 간 변동 인자들을 분리하는 능력이 후행 진단 과제에서의 성능 향상에 얼마나 기여하는가?

주요 결과

• HPU-Net은 특히 세밀한 구조적 세부 사항을 유지함으로써 표준 확률 U-Net보다 유의미하게 높은 재구성 및 샘플링 정밀도를 달성한다.
• 모델은 신경생물학적(SNEMI3D) 및 자연 영상(Cityscapes) 데이터셋 모두에서 일관된 인스턴스 분할을 생성하며, 이미지당 5개의 별도 인스턴스 ID를 학습한다.
• Cityscapes의 차량 인스턴스 분할 과제에서, 모델은 샘플 간에 모호한 영역(예: 거리/보도, 트럭/버스)을 성공적으로 뒤바꿈하여 학습된 불확실성을 반영한다.
• 가장 전역적인 척도에서만 샘플링할 경우, 모델은 굵은 해상도의 낮은 해상도 마스크를 생성한다. 반면 가장 국소적인 척도에서만 샘플링할 경우 고해상도이지만 노이즈가 많은 예측 결과를 도출하며, 이는 계층적 잠재 변수의 상호보완적 역할을 입증한다.
• 블랙아웃된 영역에서도 일관된 분할과 불확실성을 예측함으로써, 시간적 연속성 의료 영상에서의 4차원 질환 진행 모델링에의 적용 가능성을 시사한다.
• 그리드 히브리드 거리 군집 알고리즘이 효과적으로 픽셀들을 일관된 인스턴스로 그룹화하며, 후처리 과정에서 분할 경계의 아티팩트를 제거한다.