Not Just How Much, But Where: Decomposing Epistemic Uncertainty into Per-Class Contributions

2026-02-24

AI Summary

This work addresses a key limitation in existing Bayesian deep learning approaches, which rely solely on a scalar mutual information metric to quantify epistemic uncertainty and thus cannot distinguish the model’s ignorance across safety-critical versus non-critical classes. The authors propose decomposing mutual information into a per-class contribution vector \( C_k(x) = \sigma_k^{2}/(2\mu_k) \), derived from the mean and variance of class probabilities over posterior samples, with a second-order Taylor expansion ensuring that the sum of these contributions approximates mutual information. This formulation enables, for the first time, interpretable and comparable per-class epistemic uncertainty estimates, mitigates boundary suppression, facilitates fair comparison between rare and common classes, and introduces a skewness-based metric to assess approximation quality. Experiments demonstrate a 34.7% reduction in risk for critical classes in diabetic retinopathy selective prediction, state-of-the-art AUROC in out-of-distribution detection—revealing asymmetric distribution shifts invisible to mutual information—and enhanced robustness to label noise.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Other Foundations of Reasoning under Uncertainty
• Computer Vision: Adversarial Attacks & Robustness

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Security and Privacy: Large-scale security measurements
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

In safety-critical classification, the cost of failure is often asymmetric, yet Bayesian deep learning summarises epistemic uncertainty with a single scalar, mutual information (MI), that cannot distinguish whether a model's ignorance involves a benign or safety-critical class. We decompose MI into a per-class vector $C_k(x)=\sigma_k^{2}/(2\mu_k)$, with $\mu_k{=}\mathbb{E}[p_k]$ and $\sigma_k^2{=}\mathrm{Var}[p_k]$ across posterior samples. The decomposition follows from a second-order Taylor expansion of the entropy; the $1/\mu_k$ weighting corrects boundary suppression and makes $C_k$ comparable across rare and common classes. By construction $\sum_k C_k \approx \mathrm{MI}$, and a companion skewness diagnostic flags inputs where the approximation degrades. After characterising the axiomatic properties of $C_k$, we validate it on three tasks: (i) selective prediction for diabetic retinopathy, where critical-class $C_k$ reduces selective risk by 34.7\% over MI and 56.2\% over variance baselines; (ii) out-of-distribution detection on clinical and image benchmarks, where $\sum_k C_k$ achieves the highest AUROC and the per-class view exposes asymmetric shifts invisible to MI; and (iii) a controlled label-noise study in which $\sum_k C_k$ shows less sensitivity to injected aleatoric noise than MI under end-to-end Bayesian training, while both metrics degrade under transfer learning. Across all tasks, the quality of the posterior approximation shapes uncertainty at least as strongly as the choice of metric, suggesting that how uncertainty is propagated through the network matters as much as how it is measured.

Problem

• epistemic uncertainty
• safety-critical classification
• Bayesian deep learning
• asymmetric risk
• uncertainty decomposition

Innovation

• epistemic uncertainty
• per-class decomposition
• Bayesian deep learning
• mutual information
• selective prediction