[Paper Review] Ambiguous Medical Image Segmentation using Diffusion Models

Aimon Rahman, Jeya Maria Jose Valanarasu+2 more|arXiv (Cornell University)|Apr 10, 2023

Radiomics and Machine Learning in Medical Imaging7 citations

TL;DR

This paper proposes CIMD, a diffusion model-based framework for ambiguous medical image segmentation that generates multiple diverse, realistic segmentation masks from a single input by leveraging the inherent stochasticity of diffusion sampling. Unlike prior methods requiring separate prior networks, CIMD injects stochasticity at multiple hierarchical levels, achieving state-of-the-art performance in both accuracy (Dice score: 0.915) and diversity (CI score: 0.759), while preserving clinically relevant variation.

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.

Motivation & Objective

• To address the limitation of deterministic deep learning models in medical image segmentation, which produce a single mask despite inherent inter-rater variability in clinical diagnoses.
• To model the distribution of multiple plausible annotations from expert radiologists using a single probabilistic framework, rather than imitating the best individual.
• To eliminate the need for separate prior networks in c-VAE-based approaches by embedding stochasticity directly into the diffusion model’s hierarchical architecture.
• To develop a new evaluation metric (CI score) that aligns with clinical practice by measuring diversity relative to ground truth expert annotations.
• To demonstrate the effectiveness of diffusion models in ambiguous segmentation across multiple modalities: CT, ultrasound, and MRI.

Proposed method

• The proposed CIMD framework uses a conditional diffusion probabilistic model trained to reverse a forward diffusion process that gradually adds noise to segmentation masks in latent space.
• Stochastic sampling during inference naturally generates multiple plausible segmentation masks from a single input, without requiring additional inference-time sampling from a separate prior network.
• The model injects noise at multiple hierarchical feature levels, enabling diverse yet coherent predictions across different scales of the segmentation map.
• The diffusion model is trained end-to-end on multiple annotations per input image, learning the joint distribution of plausible segmentations.
• A novel loss function combines pixel-wise reconstruction loss with a contrastive loss to preserve both accuracy and diversity in generated masks.
• The method is compatible with existing diffusion-based segmentation frameworks and requires minimal architectural modification.

Experimental results

Research questions

• RQ1 Can a single diffusion model generate multiple diverse, clinically plausible segmentation masks without requiring an auxiliary prior network?
• RQ2 How does the hierarchical injection of stochasticity in a diffusion model affect the diversity and accuracy of ambiguous segmentation outputs?
• RQ3 To what extent can a diffusion model learn and reproduce the frequency distribution of expert annotations across different medical imaging modalities?
• RQ4 Does the proposed method outperform existing c-VAE-based and deterministic segmentation networks in both accuracy and diversity on real-world medical imaging datasets?
• RQ5 Can a new evaluation metric that captures both accuracy and diversity better reflect clinical relevance than standard metrics like Dice alone?

Key findings

• CIMD achieves a maximum Dice score (D_max) of 0.915 on the LIDC-IDRI dataset, significantly outperforming the next best method (Probabilistic U-Net) and demonstrating superior accuracy in matching ground truth distributions.
• The model achieves a Clinical Diversity (CI) score of 0.759 on LIDC-IDRI, indicating strong alignment with the diversity observed in expert annotations, outperforming all baselines.
• On the LIDC-IDRI dataset, CIMD reduces the Geometric Energy Distance (GED) to 0.321, outperforming DDPM-Prob-Seg (0.417) and showing better distributional alignment with ground truth.
• Qualitative results show that CIMD produces diverse, high-quality segmentations even in challenging cases such as small or ambiguous lesions in ultrasound and CT scans.
• Abllation studies confirm that the hierarchical stochastic injection is critical, as removing it leads to less diverse and less accurate outputs, even when using the same stochastic sampling process.
• The method generalizes across three modalities—CT, ultrasound, and MRI—demonstrating robustness and clinical relevance in diverse imaging contexts.