[Paper Review] Estimating Epistemic and Aleatoric Uncertainty with a Single Model

Matthew A. Chan, María J. Molina+1 more|arXiv (Cornell University)|Feb 5, 2024

Forecasting Techniques and ApplicationsDecision Sciences3 citations

TL;DR

This paper proposes Hyper-Diffusion, a single-model framework that estimates both epistemic and aleatoric uncertainty in machine learning using conditional diffusion models and hyper-networks. By sampling multiple weight sets from a hyper-network and generating predictions via a diffusion model, it achieves uncertainty estimation accuracy comparable to deep ensembles at a fraction of the training cost.

ABSTRACT

Estimating and disentangling epistemic uncertainty, uncertainty that is reducible with more training data, and aleatoric uncertainty, uncertainty that is inherent to the task at hand, is critically important when applying machine learning to high-stakes applications such as medical imaging and weather forecasting. Conditional diffusion models' breakthrough ability to accurately and efficiently sample from the posterior distribution of a dataset now makes uncertainty estimation conceptually straightforward: One need only train and sample from a large ensemble of diffusion models. Unfortunately, training such an ensemble becomes computationally intractable as the complexity of the model architecture grows. In this work we introduce a new approach to ensembling, hyper-diffusion models (HyperDM), which allows one to accurately estimate both epistemic and aleatoric uncertainty with a single model. Unlike existing single-model uncertainty methods like Monte-Carlo dropout and Bayesian neural networks, HyperDM offers prediction accuracy on par with, and in some cases superior to, multi-model ensembles. Furthermore, our proposed approach scales to modern network architectures such as Attention U-Net and yields more accurate uncertainty estimates compared to existing methods. We validate our method on two distinct real-world tasks: x-ray computed tomography reconstruction and weather temperature forecasting.

Motivation & Objective

• Address the critical need for reliable uncertainty quantification in high-stakes ML applications such as medical imaging and weather forecasting.
• Overcome the computational infeasibility of training deep ensembles for complex models, especially in generative settings.
• Disentangle epistemic uncertainty (model ignorance) from aleatoric uncertainty (inherent data variability) using a single trained model.
• Develop a scalable alternative to ensemble methods that maintains high predictive accuracy while enabling uncertainty estimation.

Proposed method

• Employ a hyper-network to generate a (pseudo-)ensemble of network weights for a conditional diffusion model at inference time.
• Use the diffusion model to sample multiple predictions from each set of generated weights, forming a predictive distribution.
• Estimate aleatoric uncertainty as the variance of the sample predictions across the diffusion model’s outputs.
• Estimate epistemic uncertainty as the variance of the mean predictions across different weight sets generated by the hyper-network.
• Train the hyper-network end-to-end to produce diverse, data-consistent weight sets that reflect uncertainty in the model’s predictions.
• Enable flexible sampling at inference by adjusting the number of weight sets (M) and predictions per set (N), mimicking ensemble behavior.

Figure 1: Block diagram of a hyper-diffusion model. A hyper-network is optimized to generate a (pseudo-)ensemble of network weights for a diffusion model which then outputs a distribution of predictions. The mean “ensemble” prediction is formed by averaging over all samples from the predictive distr

Experimental results

Research questions

• RQ1 Can a single model approximate the uncertainty estimation performance of a deep ensemble in complex inverse problems?
• RQ2 Does the proposed hyper-diffusion framework accurately disentangle epistemic and aleatoric uncertainty without sacrificing predictive accuracy?
• RQ3 How do sampling rates (M and N) affect the reliability of epistemic and aleatoric uncertainty estimates in out-of-distribution scenarios?
• RQ4 Can the method achieve competitive performance on real-world high-stakes tasks like CT reconstruction and weather forecasting while remaining computationally efficient?

Key findings

• Hyper-Diffusion achieves PSNR and SSIM scores on x-ray CT reconstruction that are marginally higher than those of MC-Dropout and comparable to DPS-UQ, despite MC-Dropout suffering from performance degradation due to inference-time dropout.
• Epistemic uncertainty estimates improve with increased sampling of weight sets (M), correctly identifying out-of-distribution features with higher uncertainty and in-distribution regions with lower uncertainty.
• Aleatoric uncertainty estimates become more uniformly distributed and stable as the number of predictions per weight set (N) increases, with irregular peaks disappearing at N ≥ 8.
• The method provides accurate uncertainty disentanglement: aleatoric uncertainty captures inherent data variability, while epistemic uncertainty reflects model uncertainty.
• Hyper-Diffusion achieves performance on par with deep ensembles in both CT reconstruction and weather forecasting, but with significantly reduced training cost.
• The framework enables scalable uncertainty estimation in high-stakes applications without the computational burden of training multiple full models.

Figure 2: Estimating aleatoric uncertainty with hyper-diffusion models. The distribution of sample variance for four hyper-diffusion models trained on datasets with varying noise levels is shown above. Notice that each distribution’s mean is approximately equal to the aleatoric variance $\sigma_{\et