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

Yunguan Fu1,20000-0002-1184-7421, Yiwen Li30000-0002-7794-9391, Shaheer U. Saeed10000-0002-5004-0663, Matthew J. Clarkson10000-0002-5565-1252, Yipeng Hu10000-0003-4902-0486
1: University College London, UK, 2: InstaDeep, UK, 3: University of Oxford, UK
Publication date: 2023/12/07
PDF · Code · arXiv


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


Image Segmentation · Diffusion Model · Recycling · Muscle Ultrasound · Abdominal CT · Prostate MR · Brain MR

Bibtex @article{melba:2023:016:fu, title = "A Recycling Training Strategy for Medical Image Segmentation with Diffusion Denoising Models", author = "Fu, Yunguan and Li, Yiwen and Saeed, Shaheer U. and Clarkson, Matthew J. and Hu, Yipeng", journal = "Machine Learning for Biomedical Imaging", volume = "2", issue = "Special Issue for Generative Models", year = "2023", pages = "507--546", issn = "2766-905X", doi = "https://doi.org/10.59275/j.melba.2023-fbe4", url = "https://melba-journal.org/2023:016" }
RISTY - JOUR AU - Fu, Yunguan AU - Li, Yiwen AU - Saeed, Shaheer U. AU - Clarkson, Matthew J. AU - Hu, Yipeng PY - 2023 TI - A Recycling Training Strategy for Medical Image Segmentation with Diffusion Denoising Models T2 - Machine Learning for Biomedical Imaging VL - 2 IS - Special Issue for Generative Models SP - 507 EP - 546 SN - 2766-905X DO - https://doi.org/10.59275/j.melba.2023-fbe4 UR - https://melba-journal.org/2023:016 ER -

2023:016 cover