# Robust semi-supervised segmentation with timestep ensembling diffusion models

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

## Key facts

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

## Placed in

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

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

- [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)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [ChestX-ray8: Hospital-scale Chest X-ray Database and Benchmarks on Weakly-Supervised Classification and Localization of Common Thorax Diseases](https://mlsanity.com/paper/1705.02315/) (2017)
- [Identifying the Best Machine Learning Algorithms for Brain Tumor Segmentation, Progression Assessment, and Overall Survival Prediction in the BRATS Challenge](https://mlsanity.com/paper/1811.02629/) (2018)
- [Label-Efficient Semantic Segmentation with Diffusion Models](https://mlsanity.com/paper/2112.03126/) (2021)
- [Contrastive learning of global and local features for medical image segmentation with limited annotations](https://mlsanity.com/paper/2006.10511/) (2020)
- [Diffusion Models for Medical Anomaly Detection](https://mlsanity.com/paper/2203.04306/) (2022)
- [Fast Unsupervised Brain Anomaly Detection and Segmentation with Diffusion Models](https://mlsanity.com/paper/2206.03461/) (2022)
- [Measurement-conditioned Denoising Diffusion Probabilistic Model for Under-sampled Medical Image Reconstruction](https://mlsanity.com/paper/2203.03623/) (2022)
- [Diffusion Deformable Model for 4D Temporal Medical Image Generation](https://mlsanity.com/paper/2206.13295/) (2022)
- [Towards performant and reliable undersampled MR reconstruction via diffusion model sampling](https://mlsanity.com/paper/2203.04292/) (2022)
- [XLSor: A Robust and Accurate Lung Segmentor on Chest X-Rays Using Criss-Cross Attention and Customized Radiorealistic Abnormalities Generation](https://mlsanity.com/paper/1904.09229/) (2019)
- [Learning Data Representations with Joint Diffusion Models](https://mlsanity.com/paper/2301.13622/) (2023)
- [Multi-Modal Hypergraph Diffusion Network with Dual Prior for Alzheimer Classification](https://mlsanity.com/paper/2204.02399/) (2022)
- [Analysing the effectiveness of a generative model for semi-supervised medical image segmentation](https://mlsanity.com/paper/2211.01886/) (2022)

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

- [Joint Diffusion models in Continual Learning](https://mlsanity.com/paper/2411.08224/) (2024)
- [Denoising Diffusion Models for Anomaly Localization in Medical Images](https://mlsanity.com/paper/2410.23834/) (2024)
- [Hybrid diffusion models: combining supervised and generative pretraining for label-efficient fine-tuning of segmentation models](https://mlsanity.com/paper/2408.03433/) (2024)
- [Learning Dynamic Collaborative Network for Semi-supervised 3D Vessel Segmentation](https://mlsanity.com/paper/2601.07377/) (2026)

## Abstract (arXiv)

Medical image segmentation is a challenging task, made more difficult by many datasets' limited size and annotations. Denoising diffusion probabilistic models (DDPM) have recently shown promise in modelling the distribution of natural images and were successfully applied to various medical imaging tasks. This work focuses on semi-supervised image segmentation using diffusion models, particularly addressing domain generalisation. Firstly, we demonstrate that smaller diffusion steps generate latent representations that are more robust for downstream tasks than larger steps. Secondly, we use this insight to propose an improved esembling scheme that leverages information-dense small steps and the regularising effect of larger steps to generate predictions. Our model shows significantly better performance in domain-shifted settings while retaining competitive performance in-domain. Overall, this work highlights the potential of DDPMs for semi-supervised medical image segmentation and provides insights into optimising their performance under domain shift.