[Paper Review] BerDiff: Conditional Bernoulli Diffusion Model for Medical Image Segmentation

Tao Chen, Chenhui Wang+1 more|arXiv (Cornell University)|Apr 10, 2023

Radiomics and Machine Learning in Medical Imaging26 references4 citations

TL;DR

This paper proposes BerDiff, a conditional Bernoulli diffusion model for medical image segmentation that uses Bernoulli noise as the diffusion kernel instead of Gaussian noise, enabling more accurate and diverse binary segmentation masks. By leveraging the stochastic reverse process, BerDiff generates multiple plausible masks highlighting salient regions, outperforming state-of-the-art methods on LIDC-IDRI and BRATS 2021 datasets with a 89.7% Dice score.

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.

Motivation & Objective

• To address the inherent ambiguity and uncertainty in medical image segmentation, especially due to unclear boundaries and multiple plausible annotations.
• To improve segmentation accuracy and diversity by leveraging the stochastic nature of diffusion models for clinical decision support.
• To develop a diffusion model tailored for discrete binary segmentation tasks, moving beyond Gaussian noise assumptions.
• To enable efficient inference through sub-sequence sampling from the reverse diffusion trajectory.
• To demonstrate superiority over Gaussian-based diffusion models and existing SOTA methods in medical segmentation.

Proposed method

• Proposes a novel conditional diffusion model using Bernoulli noise as the diffusion kernel, replacing Gaussian noise to better model discrete segmentation masks.
• Defines a Markov chain forward process where each step adds Bernoulli noise with a schedule β₁,…,βₜ, gradually transforming the ground-truth mask into pure noise.
• Uses a denoising U-Net architecture to estimate the noise in the reverse process, conditioned on the input image.
• Enables diverse sampling by reinitializing the reverse process with different random seeds for the initial Bernoulli noise and intermediate latent variables.
• Introduces sub-sequence sampling from the full reverse trajectory based on DDIM principles, accelerating inference without sacrificing quality.
• Trains the model end-to-end using a variational lower bound objective to optimize the denoising score function.

Experimental results

Research questions

• RQ1 Can Bernoulli noise as a diffusion kernel improve the performance of diffusion models in discrete binary medical image segmentation tasks?
• RQ2 How does the stochastic reverse process of a Bernoulli diffusion model generate diverse yet plausible segmentation masks for clinical reference?
• RQ3 Can sub-sequence sampling from the reverse trajectory significantly reduce inference time while preserving segmentation quality?
• RQ4 How does BerDiff compare to Gaussian-based diffusion models and other SOTA methods in terms of accuracy and diversity on real medical imaging datasets?
• RQ5 Does the use of Bernoulli noise lead to better convergence and performance, especially under limited training iterations?

Key findings

• BerDiff achieves a Dice score of 89.7% on the LIDC-IDRI dataset, outperforming previous SOTA methods including SegDiff (89.3%) and U-Net (89.2%).
• On the BRATS 2021 dataset, BerDiff surpasses both traditional U-Net and transformer-based models like Swin UNETR, demonstrating superior performance in complex tumor segmentation.
• The model generates diverse segmentation masks that highlight salient regions of interest, with saliency masks effectively capturing ground-truth structures even in ambiguous cases like small lung nodules.
• Abllation studies show that Bernoulli-based diffusion outperforms Gaussian-based diffusion, especially when training iterations are sufficient (e.g., 86,500 steps).
• Sub-sequence sampling enables faster inference by skipping intermediate steps, maintaining high-quality outputs while reducing time per sample.
• Visual comparisons show BerDiff produces more accurate segmentations on hard-to-detect regions, such as subtle tumor boundaries in MRI scans, compared to other diffusion and non-diffusion models.