Search papers, labs, and topics across Lattice.
The paper benchmarks tabular in-context learning (TICL) methods for clinical risk prediction using EHR data, finding that naive distance-based retrieval degrades performance as data heterogeneity and imbalance increase. To address this, they introduce AWARE, a task-aligned retrieval framework using supervised embedding learning and lightweight adapters. AWARE significantly improves performance, particularly under extreme imbalance, highlighting the importance of retrieval quality and alignment for TICL in clinical settings.
Naive nearest neighbor retrieval cripples tabular in-context learning for clinical risk prediction, but task-aligned retrieval unlocks significant gains, especially when data is messy.
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.
