Retrieval-aligned Tabular Foundation Models Enable Robust Clinical Risk Prediction in Electronic Health Records Under Real-world Constraints

Minh-Khoi Pham; Thang-Long Nguyen Ho; Thao Thi Phuong Dao; Tai Tan Mai; Minh-Triet Tran; Marie E. Ward; Una Geary; Rob Brennan

cs.AIApr 2, 2026

Retrieval-aligned Tabular Foundation Models Enable Robust Clinical Risk Prediction in Electronic Health Records Under Real-world Constraints

Authors: Minh-Khoi Pham, Thang-Long Nguyen Ho, Thao Thi Phuong Dao, Tai Tan Mai, Minh-Triet Tran, Marie E. Ward, Una Geary, Rob Brennan,

Organizations: School of Computing, Dublin City University, Dublin, Ireland · ADAPT Centre, Dublin City University, Dublin, Ireland · University of Science, Vietnam National University, Ho Chi Minh City, Vietnam · Department of Otolaryngology, Thong Nhat Hospital, Ho Chi Minh City, Vietnam ·

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.

Figures & tables

Panel of Figure 1 — Figure 1: Landscape of EHR benchmarks evaluated in this study. The datasets span wide variation in cohort size, feature dimensionality, outcome prevalence, and clinical domain, illustrating key challenges of (a) rarity, (b) heterogeneity, (c) cross-task, cross-institutional, multi-cohort distribution shift faced by tabular learning methods in real-world EHR settings.

Classical tabular and tree-based models
Model	Adaptation		Training		EHR Properties
	In-context	Retrieval	Pretrained	Task	Feature	Imbalance
	Learning	-augmented	Pretrained	-adaptive	Weighting	Handling
Logistic Regression	✗	✗	✗	✓	✓	✗
k-Nearest Neighbors	✗	✓	✗	✗	✗	✗
XGBoost [ 9 ]	✗	✗	✗	✓	✓	✗

Table 1: Comparison of tabular learning paradigms evaluated in this study and their adaptation characteristics for EHR-based clinical prediction.

Panel of Figure 2 — Figure 2: The challenges of raw EHR data for retrieval and the solution of task-aware alignment. (A) Raw vector is heterogeneous, sparse and task-agnostic, leading to noisy geometry and multi-task misalignment. (B) Task-aligned projection reshapes geometry for label-consistent, effective retrieval.

Figure 3: Mean performance improvement over baseline as AWARE components are progressively introduced across three retrieval-based backbones on MIMIC-IV and eICU. Results are shown for AUROC (top) and AUPRC (bottom), averaged across all clinical tasks and random seeds .

Figure 4: Aggregate performance rankings of tabular models on 12 small-medium EHR datasets, rankings are measured by average score of performance. Rankings summarize relative performance across tasks within each regime.

Panel of Figure 5 — Figure 5: Aggregate performance rankings of tabular models on large-scale EHR datasets (MIMIC-IV and eICU), measured by average (a) AUROC and (b) AUPRC across tasks within each regime.

Panel of Figure 6 — Figure 6: Performance of tabular models under three representative clinical constraints. To improve readability given the large number of models evaluated, highly competitive baselines within each paradigm are consolidated into a single curve labeled Best Classical Model and Best Deep Tabular Model , representing the strongest-performing model in each family under the respective constraint. These are compared with individual tabular in-context learning methods and their AWARE-enhanced variants (dashed lines). Performance is measured using AUPRC across all settings, while overlaid bar charts report the average relative performance gains (in %) obtained by incorporating AWARE. (a) Performance scaling with training data size. (b) Performance degradation under increasing feature heterogeneity. (c) Performance under increasing outcome imbalance.

Method	HIPE
	RDMT		CPE
	AUROC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$
XGB	$0.824_{{\color[rgb]{0.5,0.5,0.5}\pm 0.010}}$	$0.432_{{\color[rgb]{0.5,0.5,0.5}\pm 0.020}}$	$0.755_{{\color[rgb]{0.5,0.5,0.5}\pm 0.043}}$	$0.036_{{\color[rgb]{0.5,0.5,0.5}\pm 0.019}}$
CatB	$0.827_{{\color[rgb]{0.5,0.5,0.5}\pm 0.005}}$	$0.446_{{\color[rgb]{0.5,0.5,0.5}\pm 0.015}}$	$0.645_{{\color[rgb]{0.5,0.5,0.5}\pm 0.078}}$	$0.027_{{\color[rgb]{0.5,0.5,0.5}\pm 0.025}}$
LightG	$0.814_{{\color[rgb]{0.5,0.5,0.5}\pm 0.055}}$	$0.438_{{\color[rgb]{0.5,0.5,0.5}\pm 0.065}}$	$0.719_{{\color[rgb]{0.5,0.5,0.5}\pm 0.092}}$	$0.030_{{\color[rgb]{0.5,0.5,0.5}\pm 0.021}}$
TabNet	$0.827_{{\color[rgb]{0.5,0.5,0.5}\pm 0.005}}$	$0.440_{{\color[rgb]{0.5,0.5,0.5}\pm 0.009}}$	$0.558_{{\color[rgb]{0.5,0.5,0.5}\pm 0.083}}$	${0.111}_{{\color[rgb]{0.5,0.5,0.5}\pm 0.055}}$

Table 2: Performance on HIPE prediction tasks. Results are mostly reused from Pham et al. [ 56 ] . With our new inclusion of KNN, kNNPFN and TabDPT baselines and AWARE versions. Markers ( ∗∗ p

<

0.01) on AWARE models indicate they significantly outperform their vanilla (non-AWARE) counterparts (paired bootstrap, Holm-Bonferroni).

Figure 7: Overview of our designed task-aligned retrieval (AWARE) for retrieval-augmented in-context learning EHR model.

KNN
Objective	MIMIC-IV						eICU
	SEPSIS		VAP		UTI		SEPSIS		UTI
	AUROC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$
SNNL	$0.795_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.002}}}$	$0.394_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.011}}}$	$0.710_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.005}}}$	$0.233_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.011}}}$	$0.586_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.005}}}$	$0.218_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.007}}}$	$0.803_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.006}}}$	$0.325_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.018}}}$	$0.757_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.006}}}$	$0.062_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}}$
kNN L2	$0.794_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.013}}}$	$0.391_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.026}}}$	$0.709_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.015}}}$	$0.213_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.008}}}$	$0.569_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.025}}}$	$0.209_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.025}}}$	$0.795_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.023}}}$	$0.314_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.027}}}$	$0.734_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.041}}}$	$0.056_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.011}}}$
Siamese	$0.868_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.006}}}$	$0.453_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.011}}}$	$0.756_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.005}}}$	$0.235_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.014}}}$	$0.608_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.012}}}$	$0.216_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.009}}}$	$0.847_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.010}}}$	$0.332_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.017}}}$	$0.677_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.010}}}$	$0.037_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.004}}}$

Table 1: Comparison of different embedding-based retrieval objectives on large-scale ICU EHR tasks (MIMIC-IV and eICU). Performance is evaluated across models and datasets to identify the most effective objective for AWARE’s retrieval stage. Markers ( ∗ p

<

0.05, ∗∗ p

<

0.01) on a retrieval objective indicate it significantly outperforms SNNL (Full AWARE) (paired bootstrap, Holm-Bonferroni).

Figure 1: Downstream predictive performance across retrieval objectives. Distributions overlap substantially, motivating retrieval-specific evaluation.

Figure 2: Retrieval quality across methods, tasks, and context sizes

k

. SNNL achieves the best neighborhood purity and lowest contradiction rates at small

k

, while maintaining diversity comparable to SupCon.

Figure 3: Visualization of patient neighborhood structure across MIMIC-IV and eICU dataset, illustrated by projecting data into 2D space. (a) In raw feature space, proximity does not reliably correspond to predictive relevance across different clinical tasks, resulting in scattered positive examples. (b) In contrast, task-aligned embedding retrieval (AWARE) reshapes the representation space such that label-consistent examples form coherent local clusters while preserving smooth structure for downstream in-context learning. (c) Black bars illustrate task-specific feature reweighting induced by the attention-based encoder, highlighting how retrieval adapts to heterogeneous clinical variables.

KNN
Model	MIMIC-IV						eICU
	SEPSIS		VAP		UTI		SEPSIS		UTI
	AUROC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$
Full AWARE †	${\textbf{0.903}}^{**}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}}$	${\textbf{0.507}}^{**}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.006}}}$	${\textbf{0.788}}^{**}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.003}}}$	$0.250_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.005}}}$	${\textbf{0.671}}^{**}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.002}}}$	${\textbf{0.260}}^{**}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.004}}}$	${\textbf{0.971}}^{**}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}}$	${\textbf{0.711}}^{**}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.002}}}$	${\textbf{0.928}}^{**}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.002}}}$	${\textbf{0.155}}^{**}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.003}}}$
Attention Only	$\underline{0.862}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.006}}}$	$\underline{0.448}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.021}}}$	$0.763_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}}$	$\underline{0.257}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}}$	$0.629_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.007}}}$	$\underline{0.236}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}}$	$\underline{0.911}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.033}}}$	$\underline{0.492}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.112}}}$	$0.810_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.005}}}$	$0.065_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}}$
SNNL Only	$0.795_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.002}}}$	$0.394_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.011}}}$	$0.710_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.005}}}$	$0.233_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.011}}}$	$0.586_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.005}}}$	$0.218_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.007}}}$	$0.803_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.006}}}$	$0.325_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.018}}}$	$0.757_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.006}}}$	$0.062_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}}$

Table 2: Component ablation of AWARE on large-scale ICU EHR tasks (MIMIC-IV and eICU). Each row represents a configuration where only one trained component is active (added on top of the frozen PFN backbone). † Full AWARE version for KNN does not include Adapter Finetuning. Markers ( ∗ p

<

0.05, ∗∗ p

<

0.01) on Full AWARE indicate it significantly outperforms all single-component variants (paired bootstrap, Holm-Bonferroni).

KNN
Model	MIMIC-IV						eICU
	SEPSIS		VAP		UTI		SEPSIS		UTI
	AUROC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$
Full AWARE †	${\textbf{0.903}}^{**}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}}$	${\textbf{0.507}}^{**}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.006}}}$	${\textbf{0.788}}^{**}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.003}}}$	$0.250_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.005}}}$	${\textbf{0.671}}^{**}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.002}}}$	${\textbf{0.260}}^{**}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.004}}}$	${\textbf{0.971}}^{**}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}}$	${\textbf{0.711}}^{**}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.002}}}$	${\textbf{0.928}}^{**}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.002}}}$	${\textbf{0.155}}^{**}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.003}}}$
w/o Attention	$0.864_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.002}}}$	$0.440_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.003}}}$	$0.773_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}}$	$0.258_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.004}}}$	$0.632_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.002}}}$	$0.236_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.005}}}$	$0.861_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.004}}}$	$0.348_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.007}}}$	$0.827_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.003}}}$	$0.071_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.004}}}$
w/o SNNL	$\underline{0.887}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}}$	$\underline{0.486}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.006}}}$	$\underline{0.780}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}}$	$\textbf{0.266}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.005}}}$	$\underline{0.640}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.002}}}$	$\underline{0.241}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.004}}}$	$\underline{0.946}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.002}}}$	$\underline{0.604}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.007}}}$	$\underline{0.869}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}}$	$\underline{0.091}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.003}}}$

Table 3: Ablation study of AWARE on large-scale ICU EHR tasks (MIMIC-IV and eICU). Results are reported as AUROC and AUPRC (mean ± std over seeds) across clinical prediction tasks. † Full AWARE version for KNN does not include Adapter Finetuning. Markers ( ∗ p

<

0.05, ∗∗ p

<

0.01) on Full AWARE indicate it significantly outperforms all ablation variants (paired bootstrap, Holm-Bonferroni).

KNN
Model	MIMIC-IV						eICU
	SEPSIS		VAP		UTI		SEPSIS		UTI
	AUROC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$
Baseline	$0.794_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.013}}}$	$0.391_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.026}}}$	$0.709_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.015}}}$	$0.213_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.008}}}$	$0.569_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.025}}}$	$0.209_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.025}}}$	$0.795_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.023}}}$	$0.314_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.027}}}$	$0.734_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.041}}}$	$0.056_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.011}}}$
+ SNNL	$0.795_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.002}}}$	$0.394_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.011}}}$	$0.710_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.005}}}$	$0.233_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.011}}}$	$0.586_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.005}}}$	$0.218_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.007}}}$	$0.803_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.006}}}$	$0.325_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.018}}}$	$0.757_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.006}}}$	$0.062_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}}$
+ Attention	$0.843_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}}$	$\underline{0.503}_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.007}}}$	$0.736_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.002}}}$	$0.247_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.007}}}$	$0.622_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}}$	$0.246_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.003}}}$	$0.853_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.005}}}$	$0.467_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.009}}}$	$0.715_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.030}}}$	$0.048_{{{\color[rgb]{0.5,0.5,0.5}\pm 0.014}}}$

Table 4: Ablation study of AWARE and subsequent enhancements on large-scale ICU EHR tasks (MIMIC-IV and eICU). Each row progressively adds one component to the backbone. We report AUROC and AUPRC for prediction tasks (mean ± std across 3 seeds). Markers ( ∗ p

<

0.05, ∗∗ p

<

0.01) on + Adapter (Full AWARE) indicate it significantly outperforms all progressive variants (paired bootstrap, Holm-Bonferroni).

Classical and Boosting Models
Model	Regression			Classification
Model	HF $\downarrow$	OXF-PT $\downarrow$	TIT $\downarrow$	SUPPORT2 $\uparrow$	SSMI $\uparrow$	DTC $\uparrow$	GGCM $\uparrow$	AIDS $\uparrow$	ILP $\uparrow$	COVID19 $\uparrow$	Diabetes $\uparrow$	Kidney $\uparrow$
LR	$1.360_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.413_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$1.070_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.836_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.692_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.984_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.883_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.710_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.757_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.727_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.630_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.688_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$
KNN	$1.394_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.551_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$1.253_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.782_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.645_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.500_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.500_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.702_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.500_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.692_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.479_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.500_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$
CatB	$1.366_{{\color[rgb]{0.5,0.5,0.5}\pm 0.033}}$	$\underline{0.107}_{{\color[rgb]{0.5,0.5,0.5}\pm 0.013}}$	$0.942_{{\color[rgb]{0.5,0.5,0.5}\pm 0.096}}$	$0.839_{{\color[rgb]{0.5,0.5,0.5}\pm 0.006}}$	$0.690_{{\color[rgb]{0.5,0.5,0.5}\pm 0.002}}$	$0.972_{{\color[rgb]{0.5,0.5,0.5}\pm 0.007}}$	$0.851_{{\color[rgb]{0.5,0.5,0.5}\pm 0.013}}$	$0.714_{{\color[rgb]{0.5,0.5,0.5}\pm 0.004}}$	$0.741_{{\color[rgb]{0.5,0.5,0.5}\pm 0.006}}$	$0.782_{{\color[rgb]{0.5,0.5,0.5}\pm 0.005}}$	$\underline{0.676}_{{\color[rgb]{0.5,0.5,0.5}\pm 0.004}}$	$0.750_{{\color[rgb]{0.5,0.5,0.5}\pm 0.108}}$
LightG	$1.360_{{\color[rgb]{0.5,0.5,0.5}\pm 0.032}}$	$\textbf{0.092}_{{\color[rgb]{0.5,0.5,0.5}\pm 0.003}}$	$1.068_{{\color[rgb]{0.5,0.5,0.5}\pm 0.061}}$	$0.846_{{\color[rgb]{0.5,0.5,0.5}\pm 0.006}}$	$0.691_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.990_{{\color[rgb]{0.5,0.5,0.5}\pm 0.003}}$	$0.857_{{\color[rgb]{0.5,0.5,0.5}\pm 0.007}}$	$0.711_{{\color[rgb]{0.5,0.5,0.5}\pm 0.002}}$	$0.727_{{\color[rgb]{0.5,0.5,0.5}\pm 0.019}}$	$0.744_{{\color[rgb]{0.5,0.5,0.5}\pm 0.037}}$	$\textbf{0.677}_{{\color[rgb]{0.5,0.5,0.5}\pm 0.002}}$	$0.854_{{\color[rgb]{0.5,0.5,0.5}\pm 0.110}}$

Table 5: Benchmark coverage of models across small-to-medium scale EHR datasets, sourced from OpenML and UCI repositories. RMSE is reported for regression tasks, while AUROC is reported for other classification tasks. Markers ( ∗ p

<

0.05, ∗∗ p

<

0.01) on AWARE models indicate they significantly outperform all non-AWARE baselines (paired bootstrap, Holm-Bonferroni).

Classical and Boosting Models
Model	MIMIC-IV						eICU
	SEPSIS		VAP		UTI		SEPSIS		UTI
	AUROC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$
LR	$0.886_{{\color[rgb]{0.5,0.5,0.5}\pm 0.012}}$	$0.490_{{\color[rgb]{0.5,0.5,0.5}\pm 0.019}}$	$0.752_{{\color[rgb]{0.5,0.5,0.5}\pm 0.023}}$	$0.218_{{\color[rgb]{0.5,0.5,0.5}\pm 0.020}}$	$0.649_{{\color[rgb]{0.5,0.5,0.5}\pm 0.015}}$	$0.247_{{\color[rgb]{0.5,0.5,0.5}\pm 0.015}}$	$0.940_{{\color[rgb]{0.5,0.5,0.5}\pm 0.016}}$	$0.577_{{\color[rgb]{0.5,0.5,0.5}\pm 0.051}}$	$0.922_{{\color[rgb]{0.5,0.5,0.5}\pm 0.013}}$	$0.126_{{\color[rgb]{0.5,0.5,0.5}\pm 0.015}}$
KNN	$0.794_{{\color[rgb]{0.5,0.5,0.5}\pm 0.013}}$	$0.391_{{\color[rgb]{0.5,0.5,0.5}\pm 0.026}}$	$0.709_{{\color[rgb]{0.5,0.5,0.5}\pm 0.015}}$	$0.213_{{\color[rgb]{0.5,0.5,0.5}\pm 0.008}}$	$0.569_{{\color[rgb]{0.5,0.5,0.5}\pm 0.025}}$	$0.209_{{\color[rgb]{0.5,0.5,0.5}\pm 0.025}}$	$0.795_{{\color[rgb]{0.5,0.5,0.5}\pm 0.023}}$	$0.314_{{\color[rgb]{0.5,0.5,0.5}\pm 0.027}}$	$0.734_{{\color[rgb]{0.5,0.5,0.5}\pm 0.041}}$	$0.056_{{\color[rgb]{0.5,0.5,0.5}\pm 0.011}}$
CatB	$0.903_{{\color[rgb]{0.5,0.5,0.5}\pm 0.008}}$	$0.527_{{\color[rgb]{0.5,0.5,0.5}\pm 0.025}}$	$0.754_{{\color[rgb]{0.5,0.5,0.5}\pm 0.010}}$	$0.235_{{\color[rgb]{0.5,0.5,0.5}\pm 0.014}}$	$0.659_{{\color[rgb]{0.5,0.5,0.5}\pm 0.013}}$	$0.257_{{\color[rgb]{0.5,0.5,0.5}\pm 0.013}}$	$0.981_{{\color[rgb]{0.5,0.5,0.5}\pm 0.001}}$	$0.756_{{\color[rgb]{0.5,0.5,0.5}\pm 0.008}}$	$0.937_{{\color[rgb]{0.5,0.5,0.5}\pm 0.010}}$	$0.153_{{\color[rgb]{0.5,0.5,0.5}\pm 0.004}}$

Table 6: Performance on hospital-acquired infection prediction tasks on large-scale EHR datasets, including MIMIC-IV and eICU. For each task and each model, we extensively ran hyperparameters tuning for 30 rounds. Markers ( ∗ p

<

0.05, ∗∗ p

<

0.01) on AWARE models indicate they significantly outperform all non-AWARE baselines (paired bootstrap, Holm-Bonferroni).

Panel of Figure 4 — Figure 4: Training and inference efficiency trade-offs across model classes on large-scale EHR datasets. The left panel reports relative predictive ranking versus training time (measured until convergence with early stopping), while the right panel reports ranking versus per-patient inference time. Hyperparameter tuning and pretraining costs are excluded. Vanilla (non-AWARE) in-context learning models incur no task-specific training cost; AWARE-enhanced variants additionally train a lightweight retrieval encoder (and, where applicable, an adapter) per dataset/task, at substantially lower cost than end-to-end training.

Classical and Ensemble Models
Model	Lung Cancer $\uparrow$	Stroke $\uparrow$
Logistic Regression	$0.993_{{\color[rgb]{0.5,0.5,0.5}\pm 0.00}}$	$0.992_{{\color[rgb]{0.5,0.5,0.5}\pm 0.00}}$
XGBoost	$0.999_{{\color[rgb]{0.5,0.5,0.5}\pm 0.00}}$	$0.999_{{\color[rgb]{0.5,0.5,0.5}\pm 0.00}}$
Tabular In-Context Learning Models
TabPFN	$1.0_{{\color[rgb]{0.5,0.5,0.5}\pm 0.00}}$	$0.999_{{\color[rgb]{0.5,0.5,0.5}\pm 0.00}}$
TabDPT	$0.992_{{\color[rgb]{0.5,0.5,0.5}\pm 0.00}}$	$0.978_{{\color[rgb]{0.5,0.5,0.5}\pm 0.00}}$

Table 7: AUROC performance on fully synthetic EHR data generated using Synthea for two binary prediction tasks: lung cancer and stroke. Performance across model classes approaches ceiling levels due to the rule-based generative process underlying Synthea. Results are intended as an upper-bound sanity check rather than evidence of real-world generalization.

Dataset Name	#Rows	#Features	Task	Imbalance Ratio	Cohort	Feature Types
Synthetic EHR datasets
Synthea-Stroke [ 5 ]	16,055	793	Binary	–	Multi-domain	DEM, VIT, LAB, CODE
Synthea-LungCancer [ 5 ]	4,384	773	Binary	–	Oncology	DEM, VIT, LAB, CODE
Small-to-medium real-world EHR datasets
Heart Failure (Platelets Counts) [ 11 ]	299	12	Regression	–	Cardiovascular	DEM, LAB
OXF – Parkinsons Telemonitoring [ 74 ]	5,875	22	Regression	–	Neurological	DEM, PHYS, TMP

Table 8: Summary of EHR datasets used in this study, spanning diverse clinical cohorts, task types, and feature modalities. Feature type abbreviations: DEM = demographics; VIT = vital signs; LAB = laboratory results; CODE = medical codes / diagnoses; GEN = genetic / molecular features; PHYS = physiological signals; ADMIN = administrative variables; TMP = temporal features.

Hyperparameter	Range / Type
Classical Machine Learning Methods
Logistic Regression
C	Log-uniform $[10^{-5},5.0]$
penalty	Categorical $\{\ell_{2},\text{none}\}$
max_iter	Integer $[50,500]$
XGBoost

Table 9: Hyperparameter search spaces used for tuning classical machine learning and deep tabular baselines.

Tabular Foundation Models

Tabular foundation models learn reusable representations or in-context prediction skills for structured data such as spreadsheets and databases.

See all 187 papers

EveryQuery: A Promptable Foundation Model for Clinical Prediction Tasks over Electronic Health Records
Mar 9, 2026Payal Chandak, Gregory Kondas, Liat Antwarg Friedman +2EHR Foundation Models Clinical Outcome Prediction
Harvard-MIT HST · Columbia University · Harvard Medical School
A Comparative Framework for Evaluating Foundation Models on Tabular Data: A Case Study in Healthcare
Aug 26, 2026Majid Lotfian Delouee, Sjors G. J. G. In 't Veld, Martijn C. SchutTabular Foundation Models Model Selection
Distilling Tabular Foundation Models for Structured Health Data
May 18, 2026Aditya Tanna, Nassim Bouarour, Mohamed Bouadi +2Tabular Foundation Models Efficient Neural Network Inference
Lexsi Labs

Mar 9, 2026J/K move · Enter open · S save

EveryQuery: A Promptable Foundation Model for Clinical Prediction Tasks over Electronic Health Records

Payal Chandak, Gregory Kondas, Liat Antwarg Friedman, Isaac Kohane+1

Harvard-MIT HST · Columbia University · Harvard Medical School

EHR Foundation Models Clinical Outcome Prediction Clinical Prediction Longitudinal EHR Modeling Electronic Health Records

Open paper AskPDF arXiv

Autoregressive foundation models for electronic health records can make predictions for diverse clinical tasks without finetuning. These methods work by generating possible "synthetic futures" for patients, and then inferring downstream predictions over these simulated trajectories. Despite their strong capabilities, these methods have several flaws. (1) Inference is extremely computationally expensive, as each prediction requires simulating many trajectories, each with many observations. (2) Their performance is noisy due to the variance inherent to simulation, which in particular reduces efficacy on rare events, which are often of high importance clinically. (3) They are not promptable, as their only input is the patient's medical history, requiring indirect aggregation at inference time. We introduce EveryQuery, a promptable foundation model for structured EHR data. Building on the insight that clinical prediction tasks can be written in a structured query language, EveryQuery takes a query as a prompt alongside the patient's history and estimates the outcome directly. This allows EveryQuery to make predictions across diverse clinical tasks, such as classification and survival modeling, without simulation or finetuning. Across three EHR datasets, EveryQuery outperforms a competitive autoregressive baseline on clinically relevant classification and time-to-event tasks, with mean AUROC gains of 14.3% and 12.7% on MIMIC-IV and a large academic medical center and comparable performance on the smaller NWICU dataset. Its advantage is largest for rare outcomes, and it requires 580x less inference compute.

Abstract

Tabular Foundation Models

Explore similar work

EveryQuery: A Promptable Foundation Model for Clinical Prediction Tasks over Electronic Health Records

A Comparative Framework for Evaluating Foundation Models on Tabular Data: A Case Study in Healthcare

Distilling Tabular Foundation Models for Structured Health Data