[논문 리뷰] BerDiff: Conditional Bernoulli Diffusion Model for Medical Image Segmentation

Tao Chen, Chenhui Wang외 1명|arXiv (Cornell University)|2023. 04. 10.

Radiomics and Machine Learning in Medical Imaging참고 문헌 26인용 수 4

한 줄 요약

이 논문은 의료 영상 분할을 위한 조건부 베르누이 확산 모델인 BerDiff를 제안한다. 기존의 가우시안 노이즈 대신 베르누이 노이즈를 확산 커널로 사용함으로써 더 정확하고 다양한 이진 분할 마스크를 생성할 수 있다. 확률적 역확산 과정을 활용해 유의미한 영역을 강조하는 다수의 타당한 마스크를 생성하며, LIDC-IDRI 및 BRATS 2021 데이터셋에서 SOTA 기법들을 능가하여 89.7%의 Dice 스코어를 기록한다.

ABSTRACT

Medical image segmentation is a challenging task with inherent ambiguity and high uncertainty, attributed to factors such as unclear tumor boundaries and multiple plausible annotations. The accuracy and diversity of segmentation masks are both crucial for providing valuable references to radiologists in clinical practice. While existing diffusion models have shown strong capacities in various visual generation tasks, it is still challenging to deal with discrete masks in segmentation. To achieve accurate and diverse medical image segmentation masks, we propose a novel conditional Bernoulli Diffusion model for medical image segmentation (BerDiff). Instead of using the Gaussian noise, we first propose to use the Bernoulli noise as the diffusion kernel to enhance the capacity of the diffusion model for binary segmentation tasks, resulting in more accurate segmentation masks. Second, by leveraging the stochastic nature of the diffusion model, our BerDiff randomly samples the initial Bernoulli noise and intermediate latent variables multiple times to produce a range of diverse segmentation masks, which can highlight salient regions of interest that can serve as valuable references for radiologists. In addition, our BerDiff can efficiently sample sub-sequences from the overall trajectory of the reverse diffusion, thereby speeding up the segmentation process. Extensive experimental results on two medical image segmentation datasets with different modalities demonstrate that our BerDiff outperforms other recently published state-of-the-art methods. Our results suggest diffusion models could serve as a strong backbone for medical image segmentation.

연구 동기 및 목표

• 의료 영상 분할의 본질적 모호성과 불확실성, 특히 경계가 불명확하고 다수의 타당한 애너테이션 존재로 인한 문제를 해결하기 위해.
• 의료 영상 분할의 정확도와 다양성을 향상시키기 위해 확산 모델의 확률적 성질을 활용하여 임상 의사결정 지원에 기여하기 위해.
• 이산 이진 분할 작업에 특화된 확산 모델을 개발하여 가우시안 노이즈 가정을 넘어서기 위해.
• 역확산 궤적에서 부분 시퀀스 샘플링을 통해 효율적인 추론을 가능하게 하기 위해.
• 실제 의료 영상 데이터셋에서 가우시안 기반 확산 모델 및 기존 SOTA 기법들에 비해 우수한 성능을 보여주기 위해.

제안 방법

• 기존의 가우시안 노이즈를 대체하여 이산 분할 마스크를 더 잘 모델링할 수 있도록 베르누이 노이즈를 확산 커널로 사용하는 새로운 조건부 확산 모델을 제안한다.
• 각 단계에서 일정 β₁,…,βₜ에 따라 베르누이 노이즈를 추가하는 마르코프 체인 전방 과정을 정의하여, 지도 마스크가 순차적으로 순수한 노이즈로 변환되도록 한다.
• 입력 영상에 조건부로 조건부된 디노이징 U-Net 아키텍처를 사용하여 역확산 과정에서의 노이즈 추정을 수행한다.
• 초기 베르누이 노이즈 및 중간 잠재 변수에 대해 다른 난수 시드로 역확산 과정을 재초기화하여 다양한 샘플링을 가능하게 한다.
• DDIM 원리에 기반하여 전체 역확산 궤적에서 부분 시퀀스 샘플링을 도입하여 추론 속도를 향상시키면서도 품질을 유지한다.
• 변분 하한 목표함수를 사용하여 엔드 투 엔드로 모델을 훈련시켜 디노이징 스코어 함수를 최적화한다.

실험 결과

연구 질문

• RQ1 베르누이 노이즈를 확산 커널로 사용할 경우, 이산 이진 의료 영상 분할 작업에서 확산 모델의 성능 향상이 가능할까?
• RQ2 베르누이 확산 모델의 확률적 역확산 과정이 임상 참고를 위한 다양하면서도 타당한 분할 마스크를 어떻게 생성하는가?
• RQ3 역확산 궤적에서 부분 시퀀스 샘플링이 유의미하게 추론 시간을 단축시킬 수 있을까, 동시에 분할 품질은 유지되는가?
• RQ4 실제 의료 영상 데이터셋에서 BerDiff는 가우시안 기반 확산 모델 및 다른 SOTA 기법들에 비해 정확도와 다양성 측면에서 어떻게 비교되는가?
• RQ5 베르누이 노이즈 사용이 특히 훈련 반복 횟수가 제한된 상황에서도 더 나은 수렴과 성능을 이끌어낼 수 있을까?

주요 결과

• BerDiff는 LIDC-IDRI 데이터셋에서 89.7%의 Dice 스코어를 기록하여 이전 SOTA 기법들인 SegDiff(89.3%) 및 U-Net(89.2%)를 능가한다.
• BRATS 2021 데이터셋에서는 기존의 U-Net 및 스위니 UNETR와 같은 트랜스포머 기반 모델들을 초월하여 복잡한 종양 분할 작업에서 뛰어난 성능을 보였다.
• 모델은 유의미한 영역을 강조하는 다양한 분할 마스크를 생성하며, 소형 폐 결절과 같이 모호한 경우에도 정확도 있는 지도 마스크를 효과적으로 포착한다.
• 절단 실험 결과, 충분한 훈련 반복(예: 86,500단계)이 이루어진 경우, 가우시안 기반 확산보다 베르누이 기반 확산이 더 뛰어난 성능을 보였다.
• 부분 시퀀스 샘플링을 통해 중간 단계를 건너뛰어 더 빠른 추론이 가능해졌으며, 높은 품질의 출력을 유지하면서도 샘플당 시간을 줄였다.
• 시각적 비교 결과, 다른 확산 모델 및 비확산 모델 대비 MRI 스캔에서 미세한 종양 경계와 같은 어려운 영역에서도 BerDiff가 더 정확한 분할을 생성하는 것으로 나타났다.