A Recycling Training Strategy for Medical Image Segmentation with Diffusion Denoising Models-猫眼课题宝

DOI：

10.59275/j.melba.2023-fbe4

发表时间：

2023-08

期刊：

ArXiv

影响因子：

通讯作者：

中科院分区：

其他

文献类型：

作者：

Yunguan Fu;Yiwen Li;Shaheer U. Saeed;M. Clarkson;Yipeng Hu

去噪扩散模型通过生成以图像为条件的分割掩模而在图像分割中得到应用。现有的研究主要集中在调整模型结构或改善推理，如测试时的抽样策略。在这项工作中，我们专注于改进的训练策略，并提出了一种新的回收方法。在每个训练步骤中，首先在给定图像和随机噪声的情况下预测分割掩码。这个预测的掩模取代了传统的地面真实掩模，用于训练期间的去噪任务。这种方法可以被解释为通过消除对用于生成噪声样本的地面真值掩码的依赖来将训练策略与推理对齐。我们提出的方法在多个医学成像数据集上显著优于标准扩散训练，自调节和现有的回收策略：肌肉超声，腹部CT，前列腺MR和脑MR。这适用于两种广泛采用的采样策略：去噪扩散概率模型和去噪扩散隐式模型。重要的是，现有的扩散模型在推理过程中往往表现出下降或不稳定的性能，而我们的新循环始终增强或保持性能。我们首次表明，在与相同的网络架构和计算预算进行公平比较的情况下，所提出的基于循环的扩散模型与非基于扩散的监督训练取得了同等的性能。此外，通过将所提出的扩散模型和非扩散模型相结合，在所有应用中都观察到非扩散模型的显著改进，证明了这种新型训练方法的价值。本文总结了这些定量结果，并讨论了它们的价值，并在https://github.com/mathpluscode/ImgX-DiffSeg上发布了一个完全可重现的基于JAX的实现

Denoising diffusion models have found applications in image segmentation by generating segmented masks conditioned on images. Existing studies predominantly focus on adjusting model architecture or improving inference, such as test-time sampling strategies. In this work, we focus on improving the training strategy and propose a novel recycling method. During each training step, a segmentation mask is first predicted given an image and a random noise. This predicted mask, which replaces the conventional ground truth mask, is used for denoising task during training. This approach can be interpreted as aligning the training strategy with inference by eliminating the dependence on ground truth masks for generating noisy samples. Our proposed method significantly outperforms standard diffusion training, self-conditioning, and existing recycling strategies across multiple medical imaging data sets: muscle ultrasound, abdominal CT, prostate MR, and brain MR. This holds for two widely adopted sampling strategies: denoising diffusion probabilistic model and denoising diffusion implicit model. Importantly, existing diffusion models often display a declining or unstable performance during inference, whereas our novel recycling consistently enhances or maintains performance. We show for the first time that, under a fair comparison with the same network architectures and computing budget, the proposed recycling-based diffusion models achieved on-par performance with non-diffusion-based supervised training. Furthermore, by ensembling the proposed diffusion model and the non-diffusion counterpart, significant improvements to the non-diffusion models have been observed across all applications, demonstrating the value of this novel training method. This paper summarizes these quantitative results and discusses their values, with a fully reproducible JAX-based implementation, released at https://github.com/mathpluscode/ImgX-DiffSeg

A Recycling Training Strategy for Medical Image Segmentation with Diffusion Denoising Models

文献摘要

相似文献