# Investigating and Improving Latent Density Segmentation Models for Aleatoric Uncertainty Quantification in Medical Imaging

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 51 of its 93 references that are in the map.
- First posted on arXiv 31 Jul 2023.

## Builds on (most cited first, 30 of 51)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [A ConvNet for the 2020s](https://mlsanity.com/paper/2201.03545/) (2022)
- [What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?](https://mlsanity.com/paper/1703.04977/) (2017)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation](https://mlsanity.com/paper/2102.04306/) (2021)
- [Group Normalization](https://mlsanity.com/paper/1803.08494/) (2018)
- [Generating Sentences from a Continuous Space](https://mlsanity.com/paper/1511.06349/) (2015)
- [Weight Normalization: A Simple Reparameterization to Accelerate Training of Deep Neural Networks](https://mlsanity.com/paper/1602.07868/) (2016)
- [Deep Exploration via Bootstrapped DQN](https://mlsanity.com/paper/1602.04621/) (2016)
- [Wasserstein Auto-Encoders](https://mlsanity.com/paper/1711.01558/) (2017)
- [Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning](https://mlsanity.com/paper/2208.04202/) (2022)
- [A Probabilistic U-Net for Segmentation of Ambiguous Images](https://mlsanity.com/paper/1806.05034/) (2018)
- [Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks](https://mlsanity.com/paper/1807.07356/) (2018)
- [Bayesian SegNet: Model Uncertainty in Deep Convolutional Encoder-Decoder Architectures for Scene Understanding](https://mlsanity.com/paper/1511.02680/) (2015)
- [Diffusion Models for Implicit Image Segmentation Ensembles](https://mlsanity.com/paper/2112.03145/) (2021)
- [Diagnosing and Enhancing VAE Models](https://mlsanity.com/paper/1903.05789/) (2019)
- [MedSegDiff: Medical Image Segmentation with Diffusion Probabilistic Model](https://mlsanity.com/paper/2211.00611/) (2022)
- [Variational Lossy Autoencoder](https://mlsanity.com/paper/1611.02731/) (2016)
- [A Hierarchical Latent Variable Encoder-Decoder Model for Generating Dialogues](https://mlsanity.com/paper/1605.06069/) (2016)
- [Uncertainty Estimates and Multi-Hypotheses Networks for Optical Flow](https://mlsanity.com/paper/1802.07095/) (2018)
- [Inference Suboptimality in Variational Autoencoders](https://mlsanity.com/paper/1801.03558/) (2018)
- [PHiSeg: Capturing Uncertainty in Medical Image Segmentation](https://mlsanity.com/paper/1906.04045/) (2019)
- [InfoVAE: Information Maximizing Variational Autoencoders](https://mlsanity.com/paper/1706.02262/) (2017)

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

- [A Review of Bayesian Uncertainty Quantification in Deep Probabilistic Image Segmentation](https://mlsanity.com/paper/2411.16370/) (2024)
- [Harmonized Feature Conditioning and Frequency-Prompt Personalization for Multi-Rater Medical Segmentation](https://mlsanity.com/paper/2605.08210/) (2026)

## Abstract (arXiv)

Data uncertainties, such as sensor noise, occlusions or limitations in the acquisition method can introduce irreducible ambiguities in images, which result in varying, yet plausible, semantic hypotheses. In Machine Learning, this ambiguity is commonly referred to as aleatoric uncertainty. In image segmentation, latent density models can be utilized to address this problem. The most popular approach is the Probabilistic U-Net (PU-Net), which uses latent Normal densities to optimize the conditional data log-likelihood Evidence Lower Bound. In this work, we demonstrate that the PU-Net latent space is severely sparse and heavily under-utilized. To address this, we introduce mutual information maximization and entropy-regularized Sinkhorn Divergence in the latent space to promote homogeneity across all latent dimensions, effectively improving gradient-descent updates and latent space informativeness. Our results show that by applying this on public datasets of various clinical segmentation problems, our proposed methodology receives up to 11% performance gains compared against preceding latent variable models for probabilistic segmentation on the Hungarian-Matched Intersection over Union. The results indicate that encouraging a homogeneous latent space significantly improves latent density modeling for medical image segmentation.