PaperlayerView exact revision on arXiv

Source-linked AI summary

It depends: Incorporating correlations for joint aleatoric and epistemic uncertainties of high-dimensional output spaces

Leonhard F. Feiner, Manuel Nickel, Martin Menten, Laurin Lux, Rickmer Braren, Daniel Rueckert, Georgios Kaissis, Raphael Rehms, Johannes Paetzold

arXiv:2608.24518v1cs.LGcs.CV

TL;DR

The paper addresses how to jointly model aleatoric and epistemic uncertainty for high-dimensional regression without neglecting output correlations. It uses a low-rank plus diagonal covariance representation with stabilization techniques, and reports higher log-likelihood, more reliable uncertainty estimates, and tractable scaling on large outputs.

  • Problem

    Existing high-dimensional UQ methods often neglect output correlations, while exact covariance operations become infeasible as output dimensionality grows.

  • Method

    The method represents joint uncertainty with a multivariate Gaussian and low-rank plus diagonal covariance, combining aleatoric and epistemic uncertainty efficiently.

  • Results

    The approach outperforms factorized Gaussian representations, yielding higher log-likelihood, more reliable uncertainty estimates, and tractable structured uncertainty for outputs up to 10^7 dimensions.

  • Takeaways & Limitations

    The low-rank structure preserves correlated uncertainty, supports efficient sampling and log-likelihood calculations, and exposes interpretable dominant modes.

  • Takeaways & Limitations

    Evaluation is restricted to Monte Carlo Dropout, Stochastic Variational Inference, and Deep Ensemble Bayesian approximations, despite broader conceptual applicability.

Abstract

from arXiv · show

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.

1 Introduction

High-dimensional UQ must jointly represent aleatoric and epistemic uncertainty while preserving dependencies between outputs. Existing methods often neglect correlations or cannot scale beyond low-dimensional settings, motivating a globally efficient joint representation.

  • Uncertainty types: Aleatoric uncertainty represents irreducible data noise, whereas epistemic uncertainty reflects reducible uncertainty from limited model knowledge.
  • Scaling challenge: Output correlations require O(S^2) covariance entries, making standard sampling and log-likelihood operations computationally infeasible for high-dimensional outputs.
  • Why correlations matter: Ignoring dependencies between outputs can produce miscalibrated uncertainties and incoherent predictions in structured tasks such as optical flow, inpainting, and segmentation.
  • Prior limitations: Existing approaches are constrained by low-dimensional full covariances, local band sparsification, aleatoric-only LR+D models, or rank-deficient factors unsuitable for log-likelihood evaluation.
  • Contribution: Unlike a related LR+D method, the approach accounts for model uncertainty in aleatoric estimation and resolves the second-order distribution into a unified representation supporting direct log-likelihood calculation.
  • Contribution: The proposed framework jointly models aleatoric and epistemic uncertainty with LR+D covariance, preserving output correlations while reducing the costs of full covariance matrices.

2 Method

The method models high-dimensional predictive uncertainty by combining aleatoric and epistemic components in a low-rank-plus-diagonal covariance representation. Monte Carlo weight sampling, covariance decomposition, truncated SVD, and stability diagnostics make the joint covariance more tractable while preserving important correlations.

  • Joint covariance representation: Monte Carlo samples from an approximate weight posterior represent epistemic uncertainty, while the likelihood models aleatoric uncertainty.The posterior predictive distribution combines uncertainty over weights with uncertainty in predictions conditional on each sampled weight.
  • Joint covariance representation: The network models high-dimensional outputs with a multivariate Gaussian likelihood whose covariance uses a low-rank-plus-diagonal structure.The low-rank factor has rank much smaller than the number of outputs, reducing covariance complexity while retaining dominant aleatoric directions.
  • Uncertainty decomposition: The joint covariance is decomposed as the sum of aleatoric and epistemic covariance components, with the former captured by averaged covariance structure and the latter by covariance across sampled means.This multivariate covariance formulation generalizes variance-based decompositions to correlated outputs.
  • Uncertainty decomposition: Epistemic covariance is represented purely by a low-rank factor derived from centered sampled predictions, while the aleatoric diagonal remains strictly positive.Because the aleatoric diagonal is positive, the total covariance remains positive definite and invertible even when the empirical epistemic covariance is low rank.
  • Truncated SVD approximation: Truncated SVD reduces the low-rank dimension by retaining dominant variance directions and transferring dropped univariate variance to the diagonal.The approximation preserves independent variance and the most important covariances; the retained rank is selected using reconstruction error and downstream performance.
  • Stability considerations: The method addresses numerical instability arising from high output dimension, covariance inversion, and increasing low-rank rank during Monte Carlo sampling.The paper notes that condition numbers can grow for both the joint covariance and the smaller capacitance matrix used in computation.

3 Experiments

The experiments evaluate joint uncertainty modeling with LR+D covariance across four high-dimensional vision tasks and compare Bayesian, non-Bayesian, diagonal, and structured covariance variants. Results show improved predictive fit, stability, and interpretable correlated-uncertainty representations, while LR+D avoids full-covariance scaling limits.

  • Datasets and Tasks: The method is evaluated on CelebA inpainting and colorization, Flying Chairs optical flow, and NYU depth estimation.The experiments use U-Net variants with probabilistic outputs and multiple Bayesian approximation methods.
  • Baselines and Model Setup: The proposed LR+D representation combines a diagonal component with a low-rank covariance structure to capture richer output correlations.The evaluation includes diagonal, LR+D, and an expected-weights approximation alongside Bayesian methods including MCD, SVI, and deep ensembles.
  • Quantitative Results: Test log-likelihood improves consistently when epistemic uncertainty is modeled, for both diagonal and LR+D covariance parameterizations across all experiments.TLL evaluates predictive fit using the mean and covariance, while L1 and L2 assess point-prediction performance separately.
  • Quantitative Results: The unified joint multivariate representation outperforms the other tested Bayesian methods and approximations in every evaluated task.The expected-weights approximation also improves over diagonal-only baselines, but the margin varies with task complexity and output dimensionality.
  • Qualitative Results: The learned covariance factors represent hierarchical correlated uncertainty, from global coarse-scale variations to increasingly localized features.Qualitative visualizations show local and globally correlated patterns, including positive and negative relationships in colorization outputs.
  • Complexity of Covariance Parametrizations: Full covariance approaches the GPU memory limit on larger datasets, whereas LR+D supports larger matrices with more favorable scaling for log-likelihood computation.The complexity comparison reports linear memory scaling for LR+D versus quadratic scaling for full covariance until the GPU limit.
  • Stability and Truncation: Higher sample counts reduce TLL variability, while retaining more low-rank columns generally improves TLL relative to diagonal covariance.Even limited-column approximations still yield significant improvements over purely diagonal covariance, although more samples increase computation.
  • Model Stability: A fixed value of RW = 8 is adopted across tasks for stability, simplicity, and computational efficiency despite task-dependent optimal values.The diagonal update is reported to produce more robust predictions.

4 Discussion

The method jointly models aleatoric and epistemic uncertainty with a low-rank plus diagonal covariance, preserving output correlations while reducing the cost of full covariance representations. It improves uncertainty evaluation and interpretability, but remains subject to Gaussian and evaluation-scope limitations.

  • It jointly models aleatoric and epistemic uncertainty in a single framework and provides a more expressive posterior predictive approximation than traditional joint distributions.The authors associate this with improved reliability and explainability of deep-learning predictions.
  • The method uses a low-rank plus diagonal covariance to preserve vital output correlations while reducing the computational demands of full covariance matrices.It also monitors and regularizes covariance and capacitance-matrix condition numbers during training.
  • The approach outperforms the commonly used factorized Gaussian representation, with higher log-likelihood and more reliable uncertainty estimates.The low-rank structure also exposes interpretable patterns in correlated uncertainties across high-dimensional outputs.
  • The evaluation is restricted to Monte Carlo Dropout, Stochastic Variational Inference, and Deep Ensemble despite the method’s conceptual extension to other Bayesian frameworks.The restriction is attributed to simplicity and computational reasons.
  • A full covariance for 10^5 dimensions requires approximately 40GB of memory, whereas the LR+D parameterization maintains linear scaling and remains tractable up to 10^7 dimensions.The added computational overhead is presented as a trade-off for high-dimensional tasks where full covariance matrices are physically intractable.
  • The method assumes output uncertainty can be represented by a single multivariate Gaussian, which may not suit translation, rotation, or inherently multimodal tasks.In multimodal settings, the mean may average several plausible modes rather than represent one physically consistent realization.

B Reproducibility

The paper provides an official implementation and reproducibility materials, including experiment scripts, datasets, preprocessing resources, and model checkpoints.

  • The official implementation and scripts reproduce all experiments, with dataset links and preprocessing scripts included in the repository.Checkpoints for proposed and baseline models are available upon request.

C Detailed Computational Analysis and Implementation Considerations

The LR+D covariance parameterization offers efficient computation and global correlation modeling, but requires rank management and stabilization strategies that add implementation complexity.

  • Implementation: For homogeneous feature spaces, the low-rank factor can be implemented position-wise with 1 × 1 convolutions or node-wise linear layers.This applies when a constant number of channels is predicted at each grid position or graph node.
  • Computational complexity: LR+D captures global correlations while using less memory and computation than naive or Cholesky covariance representations.Sparse alternatives may restrict correlations to local or selected global structure.
  • Computational complexity: Table 3 compares covariance parametrizations by captured correlations, representation, memory, and costs for quadratic forms, determinants, and sampling.The structural size R is the bandwidth or factor rank, with R ≪S assumed.
  • Implementation: LR+D requires careful management of rank R and stabilization strategies to maintain numerical consistency in the second-order distribution.This modeling effort is presented as a compromise for scales where full-rank methods are physically impractical.

D.1 Qualitative Results

The qualitative analyses visualize predictive distributions, marginal variances, and low-rank directions across optical flow, inpainting, colorization, and depth estimation. These visualizations expose task-specific correlation structures and support entropy-based uncertainty ranking.

  • Qualitative visualizations: Figures 6–10 visualize inputs, ground truths, predictions, predictive-distribution parameters, variances, and low-rank directions across the evaluated tasks.Figure 10 additionally compares Bayesian methods and sign-inverted eigenvectors for the primary method.
  • Optical flow: Optical-flow uncertainty is represented with a colorwheel, where hue encodes direction and intensity encodes magnitude.The visualization uses the two-channel structure of optical flow.
  • Colorization: Colorization eigenvectors form a hierarchy from global color-temperature shifts to localized hair- and eye-specific features.The leading singular-value directions are typically about 10× more significant than later factors.
  • Selective prediction: Differential entropy condenses each multivariate Gaussian predictive distribution into a scalar uncertainty score for ranking predictions.The LR+D determinant is computed with the Matrix Determinant Lemma, retaining dependencies represented by the low-rank factors.
  • Selective prediction: Both diagonal and LR+D models show lower NLL risk at lower coverage, while LR+D maintains lower risk across the coverage range.This indicates that entropy predicts NLL for both parameterizations and that LR+D provides a stronger probabilistic fit in selective prediction.

D.3 Quantitative Prediction Errors

The prediction-error evaluation checks whether uncertainty improvements arise from probabilistic modeling rather than substantially different reconstruction accuracy. Errors are reported for Bayesian and non-Bayesian configurations across tasks.

  • Prediction-error evaluation: L2 and L1 prediction errors are reported across dataset–task combinations for Bayesian and non-Bayesian configurations.MCD and SVI values use means and standard deviations over five seeds, whereas DE uses a single ensemble.
  • Evaluation protocol: MCD, SVI, and non-Bayesian results are averaged over five independent model seeds, while DE uses one five-member ensemble.The single DE ensemble reflects the substantial computational cost of training multiple independent ensembles.
  • Selective prediction: Selective-prediction evaluation uses NLL risk rather than reconstruction error to assess the probabilistic fit of retained predictions.Coverage is the retained test fraction after sorting samples by predictive entropy.
  • Evaluation protocol: Bayesian models are evaluated with T = 64 weight samples.This evaluation protocol applies across the Bayesian configurations considered.

D.4 Additional Ablation Study

The ablations examine test-set variability, NPPC comparisons, training configurations, retained SVD variance, and rank choices. They reveal task-dependent trade-offs between likelihood, residual structure, stability, and computational cost.

  • Test set variability: Test-set TLL variability reflects image complexity and task ambiguity, with hard samples producing substantially lower likelihoods.High variance is associated with tasks such as MNIST and CelebA colorization, while MNIST variability also reflects discrete multimodality.
  • Comparison with NPPC: NPPC training was viable only with separate models and Gram–Schmidt projection; combined mean–uncertainty models were unstable.The NPPC loss was applied to P while diagonal covariance terms remained optimized with the LR+D loss.
  • Comparison with NPPC: LR+D-trained models consistently outperform NPPC variants in TLL, whereas NPPC reduces relative reconstruction error.The comparison indicates a metric-dependent trade-off between likelihood-based prediction and residual-structure capture.
  • Joint versus separate training: Separate uncertainty modeling outperforms joint training specifically on Flying Chairs, while joint training is superior on three of four benchmark tasks.The authors hypothesize that optical-flow inputs and predicted outputs provide direct feedback for estimating prediction errors.
  • Retaining variance of TSVD: Updating the diagonal D after SVD slightly improves average TLL, prediction consistency, and test-set variability across three dimensionality-reduction configurations.The ablation compares updating D with retaining the original diagonal.
  • Rank and dimensionality reduction: Higher numbers of learned columns generally improve results but increase training time, and RW = 32 caused training instabilities in three tasks.R denotes the resulting number of columns after dimensionality reduction.

E.1 Numerical Validation of Covariance Condition Number Bounds

MNIST inpainting validates that the proposed condition-number bounds closely enclose exact covariance condition numbers. The experiment also documents numerical artifacts and uses figure-based diagnostics for stability and conditioning.

  • Experimental setup: The MNIST inpainting setup masks 5/7 of each digit, uses S = 580 outputs, and permits exact full-covariance condition-number computation.The standard test set is used with 50,000 training and 10,000 validation examples.
  • Caveat: Boundary artifacts in MNIST uncertainty maps are attributed to the architecture’s higher spatial-padding ratio during upsampling.The smaller model uses more padding to preserve the original output resolution.
  • Qualitative results: The qualitative MNIST figure compares masked inputs, ground truth, predictive means, and marginal variance components.The displayed variance terms include Σii, Dii, and (PP ⊤)ii.
  • Conditioning diagnostics: The conditioning figures examine covariance condition-number distributions and how model parameters affect conditioning and training stability.The panels cover MNIST inpainting and CelebA colorization.
  • Condition-number validation: The estimated condition numbers remain tightly enclosed by the proposed upper and lower bounds across 1,000 sorted test images.The double-precision numerical estimates are compared with exact condition numbers, including in a magnified region.

E.2 Effect of Training Hyperparameters on Conditioning and Stability

Training stability depends strongly on LR+D conditioning and hyperparameter choices. The paper recommends staged training, conservative initialization, monitoring, and numerical safeguards to identify and avoid unstable regimes.

  • Failure modes: LR+D training can become unstable during capacitance-matrix Cholesky decomposition and inversion required for log-likelihood computation.The paper specifically identifies numerical instability for certain hyperparameter settings.
  • Hyperparameter effects: Increasing the low-rank width raises the capacitance-matrix condition number, while models with 20 or more columns consistently crash early from exploding gradients.The condition number peaks early in training and decreases after 10,000 steps.
  • Hyperparameter effects: Minimum diagonal values below 0.001 consistently cause training crashes across five random seeds by increasing the full covariance condition number.Smaller diagonal offsets improve certainty expression but worsen numerical conditioning.
  • Training protocol: A staged workflow begins with a deterministic diagonal baseline, adds conservative LR+D structure, then enables Bayesian components before iterative refinement.Suggested initial settings include R ∈[4, 10], ϵ = 0.01, and α = 1/16.
  • Robustness safeguards: Monitoring rank, diagonal-floor activity, and numerical exceptions helps identify stability boundaries and prevent covariance-factor NaNs from collapsing training.The recommended safeguards include skipping updates for batches that produce numerical exceptions.
  • Inference stability: SVD improves sampling stability for large T, whereas without SVD, lower posterior-sample counts are more likely to remain numerically stable.The trade-off is higher computational cost when using truncated SVD during sampling.
  • Covariance construction: Monte Carlo weight samples aggregate predictive means and aleatoric factors, concatenate epistemic and aleatoric components, and form Σ(x) = P(x)P(x)⊤+ D(x).Optional truncated SVD reduces the rank for computational efficiency.

G Derivations in Detail

The derivations show how LR+D covariance structure makes predictive distributions, likelihoods, inverses, determinants, and condition-number estimates tractable. SVD and matrix identities reduce the computational burden while preserving the dominant uncertainty structure.

  • Predictive distribution: The approximate posterior predictive distribution is modeled as a multivariate normal with mean µ(x) and covariance Σ(x).The derivation distinguishes the predictive distribution from the weight-conditioned normal model.
  • Likelihood computation: Determinant and inverse calculations are required for sampling and log-likelihood evaluation but are prohibitively expensive for full S × S covariance matrices.The LR+D representation is introduced to make these operations efficient.
  • Determinant: The matrix determinant lemma reduces determinant evaluation to the R × R capacitance matrix C = I_R + P⊤D^-1P.Because R ≪ S, computing the capacitance determinant is inexpensive.
  • Covariance construction: The covariance-construction algorithm samples T weights, computes means and uncertainty components, aggregates them, and combines epistemic and aleatoric factors.The resulting outputs are µ(x), D(x), P(x), and Σ(x).
  • LR+D operations: The final covariance is Σ(x) = P(x)P(x)⊤+ D(x), with optional SVD truncation applied to reduce the factor rank.The inverse uses the Woodbury identity and a small capacitance matrix.
  • Dimensionality reduction: Truncated SVD retains the largest singular values and vectors while updating the diagonal with variance removed from the low-rank approximation.This preserves essential structure with reduced dimensionality.
  • Regression objective: Gaussian regression likelihoods depend on covariance determinants, inverses, and quadratic residual terms, and are normalized by output dimensionality S.The derivation assumes predictions are Gaussian around the true value.
  • Joint uncertainty: The joint predictive covariance follows the law of total variance, combining variability within weight-conditioned predictions with variability across sampled weights.Monte Carlo integration over the proxy posterior approximates the required expectations.
Loading 2608.24518v1…