TRACE: Tracking Multimodal Clinical Trajectories for Inpatient Risk Prediction

AI summary

Overview

Research area: Clinical artificial intelligence — multimodal electronic health record (EHR) modeling for inpatient risk prediction.

Technical level: Advanced. The paper assumes familiarity with Transformer attention masking, supervised contrastive learning, and clinical database schemas (MIMIC-IV, MIMIC-III, eICU).

Scope: A daily-scale "global–local" trajectory model that organizes irregularly recorded multimodal clinical events into a temporally ordered token sequence for three inpatient prediction tasks across three public critical-care datasets.

What This Paper Is About

Hospitals record patient data unevenly: labs, medications, chest x-rays, radiology reports and discharge notes arrive on different days for different patients, and some modalities are missing entirely. Most existing multimodal models either fuse everything into one static admission-level snapshot or explicitly reconstruct the modalities that are missing. TRACE instead keeps the admission as a day-by-day trajectory, separating information that never changes during the stay (demographics, diagnosis-derived context, severity scores) from events tied to specific hospital days, and predicts in-hospital mortality, prolonged length of stay, and 30-day readmission from that trajectory.

Key Contributions

1. A day-level global–local trajectory representation. TRACE splits each admission into day-invariant global context tokens and day-specific local event tokens drawn from structured EHR events, chest x-rays, radiology reports, and discharge summaries, all encoded within a task-specific observation window.

2. Temporal global–local (TGL) attention. A custom attention mask keeps global context always accessible, preserves temporal order among local events (same-day tokens interact bidirectionally for multimodal fusion, later days cannot leak backward), and lets the [PHENO] token attend to the [TASK] token but not the reverse.

3. Missing modalities emit no tokens. Rather than imputing or reconstructing absent modalities — the strategy of M3Care, MUSE/MUSE+, and CRL-MMNAR — TRACE simply produces no token for an unavailable modality, so trajectory length varies naturally with what was actually recorded.

4. Phenotype-guided contrastive learning. Positive pairs require matching both the outcome label and a diagnosis-derived phenotype group (14 groups for MIMIC-III/MIMIC-IV, 12 for eICU), plus a phenotype-aware sampler guaranteeing at least one positive per anchor.

Main Findings

• Best AUROC in all nine settings: TRACE reports the highest AUROC across every dataset–task combination against Large-Ensemble, MultimodalPred, MM-STGNN, M3Care, MUSE+, AdaCoMed, and CRL-MMNAR. Example AUROCs: MIMIC-IV mortality 0.9804, Long LOS 0.7675, readmission 0.8371; eICU mortality 0.9490, Long LOS 0.7971, readmission 0.7584; MIMIC-III mortality 0.9770, Long LOS 0.7864, readmission 0.8306.

• Best ACC in eight of nine settings: The single exception is MIMIC-IV Long LOS, where MUSE+ reports higher ACC (0.7569 vs. TRACE's 0.7409). Gains over the strongest baseline are sometimes narrow — on eICU Long LOS, TRACE's ACC is 0.9443 versus Large-Ensemble's 0.9441.

• AUPRC and mF1 are mixed: The paper's claim is limited to AUROC and ACC. On several settings CRL-MMNAR reports higher AUPRC and mF1, for instance MIMIC-IV mortality (AUPRC 0.9061 vs. 0.8605; mF1 0.9161 vs. 0.8886).

• TGL attention mask ablation: The temporal global–local attention mask achieves the highest AUROC across all three tasks on MIMIC-IV.

• Contrastive objective ablation: Phenotype-guided contrastive learning improves mean phenotype-stratified AUROC across all nine dataset–task settings.

• Attribution localizes risk: Trajectory-level attribution identifies the specific hospital days and modalities associated with changes in predicted risk, giving an interpretable view of how predictions evolve during an admission.

• Cohort scale: Roughly 180,000 admissions and ICU stays in total — 13,819 for MIMIC-IV, 45,667 for MIMIC-III, and 119,741 for eICU (train/validation/test counts summed).

Methodology in Plain English

Each admission is broken into "global" facts that do not change (age, sex, diagnosis-derived summary, and for eICU the APACHE severity summary) and "local" events stamped with a hospital day. Structured EHR variables are turned into value-preserving text descriptors — a lab becomes something like "Hematocrit: 32.1 g/dL," a code becomes a readable concept like "Heart failure" — and all events from one day are concatenated into a single token. Chest x-rays (resized to 224×224) go through a frozen medical vision encoder; radiology reports and discharge summaries go through a frozen clinical language encoder with a 512-token limit, yielding 768-dimensional [CLS] vectors. Everything is projected into a shared space and tagged with a modality embedding and a day embedding, then serialized in a fixed order (C, E, R, D) with same-day items sorted by timestamp.

Two learnable tokens are appended: a task token for prediction and a phenotype token for the contrastive loss. A masked Transformer encodes the sequence, and the task token's state feeds a sigmoid classifier. Training combines binary cross-entropy with a supervised contrastive loss weighted by λ, with mini-batches sampled so every anchor has at least one label-and-phenotype match. Selected hyperparameters averaged λ = 0.42 and τ = 0.09 across settings, chosen per dataset–task on validation data. Models trained up to 30 epochs with AdamW and a warmup–cosine schedule; only the projection layers, trajectory encoder, and prediction/contrastive heads were updated. Cohorts excluded stays shorter than two days and used a patient-level 75:10:15 split; confidence intervals came from 2,000 bootstrap resamples.

Why This Matters

The paper argues that the temporal organization of clinical evidence, not just which modalities are fused, is an underexplored axis in inpatient prediction. Its event-driven design sidesteps a persistent engineering headache — models that require a complete multimodal grid must impute or reconstruct, which can inject fabricated signal into a record. Letting absent modalities contribute nothing is both simpler and arguably more faithful to what clinicians actually saw.

Real-world applications:

• Bed and staffing planning: early Long LOS prediction from the first 48 hours supports resource allocation.
• Deterioration triage: in-hospital mortality risk estimated from the completed admission record.
• Readmission reduction programs: discharge-time 30-day readmission scoring to target follow-up resources.
• Clinical review and audit: trajectory-level attribution points reviewers to the specific days and modalities that moved a patient's risk, supporting chart review and model governance.

Industry relevance spans hospital operations teams, EHR and clinical documentation vendors, and developers building FDA-adjacent clinical decision support, where the ability to run on whatever data exists — rather than a rigid input template — materially simplifies deployment across heterogeneous sites.

Future Directions

• Prospective validation: all evaluation here is retrospective on de-identified public datasets; whether the gains hold in live clinical deployment is not reported.
• Finer temporal granularity: the paper explicitly frames hospital-day granularity as a trade-off among event timing, sparsity, and compute cost, leaving sub-daily resolution as an open question.
• Closing the AUPRC gap: since the strongest baseline reports higher AUPRC and mF1 on several settings, calibrating or reweighting the model for precision-recall behavior under class imbalance is a natural next step.
• Generalizing beyond the ICU cohort structure: phenotype groupings were hand-mapped to 14 ICD chapter-level groups (12 for eICU); whether learned or finer-grained phenotypes improve the contrastive objective is untested here.

Target Audience

Clinical AI and healthcare ML researchers already comfortable with Transformer architectures and multimodal fusion; informatics teams working with MIMIC-IV, MIMIC-III, or eICU; and applied scientists at health systems or EHR vendors who need a risk model that tolerates incomplete, irregularly sampled multimodal records. Readers wanting a gentle introduction to clinical NLP or EHR database construction will find the paper assumes too much background.

Authors’ abstract

Accurate inpatient risk prediction requires modeling heterogeneous clinical evidence that accumulates irregularly over time. This challenge spans in-hospital mortality, prolonged length of stay, and readmission, each with task-specific observation windows. Existing multimodal approaches combine structured health record data, chest x-rays, and clinical text, but primarily focus on modality fusion or missing-modality compensation, with limited attention to temporal trajectories over an admission. We propose TRACE (TRAjectory modeling of multimodal Clinical Events), a daily-scale framework organizing clinical events into day-level trajectories, separating day-invariant global context from day-specific local events drawn from EHR data, chest x-rays, radiology reports, and discharge summaries. The trajectory is encoded using temporal attention, allowing unavailable modalities to emit no tokens. A phenotype-guided contrastive objective further organizes admissions under heterogeneous disease distributions. Across cohorts from MIMIC-IV, eICU, and MIMIC-III comprising nearly 180,000 total stays, we evaluate TRACE on mortality, length of stay, and readmission prediction. TRACE achieves the highest AUROC across all nine dataset-task settings and the highest ACC in eight, outperforming baselines. Ablation studies validate our attention masks and contrastive learning, while trajectory-level attribution localizes risk-associated hospital days and modalities.