[论文解读] 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 Imaging被引用 7

一句话总结

本文提出CIMD，一种基于扩散模型的框架，用于模糊医学图像分割，通过利用扩散采样固有的随机性，从单个输入生成多个多样化且逼真的分割掩码。与需要独立先验网络的先前方法不同，CIMD在多个分层级别注入随机性，实现了在准确率（Dice分数：0.915）和多样性（CI分数：0.759）方面的最先进性能，同时保留了临床相关的变异。

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.

研究动机与目标

• 为解决确定性深度学习模型在医学图像分割中的局限性，这些模型尽管临床诊断中存在固有的阅片者间变异性，却仅生成单一掩码。
• 使用单一概率框架建模来自专家放射科医生的多个合理标注的分布，而非模仿其中最佳个体。
• 通过将随机性直接嵌入扩散模型的分层架构中，消除基于c-VAE方法中对独立先验网络的需求。
• 开发一种新的评估指标（CI分数），通过衡量与真实专家标注的多样性相对关系，使其更符合临床实践。
• 在多种模态（CT、超声和MRI）中展示扩散模型在模糊分割中的有效性。

提出的方法

• 所提出的CIMD框架使用一个条件扩散概率模型，该模型被训练以逆转在潜在空间中逐步向分割掩码添加噪声的前向扩散过程。
• 推理期间的随机采样自然地从单个输入生成多个合理的分割掩码，而无需从独立的先验网络中进行额外的推理时间采样。
• 该模型在多个分层特征级别注入噪声，从而在分割图的不同尺度上实现多样化且一致的预测。
• 扩散模型在每个输入图像的多个标注上进行端到端训练，学习合理分割的联合分布。
• 一种新颖的损失函数结合了像素级重建损失与对比损失，以在生成的掩码中同时保留准确率和多样性。
• 该方法与现有的基于扩散模型的分割框架兼容，且仅需极少的架构修改。

实验结果

研究问题

• RQ1 单个扩散模型是否能够在无需辅助先验网络的情况下，生成多个多样化且临床合理的分割掩码？
• RQ2 在扩散模型中分层注入随机性如何影响模糊分割输出的多样性与准确率？
• RQ3 扩散模型在多大程度上能够学习并再现不同医学影像模态中专家标注的频率分布？
• RQ4 所提出的方法是否在真实世界医学影像数据集上，于准确率和多样性方面均优于现有的基于c-VAE的方法和确定性分割网络？
• RQ5 一种能同时捕捉准确率与多样性的新评估指标，是否能比仅使用Dice等标准指标更好地反映临床相关性？

主要发现

• 在LIDC-IDRI数据集上，CIMD实现了最大Dice分数（D_max）0.915，显著优于次佳方法（Probabilistic U-Net），并表现出与真实标注分布更优的匹配度。
• 在LIDC-IDRI数据集上，该模型实现了0.759的临床多样性（CI）分数，表明其与专家标注中观察到的多样性高度一致，优于所有基线方法。
• 在LIDC-IDRI数据集上，CIMD将几何能量距离（GED）降低至0.321，优于DDPM-Prob-Seg（0.417），显示出与真实标注更优的分布对齐。
• 定性结果表明，CIMD即使在超声和CT扫描中存在小病灶或模糊病灶等挑战性情况下，也能生成多样化且高质量的分割结果。
• 消融研究证实，分层随机注入至关重要，因为即使使用相同的随机采样过程，若去除该机制，输出的多样性与准确率均会下降。
• 该方法在三种模态（CT、超声和MRI）上均表现出泛化能力，展示了在多样化影像场景下的稳健性与临床相关性。