Robust semi-supervised segmentation with timestep ensembling diffusion models

Margherita Rosnati, Mélanie Roschewitz, Ben Glocker
Proceedings of the 3rd Machine Learning for Health Symposium, PMLR 225:512-527, 2023.

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 ensembling 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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v225-rosnati23a, title = {Robust semi-supervised segmentation with timestep ensembling diffusion models}, author = {Rosnati, Margherita and Roschewitz, M\'elanie and Glocker, Ben}, booktitle = {Proceedings of the 3rd Machine Learning for Health Symposium}, pages = {512--527}, year = {2023}, editor = {Hegselmann, Stefan and Parziale, Antonio and Shanmugam, Divya and Tang, Shengpu and Asiedu, Mercy Nyamewaa and Chang, Serina and Hartvigsen, Tom and Singh, Harvineet}, volume = {225}, series = {Proceedings of Machine Learning Research}, month = {10 Dec}, publisher = {PMLR}, pdf = {https://proceedings.mlr.press/v225/rosnati23a/rosnati23a.pdf}, url = {https://proceedings.mlr.press/v225/rosnati23a.html}, 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 ensembling 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.} }
Endnote
%0 Conference Paper %T Robust semi-supervised segmentation with timestep ensembling diffusion models %A Margherita Rosnati %A Mélanie Roschewitz %A Ben Glocker %B Proceedings of the 3rd Machine Learning for Health Symposium %C Proceedings of Machine Learning Research %D 2023 %E Stefan Hegselmann %E Antonio Parziale %E Divya Shanmugam %E Shengpu Tang %E Mercy Nyamewaa Asiedu %E Serina Chang %E Tom Hartvigsen %E Harvineet Singh %F pmlr-v225-rosnati23a %I PMLR %P 512--527 %U https://proceedings.mlr.press/v225/rosnati23a.html %V 225 %X 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 ensembling 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.
APA
Rosnati, M., Roschewitz, M. & Glocker, B.. (2023). Robust semi-supervised segmentation with timestep ensembling diffusion models. Proceedings of the 3rd Machine Learning for Health Symposium, in Proceedings of Machine Learning Research 225:512-527 Available from https://proceedings.mlr.press/v225/rosnati23a.html.

Related Material