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

M. M. Amaan Valiuddin; Christiaan G. A. Viviers; Ruud J. G. van Sloun; Peter H. N. de With; Fons van der Sommen

cs.CVJul 31, 2023

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

Authors: M. M. Amaan Valiuddin, Christiaan G. A. Viviers, Ruud J. G. van Sloun, Peter H. N. de With, Fons van der Sommen

Organizations: Department of Electrical Engineering, Eindhoven University of Technology, Eindhoven, 5612 AZ, The Netherlands · Phillips Research, Eindhoven, The Netherlands

Abstract

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 modify the training objective to introduce the entropy-regularized Sinkhorn Divergence in the latent space, which promotes homogeneity across all latent dimensions, effectively improving gradient-descent updates and latent space informativeness. Applied to public datasets of various clinical segmentation problems, our proposed methodology achieves the best Empirical Wasserstein Distance on every dataset compared against preceding latent variable models for probabilistic segmentation. The results indicate that encouraging a homogeneous latent space significantly improves latent density modeling for medical image segmentation.

Figures & tables

Fig. 1 : Samples of the LIDC-IDRI dataset with significant inter-observer variability where (dis-)agreement in the ground-truth masks is clearly visible.

Fig. 2 : Schematic drawing of the Sinkhorn Probabilistic U-Net (SPU-Net) training framework, which is an extension of the PU-Net introduced by Kohl et al. [ 40 ] . The ground truth is denoted as

\mathbf{Y}

, the input image as

\mathbf{X}

, and the model prediction as

\mathbf{\hat{Y}}

. During testing, only the prior density

p_{\psi}(\mathbf{z}\mid\mathbf{x})

is used to predict samples.

Dataset		Modality	Annotators	Samples
LIDC-IDRI		CT	$4$	$15096$
QUBIQ	Pancreas	CT	$2$	$702$
QUBIQ	Pancreatic lesion	CT	$2$	$206$

TABLE I : Modality, number of annotators and size for the used datasets.

Dataset	Model	$\mathbf{d=4}$		$\mathbf{d=5}$		$\mathbf{d=6}$		$\mathbf{d=7}$
Dataset	Model	$\hat{W}_{k}\downarrow$	ER $\uparrow$	$\hat{W}_{k}\downarrow$	ER $\uparrow$	$\hat{W}_{k}\downarrow$	ER $\uparrow$	$\hat{W}_{k}\downarrow$	ER $\uparrow$
LIDC-IDRI	PU-Net	$0.451\pm 0.001$	$1.68\pm 0.07$	$0.450\pm 0.003$	$1.68\pm 0.08$	$0.451\pm 0.006$	$1.92\pm 0.05$	$0.451\pm 0.002$	$1.77\pm 0.14$
	PU-Net+NF	$0.451\pm 0.004$	$2.90\pm 0.39$	$0.450\pm 0.002$	$3.31\pm 0.67$	$0.446\pm 0.003$	$2.94\pm 0.59$	$0.450\pm 0.002$	$4.00\pm 0.52$
	Ours	$0.447\pm 0.005$	$3.52\pm 0.04$	$0.441\pm 0.002$	$4.57\pm 0.03$	$\mathbf{0.440\pm 0.006}$	$5.12\pm 0.10$	$0.443\pm 0.002$	$5.59\pm 0.06$
QUBIQ Pancreas	PU-Net	$0.602\pm 0.034$	$1.76\pm 0.07$	$0.608\pm 0.024$	$1.49\pm 0.25$	$0.603\pm 0.030$	$1.52\pm 0.23$	$0.589\pm 0.002$	$1.65\pm 0.22$
	PU-Net+NF	$0.581\pm 0.004$	$2.37\pm 0.24$	$0.610\pm 0.080$	$3.15\pm 0.43$	$0.585\pm 0.021$	$2.77\pm 0.24$	$0.589\pm 0.054$	$3.77\pm 0.98$

TABLE II : Comparison of the experimented models on various datasets. The PU-Net is re-implemented from [ 40 ] and PU-Net+NF from [ 26 ] . Results are evaluations on the test set. The mean and standard deviations (threefold cross-validation) of the Empirical Wasserstein metric with cost unity minus the Intersection over Union (1-IoU) are presented. Bold indicates the best

\hat{W}_{k}

per dataset. Furthermore, the homogeneity of the latent singular values is expressed with the Effective Rank. A clear relationship between a homogeneous latent space and model performance is visible.

Model	Training time	Inference time	Memory	Parameters	FLOPs
Model	(appr. min/epoch)	(100 $\times$ , s.)	(GB)	(M)	(G)
PU-Net	1.00	2.56	3.98	6.844	119
PU-Net+NF	1.13	2.57	4.07	6.853	211
Ours	1.75	2.57	4.07	6.846	350

TABLE III : Training statistics for the used models with batch size of 32. While training costs are increased, note that our methodology adds negligible costs to inference time, memory usage and parameter count.

Panel of Figure 6 — Fig. 3 : Visualization of test set predictions. The mean and standard deviation is taken from 16 predictions for both PU-Net and our proposed model.

Panel of Figure 7 — Fig. 4 : Visualization of the latent vector mean and variances subject to various test images.

Panel of Figure 8 — Fig. 5 : Interpolation of the latent axis-aligned prior density for the LIDC-IDRI dataset (zoomed-in for viewing convenience).

Dataset	Model	$\hat{W}_{k}\downarrow$	ER $\uparrow$
LIDC-IDRI	ResNet	$0.508$	$3.48$
	ResNet (ours)	$0.435$	$3.58$
	ConvNeXt	$0.451$	$1.47$
	ConvNeXt (ours)	$0.447$	$4.55$

TABLE IV : Comparison of the PU-Net variants using different encoders with

d=6

Fig. 6 : Sequentially freezing latent dimensions (i.e., taking the mean) from smallest to largest for models trained on the LIDC-IDRI dataset with $d=6$ . The PU-Net is unaffected for the majority of the latent dimensions, while our proposed method already degrades from a single frozen dimension. This implies that the PU-Net does not use the information in the latent variances.

Panel of Figure 11 — Fig. 7 : The Effective Rank during training for the evaluated models with $d=6$ .

Panel of Figure 12 — Fig. 8 : Calibration between Effective Rank and prediction entropy for the LIDC-IDRI dataset, where $\rho$ represents the Pearson correlation coefficient.

Segmentation Uncertainty Estimation

Segmentation uncertainty estimation measures where pixel- or object-level masks may be wrong, ambiguous, or unreliable.

See all 104 papers

Delving Aleatoric Uncertainty in Medical Image Segmentation via Vision Foundation Models
Apr 13, 2026Ruiyang Li, Fang Liu, Licheng Jiao +8Medical Image Segmentation Segmentation Uncertainty Estimation
Key Laboratory of Intelligent Perception and Image Understanding of Ministry of Education, International Research Center for Intelligent Perception and Computation, Joint International Research Laboratory of Intelligent Perception and Computation, School of Artificial Intelligence, Xidian University, Xi’an 710071, China
PCA-Enhanced Probabilistic U-Net for Effective Ambiguous Medical Image Segmentation
Mar 12, 2026Xiangyu Li, Chenglin Wang, Qiantong Shen +6Medical Image Segmentation Segmentation Uncertainty Estimation
Harbin Institute of Technology, Harbin, China · Harbin Institute of Technology Zhengzhou Research Institute, China · Peking University, China +2
SegWithU: Uncertainty as Perturbation Energy for Single-Forward-Pass Risk-Aware Medical Image Segmentation
Apr 16, 2026Tianhao Fu, Austin Wang, Charles Chen +2Image Segmentation Uncertainty Quantification
University of Toronto, Toronto, ON, Canada · Vector Institute, Toronto, ON, Canada · Project Neura, Toronto, ON, Canada +4

Apr 13, 2026J/K move · Enter open · S save

Delving Aleatoric Uncertainty in Medical Image Segmentation via Vision Foundation Models

Ruiyang Li, Fang Liu, Licheng Jiao, Xinglin Xie+7

Key Laboratory of Intelligent Perception and Image Understanding of Ministry of Education, International Research Center for Intelligent Perception and Computation, Joint International Research Laboratory of Intelligent Perception and Computation, School of Artificial Intelligence, Xidian University, Xi’an 710071, China

Medical Image Segmentation Segmentation Uncertainty Estimation

Open paper AskPDF arXiv

Medical image segmentation supports clinical workflows by precisely delineating anatomical structures and lesions. However, medical image datasets medical image datasets suffer from acquisition noise and annotation ambiguity, causing pervasive data uncertainty that substantially undermines model robustness. Existing research focuses primarily on model architectural improvements and predictive reliability estimation, while systematic exploration of the intrinsic data uncertainty remains insufficient. To address this gap, this work proposes leveraging the universal representation capabilities of visual foundation models to estimate inherent data uncertainty. Specifically, we analyze the feature diversity of the model's decoded representations and quantify their singular value energy to define the semantic perception scale for each class, thereby measuring sample difficulty and aleatoric uncertainty. Based on this foundation, we design two uncertainty-driven application strategies: (1) the aleatoric uncertainty-aware data filtering mechanism to eliminate potentially noisy samples and enhance model learning quality; (2) the dynamic uncertainty-aware optimization strategy that adaptively adjusts class-specific loss weights during training based on the semantic perception scale, combined with a label denoising mechanism to improve training stability. Experimental results on five public datasets encompassing CT and MRI modalities and involving multi-organ and tumor segmentation tasks demonstrate that our method achieves significant and robust performance improvements across various mainstream network architectures, revealing the broad application potential of aleatoric uncertainty in medical image understanding and segmentation tasks.

Abstract

Segmentation Uncertainty Estimation

Explore similar work

Delving Aleatoric Uncertainty in Medical Image Segmentation via Vision Foundation Models

PCA-Enhanced Probabilistic U-Net for Effective Ambiguous Medical Image Segmentation

SegWithU: Uncertainty as Perturbation Energy for Single-Forward-Pass Risk-Aware Medical Image Segmentation