PaperlayerView exact revision on arXiv

Source-linked AI summary

Ambiguous Medical Image Segmentation using Diffusion Models

Aimon Rahman, Jeya Maria Jose Valanarasu, Ilker Hacihaliloglu, Vishal M Patel

arXiv:2304.04745v1cs.CV

TL;DR

Medical image segmentation must represent disagreement among experts, but prior AI methods largely target a single best-looking annotation rather than a distribution of collective insights. CIMD uses one diffusion model’s stochastic and hierarchical sampling to generate plausible masks without a separate inference-time prior encoder, and pairs this with the CI Score. Across CT, ultrasound, and MRI, the paper reports performance on par or better than existing ambiguous segmentation networks while preserving heterogeneous outputs.

  • Problem

    Most segmentation networks produce one mask even though diagnosticians can disagree, leaving automated methods less able to represent collective clinical insight.

  • Method

    CIMD is a diffusion-based probabilistic framework that models segmentation heterogeneity through stochastic hierarchical sampling without an additional prior-information network during inference.

  • Results

    Across three medical imaging modalities, CIMD performs on par or better than existing ambiguous segmentation networks while producing superior qualitative results.

  • Takeaways & Limitations

    The CI Score evaluates prediction accuracy and diversity in a way motivated by collective clinical assessment and sensitivity.

Abstract

from arXiv · show

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.

1. Introduction

Medical image segmentation is inherently ambiguous because diagnosticians may disagree, while most segmentation networks produce only one mask. The paper proposes CIMD, a diffusion-based model that generates plausible annotation distributions without a separate prior encoder and introduces the CI Score for clinically aligned evaluation.

  • Motivation: Deterministic segmentation networks predict a single mask despite legitimate differences among diagnosticians’ interpretations.These disagreements can produce low inter-rater agreement and make pixel-wise probabilities insufficient for representing correlated ambiguity.
  • Motivation: Clinical workflows such as case conferences and tumor boards show the value of combining multiple expert assessments.The paper frames collective intelligence as a model for safer, more accurate automated diagnosis.
  • Motivation: Existing ambiguous segmentation methods use c-VAE-style latent sampling, but deterministic models may select only the most likely hypothesis and pixel-wise uncertainty can be inconsistent.The paper identifies a need for methods that capture realistic variation across complete segmentation masks.
  • Proposed approach: CIMD uses diffusion stochasticity and hierarchical structure to model heterogeneous segmentation masks without an additional prior-information network during inference.The diffusion process learns segmentation structure and dataset ambiguity while allowing diversity control across sampling steps.
  • Evaluation: The paper introduces the CI Score because GED can reward sample diversity even when generated masks poorly match ground-truth distributions.The proposed evaluation is motivated by collective clinical assessment and sensitivity-oriented diagnosis.
  • Results: Across three medical imaging modalities, CIMD performs on par or better than existing ambiguous segmentation networks while producing superior qualitative results.The contribution is presented as preserving heterogeneous predictions without requiring an extra inference-time prior network.

2. Related Work

Prior ambiguous segmentation methods model multiple annotations through Bayesian or c-VAE-like stochasticity, but their output diversity and latent distributions can be insufficiently expressive. Hierarchical latent designs improve diversity, yet the related work identifies limitations in realism and correlation structure.

  • Ambiguous segmentation models: Approximate Bayesian inference methods can produce pixel-level variation without capturing the complex correlation structure of ground-truth distributions.This limits their ability to represent structured ambiguity in medical segmentation.
  • c-VAE-based methods: Probabilistic U-Net generates multiple hypotheses with a conditional variational autoencoder, but stochasticity injected only at the highest resolution can yield limited diversity.The backbone may ignore random latent draws when stochasticity is introduced too late.
  • c-VAE-based methods: PHi-SegNet and Hierarchical Probabilistic U-Net add hierarchical latent spaces to increase sample diversity.Their predictions still depend on variational inference and latent distributions that may not be sufficiently complex to produce realistic samples.
  • c-VAE-based methods: c-VAE-like models commonly rely on axis-aligned Gaussian latent posteriors that may be too restrictive and insufficiently expressive.The related work therefore highlights a representational limitation in modeling distributions over ambiguous annotations.

3. Proposed Method

The method conditions diffusion on the input image to generate diverse plausible segmentation masks, while auxiliary ambiguity modeling and a collective-insight metric address both prediction variation and clinical coverage.

  • Diffusion Model: Diffusion models learn a data distribution by adding Gaussian noise forward and denoising it through a learned reverse Markov process.The reverse network incrementally estimates cleaner states from noisy inputs across T timesteps.
  • Collectively Intelligent Medical Diffusion: Sampling with different Gaussian noise generates multiple plausible masks, while the hierarchical diffusion structure enables diversity control across timesteps.This stochastic process avoids requiring a separate prior distribution during inference.
  • Diffusion Model: The model concatenates the input image with noisy segmentation masks so sampling produces image-conditioned anatomical segmentations rather than arbitrary masks.At each timestep, the input image b is added to the noisy segmentation mask xb,t.
  • Collectively Intelligent Medical Diffusion: Ambiguity networks model latent distributions for ground-truth and predicted-mask ambiguity, and their KL divergence is added to the diffusion objective.The total loss combines Lsimple, Lvlb, and βLamb, with λ and β acting as regularizing weights.
  • Collective Insight Score: The CI score combines combined sensitivity, maximum Dice matching, and diversity agreement using a harmonic mean.Combined sensitivity uses unions of all ground-truth and prediction masks, while diversity agreement compares prediction and ground-truth variation.

4. Experiments

Experiments evaluate ambiguous segmentation across CT, ultrasound, and MRI using distribution-based metrics and qualitative comparisons. CIMD captures diverse, accurate masks, particularly for difficult ultrasound and small-lesion MRI cases.

  • Datasets: The experiments use LIDC-IDRI lung CT, Bone-US ultrasound, and MS-MRI datasets with multiple expert annotations.LIDC-IDRI test images have 4 expert gradings; Bone-US and MS-MRI provide additional multi-rater ambiguity settings.
  • Evaluation metrics: GED, CI, and Dmax evaluate agreement with the ground-truth distribution, with Dmax and CI emphasizing matching and diversity.Lower GED indicates better agreement, while high Dmax indicates matching and CI measures whether sample diversity matches ground-truth diversity.
  • Qualitative comparison: Qualitative comparisons sample the first 4 predicted masks against Probabilistic U-net, PHi-Seg, and available expert gradings.Bone-US examples include 1 expert and 3 novice gradings, whereas MS-MRI examples include 2 expert gradings.
  • Quantitative comparison: CIMD outperforms state-of-the-art ambiguous segmentation networks on Dmax and CI across all datasets, while matching Probabilistic U-net in GED.The evaluation uses 4 samples from each prediction distribution.
  • Qualitative comparison: CIMD performs especially well on ultrasound and captures small lesions in MRI while maintaining diversity across segmentation masks.The paper attributes this behavior to stochasticity injected at each hierarchical feature representation.

5. Discussion

The discussion uses ablations to separate CIMD’s contribution from diffusion sampling alone. Plain DDPM variants produce insufficiently diverse or poorly matched masks, whereas CIMD improves GED and CI.

  • Ablation design: The ablation compares CIMD with DDPM-det-Seg and DDPM-Prob-Seg to assess whether diffusion stochasticity alone models segmentation ambiguity.DDPM-det-Seg trains on averaged masks, while DDPM-Prob-Seg uses different masks for the same input across forward passes.
  • Ablation findings: DDPM sampling is stochastic, but its generated masks are not diverse enough and do not sufficiently resemble the ground-truth distribution.This limitation is observed for both DDPM-based segmentation variants on LIDC-IDRI.
  • Limitations: CIMD produces diverse, meaningful predictions but requires slow training and inference because of the diffusion process.The authors also identify the need for multiple radiologists’ annotations as costly and time-consuming.

6. Future work and Conclusion

The conclusion presents CIMD as a diffusion-based ambiguous segmentation network that generates multiple plausible annotations without a separate inference-time prior network. It also identifies extensions to broader vision tasks, modalities, and diffusion segmentation frameworks.

  • Conclusion: CIMD generates multiple plausible annotations from one input image by injecting stochasticity at each hierarchical level.The design requires minimal additional training and does not use a separate network to encode prior information during inference.
  • Future work: The framework can be incorporated into diffusion-based models and may be extended to general computer vision and additional medical imaging modalities.The paper specifically mentions MedSegDiff and SegDiff as possible host segmentation networks.

A. Appendix Ablation Study

Additional ablations on Bone-US and MS-MRI show that CIMD improves diffusion-model performance on both GED and CI. On LIDC-IDRI, CIMD better matches lesion presence and ground-truth variation than the DDPM variants.

  • Bone-US and MS-MRI: CIMD improves diffusion-model performance in terms of GED and CI on both Bone-US and MS-MRI.These findings are reported in the corresponding Bone-US and MS-MRI ablation tables.
  • LIDC-IDRI: On LIDC-IDRI, DDPM-det-Seg produces minimal mask changes despite stochastic sampling, while DDPM-Prob-Seg often generates empty masks.CIMD instead segments the lesion and produces masks matching the ground-truth distributions.
  • Interpretation: The ablation concludes that DDPM alone cannot model the dataset’s stochasticity, whereas CIMD captures both lesion structure and segmentation variation.The comparison uses DDPM models trained with averaged or multiple segmentation masks.

B. Appendix Network Architecture

AMN and ACN share an encoder-based architecture for modeling and controlling segmentation ambiguity. The network predicts Gaussian distributions for masks conditioned on the input image and predictions.

  • Architecture: AMN and ACN use repeated 3x3 convolution layers with 32, 64, 128, and 192 filters, each followed by ReLU and 2x2 average pooling.The encoder downsamples features using stride 2.
  • Architecture: A 1x1 convolution receives globally average-pooled feature maps and predicts the Gaussian distribution parameters.
  • Architecture: The prediction network concatenates the input image with segmentation predictions to estimate Gaussian mask distributions conditioned on the image.

C. Appendix Training details

The regularization parameter β balances the loss terms and the divergence between ACN and AMN. Its selected value avoids both noisy outputs and minimal variation.

  • Training details: β is empirically set to 0.001 for training.
  • Training details: Higher β values overwhelm the other loss terms and produce noisy outputs.
  • Training details: Lower β values ignore the KL divergence between ACN and AMN, yielding a regular diffusion model with minimal output variation.

D. Appendix Qualitative Result Analysis

CIMD consistently produces high-fidelity segmentation samples across LIDC-IDRI, Bone-US, and MS-MRI qualitative analyses. It also handles empty ultrasound annotations and fine MRI lesions more reliably than the compared baselines.

  • Average Segmentation Quality: CIMD consistently produces high-fidelity samples, whereas Prob-Unet and PHi-Seg sometimes generate incomplete or inconsistent masks.The comparison visualizes 16 samples per input image.
  • Empty Segmentation in Bone-US dataset: CIMD produces more consistent results with minimal error when Bone-US images contain no bone-surface response.Prob-Unet and PHi-Seg may segment random ultrasound contrast as bone.
  • Fine Lesion segmentation: CIMD captures the finest lesions in MS-MRI scans, which are difficult for the other networks to segment.

E. Choice of Distribution

The distribution choice controls how CIMD represents segmentation ambiguity. Although full covariance increases diversity, the axis-aligned Gaussian provides better agreement with the ground-truth distribution and qualitative outputs in this analysis.

  • Choice of Distribution: Prior approaches modeled ambiguity with diagonal-covariance multivariate Gaussians, while full covariance was hypothesized to increase sample diversity.
  • Choice of Distribution: The full covariance matrix is constructed as Σ = LLT using Cholesky decomposition, with L produced by a neural network.
  • Choice of Distribution: CIMD-FC and CIMD-AA are compared using full-covariance and axis-aligned Gaussian distributions, respectively.
  • Choice of Distribution: Samples from the full-covariance Gaussian use the reparameterization z = µ + L ∗ϵ, with ϵ drawn from N(0, I).
  • Choice of Distribution: CIMD-FC produces highly diverse but coarser outputs that are not always close to the ground-truth distribution.Its higher combined sensitivity skews the CI Score.
  • Choice of Distribution: Qualitative comparisons favor CIMD-AA over CIMD-FC and show higher-fidelity CIMD samples than Probabilistic U-net and PHi-Seg across LIDC-IDRI and Bone-US.For blank Bone-US annotations, both baselines struggle because noisy contrast resembles bone-surface response.
Loading 2304.04745v1…