A machine-learning model that represents links between many uncertain outputs reported better held-out likelihood than the alternatives tested across four computer-vision tasks. The approach combines two kinds of uncertainty - uncertainty in the model and uncertainty associated with the data - into one low-rank-plus-diagonal representation, a compact way to preserve correlations without storing a full covariance matrix.
The work is identified in its front matter as an arXiv version 1 preprint dated 25 August 2026.
A problem of many connected predictions
The study asks whether explicitly modeling correlations among outputs can make uncertainty-aware prediction more useful for high-dimensional regression while staying computationally feasible. That is the central test behind the proposed joint representation.
The proposed LR+D form combines aleatoric and epistemic uncertainty for joint downstream analysis. The experiments used U-Net variants with Monte Carlo dropout, stochastic variational inference and deep ensembles to estimate epistemic uncertainty. They were compared with non-Bayesian models and diagonal or partially approximated covariance baselines.
The main configuration used a low-rank dimension of 8, 64 sampled model weights and 576 joint columns before optional truncated singular-value decomposition compression. It used a loss weight of 0.125 and a minimum diagonal value of 0.01, with all models trained for the same number of steps.
The strongest result was likelihood
The main measure was held-out test log-likelihood, or TLL: how much probability the model assigned to outcomes it had not seen. Higher values indicate greater likelihood. The authors report that adding epistemic uncertainty improved unseen-test likelihood over non-Bayesian baselines, that an expected-weights approximation improved on diagonal-only models, and that the joint representation was superior in every evaluated task.
Across the four tasks, the proposed model recorded the best reported TLL for each of the three Bayesian configurations. With Monte Carlo dropout, the reported values were -74 +/- 13 for inpainting, 587 +/- 2 for colorization, -174 +/- 1 for optical flow and -51 +/- 10 for depth. The stochastic-variational-inference values were -174 +/- 29, 581 +/- 4, -164 +/- 6 and -174 +/- 27, respectively. Deep-ensemble values were -50, 589, -172 and -13. The columns use different scales: x100 for inpainting and x1000 for the other tasks, so the figures should not be compared across tasks as if they shared one raw scale.
The benchmark covered colorization and inpainting on CelebA-HQ, optical-flow estimation on Flying Chairs and depth estimation on NYU. CelebA-HQ used 24,183 training, 2,993 validation and 2,824 test images; Flying Chairs used 18,297, 2,287 and 2,288; and NYU used 695 training, 174 validation and 654 test samples.
Compression helped, but it did not remove trade-offs
The reason for the low-rank design is scale. The paper describes full covariance storage as growing with the square of the number of outputs, while its LR+D memory grows linearly with the retained rank. Its stated log-likelihood cost also falls from cubic growth for a full covariance to a lower-order expression involving the output size and rank. That is intended to make joint uncertainty practical for output spaces such as CelebA, which the paper gives as roughly 65,000 dimensions.
Ablation tests found that using more posterior samples reduced variability in TLL, although it increased computation. Keeping more columns after truncated singular-value decomposition generally improved TLL, while even limited retained ranks remained better than a diagonal covariance. The size of the gain depended on the task and rank.
The researchers also used differential entropy, a measure of the spread of a probability distribution, to rank predictions for selective prediction. Negative log-likelihood risk decreased as coverage decreased for both diagonal and LR+D models, with LR+D maintaining lower reported risk across the coverage range. This was a model-based diagnostic rather than an external outcome measure.
Additional comparisons showed that the joint mean-uncertainty architecture outperformed separate models on three of the four benchmark tasks, while Flying Chairs favored the separate setup. Another comparison found better TLL for LR+D-trained models but lower relative reconstruction error for NPPC variants, indicating that the preferred approach depended on which metric was used.
The method still has sharp boundaries
Numerical stability was a practical constraint. Models with 20 or more low-rank columns consistently crashed early in the reported analysis, and minimum diagonal values below 0.001 crashed across five random seeds. The authors' starting guideline was a rank between 4 and 10, a diagonal floor of 0.01 and a loss weight of 1/16; those settings are practical recommendations from the studied configurations, not universal guarantees.
The evaluation was limited to Monte Carlo dropout, stochastic variational inference, deep ensembles and vision datasets. The method assumes a single multivariate Gaussian distribution, which may be a poor fit for tasks with non-Gaussian or multimodal uncertainty. The study therefore does not establish that the approach will generalize to structured regression outside vision or outperform more sophisticated Bayesian inference methods.
The work includes an official implementation and reproduction scripts, public dataset links and preprocessing scripts, pseudocode, proofs and additional qualitative material; model checkpoints are available on request. The supplied primary text contains author affiliations but no funding statement.
Paper data and sources
Original title: It depends: Incorporating correlations for joint aleatoric and epistemic uncertainties of high-dimensional output spaces
Authors: Leonhard F. Feiner, Manuel Nickel, Martin Menten et al.
Journal/Repository: Transactions on Machine Learning Research, 05/2026
Status: Preprint, not yet peer-reviewed
First online: 2026-08-25
DOI: Not available
Original paper · Full text