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

arXiv: https://arxiv.org/abs/2303.06040 · HTML: https://mlsanity.com/paper/2303.06040/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in Medical imaging and segmentation, one of 14 directions in Diffusion models.
- Builds on 24 of its 29 references that are in the map.
- First posted on arXiv 10 Mar 2023.

## Placed in

- Diffusion models / [Medical imaging and segmentation](https://mlsanity.com/diffusion/medical-imaging-and-segmentation/)

## Builds on (most cited first, 24 of 24)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Variational Diffusion Models](https://mlsanity.com/paper/2107.00630/) (2021)
- [Segmenter: Transformer for Semantic Segmentation](https://mlsanity.com/paper/2105.05633/) (2021)
- [Label-Efficient Semantic Segmentation with Diffusion Models](https://mlsanity.com/paper/2112.03126/) (2021)
- [AMOS: A Large-Scale Abdominal Multi-Organ Benchmark for Versatile Medical Image Segmentation](https://mlsanity.com/paper/2206.08023/) (2022)
- [SegDiff: Image Segmentation with Diffusion Probabilistic Models](https://mlsanity.com/paper/2112.00390/) (2021)
- [Brain Imaging Generation with Latent Diffusion Models](https://mlsanity.com/paper/2209.07162/) (2022)
- [Diffusion Models for Implicit Image Segmentation Ensembles](https://mlsanity.com/paper/2112.03145/) (2021)
- [MedSegDiff: Medical Image Segmentation with Diffusion Probabilistic Model](https://mlsanity.com/paper/2211.00611/) (2022)
- [A Generalist Framework for Panoptic Segmentation of Images and Videos](https://mlsanity.com/paper/2210.06366/) (2022)
- [Diffusion Deformable Model for 4D Temporal Medical Image Generation](https://mlsanity.com/paper/2206.13295/) (2022)
- [Medical Diffusion: Denoising Diffusion Probabilistic Models for 3D Medical Image Generation](https://mlsanity.com/paper/2211.03364/) (2022)
- [Spot the fake lungs: Generating Synthetic Medical Images using Neural Diffusion Models](https://mlsanity.com/paper/2211.00902/) (2022)
- [Multi-Class Segmentation from Aerial Views using Recursive Noise Diffusion](https://mlsanity.com/paper/2212.00787/) (2022)
- [Bi-parametric prostate MR image synthesis using pathology and sequence-conditioned stable diffusion](https://mlsanity.com/paper/2303.02094/) (2023)
- [More unlabelled data or label more data? A study on semi-supervised laparoscopic image segmentation](https://mlsanity.com/paper/1908.08035/) (2019)

## Built on by (most cited first, 5 of 5)

- [A Comprehensive Survey of Foundation Models in Medicine](https://mlsanity.com/paper/2406.10729/) (2024)
- [A Survey of Emerging Applications of Diffusion Probabilistic Models in MRI](https://mlsanity.com/paper/2311.11383/) (2023)
- [A Recycling Training Strategy for Medical Image Segmentation with Diffusion Denoising Models](https://mlsanity.com/paper/2308.16355/) (2023)
- [ResEnsemble-DDPM: Residual Denoising Diffusion Probabilistic Models for Ensemble Learning](https://mlsanity.com/paper/2312.01682/) (2023)
- [Rethinking Diffusion Segmentation: When Does It Rely on Its Noisy State, and Does Diffusion Matter?](https://mlsanity.com/paper/2609.23967/) (2026)

## Abstract (arXiv)

Recently, denoising diffusion probabilistic models (DDPM) have been applied to image segmentation by generating segmentation masks conditioned on images, while the applications were mainly limited to 2D networks without exploiting potential benefits from the 3D formulation. In this work, we studied the DDPM-based segmentation model for 3D multiclass segmentation on two large multiclass data sets (prostate MR and abdominal CT). We observed that the difference between training and test methods led to inferior performance for existing DDPM methods. To mitigate the inconsistency, we proposed a recycling method which generated corrupted masks based on the model's prediction at a previous time step instead of using ground truth. The proposed method achieved statistically significantly improved performance compared to existing DDPMs, independent of a number of other techniques for reducing train-test discrepancy, including performing mask prediction, using Dice loss, and reducing the number of diffusion time steps during training. The performance of diffusion models was also competitive and visually similar to non-diffusion-based U-net, within the same compute budget. The JAX-based diffusion framework has been released at https://github.com/mathpluscode/ImgX-DiffSeg.