Robust semi-supervised segmentation with timestep ensembling diffusion models
This work addresses domain generalization in medical image segmentation, which is crucial for clinical applications with limited annotated data, though it appears incremental as it builds on existing diffusion model methods.
The paper tackled the problem of semi-supervised medical image segmentation under domain shift by proposing an improved ensembling scheme using diffusion models, resulting in significantly better performance in domain-shifted settings while maintaining competitive in-domain results.
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.