It depends: Incorporating correlations for joint aleatoric and epistemic uncertainties of high-dimensional output spaces
cs.LG, cs.CV
Submitted: 2026-08-25
Updated: 2026-08-25
Comments: Published in Transactions on Machine Learning Research (TMLR), 2026. 42 pages, 14 figures, 10 tables. https://openreview.net/forum?id=zw5EuUnBny
Journal ref: Transactions on Machine Learning Research, 05/2026
Code: https://github.com/LeonhardFeiner/corr-joint-ae-uq
License: http://creativecommons.org/licenses/by-nc-sa/4.0/
The gist: Uncertainty Quantification (UQ) plays a vital role in enhancing the reliability of deep learning model predictions, especially in scenarios with high-dimensional output spaces.
Terminology
Abstract
Uncertainty Quantification (UQ) plays a vital role in enhancing the reliability of deep learning model predictions, especially in scenarios with high-dimensional output spaces. This paper addresses the dual nature of uncertainty -- aleatoric and epistemic -- focusing on their joint integration in high-dimensional regression tasks. For example, in applications like medical image segmentation or restoration, aleatoric uncertainty captures inherent data noise, while epistemic uncertainty quantifies the model's confidence in unfamiliar conditions. Modeling both jointly enables more reliable predictions by reflecting both unavoidable variability and knowledge gaps, whereas modeling only one limits transparency and robustness. We propose a novel approach that approximates the resulting joint uncertainty using a low-rank plus diagonal covariance structure, capturing essential output correlations while avoiding the computational burdens of full covariance matrices. Unlike prior work, our method explicitly combines aleatoric and epistemic uncertainties into a unified second-order distribution that supports robust downstream analyses like sampling and log-likelihood evaluation. We further introduce stabilization strategies for efficient training and inference, achieving superior UQ in the tasks of image inpainting, colorization, optical flow, and depth estimation.
Sources
- Benchmarking Bayesian Deep Learning on Diabetic Retinopathy Detection Tasks
- Training VAEs Under Structured Residuals
- Sources of Uncertainty in Supervised Machine Learning -- A Statisticians' View
- Informative Priors Improve the Reliability of Multimodal Clinical Data Classification
- Variational Variance: Simple, Reliable, Calibrated Heteroscedastic Noise Variance Parameterization
- Meta Learning Low Rank Covariance Factors for Energy-Based Deterministic Uncertainty
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: Local-Consistency Optimisation for Bayesian Neural Networks
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks