[Paper Review] Diffusion Models for Implicit Image Segmentation Ensembles

Julia Wolleb, Robin Sandkühler+3 more|arXiv (Cornell University)|Dec 6, 2021

Medical Image Segmentation Techniques71 citations

TL;DR

The paper trains a diffusion probabilistic model to generate brain tumor segmentation masks and leverages stochastic sampling to form an implicit ensemble, producing uncertainty maps and improved segmentation performance on BRATS2020.

ABSTRACT

Diffusion models have shown impressive performance for generative modelling of images. In this paper, we present a novel semantic segmentation method based on diffusion models. By modifying the training and sampling scheme, we show that diffusion models can perform lesion segmentation of medical images. To generate an image specific segmentation, we train the model on the ground truth segmentation, and use the image as a prior during training and in every step during the sampling process. With the given stochastic sampling process, we can generate a distribution of segmentation masks. This property allows us to compute pixel-wise uncertainty maps of the segmentation, and allows an implicit ensemble of segmentations that increases the segmentation performance. We evaluate our method on the BRATS2020 dataset for brain tumor segmentation. Compared to state-of-the-art segmentation models, our approach yields good segmentation results and, additionally, detailed uncertainty maps.

Motivation & Objective

• Motivate uncertainty quantification in medical image segmentation and improve lesion segmentation performance.
• Propose a diffusion-model-based method that uses the input image as an atlas prior during training and sampling.
• Enable implicit ensemble generation of segmentation masks through stochastic diffusion sampling.
• Provide uncertainty maps via pixel-wise variance across multiple segmentation samples.

Proposed method

• Train a DDPM to generate segmentation masks by adding noise to ground-truth masks while incorporating the input MR image as a prior through concatenation.
• Modify the DDPM sampling equation to condition on the image prior X_t = b irst(x_b,t) and generate a mean/variance segmentation ensemble.
• Use a linear noise schedule with T=1000 steps and a U-Net architecture to predict the denoising step epsilon_theta.
• Compute pixel-wise uncertainty maps by sampling multiple segmentation masks and calculating their variance.
• Produce an ensemble by averaging multiple sampled masks and thresholding to obtain a final segmentation.
• Benchmark against nnU-Net and SegNet on BRATS2020 and report Dice, Jaccard, and HD95 metrics.

Experimental results

Research questions

• RQ1 Can diffusion models be adapted for semantic segmentation in medical images by conditioning on the corresponding image as a prior?
• RQ2 Does generating multiple segmentation samples via stochastic diffusion sampling yield improved segmentation metrics compared to single-sample predictions?
• RQ3 How do diffusion-based segmentations compare to state-of-the-art methods (nnU-Net, SegNet) on BRATS2020 in terms of accuracy and uncertainty visualization?

Key findings

• Single-sample DDPM segmentation achieves Dice 0.866 and Jaccard 0.795 with HD95 6.052.
• Ensembling 5 diffusion samples improves Dice to 0.881 and Jaccard to 0.819, with HD95 5.178.
• nnU-Net achieves Dice 0.891 and Jaccard 0.831 with HD95 5.004.
• SegNet yields Dice 0.839 and Jaccard 0.761 with HD95 7.190.
• Bayesian SegNet (MC dropout, 5 samples) yields Dice 0.838 and Jaccard 0.747 with HD95 13.707.
• Ensemble mean maps provide improved segmentation performance and enable detailed uncertainty visualization.