# Generalized Probabilistic U-Net for medical image segementation

arXiv: https://arxiv.org/abs/2207.12872 · HTML: https://mlsanity.com/paper/2207.12872/

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 28 references that are in the map.
- First posted on arXiv 26 Jul 2022.

## Builds on (most cited first, 19 of 19)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?](https://mlsanity.com/paper/1703.04977/) (2017)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [An Introduction to Variational Autoencoders](https://mlsanity.com/paper/1906.02691/) (2019)
- [Variational Dropout and the Local Reparameterization Trick](https://mlsanity.com/paper/1506.02557/) (2015)
- [Tune: A Research Platform for Distributed Model Selection and Training](https://mlsanity.com/paper/1807.05118/) (2018)
- [A Probabilistic U-Net for Segmentation of Ambiguous Images](https://mlsanity.com/paper/1806.05034/) (2018)
- [Variational Deep Embedding: An Unsupervised and Generative Approach to Clustering](https://mlsanity.com/paper/1611.05148/) (2016)
- [PHiSeg: Capturing Uncertainty in Medical Image Segmentation](https://mlsanity.com/paper/1906.04045/) (2019)
- [A Hierarchical Probabilistic U-Net for Modeling Multi-Scale Ambiguities](https://mlsanity.com/paper/1905.13077/) (2019)
- [On the Effect of Inter-observer Variability for a Reliable Estimation of Uncertainty of Medical Image Segmentation](https://mlsanity.com/paper/1806.02562/) (2018)
- [Supervised Uncertainty Quantification for Segmentation with Multiple Annotations](https://mlsanity.com/paper/1907.01949/) (2019)
- [Uncertainty quantification in medical image segmentation with normalizing flows](https://mlsanity.com/paper/2006.02683/) (2020)
- [Improving Aleatoric Uncertainty Quantification in Multi-Annotated Medical Image Segmentation with Normalizing Flows](https://mlsanity.com/paper/2108.02155/) (2021)

## Built on by (most cited first, 10 of 10)

- [Ambiguous Medical Image Segmentation using Diffusion Models](https://mlsanity.com/paper/2304.04745/) (2023)
- [FlowSDF: Flow Matching for Medical Image Segmentation Using Distance Transforms](https://mlsanity.com/paper/2405.18087/) (2024)
- [Tyche: Stochastic In-Context Learning for Medical Image Segmentation](https://mlsanity.com/paper/2401.13650/) (2024)
- [Effect of latent space distribution on the segmentation of images with multiple annotations](https://mlsanity.com/paper/2304.13476/) (2023)
- [A Review of Bayesian Uncertainty Quantification in Deep Probabilistic Image Segmentation](https://mlsanity.com/paper/2411.16370/) (2024)
- [Investigating and Improving Latent Density Segmentation Models for Aleatoric Uncertainty Quantification in Medical Imaging](https://mlsanity.com/paper/2307.16694/) (2023)
- [Diffusion Based Ambiguous Image Segmentation](https://mlsanity.com/paper/2504.05977/) (2025)
- [Uncertainty quantification for White Matter Hyperintensity segmentation detects silent failures and improves automated Fazekas quantification](https://mlsanity.com/paper/2411.17571/) (2024)
- [Towards Agnostic and Holistic Universal Image Segmentation with Bit Diffusion](https://mlsanity.com/paper/2601.02881/) (2026)
- [Neural Architecture Distributions: A New Paradigm for Stochastic Segmentation](https://mlsanity.com/paper/2606.21061/) (2026)

## Abstract (arXiv)

We propose the Generalized Probabilistic U-Net, which extends the Probabilistic U-Net by allowing more general forms of the Gaussian distribution as the latent space distribution that can better approximate the uncertainty in the reference segmentations. We study the effect the choice of latent space distribution has on capturing the uncertainty in the reference segmentations using the LIDC-IDRI dataset. We show that the choice of distribution affects the sample diversity of the predictions and their overlap with respect to the reference segmentations. For the LIDC-IDRI dataset, we show that using a mixture of Gaussians results in a statistically significant improvement in the generalized energy distance (GED) metric with respect to the standard Probabilistic U-Net. We have made our implementation available at https://github.com/ishaanb92/GeneralizedProbabilisticUNet