Robust semi-supervised segmentation with timestep ensembling diffusion models

Margherita Rosnati, Mélanie Roschewitz, Ben Glocker

Keywords: Semi-supervised learning, Medical image segmentation, Diffusion models, domain shift, generalisation

Abstract: Medical image segmentation is challenging due to limited data and annotations. Denoising diffusion probabilistic models (DDPM) show promise in modelling natural image distributions and are successfully applied in medical imaging. Our research focuses on semi-supervised image segmentation using diffusion models' latent representations and addressing domain generalisation. We found that optimal performance depends on the choice of diffusion steps and ensembling. Our model outperformed in domain-shifted settings while remaining competitive within domain, highlighting DDPMs' potential for medical image segmentation.

Submission Number: 2