Robust semi-supervised segmentation with timestep ensembling diffusion models

TL;DR

This paper introduces a semi-supervised medical image segmentation method using diffusion models, emphasizing the benefits of small diffusion steps for robustness and proposing an ensembling scheme that improves domain generalization.

Contribution

It demonstrates that smaller diffusion steps produce more robust representations and proposes an ensembling scheme leveraging both small and large steps for better domain-shift performance.

Findings

Smaller diffusion steps yield more robust latent representations.

The proposed ensembling scheme improves domain-shift segmentation performance.

The method maintains competitive in-domain performance.

Abstract

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.