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 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. By ensembling the proposed diffusion and the non-diffusion models, 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.

翻译：扩散去噪模型通过生成以图像为条件的分割掩膜，在图像分割领域得到了应用。现有研究主要集中于调整模型架构或改进推理过程，例如测试阶段的采样策略。本文聚焦于改进训练策略，提出了一种新颖的循环训练方法。在每个训练步骤中，首先根据输入图像和随机噪声预测分割掩膜。该预测掩膜将替代传统的地面真值掩膜，用于训练过程中的去噪任务。该方法可视为通过消除生成噪声样本对地面真值掩膜的依赖，使训练策略与推理过程对齐。我们提出的方法在多个医学影像数据集（肌肉超声、腹部CT、前列腺磁共振和脑部磁共振）上显著优于标准扩散训练、自条件训练及现有循环训练策略。这一优势同样适用于两种广泛采用的采样策略：去噪扩散概率模型和去噪扩散隐式模型。值得注意的是，现有扩散模型在推理过程中常出现性能下降或不稳定的现象，而我们的新型循环方法能持续增强或保持性能。实验表明，在相同网络架构和计算预算的公平对比下，基于循环训练的扩散模型达到了与非扩散监督训练相当的性能。通过集成所提扩散模型与非扩散模型，所有应用场景中非扩散模型的性能均获得显著提升，充分证明了这一新型训练方法的实用价值。本文总结了这些量化结果并探讨其意义，同时发布了基于JAX的完全可复现实现代码（https://github.com/mathpluscode/ImgX-DiffSeg）。