Abstract: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 <a class="link-external link-https" href="https://github.com/mathpluscode/ImgX-DiffSeg" rel="external noopener nofollow">this https URL</a>.

Accelerating Diffusion Models Via Pre-segmentation Diffusion Sampling for Medical Image Segmentation

Accelerating denoising diffusion probabilistic model via truncated inverse processes for medical image segmentation

Multi-Step Denoising Scheduled Sampling: Towards Alleviating Exposure Bias for Diffusion Models

Denoising Diffusions in Latent Space for Medical Image Segmentation

Edge-and-Mask Integration-Driven Diffusion Models for Medical Image Segmentation

Fast-DDPM: Fast Denoising Diffusion Probabilistic Models for Medical Image-to-Image Generation

Accelerating Diffusion Models via Early Stop of the Diffusion Process

PDDM: Poisson-Distribution Diffusion Model with Multi-Step Weighted Fusion for Medical Image Segmentation

PGDiffSeg: Prior-Guided Denoising Diffusion Model with Parameter-Shared Attention for Breast Cancer Segmentation

MedSegDiff: Medical Image Segmentation with Diffusion Probabilistic Model

Explicit-implicit priori knowledge-based diffusion model for generative medical image segmentation

BerDiff: Conditional Bernoulli Diffusion Model for Medical Image Segmentation

Cold SegDiffusion: A Novel Diffusion Model for Medical Image Segmentation

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

Analysing Diffusion Segmentation for Medical Images

Importance of Aligning Training Strategy with Evaluation for Diffusion Models in 3D Multiclass Segmentation

Diffusion model-based text-guided enhancement network for medical image segmentation

Conditional Diffusion Models for Weakly Supervised Medical Image Segmentation

Denoising Diffusion Medical Models

DiffuseExpand: Expanding dataset for 2D medical image segmentation using diffusion models

Stable Diffusion Segmentation for Biomedical Images with Single-step Reverse Process