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

arXiv: https://arxiv.org/abs/2604.01841 · HTML: https://mlsanity.com/paper/2604.01841/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 49 of its 104 references that are in the map.
- First posted on arXiv 2 Apr 2026.

## Builds on (most cited first, 30 of 49)

- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- [REPLUG: Retrieval-Augmented Black-Box Language Models](https://mlsanity.com/paper/2301.12652/) (2023)
- [Revisiting Deep Learning Models for Tabular Data](https://mlsanity.com/paper/2106.11959/) (2021)
- [TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second](https://mlsanity.com/paper/2207.01848/) (2022)
- [TabNet: Attentive Interpretable Tabular Learning](https://mlsanity.com/paper/1908.07442/) (2019)
- [Recurrent Neural Networks for Multivariate Time Series with Missing Values](https://mlsanity.com/paper/1606.01865/) (2016)
- [Med-BERT: pre-trained contextualized embeddings on large-scale structured electronic health records for disease prediction](https://mlsanity.com/paper/2005.12833/) (2020)
- [Transformers Can Do Bayesian Inference](https://mlsanity.com/paper/2112.10510/) (2021)
- [RETAIN: An Interpretable Predictive Model for Healthcare using Reverse Time Attention Mechanism](https://mlsanity.com/paper/1608.05745/) (2016)
- [TabICL: A Tabular Foundation Model for In-Context Learning on Large Data](https://mlsanity.com/paper/2502.05564/) (2025)
- [Scalable and accurate deep learning for electronic health records](https://mlsanity.com/paper/1801.07860/) (2018)
- [BEHRT: Transformer for Electronic Health Records](https://mlsanity.com/paper/1907.09538/) (2019)
- [When Do Neural Nets Outperform Boosted Trees on Tabular Data?](https://mlsanity.com/paper/2305.02997/) (2023)
- [RA-DIT: Retrieval-Augmented Dual Instruction Tuning](https://mlsanity.com/paper/2310.01352/) (2023)
- [TabArena: A Living Benchmark for Machine Learning on Tabular Data](https://mlsanity.com/paper/2506.16791/) (2025)
- [A Study of Generative Large Language Model for Medical Research and Healthcare](https://mlsanity.com/paper/2305.13523/) (2023)
- [Deep EHR: A Survey of Recent Advances in Deep Learning Techniques for Electronic Health Record (EHR) Analysis](https://mlsanity.com/paper/1706.03446/) (2017)
- [TabDPT: Scaling Tabular Foundation Models on Real Data](https://mlsanity.com/paper/2410.18164/) (2024)
- [TabM: Advancing Tabular Deep Learning with Parameter-Efficient Ensembling](https://mlsanity.com/paper/2410.24210/) (2024)
- [Why Tabular Foundation Models Should Be a Research Priority](https://mlsanity.com/paper/2405.01147/) (2024)
- [Learning to Diagnose with LSTM Recurrent Neural Networks](https://mlsanity.com/paper/1511.03677/) (2015)
- [Analyzing and Improving Representations with the Soft Nearest Neighbor Loss](https://mlsanity.com/paper/1902.01889/) (2019)
- [Zero Shot Health Trajectory Prediction Using Transformer](https://mlsanity.com/paper/2407.21124/) (2024)
- [Retrieval & Fine-Tuning for In-Context Tabular Models](https://mlsanity.com/paper/2406.05207/) (2024)
- [A Comprehensive Survey of Foundation Models in Medicine](https://mlsanity.com/paper/2406.10729/) (2024)
- [DANets: Deep Abstract Networks for Tabular Data Classification and Regression](https://mlsanity.com/paper/2112.02962/) (2021)
- [Real-TabPFN: Improving Tabular Foundation Models via Continued Pre-training With Real-World Data](https://mlsanity.com/paper/2507.03971/) (2025)
- [An Extensive Data Processing Pipeline for MIMIC-IV](https://mlsanity.com/paper/2204.13841/) (2022)
- [Context Clues: Evaluating Long Context Models for Clinical Prediction Tasks on EHRs](https://mlsanity.com/paper/2412.16178/) (2024)

## Built on by (most cited first, 4 of 4)

- [EXAONE Tabular 1.0 : Technical Report](https://mlsanity.com/paper/2608.25774/) (2026)
- [Tabular Foundation Models for Clinical Survival Analysis via Survival-Aware Adaptation](https://mlsanity.com/paper/2606.12006/) (2026)
- [Adaptation Interfaces for In-Context Tabular Foundation Models in Time-to-Event Prediction](https://mlsanity.com/paper/2609.04901/) (2026)
- [TILE: Bridging LLMs and Table Foundation Models for Tabular Learning with Strings](https://mlsanity.com/paper/2610.17065/) (2026)

## Abstract (arXiv)

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.