Reproduce “A Probabilistic U-Net for Segmentation of Ambiguous Images” on OpenResearch

About this paper

Simon A. A. Kohl, Bernardino Romera-Paredes, Clemens Meyer, Jeffrey De Fauw, Joseph R. Ledsam, Klaus H. Maier-Hein, S. M. Ali Eslami, Danilo Jimenez Rezende and 1 more

DeepMind researchers developed a Probabilistic U-Net, a generative model that learns and reproduces a distribution over diverse, consistent segmentation maps for ambiguous images. The model quantitatively outperformed baselines on medical imaging and synthetic urban scene datasets, accurately capturing the occurrence frequencies of multiple segmentation variants, including rare ones.

Key insights

• Injecting latent features at the end of the U-Net's decoder path allows for efficient sampling of multiple hypotheses by only re-evaluating a small part of the network, while enabling the U-Net core to focus on general feature extraction.
• Conditioning both the prior and posterior networks on the input image is crucial for accurately learning and capturing the true underlying distribution of segmentation variants, ensuring the model's prior distribution effectively covers the space of all possibilities.
• The Generalized Energy Distance (GED) metric provides a robust quantitative framework for evaluating generative models producing structured multi-modal outputs, effectively penalizing both inaccuracies in individual segmentations and the failure to represent existing modes.

Summary by alphaXiv.