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

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

Abstract: Uncertainty Quantification 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 Uncertainty Quantification in the tasks of image inpainting, colorization, optical flow, and depth estimation.

Submission Type: Regular submission (no more than 12 pages of main content)

Changes Since Last Submission:

- Correction of Textual Errors and LaTeX Formulas: We performed a thorough proofreading to fix minor wording and spelling mistakes. We also corrected various LaTeX formulas throughout the manuscript and appendix to ensure mathematical notation is precise and correctly formatted.

- Standardization of References: We cleaned and unified the bibliography section. All references have been updated to a consistent format regarding author lists, journal titles, and publication years, ensuring compliance with TMLR styling.

- Visual Consistency of Figures: We unified the appearance of all figures across the main body and the appendix. This includes standardizing font sizes, color schemes, and line weights to provide a cohesive visual presentation.

- De-anonymization and Repository Link: We updated the manuscript to replace the anonymous code repository link with the official, public GitHub repository link to ensure transparency and reproducibility.

- Camera-Ready Formatting: We updated the author block to include all full names, institutional affiliations, and contact information. We also added the required OpenReview forum link and ensured the document adheres to the final publication template.

Code: https://github.com/LeonhardFeiner/corr-joint-ae-uq

Assigned Action Editor: ~Jes_Frellsen1

Submission Number: 6948