Quantifying Epistemic Uncertainty in Predictions for Safer Health AI Performance Under Dataset Shifts.

Area of Science:

• Artificial Intelligence
• Machine Learning
• Health Informatics

Background:

• Out-of-distribution (OOD) data poses a significant challenge to the reliability of AI-based clinical decision support systems (CDSSs).
• Ensuring the robustness and safety of AI in healthcare requires methods to detect and manage data that differs from training distributions.

Purpose of the Study:

• To investigate the efficacy of real-time, sample-level epistemic uncertainty quantification as a safety mechanism for health AI.
• To determine if epistemic uncertainty can serve as a lightweight layer for flagging potential OOD samples, guiding model updates and human review.

Main Methods:

• A continual learning-based neural network classifier was trained on a real-world Mexican COVID-19 dataset in quarterly batches.
• Epistemic uncertainty for each prediction was estimated using Monte Carlo Dropout.
• A data-driven uncertainty threshold was established to identify potential OOD samples.

Main Results:

• Samples below the uncertainty threshold demonstrated consistently higher macro-F1 scores, indicating improved performance.
• The method effectively flagged samples that contributed to prediction errors, capturing the majority of inaccuracies.
• Performance remained largely invariant to temporal drift when using the uncertainty screening.

Conclusions:

• Real-time epistemic uncertainty screening provides a practical and efficient safety layer for health AI and CDSSs.
• This approach does not necessitate model retraining, making it suitable for dynamic healthcare environments.
• Uncertainty quantification can enhance the robustness and safety of AI systems by identifying OOD data and guiding interventions.