Source-linked AI summary
Analyzing the Role of Model Uncertainty for Electronic Health Records
Michael W. Dusenberry, Dustin Tran, Edward Choi, Jonas Kemp, Jeremy Nixon, Ghassen Jerfel, Katherine Heller, Andrew M. Dai
TL;DR
Medical prediction has serious ethical and monetary stakes, but common approaches do not address uncertainty over models. The paper studies model uncertainty with RNN ensembles and Bayesian RNNs, finding substantial patient-level disagreement despite similar dataset metrics. This variability can indicate brittle decisions and motivates examining uncertainty for individual patients.
Problem
Medical deep-learning approaches often focus on data uncertainty, leaving model uncertainty and patient-specific confidence insufficiently addressed despite the high stakes of incorrect predictions.
Source: §1 · p. 1Method
The authors study model uncertainty in predictive medicine using RNN ensembles and Bayesian RNN variants, including Bayesian RNNs with priors on subsets of parameters.
Source: §1 · p. 2, §2.4 · p. 3Results
Takeaways & Limitations
Patient-specific prediction and decision variability can help identify potentially brittle decisions and provide an opportunity to collect additional data that could reduce model uncertainty.
Source: §5 · p. 9Future work includes more specific clinically relevant decision costs, decisions accounting for both data and model uncertainty, and methods to reduce model uncertainty.
Source: §5 · p. 9
Abstract
from arXiv · showhide
In medicine, both ethical and monetary costs of incorrect predictions can be significant, and the complexity of the problems often necessitates increasingly complex models. Recent work has shown that changing just the random seed is enough for otherwise well-tuned deep neural networks to vary in their individual predicted probabilities. In light of this, we investigate the role of model uncertainty methods in the medical domain. Using RNN ensembles and various Bayesian RNNs, we show that population-level metrics, such as AUC-PR, AUC-ROC, log-likelihood, and calibration error, do not capture model uncertainty. Meanwhile, the presence of significant variability in patient-specific predictions and optimal decisions motivates the need for capturing model uncertainty. Understanding the uncertainty for individual patients is an area with clear clinical impact, such as determining when a model decision is likely to be brittle. We further show that RNNs with only Bayesian embeddings can be a more efficient way to capture model uncertainty compared to ensembles, and we analyze how model uncertainty is impacted across individual input features and patient subgroups.
1 INTRODUCTION
Medical predictions carry serious ethical and monetary stakes, yet existing deep-learning approaches often focus on data uncertainty rather than uncertainty about the model. This paper examines model uncertainty in predictive medicine, including its relationship to patient-specific predictions and clinical decisions.
Motivation
- Incorrect medical predictions can carry severe ethical and monetary costs, while individualized predictions are important in medicine.Source: §1 · p. 1
- Existing medical deep-learning approaches often focus on accurate predictive distributions but do not address model uncertainty, leaving uncertainty about individual predictions indistinguished despite high AUC-ROC.The authors identify this as a gap in approaches using a single model or a small ensemble average.Source: §1 · p. 1
Research questions
- The study examines model uncertainty in predictive medicine alongside methodological approaches including Bayesian neural networks, Monte Carlo dropout, deep ensembles, and function priors.The authors frame this work around questions about population-level metrics, calibration, patient subgroups, input features, and clinical decisions.Source: §1 · pp. 1–2
Contributions
- Models with nearly equivalent dataset-level metrics can disagree widely on individual patients’ predicted probabilities and optimal decisions, which may signal brittle decisions.The authors also report that a single model or ensemble average can mask this variability.Source: §1 · p. 2
2 BACKGROUND
The background distinguishes outcome uncertainty from uncertainty over model-induced predictive distributions, then describes calibration and two approaches for representing model uncertainty.
2.1 Data Uncertainty
- Data uncertainty concerns uncertainty about an outcome given incomplete information; binary outcomes use a Bernoulli predictive distribution, while multiclass and regression tasks can use other distributions.Source: §2.1 · p. 2
2.2 Model Uncertainty
- Model uncertainty is represented by a distribution over functions, inducing a patient-specific distribution over predictive-distribution parameters before marginalization.That distribution can quantify disagreement; the marginalized prediction combines data and model uncertainty, which cannot be separated from the marginalized distribution alone.Source: §2.2 · p. 2
2.3 Calibration
- Expected calibration error approximates calibration by taking a prediction-count-weighted average of calibration errors across bins.Perfect calibration means predictions at probability p correspond to an observed outcome frequency of p.Source: §2.3 · pp. 2–3
2.4 Deep Ensembles
- Deep ensembles train M deterministic neural networks with different random seeds and treat their predictions for one example as samples from its predictive uncertainty distribution.The paper applies this approach to RNNs for sequential patient data.Source: §2.4 · p. 3
2.5 Bayesian RNNs
- Bayesian RNNs place priors on all or a subset of model weights and use posterior draws to sample predictive outcomes and capture model uncertainty.The study uses variants with priors on different parameter subsets.Source: §2.5 · p. 3
3 MEDICAL UNCERTAINTY
The study examines model uncertainty in clinical RNNs across EHR datasets, comparing ensemble and Bayesian approaches and representing uncertainty in both predictions and optimal decisions.
3.1 Clinical Tasks
- The evaluation covers MIMIC-III mortality and discharge-diagnosis prediction and eICU mortality, with the additional dataset used to assess generalization.Source: §3.1 · p. 3
3.2 Models
- The study distinguishes data uncertainty in a patient's outcome distribution from model uncertainty across equally likely predictive distributions.Its RNNs embed and aggregate sequential and contextual patient features before producing the predictive distribution parameterization.Source: §3.2 · p. 4
3.3 Choice of Uncertainty Methods
- Ensembles vary only the random seed across replicas, while Bayesian RNNs draw prediction samples from a single model with factorized weight posteriors.Bayesian weight tensors use normal distributions with learnable means and diagonal covariance; training minimizes expected negative log likelihood plus KL regularization.Source: §3.3 · p. 4
3.4 Optimal Decisions via Sensitivity Requirements
- The sensitivity-based decision rule maximizes precision while maintaining a target sensitivity, and each ensemble model yields an optimal decision for every example.The resulting set of model decisions can be represented as a distribution whose variance represents disagreement from model uncertainty.Source: §3.4 · pp. 4–5
4 EXPERIMENTS
The experiments test whether RNN ensembles and Bayesian variants reveal patient-level uncertainty that aggregate metrics do not capture, then examine how uncertainty varies across patient subgroups and input features.
4.1 When Do We Observe Uncertainty?
- Across MIMIC-III mortality and CCS tasks, ensemble models have nearly identical dataset-level metrics, but selecting one model removes model-uncertainty information.The table reports means and standard deviations across 200 models.Source: Table 1 · p. 5, §4.1 · p. 5
- Some patients show up to 57.5% spread in mortality probabilities, with patient-specific uncertainty extending into variable optimal decisions at a 70% recall target.The decision distributions can have high variance, and predictive-distribution standard deviations vary by patient.Source: §4.1 · pp. 5–6, Figure 4 · p. 6
4.2 Comparison: Variants of Bayesian RNNs and Deterministic RNN Ensembles
- Bayesian embeddings slightly exceed ensemble AUC-PR on MIMIC-III, align qualitatively in patient uncertainty, and use 7.22 million parameters versus M × 6.16 million for the ensemble.Across both datasets, Bayesian variants are comparable to or outperform the deterministic ensemble in held-out NLL.Source: Table 2 · p. 6, §4.2 · pp. 7–8
4.1 When Do We Observe Uncertainty?
- In multiclass CCS prediction, ensemble members are relatively certain for some patients but disagree more for others.Figure 5 displays maximum-probability distributions on different x-axis scales and the corresponding predicted classes.Source: §4.1 · p. 6, Figure 5 · p. 7
4.3 Patient Subgroup Analysis
- Ensemble AUC-PR is negatively correlated between male and female patients (r = −0.442), while uncertainty and mortality both increase across age groups and are slightly higher in women.AUC-PR for the oldest adult quartile is also somewhat negatively correlated with other adults and neonates (r between −0.18 and −0.37).Source: §4.3 · p. 8
4.4 Embedding Uncertainty Analysis
- Rarer clinical-note words generally have higher Bayesian embedding uncertainty, but the relationship with frequency is nonlinear and context-specific.The analysis measures word-level uncertainty using entropy of each word’s embedding distribution.Source: §4.4 · p. 9, Figure 8 · p. 9, Table 3 · p. 9
5 CONCLUSION
The experiments show that RNN ensembles can capture patient-specific uncertainty even when their models have nearly equivalent dataset-level metrics, and that predictive disagreement can extend to optimal decisions. The paper identifies Bayesian embeddings as a way to capture uncertainty and examine individual feature values, while future work includes more specific clinical decision costs and methods to reduce uncertainty.
5 CONCLUSION
- Ensemble models with nearly equivalent dataset-level metrics can disagree substantially on some patients’ predictions and optimal decisions.Such variability can indicate a brittle decision and an opportunity to collect additional data that could reduce model uncertainty.Source: §5 · p. 9
- RNNs with only Bayesian embeddings can capture model uncertainty and support examination of uncertainty associated with individual feature values.The paper presents this as an additional benefit of the Bayesian-embedding approach.Source: §5 · p. 9
- Future work includes designing more specific clinically relevant decision costs and exploring methods to reduce model uncertainty during training and prediction.The proposed cost functions would draw on quantified medical ethics and monetary axes.Source: §5 · p. 9
A.1 Additional Metrics and Statistics
On the mortality task, the authors report no apparent correlation between deterministic-ensemble validation AUC-PR and held-out log-likelihood, likely because model differences were limited. The ensemble and Bayesian RNNs were all well-calibrated on marginalized predictions, with marginalization slightly decreasing calibration error.
A.1 Additional Metrics and Statistics
- The deterministic RNN ensemble showed no apparent correlation between validation AUC-PR and held-out log-likelihood on mortality.The authors suggest this may reflect limited differences between models; Figure 9 shows the comparison.Source: §A.1 · p. 11
- Marginalized predictions were well-calibrated across the deterministic ensemble and Bayesian RNNs on MIMIC-III mortality.Marginalization slightly decreased the deterministic ensemble’s calibration error, and the Bayesian models were comparably well-calibrated.Source: §A.1 · p. 11
A.2 Additional Training Details
The appendix reports the hyperparameter search and final model-specific settings, along with implementation and dataset-splitting details. It also reiterates that the ensemble’s validation AUC-PR and held-out log-likelihood had no apparent correlation, with limited model differences offered as a likely explanation.
A.2 Additional Training Details
- Hyperparameters were searched for the original deterministic RNN and each Bayesian model, with final model-specific values reported separately.Other deterministic ensemble members differed from the original only by random seed; Tables 5 and 6 provide the search ranges and final values.Source: §A.2 · p. 11, Table 5 · p. 12, Table 6 · p. 12
- Models were implemented in TensorFlow 2.0, trained with Adam on Nvidia V100 machines, and the MIMIC-III and eICU datasets were split 8:1:1 into train, validation, and test sets.Source: §A.2 · p. 11
- The deterministic ensemble’s validation AUC-PR showed no apparent correlation with held-out log-likelihood on mortality.The authors attribute this likely to limited differences between models; Figure 9 presents the metric comparison.Source: Figure 9 · p. 11
Related papers
This paper cites
- ADAM: A METHOD FOR STOCHASTIC OPTIMIZATION2014
- On Calibration of Modern Neural Networks2017
- TensorFlow: A system for large-scale machine learning2016
- Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning2015
- Weight Uncertainty in Neural Networks2015
- BatchEnsemble: An Alternative Approach to Efficient Ensemble and Lifelong Learning2020