Data-Driven Priors for Uncertainty-Aware Deterioration Risk Prediction with Multimodal Data

This paper introduces $\texttt{MedCertAIn}$, a novel predictive uncertainty framework that leverages data-driven priors derived from cross-modal similarities and modality-specific corruptions to significantly enhance both the performance and reliability of multimodal in-hospital risk prediction using MIMIC-IV and MIMIC-CXR datasets.

The Gist

Imagine you are a doctor in a busy hospital emergency room. You have a patient who is getting worse, and you need to decide: Is this patient going to survive the night, or do they need immediate, intensive care?

In the past, doctors relied on their experience and gut feeling. Today, we have Artificial Intelligence (AI) that can look at the patient's data and make a prediction. But here's the problem: AI is often overconfident. It might say, "This patient is fine," with 99% certainty, even when it's actually wrong. In a hospital, a wrong guess can be fatal.

This paper introduces a new AI system called MedCertAIn. Think of it not just as a doctor's assistant, but as a "Honest Assistant" who knows when they don't know the answer.

Here is how it works, broken down into simple concepts:

1. The Problem: The "Overconfident Robot"

Most current medical AI models are like a student who memorized the textbook but never took a test. When they see a question they've never seen before, they guess anyway and act like they are 100% sure.

• The Risk: If the AI is wrong and doesn't tell you, the doctor might miss a critical warning sign.
• The Goal: We need an AI that can say, "I'm not sure about this one. Please, Doctor, you take a look."

2. The Solution: Giving the AI a "Safety Net"

The authors created MedCertAIn, which uses a special mathematical trick called Bayesian Learning.

• The Analogy: Imagine a regular AI is a single person trying to solve a puzzle. MedCertAIn is like a committee of 100 slightly different people all looking at the same puzzle at the same time.
• If all 100 people agree on the solution, the AI is very confident.
• If the 100 people are arguing and can't agree, the AI knows it's uncertain. It flags the case for a human doctor to review.

3. The Secret Sauce: "Data-Driven Priors"

How does the AI learn to be humble and admit when it's unsure? Usually, you'd need a human expert to label thousands of examples as "hard cases" or "confusing cases." That takes forever and is expensive.

MedCertAIn does this automatically using two clever tricks (called "Data-Driven Priors"):

• Trick A: The "Distorted Mirror" (Data Corruption)
  Imagine you show the AI a picture of a patient's heart scan. Then, you take that picture and do weird things to it: you flip it upside down, add static noise, or cut off a corner.

  • Why? If the AI gets confused by these "broken" versions, it learns that when data looks messy or weird, it should be less confident. It teaches the AI to recognize "garbage" data.

• Trick B: The "Mismatched Puzzle" (Cross-Modal Similarity)
  The AI looks at two types of data for every patient:

  1. Time-Series: Numbers from monitors (heart rate, blood pressure over time).
  2. Images: X-ray photos of the chest.
  • The Trick: The AI checks if the numbers and the picture tell the same story. If the heart rate says "everything is fine" but the X-ray looks terrible, the AI realizes there is a conflict. It learns that when the data sources disagree, it should raise a red flag and say, "I'm confused, human, please check this."

4. The Result: A Smarter Workflow

When MedCertAIn is tested on real hospital data (from the MIMIC database), it does two amazing things:

1. It predicts better: It is more accurate at spotting patients who are at risk of dying than standard AI models.
2. It knows when to quit: It successfully identifies the "tricky" cases where it is likely to be wrong.

The Real-World Impact:
Instead of the AI trying to do everything, it acts like a filter.

• Confident cases: The AI says, "I'm 95% sure this patient is safe," and the doctor moves on to the next patient.
• Uncertain cases: The AI says, "I'm only 60% sure. This looks weird. Doctor, please review this one."

This saves the doctors' time (they don't have to check every single patient) and saves lives (they focus their attention on the patients the AI is worried about).

Summary

MedCertAIn is a new type of medical AI that combines X-rays and vital signs to predict patient health. But its superpower isn't just guessing right; it's knowing when it's wrong. By teaching itself to recognize confusing data and conflicting information, it acts as a reliable partner that knows when to step back and let the human doctor take the wheel. This makes AI safer and more trustworthy for high-stakes medical decisions.

Technical Summary

1. Problem Statement

The integration of machine learning into clinical decision support systems requires not only high predictive accuracy but also reliable uncertainty quantification (UQ). In safety-critical healthcare settings, models must be able to express when they are uncertain about a prediction, allowing clinicians to defer difficult cases for human review.

Current challenges include:

• Lack of Reliable UQ: Standard deterministic models cannot estimate uncertainty, and existing Bayesian methods often rely on uninformative priors (e.g., standard Gaussian) that fail to capture complex data distributions.
• Multimodal Complexity: Clinical data is inherently multimodal (e.g., time-series EHR data and Chest X-rays). Existing UQ research focuses heavily on unimodal settings (mostly imaging). Fusing different modalities introduces new sources of uncertainty, such as modality disagreement, temporal misalignment, and missing data.
• Scalability: Current methods often require manual annotation of "high-uncertainty" cases by domain experts, which is not scalable for real-world deployment.

2. Methodology: MedCertAIn

The authors propose MedCertAIn, a multimodal uncertainty-aware framework that combines Bayesian learning with data-driven priors. The core innovation is the construction of "context sets" (samples representing high uncertainty) without manual labels, which are used to regularize the model's learning objective.

A. Architecture

• Backbone: The framework uses MedFuse, a multimodal fusion architecture.
  • EHR Encoder: A two-layer LSTM processes clinical time-series data.
  • CXR Encoder: A ResNet-34 processes chest X-ray images.
  • Fusion: Features are concatenated and passed to a classifier for in-hospital mortality prediction.
• Stochastic Extension: The deterministic backbone is converted into a Stochastic Neural Network using Variational Inference (VI). Instead of fixed weights, parameters are treated as random variables with a posterior distribution approximated via Gaussian mean-field VI.

B. Data-Driven Priors (The Core Contribution)

To overcome the limitations of uninformative priors, MedCertAIn constructs a context-conditioned prior $p(\theta | X_c)$ using two label-free strategies to identify "high-uncertainty" samples ($X_c$):

1. Modality-Specific Data Perturbations:

   • EHR: Controlled corruptions including dropping initial sequence segments, adding Gaussian noise, and reversing time dimensions.
   • CXR: Seven standard augmentations (random crop, flips, blur, solarization, inversion, color jitter).
   • Goal: Simulate distribution shifts and data noise to force the model to learn uncertainty in corrupted regions.

2. Cross-Modal Similarity in Latent Space:

   • A self-supervised ConVIRT model (trained with InfoNCE loss) generates latent representations for EHR and CXR.
   • Selection Criteria: Samples where the cosine distance between the EHR and CXR latent representations is high (indicating modality mismatch or disagreement) are selected.
   • Goal: Identify cases where the modalities provide conflicting information, which is a strong signal of predictive uncertainty.

C. Training Objective

The variational objective function is augmented with an uncertainty regularization term. The model is trained to maximize the expected log-likelihood on the standard training set while simultaneously learning to assign high uncertainty (low confidence) to the context set $X_c$.

The objective function $F(q_\Theta)$ is:
$$F(q_\Theta) \approx \mathbb{E}_{q_\Theta}[\log p(y_D | x_D, \Theta)] - \text{KL}(q_\Theta \parallel p_\Theta) + \mathbb{E}_{q_\Theta}[\mathbb{E}_{X_c}[\log \tilde{p}(0 | X_c, \Theta)]]$$

• Term 1: Standard log-likelihood (fitting the data).
• Term 2: KL Divergence (regularization against the prior).
• Term 3: Uncertainty Regularization: Encourages the model to predict a specific "uncertain" distribution (modeled as a Dirac delta at 0) for the context set $X_c$.

D. Evaluation Strategy

• Selective Prediction: The model ranks predictions by Shannon Entropy. Clinicians can set a threshold to "reject" (defer) predictions with high uncertainty.
• Metrics: Standard AUROC/AUPRC are reported alongside Selective AUROC and Selective AUPRC, which measure performance only on the subset of cases the model chooses to predict (i.e., low-uncertainty cases).

3. Key Contributions

1. MedCertAIn Framework: A novel multimodal uncertainty-aware framework that augments existing architectures with principled uncertainty estimates without requiring architectural changes to the fusion module.
2. Label-Free Prior Construction: A flexible pipeline to automatically construct high-uncertainty context sets using data corruptions and self-supervised latent-space divergence, eliminating the need for expert-annotated uncertainty datasets.
3. Empirical Validation: Comprehensive evaluation on the MIMIC-IV (EHR) and MIMIC-CXR (Chest X-ray) datasets for in-hospital mortality prediction.
4. Open Source: Implementation in JAX with public code availability to promote research in uncertainty-aware multimodal modeling.

4. Results

The framework was evaluated against deterministic baselines (MedFuse, DrFuse) and a stochastic baseline with an uninformative Gaussian prior.

• Predictive Performance: MedCertAIn achieved state-of-the-art performance with an AUROC of 0.835 and AUPRC of 0.498, outperforming the deterministic MedFuse baseline (AUROC 0.819, AUPRC 0.466).
• Uncertainty & Selective Prediction: The most significant gains were in selective prediction metrics:
  • Selective AUROC: 0.857 (vs. 0.659 for MedFuse, a 30% improvement).
  • Selective AUPRC: 0.599 (vs. 0.203 for MedFuse, a 195% improvement).
  • This indicates MedCertAIn is exceptionally good at identifying and deferring its own uncertain predictions, thereby increasing the reliability of the predictions it does make.
• Ablation Studies: Combining both corruption strategies and cross-modal similarity (MedCertAIn I) yielded the best results, proving that data-driven priors are superior to random corruptions or similarity selection alone.
• Subgroup Analysis: The model consistently outperformed baselines across different age groups and sexes, demonstrating robustness in diverse patient populations.

5. Significance

• Clinical Safety: By enabling selective prediction, MedCertAIn allows clinical workflows to automatically flag low-confidence cases for human review, reducing the risk of AI-induced errors in critical care.
• Scalability: The label-free approach to constructing uncertainty priors makes the method applicable to diverse clinical tasks and modalities without the bottleneck of expert annotation.
• Multimodal UQ: It addresses a critical gap in the literature by demonstrating how to effectively quantify uncertainty in multimodal settings, where modality disagreement is a key source of risk.
• Trustworthy AI: The results suggest that Bayesian multimodal learning with data-driven priors is a viable path toward deploying robust, calibrated, and trustworthy AI tools in high-stakes hospital environments.