[논문 리뷰] Ambiguous Medical Image Segmentation using Diffusion Models

Aimon Rahman, Jeya Maria Jose Valanarasu외 2명|arXiv (Cornell University)|2023. 04. 10.

Radiomics and Machine Learning in Medical Imaging인용 수 7

한 줄 요약

이 논문은 확산 모델 기반 프레임워크인 CIMD를 제안한다. 이는 확산 샘플링의 본질적 확률적 특성을 활용해 단일 입력에서 다수의 다양한 실재성 있는 분할 마스크를 생성함으로써 모호한 의료 영상 분할 문제를 해결한다. 기존 방법과 달리 별도의 사전 네트워크가 필요로 하지 않으며, CIMD는 다중 계층적 수준에서 확률성을 주입함으로써 최신 기술 수준의 성능을 달성한다. 정확도(Dice 점수: 0.915)와 다양성(CI 점수: 0.759) 측면에서 모두 뛰어난 성능을 보이며, 임상적으로 관련성이 있는 변동성을 유지한다.

ABSTRACT

Collective insights from a group of experts have always proven to outperform an individual's best diagnostic for clinical tasks. For the task of medical image segmentation, existing research on AI-based alternatives focuses more on developing models that can imitate the best individual rather than harnessing the power of expert groups. In this paper, we introduce a single diffusion model-based approach that produces multiple plausible outputs by learning a distribution over group insights. Our proposed model generates a distribution of segmentation masks by leveraging the inherent stochastic sampling process of diffusion using only minimal additional learning. We demonstrate on three different medical image modalities- CT, ultrasound, and MRI that our model is capable of producing several possible variants while capturing the frequencies of their occurrences. Comprehensive results show that our proposed approach outperforms existing state-of-the-art ambiguous segmentation networks in terms of accuracy while preserving naturally occurring variation. We also propose a new metric to evaluate the diversity as well as the accuracy of segmentation predictions that aligns with the interest of clinical practice of collective insights.

연구 동기 및 목표

• 의료 영상 분할에서 결정론적 딥 러닝 모델의 한계를 해결하기 위해, 임상 진단에서의 본질적 인식자 간 변동성에도 불구하고 단일 마스크만 생성하는 문제를 해결한다.
• 최고의 개인을 모방하는 대신, 단일 확률적 프레임워크를 사용해 전문 방사선의사가 제시하는 다수의 타당한 애너테이션 분포를 모델링한다.
• c-VAE 기반 접근법에서 요구되는 별도의 사전 네트워크가 필요 없도록, 확률성을 직접 확산 모델의 계층적 아키텍처에 통합한다.
• 임상 실무와 부합하는 평가 지표(CI 점수)를 새롭게 개발하여, 지표 마스크에 대한 다양성을 지상 진술 전문가 애너테이션과 비교해 측정한다.
• CT, 초음파, MRI를 포함한 다양한 모odalities에서 모호한 분할 문제에 대해 확산 모델의 효과성을 입증한다.

제안 방법

• 제안된 CIMD 프레임워크는 잠재 공간에서 분할 마스크에 점차 노이즈를 추가하는 전방 확산 과정을 역으로 수행하도록 훈련된 조건부 확산 확률 모델을 사용한다.
• 추론 시 확률적 샘플링을 통해 별도의 사전 네트워크에서 추가 추론 단계를 거치지 않아도 단일 입력에서 다수의 타당한 분할 마스크를 자연스럽게 생성한다.
• 모델은 다중 계층적 특징 수준에서 노이즈를 주입함으로써 분할 맵의 다양한 스케일에서 다양하면서도 일관된 예측을 가능하게 한다.
• 확산 모델은 입력 이미지당 다수의 애너테이션에서 끝내기로 훈련되어, 타당한 분할의 결합 분포를 학습한다.
• 새로운 손실 함수는 픽셀 단위 재구성 손실과 대비 손실을 조합하여 생성된 마스크의 정확성과 다양성을 모두 유지한다.
• 기존의 확산 기반 분할 프레임워크와 호환되며, 아키텍처 수정이 최소한으로 필요하다.

실험 결과

연구 질문

• RQ1 보조 사전 네트워크가 필요 없이 단일 확산 모델이 별도의 사전 네트워크 없이 다수의 다양한 임상적으로 타당한 분할 마스크를 생성할 수 있는가?
• RQ2 확산 모델에서 계층적 방식으로 확률성을 주입할 경우, 모호한 분할 출력의 다양성과 정확도에 어떤 영향을 미치는가?
• RQ3 확산 모델이 다양한 의료 영상 모달리티에서 전문가 애너테이션의 빈도 분포를 어느 정도 학습하고 재현할 수 있는가?
• RQ4 제안된 방법이 실제 의료 영상 데이터셋에서 기존의 c-VAE 기반 및 결정론적 분할 네트워크보다 정확도와 다양성 측면에서 모두 뛰어나게 성능을 발휘하는가?
• RQ5 정확도와 다양성을 동시에 반영하는 새로운 평가 지표는 기존의 Dice 점수만을 사용하는 지표보다 임상적 관련성을 더 잘 반영할 수 있는가?

주요 결과

• CIMD는 LIDC-IDRI 데이터셋에서 최대 Dice 점수(D_max) 0.915를 기록하여, 다음으로 우수한 성능을 보인 방법(Probabilistic U-Net)을 크게 앞서며 지상 진술 분포와의 일치도에서 뛰어난 정확도를 입증한다.
• LIDC-IDRI 데이터셋에서 CIMD는 임상적 다양성(CI) 점수 0.759를 기록하여 전문가 애너테이션에서 관찰되는 다양성과 강력한 일치를 보이며, 모든 베이스라인을 초월한다.
• LIDC-IDRI 데이터셋에서 CIMD는 기하학적 에너지 거리(GED)를 0.321로 줄여, DDPM-Prob-Seg(0.417)를 능가하며 지상 진술과의 분포 일치도가 뛰어나다.
• 정성적 결과에서는 초음파 및 CT 스캔에서 작은 또는 모호한 병변이 있는 도전적인 케이스에서도 CIMD가 다양한 고품질 분할을 생성함을 보여준다.
• 단절 실험 결과는 계층적 확률성 주입이 핵심임을 확인한다. 동일한 확률적 샘플링 과정을 사용하더라도 이를 제거할 경우 다양성과 정확도가 떨어지는 결과를 보인다.
• 이 방법은 CT, 초음파, MRI 세 가지 모달리티에 걸쳐 일반화되며, 다양한 영상 환경에서의 강건성과 임상적 관련성을 입증한다.