Source-linked AI summary
Explainable Uncertainty Estimation for Reliable Medical AI
Li Rong Wang, Jamie Duell, Xinran Xu, Thomas C. Henderson, Yu Yue Hew, Pik Wan Erica Chiang, Xiao Wei Alstar Ang, Bingwen Eugene Fan, Xiuyi Fan
TL;DR
Medical AI adoption is limited by trust, and existing uncertainty and explanation methods do not reveal which features drive uncertain predictions. The paper introduces egRUE, which integrates Expected Gradients with reconstruction uncertainty and decomposes uncertainty feature-wise. Across four healthcare datasets, egRUE improves several reliability measures over existing methods, while a medical-expert study finds that its explanations improve calibrated trust over uncertainty scores alone.
Problem
AI deployment in healthcare is constrained by limited trust, while uncertainty estimation and XAI are treated separately without feature-level explanations of prediction uncertainty.
Method
egRUE integrates Expected Gradients feature attributions into Reconstruction Uncertainty Estimate and decomposes uncertainty into feature-wise contributions.
Results
151–330% gains in reliability and 4–53% gains in misclassification detection over Deep Ensembles were reported, alongside improved calibrated trust from explanations.
Takeaways & Limitations
egRUE clarifies when predictions may be unreliable and which features drive uncertainty, supporting more calibrated medical decision-making.
Abstract
from arXiv · showhide
Artificial intelligence has strong potential to support clinical decision-making, yet its adoption in healthcare remains limited due to a lack of trust. Uncertainty estimation can signal unreliable predictions, and explainable AI (XAI) can clarify how predictions are made but existing methods treat them separately, providing no feature-level insight into why a prediction is uncertain or which tests to prioritize to reduce it. To address this gap, we propose explainable uncertainty estimation, which unifies uncertainty estimation and XAI to both quantify uncertainty and explain feature-level contributions. We introduce the Expected Gradients Reconstruction Uncertainty Estimate (egRUE), which incorporates prediction explanations into its uncertainty computation and decomposes uncertainty into feature-wise contributions. We prove theoretical properties of egRUE and show through experiments that it improves reliability and interpretability compared to existing methods. A user study with medical experts further demonstrates that egRUE's explanations improve calibrated trust over uncertainty scores alone, increasing confidence in correct predictions and reducing confidence in incorrect ones. By combining prediction uncertainty with feature-level explanations, egRUE strengthens decision-making support in safety-critical healthcare settings, clarifying both when predictions may be unreliable and which features drive that uncertainty.
I. INTRODUCTION
Clinical AI adoption is limited by trust, while uncertainty estimation and XAI separately leave unclear why predictions are uncertain. The paper therefore combines feature attribution with uncertainty estimation in egRUE.
- AI shows promise for clinical decision-making, but deployment remains constrained by a lack of trust.
- Uncertainty estimation quantifies prediction uncertainty, whereas XAI explains feature importance; existing methods generally do not attribute uncertainty to specific sources.
- In a lung cancer screening example, egRUE attributed high uncertainty of 0.48 to the feature Smoker for an incorrect positive prediction.
- Combining uncertainty estimation and XAI can clarify both model uncertainty and the features contributing to it.
- The work introduces egRUE, evaluates theoretical properties and four real-world datasets, and studies calibrated trust with practicing clinicians.
- Expected Gradients provides feature attributions by averaging gradients over multiple baseline inputs, supporting the explanatory component of egRUE.
III. METHODOLOGY
The methodology builds explainable uncertainty estimation on reconstruction-based uncertainty and feature attribution. egRUE estimates uncertainty with attribution-informed weighting and decomposes it into feature-wise contributions.
- The framework comprises egRUE, an XAI-informed uncertainty estimate, and egRUE Explanations, its feature-wise uncertainty decomposition.
- A. Uncertainty Estimation: Uncertainty estimation maps an instance to a real value quantifying uncertainty for a prediction function applied to that instance.
- 1) RUE: RUE uses an encoder, predictor head, and decoder, with decoder reconstruction error serving as the uncertainty estimate.
- 1) RUE: RUE assumes similar-to-training instances are predicted and reconstructed more accurately than dissimilar instances.
2) egRUE:
egRUE extends reconstruction-based uncertainty estimation by weighting feature-wise reconstruction errors according to Expected Gradients feature attributions, producing targeted uncertainty estimates.
- egRUE: RUE assumes that each feature contributes equally to prediction uncertainty, whereas egRUE assigns greater weight to important features.The method is motivated by the hypothesis that inability to reconstruct important features indicates higher uncertainty.
- egRUE: egRUE extends RUE with feature-attribution values from Expected Gradients using the same prediction model.Feature-wise reconstruction errors are formally defined before computing the extended estimate.
- Expected Gradients adaptation: Expected Gradients is adapted by using the RUE encoder and predictor head, training-distribution samples as baselines, and the input instance as the target.These modifications define the feature-attribution values used by egRUE.
- Uncertainty computation: egRUE normalizes absolute Expected Gradients attributions into feature-importance weights and applies them to feature-wise reconstruction errors.The resulting uncertainty is a feature-importance-weighted sum rather than an unweighted reconstruction-error aggregate.
B. Uncertainty Explanation
The paper explains uncertainty by decomposing aggregate estimates into feature-wise contributions and evaluates both estimate quality and explanation informativeness across four medical datasets.
- Uncertainty Explanation: Uncertainty explanations map an input instance to feature-wise values that quantify each feature’s contribution to prediction uncertainty.This formalizes explanations as a mapping from R^J to R^J.
- Uncertainty Explanation: RUE explanations use feature-wise reconstruction errors, while egRUE explanations multiply those errors by feature-importance values.Both explanations follow from the aggregate structure of their corresponding uncertainty estimates.
- Evaluation: The evaluation measures reliability, misclassification detection, selective prediction, robustness to false negatives, out-of-distribution detection, and explanation informativeness.Explanation informativeness is assessed qualitatively and quantitatively using Remove and Retrain (ROAR).
- Results: Experiments on four medical diagnosis datasets show egRUE consistently outperforming existing methods across reliability, misclassification detection, selective prediction, and robustness to false negatives.Relative to Deep Ensembles, reported gains are 151–330% in reliability, 4–53% in misclassification detection, 32–77% in selective prediction, and 78–98% in robustness to false negatives; OOD detection improves by 11-97%.
B. Uncertainty Estimation Baselines
The study compares egRUE with established uncertainty baselines using metrics for reliability, selectivity, robustness to false negatives, and OOD detection, alongside qualitative explanation comparisons.
- Baselines: Compared baselines include Entropy, Bayesian Neural Networks, Monte Carlo Dropout, Deep Ensembles, Deep Evidential Classification, Posterior Networks, and Gaussian Processes.These methods provide reference points for evaluating egRUE uncertainty estimates.
- Metrics: The evaluation metrics are correlation, misclassification-prediction AUROC, AURC, σ-risk, and OOD-detection AUROC.They assess reliability, selective prediction, robustness to false negatives, and separation of in-distribution from OOD instances.
- Table I results: egRUE consistently performs well across reliability, selectivity, and robustness to false negatives on all four datasets.It achieves the highest correlation on Lung Cancer, OCTMNIST, and BloodMNIST, and the second highest on Colorectal.
- Explanation comparison: In mispredicted OCTMNIST and BloodMNIST examples, egRUE focuses on edema or an atypical nucleus, while EG and RUE produce broader or noisier regions.The figure contrasts prediction explanations from EG with uncertainty explanations from RUE and egRUE.
- OOD detection: egRUE achieves the highest OOD-detection AUROC on both OCTMNIST and BloodMNIST across OOD datasets with varying similarity to training data.This result is reported in Table II for image datasets.
2) Uncertainty Explanation:
egRUE produces uncertainty explanations that identify clinically meaningful features and align more closely with true feature-level uncertainty than RUE. Across evaluations, it improves uncertainty-estimate performance and supports more calibrated trust among medical experts.
- Explanation quality: egRUE highlights clinically meaningful uncertain features, whereas EG and RUE produce noisy or nonspecific explanations.Examples include edema-associated regions and atypical nuclear morphology in blood-cell images.
- Explanation quality: egRUE generally yields more stable and higher ROAR AUC than RUE after uncertain features are removed and the model is retrained.AUC was used instead of accuracy because the datasets were class-imbalanced.
- Uncertainty-estimate performance: egRUE outperforms RUE across four datasets by incorporating feature-attribution weights into uncertainty estimation.The comparison covers reliability, selectivity, and robustness to false negatives.
- Explanation quality: On BloodMNIST, egRUE achieves the highest ROAD score among evaluated attribution methods paired with RUE.ROAD progressively removes features ranked by uncertainty and measures the area under the accuracy–removal curve.
- Computational efficiency: egRUE requires 0.002s per instance, faster than DE (0.0103s), MCD (0.0125s), and BNN (0.0126s), while Entropy and DEC are slightly faster.Entropy requires 0.0013s and DEC 0.0014s per instance.
- User study: With uncertainty scores alone, experts agreed with incorrect model predictions in 40% of responses; explanations improved calibrated trust for correct and incorrect predictions.Confidence increased for correct predictions and decreased for incorrect predictions, while effects for a third case were mixed.
VI. THEORETICAL EVALUATION
The theoretical evaluation establishes bounds connecting prediction error with RUE and egRUE uncertainty, then verifies formal properties of their uncertainty explanations. egRUE inherits error-bounding guarantees while satisfying implementation invariance, sensitivity, and consistency.
- Uncertainty bounds: Both RUE and egRUE provide bounds that relate model prediction error to their uncertainty estimates.These bounds are presented as necessary for aligning model confidence with actual performance.
- Uncertainty bounds: RUE is lower-bounded by egRUE and upper-bounded by egRUE multiplied by a constant.The relation follows because RUE sums feature-wise reconstruction errors, whereas egRUE weights them by normalized feature importance.
- Uncertainty bounds: Theorem 2 states that positive constants bound prediction error below and above using egRUE uncertainty for every instance.The result extends the RUE error bound through the relation between RUE and egRUE.
- Attribution properties: For functionally equivalent prediction models, implementation invariance requires identical uncertainty explanations.Theorem 3 establishes this property for ζR, and Theorem 4 extends it to ζe.
- Attribution properties: Consistency requires feature attributions to preserve relative feature importance across prediction models.The property prevents a feature receiving greater predictive contribution from receiving a decreased uncertainty attribution.
- Attribution properties: egRUE’s uncertainty attribution satisfies implementation invariance, sensitivity, and consistency.Sensitivity assigns zero attribution to features on which the prediction function does not depend.
VII. CONCLUSION
The paper proposes egRUE to quantify prediction uncertainty while explaining feature-level contributions. Theory, experiments, and a medical-expert user study support its reliability, interpretability, and calibrated-trust benefits.
- Conclusion: egRUE integrates Expected Gradient scores into RUE to quantify uncertainty and explain how individual features contribute to it.The method is designed to address trust limitations affecting medical AI adoption.
- Conclusion: Across four healthcare datasets, egRUE improves reliability, misclassification detection, selective prediction, robustness to false negatives, and out-of-distribution detection compared with state-of-the-art methods.It also provides explanations that align more closely with true feature-level uncertainties than RUE.
- Conclusion: A medical-expert user study found that egRUE explanations improve calibrated trust by increasing confidence in correct predictions and reducing confidence in incorrect ones.The study compares explanations with uncertainty scores alone.
APPENDIX
The appendix details assumptions and recursive loss arguments used to prove RUE’s prediction-error bounds. It defines layer-wise errors and losses and records the supporting inequalities and proof conclusion.
- Experimental protocol: The experiments were repeated 10 times for tabular data and 5 times for image data, with means and standard deviations reported.Architecture, hyperparameters, seeds, and computing infrastructure are documented in the code repository.
- Proof assumptions: Theorem 1 assumes an L-layer MLP with linear activations, weight matrices, biases, an encoder, and a predictor head.Layers 1 through M form the encoder, while layers M+1 through L form the predictor head.
- Proof construction: Layer-wise error is defined as a vector, and layer-wise loss is its L1 norm at each layer.The proof relates layer-wise loss at the encoder output to prediction and reconstruction error.
- Proof construction: Lemma 1 establishes recursive upper bounds on layer-wise losses using positive constants p and q.The bounds are derived through per-output losses, maximum weights, triangle inequalities, and recursive application across layers.
C. Proof of Theorem 3
The proof establishes implementation invariance for ζR by tying it to feature-wise reconstruction errors produced by the egRUE architecture.
- ζR returns feature-wise reconstruction errors from the egRUE network architecture.
- Functionally equivalent networks receive the same input x and produce the same output (g, fψ ◦fϕ(x)).
- Therefore, ζR preserves reconstruction error for the same inputs and outputs regardless of implementation.
D. Proof of Theorem 4
The proof argues that egRUE inherits implementation invariance, sensitivity, and consistency from gradient-based explanations and feature-wise reconstruction errors.
- Theorem 3 supplies implementation invariance for the feature-wise reconstruction error ζR.
- Expected Gradients satisfies implementation invariance because it reformulates Integrated Gradients as an expectation over baselines.
- For egRUE, gradients through the predictor head fψ provide explanations that remain unchanged across equivalent implementations with the same input and output.
- Sensitivity holds because a feature xj irrelevant to fψ receives a zero returned value.
- Consistency follows because Expected Gradients reflects changes in prediction probability, while ζe normalizes that reflection for the same input and ζR.