EHR-RAG: Bridging Long-Horizon Structured Electronic Health Records and Large Language Models via Enhanced Retrieval-Augmented Generation

TL;DR

This work tackles the challenge of predicting clinical outcomes from long-horizon structured EHRs under fixed-context limits by introducing EHR-RAG, a retrieval-augmented framework with three novel components: Event- and Time-Aware Hybrid Retrieval (ETHER) to preserve event structure and temporal dynamics, Adaptive Iterative Retrieval (AIR) to progressively expand evidence coverage, and Dual-Path Evidence Retrieval and Reasoning (DER) to evaluate both factual and counterfactual evidence. The approach systematically constructs a compact, task-relevant evidence context and leverages dual-path hypothesis evaluation to improve robustness against bias and spurious correlations. Empirical results on the EHRSHOT benchmark across four long-horizon tasks show consistent Macro-F1 gains (average $10.76\%$) over strong LLM baselines and demonstrate robustness across multiple LLM backbones and data regimes. The findings highlight the practical value of structured, evidence-grounded retrieval for clinical prediction with longitudinal EHR data and establish a foundation for future scalable reasoning over long healthcare trajectories.

Abstract

Electronic Health Records (EHRs) provide rich longitudinal clinical evidence that is central to medical decision-making, motivating the use of retrieval-augmented generation (RAG) to ground large language model (LLM) predictions. However, long-horizon EHRs often exceed LLM context limits, and existing approaches commonly rely on truncation or vanilla retrieval strategies that discard clinically relevant events and temporal dependencies. To address these challenges, we propose EHR-RAG, a retrieval-augmented framework designed for accurate interpretation of long-horizon structured EHR data. EHR-RAG introduces three components tailored to longitudinal clinical prediction tasks: Event- and Time-Aware Hybrid EHR Retrieval to preserve clinical structure and temporal dynamics, Adaptive Iterative Retrieval to progressively refine queries in order to expand broad evidence coverage, and Dual-Path Evidence Retrieval and Reasoning to jointly retrieves and reasons over both factual and counterfactual evidence. Experiments across four long-horizon EHR prediction tasks show that EHR-RAG consistently outperforms the strongest LLM-based baselines, achieving an average Macro-F1 improvement of 10.76%. Overall, our work highlights the potential of retrieval-augmented LLMs to advance clinical prediction on structured EHR data in practice.

Figures (8)

• Figure 1: Overview of the EHR-RAG framework for long-horizon clinical prediction. Compared with Direct Generation and Vanilla RAG, our framework explicitly addresses context truncation and incomplete retrieval. It integrates (a) Event- and Time-aware Hybrid Retrieval (ETHER), (b) Adaptive Iterative Retrieval (AIR), and (c) Dual-Path Factual and Counterfactual Reasoning (DER), ensuring retrieval quality, robustness, and completeness for reliable clinical decision-making.
• Figure 2: Macro-F1 performance comparison between EHR-RAG, vanilla RAG, and traditional machine learning (ML) baselines under varying amounts of labeled training data across four tasks. The x-axis denotes the number of training samples per class, while dashed horizontal lines indicate the zero-shot performance of LLM-based methods. EHR-RAG consistently matches or outperforms ML baselines, particularly in low-resource settings.
• Figure 3: Prompt template used for LLM-based clinical prediction baselines. Blue text indicates input variables.
• Figure 4: Prompt template used for clinical indicator selection in the Event- and Time-Aware Hybrid EHR Retrieval component of EHR-RAG. Blue text denotes input variables.
• Figure 5: Prompt template used for query refinement in the Adaptive Iterative Retrieval component of EHR-RAG. Blue text denotes input variables.