Retrieval-aligned Tabular Foundation Models Enable Robust Clinical Risk Prediction in Electronic Health Records Under Real-world Constraints
cs.AI
Submitted: 2026-04-02
Updated: 2026-08-25
Comments: Not peer-reviewed. Under revision
DOI: 10.21203/rs.3.rs-9085469/v1
License: http://creativecommons.org/licenses/by/4.0/
The gist: Clinical prediction from structured electronic health records (EHRs) is challenging due to high dimensionality, heterogeneity, class imbalance, and distribution shift.
Terminology
Abstract
Clinical prediction from structured electronic health records (EHRs) is challenging due to high dimensionality, heterogeneity, class imbalance, and distribution shift. While tabular in-context learning (TICL) and retrieval-augmented methods perform well on generic benchmarks, their behavior in clinical settings remains unclear. We present a multi-cohort EHR benchmark comparing classical, deep tabular, and TICL models across varying data scale, feature dimensionality, outcome rarity, and cross-cohort generalization. PFN-based TICL models are sample-efficient in low-data regimes but degrade under naive distance-based retrieval as heterogeneity and imbalance increase. We propose AWARE, a task-aligned retrieval framework using supervised embedding learning and lightweight adapters. AWARE improves AUPRC by up to 12.2% under extreme imbalance, with gains increasing with data complexity. Our results identify retrieval quality and retrieval-inference alignment as key bottlenecks for deploying tabular in-context learning in clinical prediction.
Sources
- EHRMamba: Towards Generalizable and Scalable Foundation Models for Electronic Health Records
- Analyzing and Improving Representations with the Soft Nearest Neighbor Loss
- Count-Based Approaches Remain Strong: A Benchmark Against Transformer and LLM Pipelines on Structured EHR
- Real-TabPFN: Improving Tabular Foundation Models via Continued Pre-training With Real-World Data
- TabR: Tabular Deep Learning Meets Nearest Neighbors in 2023
- Longitudinal Progression Prediction of Alzheimer's Disease with Tabular Foundation Model
- TabM: Advancing Tabular Deep Learning with Parameter-Efficient Ensembling
- TabArena: A Living Benchmark for Machine Learning on Tabular Data
- A Comprehensive Survey of Foundation Models in Medicine
- An Extensive Data Processing Pipeline for MIMIC-IV
- Large Language Models are Powerful Electronic Health Record Encoders
- TabTransformer: Tabular Data Modeling Using Contextual Embeddings
- Learning to Diagnose with LSTM Recurrent Neural Networks
- Transformers Can Do Bayesian Inference
- IM-Context: In-Context Learning for Imbalanced Regression Tasks
- CORE-BEHRT: A Carefully Optimized and Rigorously Evaluated BEHRT
- CEHR-XGPT: A Scalable Multi-Task Foundation Model for Electronic Health Records
- Scalable In-Context Learning on Tabular Data via Retrieval-Augmented Large Language Models
- Revisiting Nearest Neighbor for Tabular Data: A Deep Tabular Baseline Two Decades Later
Related papers
- MAVEN-T: Reinforced Heterogeneous Distillation for Real-Time Multi-Agent Trajectory Prediction
- Model Discovery Agent: LLM-assisted Bayesian experiment design for data-efficient discovery of mechanistic world models
- Agents' Last Exam
- The Clinician's Veto: Navigating Trust, Liability, and Uncertainty in Autonomous AI Prescribing
- MindHelper: Closed-Loop Embodied Mental-State Reasoning for Precision Intervention
- Incumbent Advantage: Brand Bias and Cognitive Manipulation Dynamics in LLM Recommendation Systems