📊 statistics

Cooperative Variance Estimation and Bayesian Neural Networks for Disentangling Aleatoric and Epistemic Uncertainties

This paper proposes a cooperative training framework that integrates a variance estimation network with a Bayesian neural network to effectively disentangle aleatoric and epistemic uncertainties while improving mean estimation across diverse datasets.

Original authors: Jiaxiang Yi, Miguel A. Bessa

Published 2026-05-29

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jiaxiang Yi, Miguel A. Bessa

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to predict the weather. You have a model that says, "Tomorrow will be 75°F." But how sure is the model?

In the world of AI, there are two main reasons why a model might be unsure:

The "Messy Data" Problem (Aleatoric Uncertainty): The thermometer is broken, or the wind is blowing wildly. No matter how smart your model is, you can't predict the exact temperature because the data itself is noisy. This is like trying to guess the exact outcome of rolling a die; the noise is built into the process.
The "I've Never Seen This Before" Problem (Epistemic Uncertainty): You ask your model about a weather pattern it has never seen in its training data. The model is unsure because it lacks knowledge, not because the data is noisy. This is like asking a chef who only cooks Italian food to guess the taste of a dish from a cuisine they've never studied.

The Problem with Current AI
Most AI models are like overconfident students. They give you a single answer (the mean) but act like they know everything. If they are wrong, they don't tell you.

Some models try to admit they are unsure about the "messy data," but they get confused and often overfit (memorize the noise), leading to bad predictions.
Other models (Bayesian Neural Networks) try to admit they are unsure about their "lack of knowledge," but they are incredibly hard to train and often get stuck in a mathematical rut.

Trying to train a model to do both at the same time is like asking a student to solve a math problem while simultaneously trying to figure out how messy their pencil is. The two tasks fight each other, and the result is usually a mess.

The Solution: A Cooperative Team (VeBNN)
The authors of this paper propose a new way to train AI called VeBNN (Variance estimation Bayesian Neural Network). Instead of forcing one model to do everything at once, they break the job into three distinct steps, like a relay race where each runner specializes in one leg.

Step 1: The "Mean" Runner
First, they train a simple model just to find the average answer (the mean). They pretend the data is perfectly clean for a moment. This gives them a solid baseline prediction without getting distracted by the noise.

Step 2: The "Noise" Detective
Next, they take that solid baseline and train a second, specialized model just to look at the mistakes the first model made. This "Noise Detective" learns to map out where the data is messy and where it is clean. It learns the Aleatoric Uncertainty (the irreducible noise). Crucially, it does this without trying to change the average prediction, so it doesn't get confused.

Step 3: The "Knowledge" Expert
Finally, they train a complex Bayesian Neural Network. This is the expert that knows how to handle "I don't know" situations. But here is the trick: they give this expert the map of the "noise" created by the detective in Step 2. Because the expert already knows exactly where the data is messy, it can focus entirely on figuring out where it lacks knowledge. This allows it to accurately measure Epistemic Uncertainty.

Why This Works (The Analogy)
Think of it like a team of doctors diagnosing a patient:

Doctor A measures the patient's temperature and heart rate (the Mean).
Doctor B looks at the patient's history and notes that the thermometer is shaky and the patient is sweating (the Aleatoric Noise).
Doctor C (the Bayesian expert) uses the measurements from A and the context from B to decide: "Is this a rare disease I've never seen?" (The Epistemic Uncertainty).

If you ask Doctor C to do all three jobs at once, they might get overwhelmed. But by having them work cooperatively in sequence, the diagnosis is much more accurate.

What They Found
The authors tested this method on various datasets, including:

Standard math problems.
Image recognition tasks.
A special dataset they created involving material science (predicting how metal deforms under stress). In this specific case, they actually knew the true amount of noise in the data, allowing them to prove their method was correct.

The Results

Better Accuracy: The models predicted the average outcome better than previous methods.
True Disentanglement: The model successfully separated "messy data" from "lack of knowledge." In previous methods, these two uncertainties often got mixed up, making the model either too confident or too scared.
Robustness: The method worked well even when the data was scarce or when the model was asked to predict things outside its training range (extrapolation).

In Summary
This paper introduces a "divide and conquer" strategy for AI uncertainty. By training three specialized networks in a specific order, the authors created a system that can tell you not just what will happen, but how sure it is about the data's noise versus its own lack of knowledge. It's a simpler, more robust way to make AI admit what it doesn't know.

Technical Summary: Cooperative Variance Estimation and Bayesian Neural Networks for Disentangling Aleatoric and Epistemic Uncertainties

Problem Statement

Real-world data inherently contains two distinct types of uncertainty: aleatoric uncertainty (data noise arising from imperfect measurements or stochastic generation processes) and epistemic uncertainty (model uncertainty arising from limited data or model capacity). While mean-variance estimation (MVE) networks can learn aleatoric uncertainty, they often suffer from overfitting and fail to predict epistemic uncertainty. Conversely, Bayesian Neural Networks (BNNs) can predict epistemic uncertainty but are notoriously difficult to train, particularly when attempting to simultaneously model heteroscedastic (input-dependent) aleatoric noise.

Existing approaches to jointly train MVE networks often encounter optimization challenges. The Negative Log-Likelihood (NLL) loss function used for joint training creates a non-convex landscape with saddle points. Specifically, the gradient with respect to the variance contains high-order terms in the denominator, leading to imbalances where the network may either overfit the noise (driving variance to zero) or fail to converge (driving variance to infinity). Furthermore, attempts to disentangle uncertainties using ensemble methods or MC-Dropout on MVE networks have been shown to yield unreliable results, often conflating the two uncertainty types.

Methodology: VeBNN

The authors propose VeBNN (Variance estimation Bayesian Neural Network), a cooperative training strategy that sequentially decouples the estimation of mean, aleatoric uncertainty, and epistemic uncertainty. The method proceeds in three distinct steps:

Mean Network Training (Step 1): A deterministic network is trained to estimate the mean ( $\mu$ ) assuming constant (homoscedastic) aleatoric uncertainty. This simplifies the optimization landscape, allowing the network to find a robust point estimate of the mean without the interference of variance gradients.
Variance Estimation Network Training (Step 2): A separate network is trained to predict the aleatoric uncertainty ( $\sigma^2_a$ ) based on the residuals of the fixed mean estimate from Step 1. Crucially, instead of using the standard Gaussian NLL loss (which suffers from gradient imbalance), the authors propose modeling the squared residuals using a Gamma distribution. This leads to a Gamma-based NLL loss where the gradients with respect to the network parameters do not contain high-order denominator terms, ensuring stable optimization. The aleatoric variance is then derived as the expected value of this Gamma distribution.
Bayesian Neural Network Training (Step 3): A BNN is trained with a "warm start" using the mean estimate from Step 1 and the fixed aleatoric variance estimate from Step 2. Because the aleatoric noise is now treated as a known, fixed quantity during this stage, the BNN focuses solely on learning the posterior distribution of the weights to capture epistemic uncertainty.

The final prediction for a new input $x'$ decomposes the total uncertainty into:

Predictive Mean: The expectation of the BNN output.
Aleatoric Uncertainty: The fixed variance estimate from Step 2.
Epistemic Uncertainty: The variance of the BNN's predictive mean across the posterior distribution.

The authors note that the process can iterate between Steps 2 and 3 to refine the disentanglement, though empirical results suggest convergence is often achieved with a single iteration ( $K=1$ or $K=2$ ).

Key Contributions

Cooperative Training Strategy: The paper introduces a sequential training framework that avoids the coupled optimization problems inherent in joint MVE training. By separating the tasks, the method mitigates the non-convexity and gradient imbalance issues associated with joint mean-variance estimation.
Gamma Likelihood for Variance Estimation: The authors propose using a Gamma likelihood to model squared residuals for variance estimation. This avoids the numerical instability of the Gaussian NLL loss and provides a theoretically grounded, convex optimization problem for the variance network.
Effective Disentanglement: The method successfully separates aleatoric and epistemic uncertainties, a task where many existing methods (including deep ensembles and MC-Dropout applied to MVE) fail to provide reliable separation.
Scalability and Implementation: The approach is straightforward to implement using standard architectures (MLPs, GRUs, ResNets) and inference methods (pSGLD, Ensembles, BBB, MC-Dropout), demonstrating that cooperative training can enhance even "weak" inference methods like BBB.

Experimental Results

The authors evaluated VeBNN across five categories of datasets:

One-dimensional illustrative examples: Demonstrated that VeBNN avoids the overconfident variance predictions and poor mean estimation seen in joint MVE training, particularly in extrapolation regions.
UCI Regression Datasets: VeBNN consistently improved Test Log-Likelihood (TLL) and Root Mean Square Error (RMSE) across various inference methods (Ensembles, MC-Dropout, BBB, pSGLD) compared to their joint MVE counterparts. Notably, VeBNN (BBB) achieved competitive performance, whereas joint MVE (BBB) often failed to converge.
Large-scale Image Regression: Under real-world distribution shifts, VeBNN (Ensembles) achieved higher TLL and better test coverage than baselines, demonstrating a stronger ability to handle shifted data.
Material Plasticity Law Discovery (Known Ground Truth): Using a custom dataset where ground-truth aleatoric uncertainty is known (via multiple physical simulations), VeBNN (pSGLD) outperformed all baselines in RMSE, Wasserstein distance (measuring the closeness of estimated aleatoric uncertainty to ground truth), and Epistemic TLL. The method correctly identified that aleatoric uncertainty remained stable while epistemic uncertainty increased in out-of-distribution regions.
APPA-REAL Dataset: VeBNN achieved the best performance across all metrics (RMSE, WA, Epistemic TLL) compared to MVE baselines.

Ablation studies confirmed that the Gamma loss is superior to Gaussian NLL or $\beta$ -NLL for the variance network and that the method converges rapidly, often requiring only one additional iteration beyond the initial training.

Significance and Claims

The paper claims that the proposed cooperative training strategy offers a robust, efficient, and scalable solution for uncertainty quantification in deep learning. By decoupling the estimation of aleatoric and epistemic uncertainties, the method overcomes the optimization hurdles that have limited the widespread adoption of Bayesian heteroscedastic regression.

The authors emphasize that while training BNNs remains more challenging than training deterministic networks, their approach facilitates this process by providing a high-quality warm start (the mean) and a fixed aleatoric noise estimate. They state that this method is particularly valuable for safety-critical applications where distinguishing between reducible (epistemic) and irreducible (aleatoric) uncertainty is essential for active learning and risk-sensitive decision-making. The paper concludes that VeBNN is a practical, implementable alternative to current state-of-the-art methods, capable of improving mean estimation and uncertainty disentanglement across diverse model architectures and inference techniques.