# Modeling Multimodal Aleatoric Uncertainty in Segmentation with Mixture of Stochastic Experts

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

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Builds on 43 of its 66 references that are in the map.
- First posted on arXiv 14 Dec 2022.

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

- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Image-to-Image Translation with Conditional Adversarial Networks](https://mlsanity.com/paper/1611.07004/) (2016)
- [The Cityscapes Dataset for Semantic Urban Scene Understanding](https://mlsanity.com/paper/1604.01685/) (2016)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation](https://mlsanity.com/paper/1802.02611/) (2018)
- [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)
- [V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation](https://mlsanity.com/paper/1606.04797/) (2016)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Concrete Problems in AI Safety](https://mlsanity.com/paper/1606.06565/) (2016)
- [Aleatoric and Epistemic Uncertainty in Machine Learning: An Introduction to Concepts and Methods](https://mlsanity.com/paper/1910.09457/) (2019)
- [Evidential Deep Learning to Quantify Classification Uncertainty](https://mlsanity.com/paper/1806.01768/) (2018)
- [Image Segmentation Using Deep Learning: A Survey](https://mlsanity.com/paper/2001.05566/) (2020)
- [Predictive Uncertainty Estimation via Prior Networks](https://mlsanity.com/paper/1802.10501/) (2018)
- [Beyond temperature scaling: Obtaining well-calibrated multiclass probabilities with Dirichlet calibration](https://mlsanity.com/paper/1910.12656/) (2019)
- [Toward Multimodal Image-to-Image Translation](https://mlsanity.com/paper/1711.11586/) (2017)
- [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)
- [Deep Unsupervised Clustering with Gaussian Mixture Variational Autoencoders](https://mlsanity.com/paper/1611.02648/) (2016)
- [Variational Deep Embedding: An Unsupervised and Generative Approach to Clustering](https://mlsanity.com/paper/1611.05148/) (2016)
- [Learning with a Wasserstein Loss](https://mlsanity.com/paper/1506.05439/) (2015)
- [PHiSeg: Capturing Uncertainty in Medical Image Segmentation](https://mlsanity.com/paper/1906.04045/) (2019)
- [Label Distribution Learning](https://mlsanity.com/paper/1408.6027/) (2014)
- [A Variational U-Net for Conditional Appearance and Shape Generation](https://mlsanity.com/paper/1804.04694/) (2018)
- [The Cramer Distance as a Solution to Biased Wasserstein Gradients](https://mlsanity.com/paper/1705.10743/) (2017)
- [Reverse KL-Divergence Training of Prior Networks: Improved Uncertainty and Adversarial Robustness](https://mlsanity.com/paper/1905.13472/) (2019)
- [Stochastic Segmentation Networks: Modelling Spatially Correlated Aleatoric Uncertainty](https://mlsanity.com/paper/2006.06015/) (2020)
- [Learning in an Uncertain World: Representing Ambiguity Through Multiple Hypotheses](https://mlsanity.com/paper/1612.00197/) (2016)
- [Learning with minibatch Wasserstein : asymptotic and gradient properties](https://mlsanity.com/paper/1910.04091/) (2019)

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

- [Stochastic Segmentation with Conditional Categorical Diffusion Models](https://mlsanity.com/paper/2303.08888/) (2023)
- [A Review of Bayesian Uncertainty Quantification in Deep Probabilistic Image Segmentation](https://mlsanity.com/paper/2411.16370/) (2024)
- [Extracting Uncertainty Estimates from Mixtures of Experts for Semantic Segmentation](https://mlsanity.com/paper/2509.04816/) (2025)
- [Flow Stochastic Segmentation Networks](https://mlsanity.com/paper/2507.18838/) (2025)
- [DiffOSeg: Omni Medical Image Segmentation via Multi-Expert Collaboration Diffusion Model](https://mlsanity.com/paper/2507.13087/) (2025)
- [Sequential Amodal Segmentation via Cumulative Occlusion Learning](https://mlsanity.com/paper/2405.05791/) (2024)
- [Adaptive Individual Uncertainty under Out-Of-Distribution Shift with Expert-Routed Conformal Prediction](https://mlsanity.com/paper/2510.15233/) (2025)
- [Deterministic Mode Proposals: An Efficient Alternative to Generative Sampling for Ambiguous Segmentation](https://mlsanity.com/paper/2603.20191/) (2026)
- [Information-Theoretic Optimization for Task-Adapted Compressed Sensing Magnetic Resonance Imaging](https://mlsanity.com/paper/2604.12709/) (2026)
- [Neural Architecture Distributions: A New Paradigm for Stochastic Segmentation](https://mlsanity.com/paper/2606.21061/) (2026)
- [SARFA: Segment Anything with Radiomic Feature Alignment](https://mlsanity.com/paper/2607.13323/) (2026)
- [KANResDiff: Learning Local Residual Diffusion via Kolmogorov-Arnold Network for Ambiguous Medical Image Segmentation](https://mlsanity.com/paper/2608.11617/) (2026)
- [Latent-to-Latent Flow for Volumetric Stochastic Segmentation](https://mlsanity.com/paper/2609.07460/) (2026)

## Abstract (arXiv)

Equipping predicted segmentation with calibrated uncertainty is essential for safety-critical applications. In this work, we focus on capturing the data-inherent uncertainty (aka aleatoric uncertainty) in segmentation, typically when ambiguities exist in input images. Due to the high-dimensional output space and potential multiple modes in segmenting ambiguous images, it remains challenging to predict well-calibrated uncertainty for segmentation. To tackle this problem, we propose a novel mixture of stochastic experts (MoSE) model, where each expert network estimates a distinct mode of the aleatoric uncertainty and a gating network predicts the probabilities of an input image being segmented in those modes. This yields an efficient two-level uncertainty representation. To learn the model, we develop a Wasserstein-like loss that directly minimizes the distribution distance between the MoSE and ground truth annotations. The loss can easily integrate traditional segmentation quality measures and be efficiently optimized via constraint relaxation. We validate our method on the LIDC-IDRI dataset and a modified multimodal Cityscapes dataset. Results demonstrate that our method achieves the state-of-the-art or competitive performance on all metrics.