Cooperative Bayesian and variance networks disentangle aleatoric and epistemic uncertainties

TL;DR

The paper addresses the challenge of splitting aleatoric (data) uncertainty from epistemic (model) uncertainty in regression tasks. It introduces a cooperative framework that sequentially trains a mean network, a variance estimation network for aleatoric noise, and a Bayesian neural network that updates the mean and captures epistemic uncertainty, ensuring explicit disentanglement. The approach uses a Gaussian likelihood with heteroscedastic variance for the mean, a Gamma-based residual model for aleatoric variance, and pSGLD-based Bayesian inference to obtain a predictive distribution composed of the mean, aleatoric, and epistemic components. Across 18 datasets, including a synthetic plasticity problem with known aleatoric noise, the method consistently improves mean predictions and provides well-calibrated uncertainty estimates, outperforming end-to-end baselines and enabling robust uncertainty quantification for decision-making and active learning.

Abstract

Real-world data contains aleatoric uncertainty - irreducible noise arising from imperfect measurements or from incomplete knowledge about the data generation process. Mean variance estimation (MVE) networks can learn this type of uncertainty but require ad-hoc regularization strategies to avoid overfitting and are unable to predict epistemic uncertainty (model uncertainty). Conversely, Bayesian neural networks predict epistemic uncertainty but are notoriously difficult to train due to the approximate nature of Bayesian inference. We propose to cooperatively train a variance network with a Bayesian neural network and demonstrate that the resulting model disentangles aleatoric and epistemic uncertainties while improving the mean estimation. We demonstrate the effectiveness and scalability of this method across a diverse range of datasets, including a time-dependent heteroscedastic regression dataset we created where the aleatoric uncertainty is known. The proposed method is straightforward to implement, robust, and adaptable to various model architectures.

Figures (19)

• Figure 1: Illustration of the proposed cooperative training of a mean network, a variance network, and a Bayesian neural network for disentangling aleatoric and epistemic uncertainties. The top left figure shows the unseen ground truth mean (thick dashed gray line) and aleatoric uncertainty (credible interval within the thin dashed gray lines), as well as the respective data for training (magenta crosses). The method starts by training the mean network to only estimate the mean (green solid line in Step 1). Then, without updating the mean estimate, a variance network is trained to only predict aleatoric uncertainty (orange credible interval in Step 2). Subsequently, considering this aleatoric uncertainty estimation, a Bayesian neural network is trained to obtain an updated mean and corresponding epistemic uncertainty (solid blue line for the new mean, and shaded blue credible interval for the epistemic uncertainty in Step 3). If needed, the method can iterate between steps 3 and 2 to improve the disentanglement of uncertainties. Note the disentanglement of uncertainties together with the improvement of the mean estimation away from the data support ($x<0$ and $x>10$ in Step 3).
• Figure 2: Heteroscedastic regression by our method (right) compared to existing end-to-end training methods (left) for each inference type.
• Figure 3: Accuracy metrics (RMSE $\downarrow$, Epistemic TLL $\uparrow$, and WA $\downarrow$) obtained for the plasticity law discovery dataset considering a training set with different number $N$ of training sequences (history-dependent paths), where each training sequence has 100 points -- an example of a typical input and heteroscedastic output path is shown in \ref{['fig:plasticity_data_noise']} of \ref{['sec:des_plasticity_law']}. The Wasserstein distance (WA) represents the closeness of the estimated aleatoric uncertainty distribution to the ground truth distribution. Note that the MVE (MC-Dropout) and Evidential methods have large values of WA, and part of their curves is cut off. All metrics result from repeating the training of each method 5 times by resampling points randomly from the training datasets.
• Figure 4: Predictions of different methods on plasticity law discovery dataset. We randomly pick one test point from the dataset and show the entire third component $\mathbf{y}_{2}$ of 100-time steps under the 50 and 800 training sequences respectively.
• Figure 5: Trajectory of all essential components with respect to the iteration $K$ for 800 training sequences in the plasticity law discovery problem. In the figures, "Init" represents the initialization of training the mean network only as described in \ref{['sec: mean_training']}, and different curves are realizations under different seeds to restart the procedure, as well as using new samples of training sequences in the dataset.

Theorems & Definitions (1)

• proof