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、前列腺MR和脑MR）上显著优于标准扩散训练、自条件化及现有回收策略。这一优势在两种广泛采用的采样策略——去噪扩散概率模型与去噪扩散隐模型——中均成立。值得注意的是，现有扩散模型在推理时常出现性能下降或不稳定现象，而本研究的回收方法能持续增强或保持性能。实验表明，在相同的网络架构与计算预算下进行公平比较，基于回收的扩散模型可达到与非扩散监督训练相当的性能。通过集成所提扩散模型与非扩散模型，所有应用场景中非扩散模型均获得显著提升，凸显了该新型训练方法的价值。本文汇总了这些量化结果，讨论了其意义，并提供了完全可复现的JAX实现代码（发布于https://github.com/mathpluscode/ImgX-DiffSeg）。