[论文解读] 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 Imaging参考文献 26被引用 4

一句话总结

本文提出BerDiff，一种用于医学图像分割的条件伯努利扩散模型，采用伯努利噪声作为扩散核而非高斯噪声，从而实现更精确且多样的二值分割掩码。通过利用随机反向过程，BerDiff可生成多个合理的掩码，突出显示显著区域，在LIDC-IDRI和BRATS 2021数据集上取得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 使用伯努利噪声是否能带来更好的收敛性与性能表现，尤其是在训练迭代次数有限的情况下？

主要发现

• 在LIDC-IDRI数据集上，BerDiff取得89.7%的Dice分数，优于先前SOTA方法，包括SegDiff（89.3%）和U-Net（89.2%）。
• 在BRATS 2021数据集上，BerDiff超越了传统U-Net及基于Transformer的模型（如Swin UNETR），在复杂肿瘤分割任务中表现更优。
• 该模型可生成多样化的分割掩码，有效突出感兴趣区域，其显著性掩码即使在小肺结节等模糊情况下也能准确捕捉真实结构。
• 消融实验表明，基于伯努利的扩散方法优于基于高斯的扩散方法，尤其在训练迭代足够多时（如86,500步）表现更优。
• 子序列采样通过跳过中间步骤实现更快推理，同时保持高质量输出。
• 可视化对比显示，与其它扩散及非扩散模型相比，BerDiff在MRI扫描中难以检测区域（如细微肿瘤边界）的分割结果更精确。