Conceptio › Archive › arXiv CS

arXiv CSopen access

EHR-RAGp: Retrieval-Augmented Prototype-Guided Foundation Model for Electronic Health Records

2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

neural-networks

machine learning, deep learning, neural networks

EHR-RAGp: Retrieval-Augmented Prototype-Guided Foundation Model for Electronic Health Records

arXiv:2605.12335v1 [cs.IR] 12 May 2026

Saeed Shurrab1,2 , Mariam Al-Omari2 , Dana El Samad2 , Farah E. Shamout1,2 New York University1 New York University Abu Dhabi2

Abstract Electronic Health Records (EHR) contain rich longitudinal patient information and are widely used in predictive modeling applications. However, effectively leveraging historical data remains challenging due to long trajectories, heterogeneous events, temporal irregularity, and the varying relevance of past clinical context. Existing approaches often rely on fixed windows or uniform aggregation, which can obscure clinically important signals. In this work, we introduce EHR-RAGp, a retrieval-augmented foundation model that dynamically integrates the most relevant patient history across diverse clinical event types. We propose a prototype-guided retrieval module that acts as an alignment mechanism and estimates the relevance of retrieved historical chunks with respect to a given prediction task, guiding the model towards the most informative context. Across multiple clinical prediction tasks, EHR-RAGp consistently outperforms state-of-the-art EHR foundation model and transformer-based baselines. Furthermore, integrating EHR-RAGp with existing clinical foundation models yields substantial performance gains. Overall, EHR-RAGp provides a scalable and efficient framework for leveraging long-range clinical context to improve downstream performance.

Introduction

Healthcare systems rely extensively on Electronic Health Record (EHR) databases to store diverse clinical events collected during patient encounters, such as demographics, diagnoses, vital signs, lab results, procedures and medications [1, 2]. Recent advances in machine learning have enabled the adaptation of foundation models for EHR data, where patient trajectories are represented as chronologically ordered sequences of events, in a manner that mimics the structures of sentences in natural language [3, 4, 5, 6]. However, unlike natural language, clinical events are inherently irregular, sparse, fragmented, and span multiple episodes of care that accumulate over months or years [7]. Existing EHR foundation models are designed under specific assumptions, including: (i) information present only within a predefined time window for a specific patient encounter is relevant to predict the outcome of interest [8, 9]; and (ii) a preselected subset of events or features, such as lab test results, vital sign measurements, diagnoses, or combinations thereof, is sufficient to predict patient outcomes [10, 11, 12, 13, 14]. This often stems from computational constraints and the difficulty of representing long-range irregular patient trajectories. Such constraints could also lead to the omission of clinically important historical events. From a clinical perspective, the patient’s state is almost always shaped by their accumulated longitudinal history [15, 16]. This includes slow-moving chronic patterns and cross-encounter interactions that clinicians routinely examine when making decisions. Such a central gap highlights a fundamental limitation in current EHR modeling paradigms, namely the inability to integrate and reason over a patient’s full history within a unified representation learning framework. Filling this gap requires a new outlook on EHR representation learning. Rather than limiting the patient’s predictive context to a predefined time window, a model must be equipped to access, retrieve, Preprint.

Raw EHR Database

Patient 1 OUTP GAP L1

Vitals (V) Labs (L) Medications (M) Procedures (P) Diagnosis (D) - - - - - - -

Outpatient (OUTP)

Patient 2

Hospital (HOSP)

Patient 3

OUTP GAP L1

ED L3

M1 M2

OUTP GAP D1

M1 M2 M3

HOSP

ICU

HOSP

Concept L1

Values

null

null null null null

Time

Δt Δt

Δt

Δt Δt Δt Δt Δt Δt Δt

Stage

Visit

Type

EHR-RAG End-to-End Fine-tuning

HOSP

L2 T

HOSP

f π(q)

OUTP GAP

Embed query segment

PGR

Top-M

P2 null Δt 4 2 3

ICU M3 M4 null null HOSP Embed Mt3 M4 ΔP2t Δ retrieved null null 4 ICU 4 null chunks Mt3 M4 Δ2 t ΔP22t Δ 54 null 54 null 4 null 2 Δ2 t Δ2 t Δ t 4 54 3 54 2 2 2 5 3 5

π(ci)

Weigh and Fuse

ICU P2 M3 M4 null null null Qp Δ t Δ t Δ t 4 4 4 2 2 2 5 3 5

Embedded & aggregated to obtain a single event representation e

Masked Pretraining with Chunked Trajectories

HOSP

ED L1

HOSP

ED V1

Intensive Care Unit (ICU)

ED L1

Patient 4

Chunking & Storage of EHR vector database

Clinical Events Grouping & Encoding across Patient Timeline

Structured Patient Timeline (MEDS Format)

Emergency (ED)

zfused

g wi

f ci

Figure 1: Overview of the EHR-RAGp, including chunking of patient trajectories to build the EHR vector database (top panel), masked pre-training of the backbone encoder with chunked trajectories (bottom left), and end-to-end fine-tuning of the final architecture (bottom right). and integrate information scattered across the patient’s complete trajectory. The central assumption motivating this work is that retrieval provides a scalable mechanism for incorporating long-range historical context without forcing the model to encode the entire trajectory within its fixed internal state. Given the remarkable success of Retrieval-Augmented Generation [17] (RAG) techniques in Natural Language Processing (NLP) and their ability to leverage external knowledge beyond the training data of the base Large Language Model (LLM) [18, 19, 20, 21], the same concept can be applied for retrieving and integrating historical EHR data. In other words, by considering the patient’s historical events as retrievable pieces of information, an EHR foundation model can dynamically condition on what matters the most for the patient’s current state. This perspective shifts longitudinal modeling from a purely sequential problem into a retrieval-augmented one. Hence, in this work, we propose EHR-RAGp, a retrieval-based foundation model designed to reason over the full breadth of a patient’s longitudinal EHR history. The main contributions of the work are summarized as follows: 1. We propose a new framework for constructing a vector database of longitudinal EHR patient trajectories using multi-granular chunking strategies, including event-based, time-based, visit-level, and care-stage segmentation, enabling flexible representation of heterogeneous events, temporal structure, and care progression. 2. We introduce a prototype-guided retrieval mechanism that acts as an alignment operator, estimating the relevance of retrieved historical patient chunks with respect to the prediction task, guiding the model to integrate the most informative long-range clinical context. 3. We present EHR-RAGp, an EHR-native retrieval-augmented foundation model trained on large-scale real-world EHR data. Extensive evaluation across multiple clinical prediction tasks demonstrates that EHR-RAGp consistently outperforms multiple state-of-the-art models. 4. By training and optimizing all baselines using an expanded and unified feature set, we establish new benchmark results on four clinical prediction tasks and make our code publicly available at: https://github.com/nyuad-cai/EHR-RAGp.

Related Work

EHR Foundation Models. Traditional EHR modeling approaches process patient trajectories as sequences of discrete clinical events collected over time. Some models primarily consider the occurrence of certain events, while others process the occurrence of the clinical event and its measured value. Early techniques utilized recurrent models such as Long-Short Term Memory (LSTM) [22] or Gated Recurrent Unit (GRU) networks [23] to capture long-term dependencies [24, 25, 16, 26]. 2

The emergence of transformer-based LLMs [27, 28] paved the way for more advanced EHR foundation models. For example, encoder-based models such as BEHRT [29], Hi-BEHRT [30], CEHRBERT [31], and Med-BERT [32] introduced large-scale pretraining over longitudinal sequences of medical events to enhance the contextualization of patient trajectories across visits. Other decoderbased transformer models such as ETHOS [33], CEHR-GPT [34], Curiosity [35], and EHR2Path [36] focused on the development of generative EHR foundation models for clinical prediction tasks. The common link between all these models is that they mimic standard language models by learning contextualized medical code vocabularies as tokens. Another category of EHR foundation models shifts away from code-based modeling towards textbased modeling by treating each medical code as a textual description. Hence, the patient timeline constitutes a set of textual descriptions that can be processed directly by the LLM [12, 13, 37]. Despite the remarkable progress, these models continue to inherit core limitations from their language-model roots. They operate on truncated encounter windows, compress long and irregular patient histories into fixed-length sequences, and rely on restricted types of clinical events. While some methods have explored the use of longer context windows to accommodate extended trajectories [35, 12, 13], their vocabularies are heavily tokenized for medical codes and numerical values, or textualized and dependent on the base model’s reasoning capabilities. As a result, substantial portions of a patient’s longitudinal trajectory remain inaccessible in downstream prediction tasks. This emphasizes the need for scalable frameworks that are capable of integrating and reasoning over complete EHR trajectories. Retrieval-Augmented Modeling. RAG has emerged as a powerful technique in NLP for extending large language models capabilities by allowing them to condition on external, non-parametric knowledge sources. The original RAG framework [17] demonstrated how combining dense retrieval with generative transformers improves factual correctness and reduces hallucinations. The concept was later extended by models such as RETRO [38], FLARE [39], and REPLUG [40]. These architectures share a common principle: retrieval allows the model to incorporate long-range or domain-specific information without encoding the entire knowledge base within its internal parameters. The success of the retrieval paradigm has also influenced healthcare, particularly in clinical question answering [41, 42], evidence grounding [43, 44], biomedical knowledge integration [45, 46], phenotyping and cohort identification [47], and many other applications related to medical textual data [48, 49]. In the context of EHR modeling, RAG has been primarily explored as a mechanism for integrating external knowledge into predictive models, such as EMERGE [50], RAM-EHR [51], REALM [52], and KAMELEON [53]. To the best of our knowledge, REMED [54] is our most closely related work. Its core idea is to retrieve the most relevant subset of medical codes for a given outcome prediction task, to overcome manual selection of clinical events or features. Overall, recent advancements highlight the potential of retrieval-augmented modeling for EHR, but they also reveal a clear unmet need for mechanisms that can dynamically access and integrate a patient’s entire structured history.

Methodology

3.1

Overview

EHR-RAGp is a retrieval-augmented foundation model for reasoning over longitudinal EHR trajectories (Figure 1). The framework encodes structured clinical events, organizes them into retrievable chunks, retrieves relevant historical segments through prototype-guided alignment, and integrates them with a task-specific query for downstream prediction. Here , we describe the core architectural components. 3.2

Preliminaries

Let D = {Hp }P p=1 represent a large-scale EHR dataset consisting of patient-level longitudinal records for a cohort of P patients. For a given patient p, the longitudinal medical history is represented as: Hp = {Vp(1) , Vp(2) , . . . , Vp(Np ) }, where Np denotes the total number of visits for the respective patient (see Figure S1). The n-th visit (n) Vp consists of a set of clinical events: n Vp(n) = {(cnk , vnk , tnk )}K k=1 ,

where cnk denotes a clinical event such as a diagnosis, procedure, laboratory test, vital sign, or medication administration, vnk is an optional associated value such as a lab result, medication dosage, 3

or vital sign measurement, and tnk is the exact timestamp at which the event occurred. Events within a visit may vary in type, granularity, and frequency, reflecting the inherently irregular and heterogeneous nature of structured EHR data. (τ )

For a given index visit, denoted by Vp , the goal is to predict a target patient outcome based on a (τ ) true label yp . Hence, the proposed system role is to use the index visit information in retrieving the (τ ) (τ ) most relevant information. Therefore, we define a query segment Qp ⊆ Vp , which represents the portion of the current visit that is used by the model as its immediate input context. The query segment may correspond to only a portion of the full visit, as the entire visit can exceed the maximum input length supported by the model. Therefore, we restrict the query to a fixed-size segment that fits within the model’s context window. All events occurring before the query segment, either from the same visit or from any earlier visits, form the historical context: Hp(<τ ) = {Vp(1) , . . . , Vp(τ −1) , Vp(τ <q) }, (τ <q)

where Vp denotes the prefix of the last visit preceding the query (see Figure S2). Hence, the model aims to learn a mapping function: ) (<τ ) γ : Q(τ → ŷp(τ ) . p , Hp A naive formulation of this problem would require the model to encode the full patient history in its fixed internal state, which is a significant challenge due to computational constraints. Hence, the goal of our proposed framework is to enable the model to retrieve and integrate the most relevant (<τ ) (τ ) information from the patient’s history Hp based on the query segment Qp . 3.3

Event Representation Learning

EHR data contains heterogeneous clinical events with varying semantics, numerical values, temporal patterns, and contextual metadata. To process these events within a unified architecture, we propose mapping each event (cnk , vnk , tnk ) and its associated metadata into a single representation enk ∈ Rd , where d is the vector dimension. This embedding considers six key components: clinical concept ucnk , numeric value vnk , time tnk , visit-order rvisit(n) , care stage sstage(n) , and event-type wtype(nk) embeddings as shown below. Detailed information pertaining to data encoding is provided in Appendix B. The unified event representation is then constructed by summation following [10, 29]: enk = (ucnk + vnk + tnk + rvisit(n) + sstage(n) + wtype(nk) ) ∈ Rd This unified representation enables encoding of heterogeneous events, longitudinal ordering, and care-setting context within a single embedding space. 3.4

Chunking

After representing individual events, EHR-RAGp partitions the patient timeline into coherent chunks that serve as retrievable units. We propose four chunking methods, namely, event-based, time-based, visit-level, and care-stage chunking (Figure S3). Each chunking method captures a different aspect of longitudinal patient structure, enabling retrieval over heterogeneous temporal granularities and clinical contexts. Event-based chunking groups a fixed number of consecutive events, time-based chunking partitions the timeline into fixed temporal windows, visit-level chunking organizes events according to hospital admissions, and care-stage chunking separates events based on transitions between clinical care settings, such as outpatient, emergency department, hospital ward, and ICU. 3.5

Prototype-Guided Retrieval (τ )

Given a query segment Qp , EHR-RAGp retrieves relevant historical segments from the patient’s (<τ ) longitudinal history Hp using a two-stage procedure: (1) coarse-grained semantic retrieval to identify candidate chunks, followed by (2) prototype-guided alignment that estimates the relevance of each retrieved chunk via agreement in the learned prototype assignment space, to enable weighted integration of context for the downstream task. 4

(τ )

Candidate Retrieval. The query segment Qp is first encoded using a frozen pretrained retriever R, producing a fixed-dimensional query representation: ) d q ′ = R(Q(τ p )∈R .

The retriever, R, is pretrained using masked language modeling. Further details are provided in (<τ ) Appendix C. Using the same retriever R, each historical chunk hi ∈ Hp is embedded as c′i = R(hi ) ∈ Rd . Candidate retrieval is then performed by computing cosine similarity between the query embedding q ′ and historical chunk embeddings {c′i } stored in precomputed similarity index I. We then retrieve the top-M most similar chunks: {h1 , h2 , . . . , hM } ⊂ Hp(<τ ) . (τ )

Query and Candidate Encoding. Upon retrieval, the query segment Qp and the top-M chunks {h1 , h2 , . . . , hM } are passed to a shared backbone encoder f to obtain q and {ci }, respectively. ) d q = f (Q(τ ci = f (hi ) ∈ Rd . p )∈R , Using the same encoder f for queries and historical chunks ensures that both are represented in a common embedding space, enabling consistent comparison and alignment between the query and retrieved context.

Prototype-Guided Alignment. EHR-RAGp maintains a learnable set of L prototype vectors, which serve as latent clinical centroids that organize a patient’s longitudinal history into distinct semantic modes: P = {p1 , p2 , . . . , pL }, pl ∈ Rd , These prototypes serve as latent anchors that organize the representation space into distinct modes, enabling structured comparison between query and historical segments. Rather than being predefined, they are learned end-to-end, allowing the model to capture recurring patterns in clinical trajectories and facilitate alignment between the query and retrieved context with respect to the downstream task. For a chunk embedding x ∈ Rd , the model computes a probability distribution over the L prototypes: ⊤ xP π(x) = softmax ∈ RL , Tπ where x corresponds to either to the query q or a candidate chunk ci , and Tπ is a temperature parameter that controls the sharpness of the distribution [55]. In practice, we use different temperatures for the query and history branches, denoted by Tq and Th , respectively, with Tq < Th . This asymmetric design reflects the roles of the two branches: the query must identify a focused subset of relevant patterns for the prediction task, while historical chunks represent diverse clinical contexts that may only partially match the query. Accordingly, we obtain πq = π(q),

πi = π(ci ), i = 1, . . . , M.

πq represents the prototype assignment distribution of the query, while each πi represents the assignment of candidate chunk ci to the set of learned prototypes L. To measure how well a candidate history chunk aligns with the query in the prototype space, EHR-RAGp measures the agreement between their prototype distributions using a cross-entropy-based alignment score: αi = −

L X

(l)

πq(l) log πi ,

l=1

Where α ∈ [0, ∞). This formulation measures how well the candidate distribution πi explains the query distribution πq , treating the query as a reference. Lower negative cross-entropy values indicate stronger alignment, meaning that the candidate assigns high probability mass to the prototypes emphasized by the query, whereas higher values indicate weaker alignment. In addition, this distributional alignment enables smooth and differentiable matching in the prototype space without requiring hard assignment to a single prototype. 5

Prototype-Guided Weighting. Given the prototype alignment scores {αi }M i=1 , EHR-RAGp converts them into normalized importance weights using a temperature-controlled softmax: −αi wi = softmax , Ts where Ts is a temperature parameter that controls the smoothness of the weighting distribution. EHR-RAGp adopts a fully differentiable soft weighting mechanism, where all retrieved chunks contribute to the final prediction with varying degrees of importance. This design allows the model to preserve distributed clinical evidence across multiple chunks, which is particularly important in EHR data where relevant signals are often weak and spread over time. The resulting weights {wi } are used to modulate the contribution of each retrieved chunk during the subsequent fusion stage Retrieval-Augmented Fusion. The query representation q and the retrieved historical chunks (<τ ) Rp = {c1 , . . . , cn } are combined through a prototype-guided weighting mechanism. Each retrieved chunk is modulated by its corresponding importance weight wi : zfused = q, w1 c1 , w2 c2 , . . . , wn cn . This sequence is then processed by a transformer encoder g(·) followed by a a linear classifier h(·) to yield the final prediction: ŷp(τ ) = h (g(zfused )) . Training Objective. EHR-RAGp is trained end-to-end using a supervised binary cross-entropy loss associated with the target prediction task, while keeping the retrieval encoder, R, frozen and updating the backbone encoder, f , the prototype parameters, and the fusion module. The overall objective combines the task loss with a prototype usage regularization term that encourages balanced utilization of the prototype space across both query and history representations: L = l yp(τ ) , ŷp(τ ) − λu (H(π̄q ) + H(π̄h )) , PB P (b) (b) where π̄q = B1 b=1 πq and π̄h = P 1Kb b,k πhk denote the average prototype assignment b distributions over a batch for the query and historical chunks, respectively, H(·) is the entropy function, and λu > 0 determines the strength of the regularization term. This regularization encourages the model to avoid degenerate solutions in which only a small subset of prototypes are used.

Experiments

Dataset & Prediction Tasks. We evaluate Table 1: Summary of data splits used for downour proposed framework on multiple tasks using stream evaluation. the large-scale real-world MIMIC-IV dataset Training Val Test Total (V3.1) [56]. MIMIC-IV consists of heteroge- Statistic 34, 887 4, 984 9, 968 49, 839 neous EHR data gathered from patients admit- # Patients # ICU stays 42, 679 6, 133 12, 363 61175 ted to the Beth Israel Deaconess Medical Center # ICU events 231M 33M 67M 332M between 2008-2022. It includes clinical data Mean LOS (days) 4.164 4.160 4.170 − collected from 364,627 patients associated with Percentage (%) 70 10 20 100 546,028 hospital admissions and 94,458 ICU stays [57]. It consists of two main modules: hosp and icu. While the former covers general admission data such as labs, medications, and microbiology events, the latter covers ICU-level charted data such as chartevents, infusions, and fluid outputs. We consider all tables that hold clinically relevant information and exclude those tables holding operational information such as provider information. Table 1 summarizes the data statistics with further information provided in Appendix B. We assess EHR-RAGp for four risk prediction tasks: (1) in-hospital mortality, (2) ICU-readmission within 30 days, (3) long length of stay (7 days or more), and (4) 1-year mortality post-discharge. The selected tasks span both short-term and long-term clinical outcomes and reflect common benchmarks in EHR-based modeling. Together, they assess the model’s ability to capture acute risk, longitudinal disease progression, and healthcare utilization by leveraging structured patient histories. Detailed information pertaining to the tasks and data preprocessing are provided in Appendix C.2. 6

Table 2: Performance comparison between EHR-RAGp and all baselines. Best results are marked in bold. Second best results are underlined. Model

ICU-Readmit 30d AUROC (CI) AUPRC (CI)

In-hospital Mortality AUROC (CI) AUPRC (CI)

Long LOS 7d AUROC (CI) AUPRC (CI)

1YR Mortality AUROC (CI) AUPRC (CI)

DescEmb [58] DescEmb*‡ [58] GenHPF [13] Med-BERT [32] CEHR-BERT [31] BEHRT [29] Hi-BEHRT [30]

0.660 (0.634, 0.684) 0.632 (0.605, 0.658) 0.705 (0.681, 0.728) 0.715 (0.693, 0.739) 0.718 (0.694, 0.739) 0.737 (0.714, 0.761) 0.596 (0.571, 0.621)

0.096 (0.080, 0.118) 0.098 (0.080, 0.124) 0.134 (0.108, 0.167) 0.106 (0.091, 0.125) 0.133 (0.110, 0.163) 0.133 (0.112, 0.162) 0.055 (0.048, 0.067)

0.933 (0.927, 0.940) 0.911 (0.903, 0.918) 0.932 (0.926, 0.939) 0.907 (0.899, 0.915) 0.933 (0.927, 0.940) 0.925 (0.918, 0.932) 0.897 (0.888, 0.904)

0.849 (0.840, 0.858) 0.815 (0.804, 0.825) 0.850 (0.841, 0.859) 0.802 (0.790, 0.814) 0.838 (0.829, 0.848) 0.831 (0.821, 0.841) 0.775 (0.761, 0.787)

0.480 (0.455, 0.505) 0.403 (0.378, 0.427) 0.482 (0.456, 0.510) 0.387 (0.364, 0.411) 0.447 (0.423, 0.472) 0.439 (0.417, 0.463) 0.355 (0.332, 0.380)

0.766 (0.755, 0.777) 0.740 (0.728, 0.752) 0.781 (0.770, 0.793) 0.749 (0.738, 0.760) 0.778 (0.767, 0.789) 0.771 (0.760, 0.783) 0.700 (0.688, 0.712)

0.329 (0.308, 0.351) 0.304 (0.283, 0.327) 0.345 (0.324, 0.369) 0.305 (0.285, 0.327) 0.342 (0.321, 0.365) 0.337 (0.315, 0.361) 0.254 (0.237, 0.273)

EHRMamba [10] ModernBERT [59]

0.725 (0.702, 0.749) 0.740 (0.715, 0.762)

0.130 (0.109, 0.160) 0.137 (0.114, 0.167)

0.925 (0.918, 0.932) 0.937 (0.931, 0.943)

0.865 (0.855, 0.874) 0.869 (0.86, 0.879)

0.578 (0.552, 0.605) 0.597 (0.573, 0.624)

0.757 (0.746, 0.769) 0.807 (0.797, 0.817)

0.303 (0.283, 0.325) 0.370 (0.349, 0.395)

0.476 (0.453, 0.501) 0.546 (0.521, 0.570) 0.559 (0.534, 0.585) 0.568 (0.542, 0.593)

0.781 (0.770, 0.791) 0.789 (0.778, 0.800) 0.787 (0.777, 0.798) 0.798 (0.788, 0.808)

0.345 (0.325, 0.369) 0.356 (0.336, 0.382) 0.355 (0.332, 0.377) 0.364 (0.342, 0.387)

REMed [54] 0.535 (0.510, 0.561) 0.044 (0.039, 0.052) 0.867 (0.857, 0.876) 0.468 (0.439, 0.499) 0.800 (0.790, 0.811) 0.363 (0.342, 0.386) 0.622 (0.608, 0.636) Vanilla EHR-RAGp 0.742 (0.719, 0.764) 0.146 (0.121, 0.177) 0.939 (0.933, 0.945) 0.710 (0.687, 0.732) 0.879 (0.868, 0.888) 0.618 (0.592, 0.642) 0.814 (0.804, 0.824) EHR-RAGp (Ours) 0.747 (0.724, 0.768) 0.156 (0.128, 0.189) 0.940 (0.933, 0.945) 0.716 (0.693, 0.738) 0.885 (0.876, 0.893) 0.628 (0.603, 0.652) 0.821 (0.811, 0.831) ‡ DescEmb represents DescEmb BERT-FT variant, whereas DescEmb* represents DescEmb CLS-FT variant. For more information, see Appendix D.

0.185 (0.175, 0.199) 0.381 (0.359, 0.405) 0.396 (0.372, 0.422)

Clinical baselines 0.693 (.669, 0.716) 0.631 (0.606, 0.655) 0.706 (0.682, 0.730) 0.609 (0.581, 0.635) 0.700 (0.677, 0.724) 0.674 (0.648, 0.697) 0.585 (0.557, 0.611)

Long context baselines 0.640 (0.614, 0.670) 0.713 (0.688, 0.735)

Transformer-based baselines RoBERTa [60] LongFormer [61] BigBird [62] RoFormer [63]

0.734 (0.710, 0.756) 0.736 (0.713, 0.760) 0.724 (0.700, 0.747) 0.739 (0.716, 0.760)

0.141 (0.118, 0.173) 0.123 (0.105, 0.153) 0.133 (0.109, 0.164) 0.142 (0.116, 0.173)

0.925 (0.917, 0.932) 0.922 (0.914, 0.929) 0.931 (0.925, 0.938) 0.935 (0.928, 0.942)

0.671 (0.644, 0.697) 0.645 (0.615, 0.673) 0.690 (0.667, 0.714) 0.701 (0.679, 0.724)

0.842 (0.833, 0.851) 0.855 (0.846, 0.864) 0.864 (0.855, 0.874) 0.872 (0.863, 0.880)

Retrieval-based models

Table 3: Performance gains of existing EHR baseline with EHR-RAGp. ICU-Readmit 30d AUROC (CI) AUPRC (CI)

In-hospital Mortality AUROC (CI) AUPRC (CI)

Long LOS 7d AUROC (CI) AUPRC (CI)

1YR Mortality AUROC (CI) AUPRC (CI)

Med-BERT

× √

0.715 (0.693, 0.739) 0.725 (0.703, 0.748)

0.106 (0.091, 0.125) 0.121 (0.103, 0.147)

0.907 (0.899, 0.915) 0.929 (0.922, 0.936)

0.609 (0.581, 0.635) 0.683 (0.658, 0.709)

0.802 (0.790, 0.814) 0.861 (0.851, 0.870)

0.387 (0.364, 0.411) 0.586 (0.561, 0.610)

0.749 (0.738, 0.760) 0.773 (0.761, 0.784)

0.305 (0.285, 0.327) 0.307 (0.289, 0.328)

CEHR-BERT

× √

0.718 (0.694, 0.739) 0.745 (0.724, 0.767)

0.133 (0.110, 0.163) 0.167 (0.137, 0.203)

0.933 (0.927, 0.940) 0.939 (0.933, 0.946)

0.700 (0.677, 0.724) 0.717(0.694, 0.741)

0.838 (0.829, 0.848) 0.876 (0.867, 0.885)

0.447 (0.423, 0.472) 0.604 (0.580, 0.627)

0.778 (0.767, 0.789) 0.802 (0.791, 0.811)

0.342 (0.321, 0.365) 0.372 (0.350, 0.396)

BEHRT

× √

0.737 (0.714, 0.761) 0.743 (0.722, 0.765)

0.133 (0.112, 0.162) 0.135 (0.112, 0.166)

0.925 (0.918, 0.932) 0.932 (0.925, 0.939)

0.674 (0.648, 0.697) 0.680 (0.656, 0.704)

0.831 (0.821, 0.841) 0.871 (0.862, 0.879)

0.439 (0.417, 0.463) 0.599 (0.575, 0.623)

0.771 (0.760, 0.783) 0.792 (0.783, 0.803)

0.337 (0.315, 0.361) 0.361 (0.338, 0.385)

Model

EHR-RAGp

Implementation Details. We use RoFormer base [63] as the backbone encoder for EHR-RAGp. In the first stage, we pretrain the encoder using an MLM objective [60]. For fairness of comparison, all code-based baselines in our evaluation are pretrained using the same objective and data splits. The pretrained EHR-RAGp encoder is subsequently reused as the embedding model for constructing and querying the vector database. In the second stage, we fine-tune the full EHR-RAGp model end-to-end for the downstream prediction tasks in a supervised setting. We conduct extensive hyperparameter tuning using Bayesian optimization for all baselines and EHR-RAGp models. Detailed optimization settings and training schedules are provided in Appendix C. We evaluate all models using the Area Under the Operator Receiver Curve (AUROC) and the Area Under the Precession-Recall (AUPRC) and report them along with %95 confidence interval (CI) using bootstrapping [64]. Baselines. We consider a diverse set of baselines spanning clinical EHR models, general transformer architectures, long-context models, retrieval-based methods, and LLMs. Clinical baselines include existing EHR-specific foundation models such as DescEmb [58], GenHPF [13], Med-BERT [32], CEHR-BERT [31], BEHRT [29], and Hi-BEHRT [30]. Since these models were originally trained on different subsets of clinical features, we adapt all baselines using a shared vocabulary and EHR embedding layer to ensure unified input representations and fair comparison. We further include general transformer-based encoders adapted to EHR data, including RoBERTa [60], Longformer [61], BigBird [62], and RoFormer [63]. For long-context modeling, we consider EHRMamba [10], designed for sequences up to 2048 events, and ModernBERT [59], supporting contexts up to 8192 tokens. Retrieval baselines include REMed [54] and a vanilla RAG variant without prototype-guided refinement. Finally, we evaluate both general-purpose (Qwen2.5-7B [65], Mistral-7B [66]) and medically adapted LLMs (MedGemma-1.5-4B [67], BioMistral-7B [68]) in a zero-shot setting. Additional details are provided in Appendix D.

Main Results

Overall Performance Results. Table 2 reports the downstream performance across four clinical prediction tasks. EHR-RAGp consistently achieves the best results across all metrics, outperforming strong baselines from multiple model categories. For ICU readmission, EHR-RAGp attains an AUROC of 0.747 and AUPRC of 0.156, improving over the strongest transformer baseline RoFormer and 7

12.5

Care stage OUTP ED HOSP ICU Prototype

10.0 7.5

UMAP-2

5.0

2.5 0.0

Care stage OUTP ED HOSP ICU Prototype 5

2.5 5.0

5 UMAP-1

7.5

(a) In-hospital mortality

5 UMAP-1

(b) Long length-of-stay (7 days)

Figure 2: UMAP projections of history embeddings and prototypes for (a) long length of stay and (b) in-hospital mortality tasks. long-context model ModernBERT. For in-hospital mortality, EHR-RAGp achieves the best performance (AUROC 0.940, AUPRC 0.716), surpassing ModernBERT (AUPRC 0.713) and CEHR-BERT (AUPRC 0.700). For long length-of-stay, EHR-RAGp reaches AUROC 0.885 and AUPRC 0.628, outperforming both long-context baselines, such as EHRMamba (AUROC 0.865, AUPRC 0.578), and transformer models like RoFormer (AUROC 0.872, AUPRC 0.568). Similarly, for 1-year mortality, it achieves an AUROC of 0.821 and AUPRC 0.396, exceeding RoFormer (AUROC 0.798, AUPRC 0.364) and ModernBERT (AUROC 0.807, AUPRC 0.370). Compared to clinical foundation models, the gains are substantial across all tasks. While the vanilla version without prototypes consistently underperforms the full model, it still outperforms other baselines almost on all tasks. Overall, the results demonstrate that EHR-RAGp provides consistent improvements over all baselines, indicating effective utilization of long-range clinical context beyond increasing context length alone. Table S5 shows that LLM baselines perform substantially worse across all tasks. EHR-RAGp Improves Clinical Foundation Models. Table 3 evaluates whether EHR-RAGp can improve existing EHR foundation models when used as a retrieval-augmented extension. Across Med-BERT, CEHR-BERT, and BEHRT, adding EHR-RAGp consistently improves performance across all four tasks. For Med-BERT, the largest gains appear in long LOS prediction, where AUPRC increases from 0.387 to 0.586, and in-hospital mortality, where AUPRC improves from 0.609 to 0.683. For CEHR-BERT, EHR-RAGp improves ICU readmission AUPRC from 0.133 to 0.167, long LOS AUPRC from 0.447 to 0.604, and 1-year mortality AUPRC from 0.342 to 0.372. Similar gains are observed for BEHRT, especially for long LOS, where AUPRC increases from 0.439 to 0.599. These results show that EHR-RAGp is not only a standalone model, but also a general retrieval-augmentation framework that can strengthen existing EHR foundation models. Latent Representation Structure. Figure 2 presents UMAP projections of history embeddings for two tasks: in-hospital mortality and long length-of-stay. Across both tasks, the embedding space exhibits a consistent global structure, with clear separation between high-acuity (ICU) and loweracuity (HOSP, ED, OUTP) regions, indicating that the model organizes patient history primarily along clinical severity. Additionally, task-specific variations are evident in the internal geometry of clusters, suggesting that the learned representations adapt to the prediction objective. Notably, prototype centers (black stars) are distributed across dense regions of the manifold, covering both dominant ICU clusters and smaller peripheral structures, which supports their role in capturing diverse modes of patient trajectories. Overall, the visualization highlights that the model learns a structured and task-sensitive representation space while preserving clinically meaningful organization across care stages. See Figure S4 for UMAP plots of ICU-Readmit 30d and 1YR Mortality tasks. To further analyze retrieval behavior and prediction dynamics, we provide qualitative analysis in Appendix E.3, illustrating prototype usage, chunk relevance weighting, and care-stage contributions at inference.

Table 5: Impact of choice of chunking strategy on model performance across all tasks. Model

ICU-Readmit 30d AUROC (CI) AUPRC (CI)

In-hospital Mortality AUROC (CI) AUPRC (CI)

Long LOS 7d AUROC (CI) AUPRC (CI)

1YR Mortality AUROC (CI) AUPRC (CI)

Event-Based Time-Based Visit-Level Care-Stage

0.747 (0.724, 0.768) 0.745 (0.722, 0.766) 0.738 (0.715, 0.762) 0.743 (0.720, 0.764)

0.939 (0.933, 0.945) 0.937 (0.930, 0.943) 0.939 (0.933, 0.945) 0.940 (0.933, 0.945)

0.884 (0.875, 0.891) 0.884 (0.876, 0.893) 0.885 (0.876, 0.893) 0.884 (0.875, 0.893)

0.817 (0.807, 0.827) 0.807 (0.797, 0.817) 0.821 (0.811, 0.831) 0.820 (0.811, 0.829)

0.156 (0.128, 0.189) 0.152 (0.126, 0.185) 0.153 (0.125, 0.184) 0.150 (0.124, 0.182)

0.710 (0.684, 0.732) 0.706 (0.683, 0.730) 0.715 (0.692, 0.737) 0.716 (0.693, 0.738)

0.611 (0.586, 0.634) 0.623 (0.598, 0.647) 0.628 (0.603, 0.652) 0.618 (0.590, 0.645)

0.390 (0.365, 0.416) 0.377 (0.353, 0.402) 0.396 (0.372, 0.422) 0.393 (0.369, 0.418)

Ablations

Effect of Regularization. Table 4 evaluates the Table 4: Ablation of regularization term with impact of prototype regularization (λu ) on both different number of prototypes (#p). representation behavior and downstream perforQuery History Performance AUPRC mance. Without regularization, prototype assign- #p λ P max Entropy P Max Entropy AUROC Long LOS 7d ments are highly peaked for queries (e.g., Pmax = × 0.887 0.663 0.094 4.048 0.866 0.579 64 ✓ 0.084 4.096 0.027 4.156 0.884 0.618 0.887 for Long LOS), indicating collapse to a few ICU-Readmit 30d dominant prototypes, while history assignments × 0.716 1.578 0.030 6.099 0.730 0.118 0.156 remain sparse and less informative. Enabling reg- 512 ✓ 0.014 6.198 0.004 6.232 0.747 ularization significantly increases entropy and reduces Pmax for both query and history, leading to more distributed and balanced prototype usage. This results in consistent performance gains across tasks, with Long LOS AUPRC improving from 0.579 to 0.618, and ICU readmission AUPRC from 0.118 to 0.156. Overall, the results show that regularization prevents prototype collapse, promotes better utilization of the prototype space, and improves performance. See Appendix E.4 for further analysis on the usage regularization λu . u

(0.857, 0.876)

(0.552, 0.605)

(0.875, 0.893)

(0.590, 0.645)

(0.709, 0.752)

(0.098, 0.144)

(0.724, 0.768)

(0.128, 0.189)

Effect of Chunking Strategies. Table 5 presents an ablation over chunking strategies. Overall, performance is consistent across all strategies, indicating that EHR-RAGp is robust to how patient history is segmented. Event-based chunking achieves the best results for ICU readmission (AUROC 0.747, AUPRC 0.156), while care-stage chunking performs best for in-hospital mortality (AUROC 0.940, AUPRC 0.716). Visit-level chunking yields the strongest performance for long-term outcomes, including long LOS (AUPRC 0.628) and 1-year mortality (AUPRC 0.396). Time-based chunking performs competitively across all tasks but does not outperform other strategies. These results suggest that no single chunking scheme dominates universally, and different chunking methods capture complementary aspects of patient trajectories depending on the prediction task. Effect of Prototypes Count. Table 6 shows Table 6: Ablation of number of prototypes. the effect of varying the number of prototypes. ICU-Readmit 30d In-hospital Mortality Increasing the number of prototypes consis- #p AUROC AUPRC AUROC AUPRC tently improves performance for ICU readmis- 64 0.728 (0.705, 0.751) 0.132 (0.109, 0.161) 0.937 (0.931, 0.943) 0.704 (0.680, 0.726) 128 0.736 (0.711, 0.757) 0.144 (0.121, 0.176) 0.937 (0.930, 0.943) 0.706 (0.683, 0.730) sion, with AUROC rising from 0.728 to 0.743, 256 0.739 (0.715, 0.761) 0.145 (0.121, 0.179) 0.934 (0.928, 0.94) 0.695 (0.672, 0.721) with 64 and 512 prototypes, and AUPRC from 512 0.743 (0.720, 0.764) 0.150 (0.124, 0.182) 0.935 (0.928, 0.941) 0.697 (0.673, 0.723) 0.132 to 0.150, indicating that a richer prototype set better captures fine-grained structure in patient trajectories. In contrast, performance on in-hospital mortality remains relatively stable across configurations, with marginal variations around AUROC and AUPRC. Overall, this suggests that the effective number of prototypes is also task dependent.

Discussion & Conclusion

In this work, we introduce EHR-RAGp, a retrieval-augmented, prototype-guided foundation model designed to reason over heterogeneous EHR trajectories by selectively incorporating relevant patient history. By framing longitudinal EHR modeling as a retrieval problem, EHR-RAGp addresses the limitations of fixed-window approaches and enables scalable integration of long-range clinical context. Across multiple downstream tasks, our approach consistently outperforms strong baselines as well as boosting existing clinical foundation models when used as an extension. Beyond quantitative improvements, our qualitative analyses further show that EHR-RAGp learns structured latent representations and dynamically reweighs retrieved patient history according to downstream relevance rather than retrieval similarity alone. Together, these results support the central claim of this work: increasing the effective medical context available to prediction models improves performance [6, 10, 30]. Limitations. While EHR-RAGp shows strong potential, some limitations remain to be investigated in future research. First, experiments are conducted on a single EHR dataset, and evaluation on different 9

datasets is needed to assess generalizability across different patient populations and coding systems. Second, experiments were conducted under a limited set of retrieval, and chunking configurations (e.g., chunk size of 256, 6-hour temporal windows, and Top-M of 24), further investigation of larger contexts and alternative chunking granularities is required. Third, although prototype-guided alignment improves performance and organizes the representation space into latent modes, the clinical interpretability of individual prototypes remains limited and requires deeper analysis with expert validation. Finally, while EHR-RAGp is designed on purpose for event-based EHR data, incorporation of other clinical modalities may provide complementary context for retrieval and prediction.

References [1] Nurul Athirah Nasarudin, Fatma Al Jasmi, Richard O Sinnott, Nazar Zaki, Hany Al Ashwal, Elfadil A Mohamed, and Mohd Saberi Mohamad. A review of deep learning models and online healthcare databases for electronic health records and their use for health prediction. Artificial Intelligence Review, 57(9):249, 2024. [2] Cao Xiao, Edward Choi, and Jimeng Sun. Opportunities and challenges in developing deep learning models using electronic health records data: a systematic review. Journal of the American Medical Informatics Association, 25(10):1419–1428, 2018. [3] Michael Wornow, Yizhe Xu, Rahul Thapa, Birju Patel, Ethan Steinberg, Scott Fleming, Michael A Pfeffer, Jason Fries, and Nigam H Shah. The shaky foundations of large language models and foundation models for electronic health records. npj digital medicine, 6(1):135, 2023. [4] Lingyao Li, Jiayan Zhou, Zhenxiang Gao, Wenyue Hua, Lizhou Fan, Huizi Yu, Loni Hagen, Yongfeng Zhang, Themistocles L Assimes, Libby Hemphill, et al. A scoping review of using large language models (llms) to investigate electronic health records (ehrs). arXiv preprint arXiv:2405.03066, 2024. [5] Weijieying Ren, Jingxi Zhu, Zehao Liu, Tianxiang Zhao, and Vasant Honavar. A comprehensive survey of electronic health record modeling: From deep learning approaches to large language models. arXiv preprint arXiv:2507.12774, 2025. [6] Michael Wornow, Suhana Bedi, Miguel Angel Fuentes Hernandez, Ethan Steinberg, Jason Alan Fries, Christopher Ré, Sanmi Koyejo, and Nigam H Shah. Context clues: Evaluating long context models for clinical prediction tasks on ehrs. arXiv preprint arXiv:2412.16178, 2024. [7] Shaojie Zhong, Li Rong Wang, Zhuoxuan Zhan, Yih Yng Ng, and Xiuyi Fan. A hybrid approach for irregular-time series prediction using electronic health records: an intensive care unit mortality case study. ACM Transactions on Computing for Healthcare, 6(4):1–33, 2025. [8] Michael Wornow, Rahul Thapa, Ethan Steinberg, Jason Fries, and Nigam Shah. Ehrshot: An ehr benchmark for few-shot evaluation of foundation models. Advances in Neural Information Processing Systems, 36:67125–67137, 2023. [9] Shaza Elsharief, Saeed Shurrab, Baraa Al Jorf, L Julián Lechuga López, and Farah E Shamout. Medmod: Multimodal benchmark for medical prediction tasks with electronic health records and chest x-ray scans. Proceedings of Machine Learning Research, 287:1–23, 2025. [10] Adibvafa Fallahpour, Mahshid Alinoori, Wenqian Ye, Xu Cao, Arash Afkanpour, and Amrit Krishnan. Ehrmamba: Towards generalizable and scalable foundation models for electronic health records. arXiv preprint arXiv:2405.14567, 2024. [11] Mikkel Odgaard, Kiril Vadimovic Klein, Sanne Møller Thysen, Espen Jimenez-Solem, Martin Sillesen, and Mads Nielsen. Core-behrt: A carefully optimized and rigorously evaluated behrt. arXiv preprint arXiv:2404.15201, 2024. [12] Kyunghoon Hur, Jungwoo Oh, Junu Kim, Jiyoun Kim, Min Jae Lee, Eunbyeol Cho, Seong-Eun Moon, Young-Hak Kim, and Edward Choi. Unihpf: Universal healthcare predictive framework with zero domain knowledge. arXiv preprint arXiv:2211.08082, 2022. 10

[13] Kyunghoon Hur, Jungwoo Oh, Junu Kim, Jiyoun Kim, Min Jae Lee, Eunbyeol Cho, SeongEun Moon, Young-Hak Kim, Louis Atallah, and Edward Choi. Genhpf: General healthcare predictive framework for multi-task multi-source learning. IEEE Journal of Biomedical and Health Informatics, 28(1):502–513, 2023. [14] Zhichao Yang, Avijit Mitra, Weisong Liu, Dan Berlowitz, and Hong Yu. Transformehr: transformer-based encoder-decoder generative model to enhance prediction of disease outcomes using electronic health records. Nature communications, 14(1):7857, 2023. [15] Leonard E Braitman and Frank Davidoff. Predicting clinical states in individual patients. Annals of Internal Medicine, 125(5):406–412, 1996. [16] Huilong Duan, Zhoujian Sun, Wei Dong, Kunlun He, and Zhengxing Huang. On clinical event prediction in patient treatment trajectory using longitudinal electronic health records. IEEE Journal of Biomedical and Health Informatics, 24(7):2053–2063, 2019. [17] Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al. Retrieval-augmented generation for knowledge-intensive nlp tasks. Advances in neural information processing systems, 33:9459–9474, 2020. [18] Shangyu Wu, Ying Xiong, Yufei Cui, Haolun Wu, Can Chen, Ye Yuan, Lianming Huang, Xue Liu, Tei-Wei Kuo, Nan Guan, et al. Retrieval-augmented generation for natural language processing: A survey. arXiv preprint arXiv:2407.13193, 2024. [19] Hao Yu, Aoran Gan, Kai Zhang, Shiwei Tong, Qi Liu, and Zhaofeng Liu. Evaluation of retrieval-augmented generation: A survey. In CCF Conference on Big Data, pages 102–120. Springer, 2024. [20] Penghao Zhao, Hailin Zhang, Qinhan Yu, Zhengren Wang, Yunteng Geng, Fangcheng Fu, Ling Yang, Wentao Zhang, Jie Jiang, and Bin Cui. Retrieval-augmented generation for ai-generated content: A survey. arXiv preprint arXiv:2402.19473, 2024. [21] Xu Zheng, Ziqiao Weng, Yuanhuiyi Lyu, Lutao Jiang, Haiwei Xue, Bin Ren, Danda Paudel, Nicu Sebe, Luc Van Gool, and Xuming Hu. Retrieval augmented generation and understanding in vision: A survey and new outlook. arXiv preprint arXiv:2503.18016, 2025. [22] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):1735–1780, 1997. [23] Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. arXiv preprint arXiv:1409.1259, 2014. [24] Zhengping Che, Sanjay Purushotham, Kyunghyun Cho, David Sontag, and Yan Liu. Recurrent neural networks for multivariate time series with missing values. Scientific reports, 8(1):6085, 2018. [25] Hrayr Harutyunyan, Hrant Khachatrian, David C Kale, Greg Ver Steeg, and Aram Galstyan. Multitask learning and benchmarking with clinical time series data. Scientific data, 6(1):96, 2019. [26] Nasir Hayat, Krzysztof J Geras, and Farah E Shamout. Medfuse: Multi-modal fusion with clinical time-series data and chest x-ray images. In Machine Learning for Healthcare Conference, pages 479–503. PMLR, 2022. [27] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume 1 (long and short papers), pages 4171–4186, 2019. [28] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019. 11

[29] Yikuan Li, Shishir Rao, José Roberto Ayala Solares, Abdelaali Hassaine, Rema Ramakrishnan, Dexter Canoy, Yajie Zhu, Kazem Rahimi, and Gholamreza Salimi-Khorshidi. Behrt: transformer for electronic health records. Scientific reports, 10(1):7155, 2020. [30] Yikuan Li, Mohammad Mamouei, Gholamreza Salimi-Khorshidi, Shishir Rao, Abdelaali Hassaine, Dexter Canoy, Thomas Lukasiewicz, and Kazem Rahimi. Hi-behrt: hierarchical transformer-based model for accurate prediction of clinical events using multimodal longitudinal electronic health records. IEEE journal of biomedical and health informatics, 27(2):1106–1117, 2022. [31] Chao Pang, Xinzhuo Jiang, Krishna S Kalluri, Matthew Spotnitz, RuiJun Chen, Adler Perotte, and Karthik Natarajan. Cehr-bert: Incorporating temporal information from structured ehr data to improve prediction tasks. In Machine Learning for Health, pages 239–260. PMLR, 2021. [32] Laila Rasmy, Yang Xiang, Ziqian Xie, Cui Tao, and Degui Zhi. Med-bert: pretrained contextualized embeddings on large-scale structured electronic health records for disease prediction. NPJ digital medicine, 4(1):86, 2021. [33] Pawel Renc, Michal K Grzeszczyk, Nassim Oufattole, Deirdre Goode, Yugang Jia, Szymon Bieganski, Matthew BA McDermott, Jaroslaw Was, Anthony E Samir, Jonathan W Cunningham, et al. Foundation model of electronic medical records for adaptive risk estimation. GigaScience, 14:giaf107, 2025. [34] Chao Pang, Xinzhuo Jiang, Nishanth Parameshwar Pavinkurve, Krishna S Kalluri, Elise L Minto, Jason Patterson, Linying Zhang, George Hripcsak, Gamze Gürsoy, Noémie Elhadad, et al. Cehr-gpt: Generating electronic health records with chronological patient timelines. arXiv preprint arXiv:2402.04400, 2024. [35] Shane Waxler, Paul Blazek, Davis White, Daniel Sneider, Kevin Chung, Mani Nagarathnam, Patrick Williams, Hank Voeller, Karen Wong, Matthew Swanhorst, et al. Generative medical event models improve with scale. arXiv preprint arXiv:2508.12104, 2025. [36] Chantal Pellegrini, Ege Özsoy, David Bani-Harouni, Matthias Keicher, and Nassir Navab. From ehrs to patient pathways: Scalable modeling of longitudinal health trajectories with llms. arXiv preprint arXiv:2506.04831, 2025. [37] Simon A Lee, Sujay Jain, Alex Chen, Kyoka Ono, Arabdha Biswas, Ákos Rudas, Jennifer Fang, and Jeffrey N Chiang. Clinical decision support using pseudo-notes from multiple streams of ehr data. npj Digital Medicine, 8(1):394, 2025. [38] Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George Bm Van Den Driessche, Jean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, et al. Improving language models by retrieving from trillions of tokens. In International conference on machine learning, pages 2206–2240. PMLR, 2022. [39] Zhengbao Jiang, Frank F Xu, Luyu Gao, Zhiqing Sun, Qian Liu, Jane Dwivedi-Yu, Yiming Yang, Jamie Callan, and Graham Neubig. Active retrieval augmented generation. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 7969–7992, 2023. [40] Weijia Shi, Sewon Min, Michihiro Yasunaga, Minjoon Seo, Richard James, Mike Lewis, Luke Zettlemoyer, and Wen-tau Yih. Replug: Retrieval-augmented black-box language models. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pages 8371–8384, 2024. [41] Walid Saba, Suzanne Wendelken, and James Shanahan. Question-answering based summarization of electronic health records using retrieval augmented generation. arXiv preprint arXiv:2401.01469, 2024. [42] Jiwoong Sohn, Yein Park, Chanwoong Yoon, Sihyeon Park, Hyeon Hwang, Mujeen Sung, Hyunjae Kim, and Jaewoo Kang. Rationale-guided retrieval augmented generation for medical question answering. In Proceedings of the 2025 Conference of the Nations of the Americas 12

Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pages 12739–12753, 2025. [43] Matthew Lewis, Samuel Thio, Richard JB Dobson, and Spiros Denaxas. Grounding large language models in clinical evidence: A retrieval-augmented generation system for querying uk nice clinical guidelines. arXiv preprint arXiv:2510.02967, 2025. [44] Shuyue Jia, Subhrangshu Bit, Varuna H Jasodanand, Yi Liu, and Vijaya B Kolachalama. Agentic memory-augmented retrieval and evidence grounding for medical question-answering tasks. medRxiv, pages 2025–08, 2025. [45] Yichun Feng, Jiawei Wang, Ruikun He, Lu Zhou, and Yixue Li. A retrieval-augmented knowledge mining method with deep thinking llms for biomedical research and clinical support. GigaScience, 14:giaf109, 2025. [46] Hui Feng, Yuntzu Yin, Emiliano Reynares, and Jay Nanavati. Ontologyrag: better and faster biomedical code mapping with retrieval-augmented generation (rag) leveraging ontology knowledge graphs and large language models. In International Workshop on Knowledge-Enhanced Information Retrieval, pages 71–86. Springer, 2025. [47] Angelo Ziletti and Leonardo D’Ambrosi. Generating patient cohorts from electronic health records using two-step retrieval-augmented text-to-sql generation. arXiv preprint arXiv:2502.21107, 2025. [48] Mohamed Abo El-Enen, Sally Saad, and Taymoor Nazmy. A survey on retrieval-augmentation generation (rag) models for healthcare applications. Neural Computing and Applications, 37(33):28191–28267, 2025. [49] Lameck Mbangula Amugongo, Pietro Mascheroni, Steven Brooks, Stefan Doering, and Jan Seidel. Retrieval augmented generation for large language models in healthcare: A systematic review. PLOS Digital Health, 4(6):e0000877, 2025. [50] Yinghao Zhu, Changyu Ren, Zixiang Wang, Xiaochen Zheng, Shiyun Xie, Junlan Feng, Xi Zhu, Zhoujun Li, Liantao Ma, and Chengwei Pan. Emerge: Enhancing multimodal electronic health records predictive modeling with retrieval-augmented generation. In Proceedings of the 33rd ACM International Conference on Information and Knowledge Management, pages 3549–3559, 2024. [51] Ran Xu, Wenqi Shi, Yue Yu, Yuchen Zhuang, Bowen Jin, May Dongmei Wang, Joyce Ho, and Carl Yang. Ram-ehr: Retrieval augmentation meets clinical predictions on electronic health records. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pages 754–765, 2024. [52] Yinghao Zhu, Changyu Ren, Shiyun Xie, Shukai Liu, Hangyuan Ji, Zixiang Wang, Tao Sun, Long He, Zhoujun Li, Xi Zhu, et al. Realm: Rag-driven enhancement of multimodal electronic health records analysis via large language models. arXiv preprint arXiv:2402.07016, 2024. [53] Rituparna Datta, Jiaming Cui, Zihan Guan, Vishal G Reddy, Joshua C Eby, Gregory Madden, Rupesh Silwal, and Anil Vullikanti. Improving hospital risk prediction with knowledgeaugmented multimodal ehr modeling. arXiv preprint arXiv:2508.01970, 2025. [54] Junu Kim, Chaeeun Shim, Bosco Seong Kyu Yang, Chami Im, Sung Yoon Lim, Han-Gil Jeong, and Edward Choi. General-purpose retrieval-enhanced medical prediction model using near-infinite history. arXiv preprint arXiv:2310.20204, 2023. [55] Mahmoud Assran, Mathilde Caron, Ishan Misra, Piotr Bojanowski, Florian Bordes, Pascal Vincent, Armand Joulin, Mike Rabbat, and Nicolas Ballas. Masked siamese networks for label-efficient learning. In European conference on computer vision, pages 456–473. Springer, 2022. [56] Alistair Johnson, Lucas Bulgarelli, Tom Pollard, Steven Horng, Leo Anthony Celi, and Roger Mark. Mimic-iv. PhysioNet. Available online at: https://physionet. org/content/mimiciv/1.0/(accessed August 23, 2021), pages 49–55, 2020. 13

[57] Alistair EW Johnson, Lucas Bulgarelli, Lu Shen, Alvin Gayles, Ayad Shammout, Steven Horng, Tom J Pollard, Sicheng Hao, Benjamin Moody, Brian Gow, et al. Mimic-iv, a freely accessible electronic health record dataset. Scientific data, 10(1):1, 2023. [58] Kyunghoon Hur, Jiyoung Lee, Jungwoo Oh, Wesley Price, Younghak Kim, and Edward Choi. Unifying heterogeneous electronic health records systems via text-based code embedding. In Conference on Health, Inference, and Learning, pages 183–203. PMLR, 2022. [59] Benjamin Warner, Antoine Chaffin, Benjamin Clavié, Orion Weller, Oskar Hallström, Said Taghadouini, Alexis Gallagher, Raja Biswas, Faisal Ladhak, Tom Aarsen, et al. Smarter, better, faster, longer: A modern bidirectional encoder for fast, memory efficient, and long context finetuning and inference. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 2526–2547, 2025. [60] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692, 2019. [61] Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020. [62] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283–17297, 2020. [63] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063, 2024. [64] Marie-Therese Puth, Markus Neuhäuser, and Graeme D Ruxton. On the variety of methods for calculating confidence intervals by bootstrapping. Journal of Animal Ecology, 84(4):892–897, 2015. [65] An Yang, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, Guanting Dong, Haoran Wei, Huan Lin, Jialong Tang, Jialin Wang, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Ma, Jin Xu, Jingren Zhou, Jinze Bai, Jinzheng He, Junyang Lin, Kai Dang, Keming Lu, Keqin Chen, Kexin Yang, Mei Li, Mingfeng Xue, Na Ni, Pei Zhang, Peng Wang, Ru Peng, Rui Men, Ruize Gao, Runji Lin, Shijie Wang, Shuai Bai, Sinan Tan, Tianhang Zhu, Tianhao Li, Tianyu Liu, Wenbin Ge, Xiaodong Deng, Xiaohuan Zhou, Xingzhang Ren, Xinyu Zhang, Xipin Wei, Xuancheng Ren, Yang Fan, Yang Yao, Yichang Zhang, Yu Wan, Yunfei Chu, Yuqiong Liu, Zeyu Cui, Zhenru Zhang, and Zhihao Fan. Qwen2 technical report. arXiv preprint arXiv:2407.10671, 2024. [66] Albert Q Jiang, A Sablayrolles, A Mensch, C Bamford, D Singh Chaplot, Ddl Casas, F Bressand, G Lengyel, G Lample, L Saulnier, et al. Mistral 7b. arxiv. arXiv preprint arXiv:2310.06825, 10:3, 2023. [67] Andrew Sellergren, Chufan Gao, Fereshteh Mahvar, Timo Kohlberger, Fayaz Jamil, Madeleine Traverse, Alberto Tono, Bashir Sadjad, Lin Yang, Charles Lau, et al. Medgemma 1.5 technical report. arXiv preprint arXiv:2604.05081, 2026. [68] Yanis Labrak, Adrien Bazoge, Emmanuel Morin, Pierre-Antoine Gourraud, Mickael Rouvier, and Richard Dufour. Biomistral: A collection of open-source pretrained large language models for medical domains, 2024. [69] Bert Arnrich, Edward Choi, Jason Alan Fries, Matthew BA McDermott, Jungwoo Oh, Tom Pollard, Nigam Shah, Ethan Steinberg, Michael Wornow, and Robin van de Water. Medical event data standard (meds): Facilitating machine learning for health. In ICLR 2024 Workshop on Learning from Time Series For Health, pages 03–08, 2024. [70] Matthew BA McDermott, Justin Xu, Teya S Bergamaschi, Hyewon Jeong, Simon A Lee, Nassim Oufattole, Patrick Rockenschaub, Kamilė Stankevičiūtė, Ethan Steinberg, Jimeng Sun, et al. Meds: Building models and tools in a reproducible health ai ecosystem. In Proceedings of 14

the 31st ACM SIGKDD Conference on Knowledge Discovery and Data Mining V. 2, pages 6243–6244, 2025. [71] Seyed Mehran Kazemi, Rishab Goel, Sepehr Eghbali, Janahan Ramanan, Jaspreet Sahota, Sanjay Thakur, Stella Wu, Cathal Smyth, Pascal Poupart, and Marcus Brubaker. Time2vec: Learning a vector representation of time. arXiv preprint arXiv:1907.05321, 2019. [72] Hong Meng, Liang Guo, Yucheng Pan, Bin Kong, Wei Shuai, and He Huang. Machine learning based clinical prediction model for 1-year mortality in sepsis patients with atrial fibrillation. Heliyon, 10(21), 2024. [73] Beatriz Nistal-Nuño. Developing machine learning models for prediction of mortality in the medical intensive care unit. Computer Methods and Programs in Biomedicine, 216:106663, 2022. [74] Marzyeh Ghassemi, Tristan Naumann, Finale Doshi-Velez, Nicole Brimmer, Rohit Joshi, Anna Rumshisky, and Peter Szolovits. Unfolding physiological state: Mortality modelling in intensive care units. In Proceedings of the 20th ACM SIGKDD international conference on Knowledge discovery and data mining, pages 75–84, 2014. [75] Min Zhang and Tsung-Ting Kuo. Early prediction of long hospital stay for intensive care units readmission patients using medication information. Computers in biology and medicine, 174:108451, 2024. [76] Khalid Alghatani, Nariman Ammar, Abdelmounaam Rezgui, and Arash Shaban-Nejad. Predicting intensive care unit length of stay and mortality using patient vital signs: machine learning model development and validation. JMIR medical informatics, 9(5):e21347, 2021. [77] Huiling Hu, Jiashuai Li, Hui Ge, Bilin Wu, Tingting Feng, Xue Wu, and Xuanna Wu. Prognostic models for unplanned intensive care unit readmission risk prediction: A systematic review and meta-analysis based on hsroc model. Nursing in critical care, 30(2):e13306, 2025. [78] Talen Chen, Samaneh Madanian, David Airehrour, and Marianne Cherrington. Machine learning methods for hospital readmission prediction: systematic analysis of literature. Journal of Reliable Intelligent Environments, 8(1):49–66, 2022. [79] Alex GC de Sá, Daniel Gould, Anna Fedyukova, Mitchell Nicholas, Lucy Dockrell, Calvin Fletcher, David Pilcher, Daniel Capurro, David B Ascher, Khaled El-Khawas, et al. Explainable machine learning for icu readmission prediction. arXiv preprint arXiv:2309.13781, 2023. [80] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017. [81] Leslie N Smith. Cyclical learning rates for training neural networks. In 2017 IEEE winter conference on applications of computer vision (WACV), pages 464–472. IEEE, 2017. [82] Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. Advances in neural information processing systems, 33:21271–21284, 2020. [83] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In First conference on language modeling, 2024.

Broader Impact

Our work presents a step forward toward retrieval-augmented foundation models for structured EHR data. By enabling reasoning over a patient’s full longitudinal history, our approach may improve the performance of clinical prediction models and decision-support systems, with downstream applications in patient risk stratification, resource allocation, and personalized care. Such advances could support clinicians in managing complex medical histories and contribute to improved healthcare outcomes. At the same time, models trained on sensitive health data raise important ethical considerations, including risks related to privacy, data bias, and inappropriate reliance on automated predictions. While our experiments use de-identified data and are intended for research purposes, real-world deployment would require careful validation, governance, and safeguards to ensure fairness, transparency, and patient safety. We encourage continued interdisciplinary collaboration to responsibly guide the development and use of retrieval-augmented EHR foundation models.

Data Preprocessing

B.1

Cohort construction

Pretraining Cohort: We first convert the MIMIC-IV dataset into the Medical Event Data Standard (MEDS) format [69, 70], representing each patient history as a chronologically ordered sequence of medical events covering all patient visits within a single sequence. We then filter the dataset to construct the pretraining cohort. Specifically, we exclude patients listed in the patients table who do not have any valid hospital or ICU admissions. Next, we examine the total number of events within each visit and exclude visits containing fewer than 10 events. At this stage, we do not apply any age-based filtering, as this cohort is intended solely for pretraining. We further exclude all patients included in the test set from the pretraining cohort to prevent data leakage during evaluation. The resulting cohort consists of 199,012 patients with 463,436 hospital admissions, 80,005 ICU stays, and 496,199,673 medical events. Downstream Cohort: For downstream tasks, we construct the dataset using ICU stays only and exclude patients with hospital admissions but no ICU stays. We follow the inclusion and exclusion criteria proposed by [25] to define the downstream cohort. First, we include all adult patients with an age at admission of ≥ 18 years. Second, we retain patients with a single ICU stay within the same hospital admission. Third, we include only patients whose admission and discharge ICU units are identical. Finally, we examine the length of each valid ICU stay and retain only those stays lasting at least 24 hours. This process results in a downstream cohort comprising 49,839 patients with 61,175 ICU stays. B.2

Data Preprocessing

We perform comprehensive preprocessing to reduce irregularities and transform the data into a standardized format. First, MIMIC-IV contains ICD diagnosis and procedure codes from both versions 9 and 10. To unify all codes under a single coding system, we map all ICD-9 diagnosis and procedure codes to ICD-10. We use the General Equivalence Mappings (GEMS1 ) released in 2018, which is the most recent version available. Second, we examine medication events and remove any medications appearing in the patient timeline with ambiguous or unclear names. For laboratory tests, we first identify all lab entries listed in the MIMIC-IV file d_labitems_to_loinc.csv that are marked as non-laboratory measurements and filter them out. Subsequently, we modify lab event naming to incorporate the lab item_id, specimen fluid, and test name. This naming convention preserves the original structure of the MIMIC-IV dataset, increases the expressiveness of laboratory events, and reduces semantic ambiguity. For ICU data, we modify event names to explicitly reflect the source table. For example, events originating from the chartevents.csv table are prefixed with ICU-CHART. We further include the corresponding item_id as a secondary identifier, followed by the event-specific name. Representative examples are provided in Table S1. 1 https://www.cms.gov/medicare/coding-billing/icd-10-codes/icd-10-cm-icd-10-pcs-gem-archive

Next, we handle outliers in numeric values using the MEDS-transform Python package by computing dataset-wide statistics for each numerical event and excluding values that lie beyond three standard deviations from the mean. Finally, we normalize all remaining numeric values using MEDS-transform by recomputing normalization statistics after outlier removal. B.3

Patient Timeline Construction

Figure S1 depicts a sample patient timeline segmented into distinct care stages and covering all relevant clinical information. To construct a patient timeline, we first examine each patient sequence and identify hospital admissions (HOSP) using the hadm_id field, as well as ICU stays (ICU) using the icustay_id field, which are the only care stages in MIMIC-IV associated with unique identifiers. During this process, we observe that certain medical events, such as laboratory tests and microbiology orders, appear in the patient timeline without an explicit reference to any hospital admission. To handle these events without excluding them, we assess their temporal proximity to admissions present in the patient timeline. Based on the time difference, we associate each such event with the nearest hospital admission by assigning the corresponding hadm_id. Events occurring within 24 hours of the nearest hospital admission are labeled as emergency department (ED) events, while events occurring within 30 days are treated as outpatient (OUTP) events. To represent time gaps between consecutive visits, we insert artificial time tokens following prior work [31, 10], denoted as GAP, as illustrated in Figure S1. These tokens encode temporal intervals spanning weeks, months, or years, depending on the elapsed time between visits. A standard patient timeline primarily consists of hospital events including laboratory measurements, medications, microbiology tests, ICD procedure codes, ICD diagnosis codes, and DRG codes, as well as ICU events such as chart events, infusions, procedures, and fluid outputs. In addition, administrative events are included for both hospital and ICU stays to indicate boundaries, locations, and care types. To explicitly preserve care-stage boundaries within the timeline, we introduce special boundary tokens, including OUTPATIENT-START, OUTPATIENT-END, EMERGENCY-START, and EMERGENCY-END for OUTP and ED stages, respectively. Demographic attributes such as gender and race are treated as static events without timestamps and are prepended to the beginning of each sequence. Patient age is computed at the start of each stay and represented as a dedicated event token, AGE-AT-ADMISSION, with an associated numeric value. If a patient dies during or after discharge, a special timestamped token, MEDS_DEATH, is appended to the sequence. We also include special tokens required for language modeling objectives, such as PAD, MASK, CLS, and UNK, where applicable. Table S1 summarizes the different event types present in the patient timeline along with their frequencies. For events associated with numeric values, such as laboratory tests and ICU chart events, we attach the corresponding numeric measurements. For events without numeric values, including PROCEDURE-ICD and DIAGNOSIS-ICD, we associate a learnable null parameter. To encode temporal information, we compute the time difference between consecutive events in minutes and attach this value to each event. Because time gaps may vary substantially both within and across visits, we scale time deltas using the transformation defined in Equation 1, which maps values to the range [0, 1] and prevents extreme magnitudes. Temporal representations are then processed using Time2Vec [71] during training. ′

∆ti =

log(1 + ∆ti ) log(∆tmax )

(1)

where ∆ti is the consecutive time difference at event i, and ∆max is the maximum time difference present in the dataset. To further contextualize each event within the patient timeline, we introduce three additional categorical representations, care stage, visit order, and event type, as illustrated in the bottom rows of Figure S1. The care stage embedding explicitly encodes the clinical context in which each event occurs, distinguishing between outpatient (OUTP), emergency department (ED), inpatient hospitalization (HOSP), ICU stay (ICU), and artificial gap (GAP) events. This representation enables the model to differentiate identical medical events occurring under different clinical settings, which often carry distinct semantic meanings. 17

Visit 1

OUTP GAP

Visit 2

GAP

HOSP

ICU

HOSP

M3 M4

Concept

Value

null

null null null null null

null null null null null null null null

Time

Δt Δt

Δt

Δt Δt Δt Δt Δt Δt Δt

Δt

Δt Δt Δt Δt

Stage

Visit

Type

M1 M2

Δt Δt Δt Δt Δt Δt

Figure S1: Sample patient timeline consisting of two visits. We illustrate the main components constituting the patient timeline in our implementation including medical events, numeric values, time values, care stage representation, visit order representation, and event type representation. Abbreviations, OUTP: Outpatient, ED: Emergency Department, HOSP: Hospital Admission, ICU: Incentive Care Unit, L: Labs, T: TIME-GAP, P: Procedure, M Medication, D: Diagnosis, xa numeric value. This is a sample timeline created randomly for illustration purposes and does not reflect real data. In addition, we assign a visit order index to each event, indicating the chronological visit number to which it belongs. The visit order is incremented across care episodes and reset for time gap tokens, allowing the model to reason over longitudinal disease progression across multiple encounters rather than treating the timeline as a flat sequence. Finally, the event type embedding encodes the high-level category of each event, such as laboratory tests, medications, diagnoses, procedures, administrative markers, or special tokens. Below, we provide a formal definition of each component in that define a complete event: 1. Concept Embedding: Each clinical concept cnk (e.g., diagnosis, lab test, procedure, medication) is associated with a learnable semantic embedding: ucnk ∈ Rd . 2. Value Embedding: For events with a numerical measurement vnk , the value is normalized and projected, such that vnk = θ(vnk ) ∈ Rd , and θ(·) is a multi-layer perceptron. Events without values are assigned a learnable (null-value) parameter. 3. Time Encoding: Each timestamp tnk is mapped into a temporal embedding representing the local temporal information as time deltas between consecutive events, such that tnk ∈ Rd . 4. Visit Order Embedding: To capture patient-level longitudinal ordering, each event inherits a visit index embedding: rvisit(n) ∈ Rd , where the indices are encoded as continuous learnable positional vectors. This allows the model to distinguish between early and late episodes in the patient trajectory (global temporal representation). 5. Care Stage Embedding: Each visit belongs to a clinical stage (e.g., outpatient, emergency, inpatient, ICU). We assign a learned embedding: sstage(n) ∈ Rd to enable the model to incorporate care-context information, e.g., ICU vs. emergency events carry different clinical semantics. 6. Type Embedding: Each event also receives a higher-level event type identifier that encodes the data source from which it originates (ICU chart event, ICU fluid output, etc.): wtype(nk) ∈ Rd . This provides local structural information that allows the model to differentiate heterogeneous event types. By jointly modeling care stage, visit order, and event type, the patient timeline representation preserves structural, temporal, and semantic distinctions across events, facilitating more expressive and context-aware representation learning over heterogeneous EHR sequences. B.4

Query/History Definition

For each prediction instance in the downstream cohort, we decompose the patient timeline into a query segment and a corresponding history segment to enable retrieval-based conditioning. The query represents the patient’s current clinical context at the time of prediction and consists of events observed within a task-specific observation window. For details on the observation windows used for each task, we refer the reader to Section C.2 and Table S2. All events occurring strictly before the 18

Table S1: Summary of the various events present in the patient timeline. We summarize all events present in the patient timeline and list the count along with representative examples. The == sign indicates that event type and example are identical. Event type

# of unique events

Example

Administrative Events OUTPATIENT-START OUTPATIENT-END EMERGENCY-START EMERGENCY-END ADMISSION-AT-HOSPITAL ADMISSION-AT-ICU ADMISSION-LOCATION ADMISSION-TYPE AGE_AT_ADMISSION DISCHARGE-FROM-HOSPITAL DISCHARGE-FROM-ICU DISCHARGE-LOCATION

1 1 1 1 1 1 12 9 1 1 1 14

Gender Race MEDS_DEATH

2 11 1

== == == == == == ADMISSION-LOCATION//TRANSFER FROM HOSPITAL ADMISSION-TYPE//EW EMER == == == DISCHARGE-LOCATION//HOME Static Events GENDER//F RACE//HISPANIC == Medical Events

LAB MEDICATION MICROBIOLOGY PROCEDURE-ICD DIAGNOSIS-ICD DRG

851 7486 169 13711 21135 1087

ICU-CHART ICU-FLUID-OUTPUT ICU-PROCEDURE ICU-INFUSION

2311 71 151 327

LAB//51237//Blood//INR(PT) MEDICATION//heparin MICROBIOLOGY//90039//URINE CULTURE PROCEDURE-ICD//3E0G76Z DIAGNOSIS-ICD//K429 DRG//APR//228 ICU Events ICU-CHART//220228//Hemoglobin ICU-FLUID-OUTPUT//226559//Foley ICU-PROCEDURE//225752//Arterial Line Norepinephrine infusion Special Tokens

PAD MASK CLS UNK

1 1 1 1

Week tokens: W1-W3 Month tokens: M1-M12 Year token: 1Y-Y+

3 12 1

== == == == Time Tokens TIME-GAP//1-W TIME-GAP//11-M TIME-GAP//1-Y+

query window constitute the patient history and are segmented into fixed-length units that serve as retrievable candidates. Figure S2 illustrates the query/history definitions for the various prediction windows considered in this work. This formulation mirrors real-world clinical reasoning, in which decisions are made based on the present encounter while selectively consulting relevant past visits. By explicitly separating query and history, the model is encouraged to retrieve and integrate only the most informative historical context, rather than uniformly encoding the entire past, enabling scalable reasoning over long and heterogeneous patient trajectories.

History

OUTP GAP L1

ED L1

GAP

HOSP L3

24/48 Hrs

ICU

Query L1

M1 M2

(a)

OUTP GAP L1

ED L1

M3 M4

History

GAP

HOSP

ICU L1

M1 M2

Query M3 M4

(b)

Figure S2: Illustration of query and history definition (a): Query/history boundaries for tasks with time-bounded prediction window. (b) Query/history boundaries task with entire stay prediction window.

History

OUTP GAP

GAP

HOSP

ICU

Query

M1 M2

M3 M4

OUTP GAP L1 t1

...

... (a)

ICU

M1 M2

...

GAP

HOSP

...

Query

M3 M4

(b)

OUTP GAP

GAP

HOSP

ICU

Query

M1 M2

M3 M4

C1 (c)

OUTP GAP

GAP

HOSP

ICU

Query

M1 M2

M3 M4

...

(d)

Figure S3: Visual illustration of the chunking strategies employed with EHR-RAGp. (a) Event-based chunking. (b) Time-based chunking. (c) Visit-level chunking. (d) Care-stage chunking.

B.5

Chunking

Structured EHR trajectories are inherently irregular, with clinical events occurring at variable temporal resolutions and across diverse care settings. To organize longitudinal patient history into retrievable units, EHR-RAGp partitions the timeline into clinically coherent chunks. We investigate four chunking strategies: event-based, time-based, visit-level, and care-stage chunking. These strategies provide complementary views of patient history by capturing different temporal and clinical granularities. Figure S3 visually illustrates these chunking strategies. Below, we formally outline our proposed chunking strategies: • Event-Based: Events are segmented into fixed-size groups based on event count. Each chunk contains a predefined number of embedded clinical events. Consecutive chunks overlap by a fixed number of events to preserve continuity across chunk boundaries. • Time-Based: Events are segmented into fixed time windows, for example 6-hour, 12-hour, 24-hour blocks. Each chunk contains all events occurring in its temporal window. • Visit-level: Each hospital visit, admission, or encounter is treated as a candidate chunk. This aligns with natural clinical boundaries: emergency department stays, inpatient admissions, and outpatient encounters. • Care-Stage: Events are grouped by clinical stage (outpatient, emergency, inpatient, ICU), ensuring that each chunk reflects a coherent care stage. In cases where the chunk length is shorter than the context window, we pad the sequence to match the context length. In cases where the chunk length is greater than the context length, we recursively impose Event-Based Chunking with proper metadata handling. 20

Experimental Setup

C.1

Data Splitting

To ensure a fair evaluation and prevent information leakage across training stages, we adopt a strict data splitting strategy. All data splits are performed at the patient level via subject_id field rather than the admission or visit level to ensure that no clinical information from the same patient appears in more than one split. For downstream prediction tasks, we partition the eligible ICU cohort into mutually exclusive training, validation, and test sets with proportions of 70%, 10%, and 20%, respectively. The test set is held out entirely and is not used during any stage of model development, including pretraining, hyperparameter tuning, or early stopping. This guarantees that the pretrained representations do not indirectly encode information from evaluation patients. Instead, the validation set is used exclusively for model selection and hyperparameter optimization. Further, for pretraining we hold out 5% as validation set to examine the pretraining quality. This strict separation across patients and training stages ensures that reported performance accurately reflects the model’s ability to generalize to unseen patients and supports reliable comparison across baselines. C.2

Downstream Tasks • In-Hospital Mortality: This task aims to predict whether a patient will die during the current hospital admission as a binary classification problem, using information available up to a predefined observation window. It is a widely used benchmark for evaluating clinical risk prediction models. In our setup, we condition the in-hospital mortality task on the first 48 hours spent in the ICU, following [25, 26, 9]. • 1-Year Mortality Post-Discharge: This task focuses on predicting whether a patient will die within one year following hospital discharge or no. It captures longer-term outcomes and reflects the model’s ability to leverage both acute and chronic patterns in a patient’s longitudinal history. We consider the entire ICU stay data as a valid prediction window [72, 73, 74]. For patients with multiple ICU stay and positive label of mortality within 1-year, we only assign the positive label to all stays that are within one year from the mortality date. • Long Length of Stay: This task predicts whether a patient’s hospital stay will exceed a predefined duration threshold. It serves as a proxy for resource utilization and clinical complexity and is commonly used to assess a model’s ability to identify patients at risk of prolonged hospitalization. We set the length-of-stay threshold to 7 days following [6, 8, 54] and define the prediction window as the first 24 hours spent in the ICU, following [75, 76]. • ICU Readmission: This task aims to predict whether a patient will be readmitted to the intensive care unit within a specified time window after ICU discharge. It is an important indicator of care quality and patient instability and requires effective modeling of prior clinical trajectories. We set the readmission window to 30 days post-discharge following [77, 8, 78] and use the entire ICU stay as the prediction window, following [8, 79]. Table S2: Overview of the downstream tasks and the query/ history definition for each task

Task

Label type

Prediction Window

Reference point

Query Window

History Chunk

Binary Binary Binary Binary

48 hours since ICU admission 24 hours since ICU admission Entire ICU stay Entire ICU stay

ICU admission ICU admission ICU discharge ICU discharge

Last L event ≤ ICU admission + 48 hours Last L event ≤ ICU admission + 24 hours Last L event ≤ ICU discharge Last L event ≤ ICU discharge

All events < query start All events < query start All events < query start All events < query start

In-hospital mortality Long Length of stay-7 Days ICU Readmission-30 Days 1 Year Mortality

Table S3: Label distribution per data split for downstream prediction tasks. Task In-hospital mortality Long Length of stay-7 Days ICU Readmission-30 Days 1 Year Mortality

Train (stay)

Validation (stay)

Test (stay)

N Pos (%)

N Neg (%)

N Pos (%)

N Neg (%)

N Pos (%)

Neg (%)

4, 445 (10.415%) 5, 810(13.613%) 1, 558 (3.651%) 5, 741 (13.452%)

38, 234 (89.585%) 36, 869(86.387%) 41, 121 (96.349%) 36, 938 (86.548%)

621 (10.126%) 835 (13.615%) 230 (3.750%) 842 (13.729%)

5, 512 (89.874%) 5, 298 (86.385%) 5, 903 (96.250%) 5, 291 (86.271%)

1, 210 (9.787%) 1, 637 (13.241%) 480 (3.883%) 1, 654 (13.379%)

11, 153 (90.213%) 10, 726 (86.759%) 11, 883 (96.117%) 10, 709 (86.621%)

C.3

EHR-RAGp Architecture

EHR-RAGp is built on a RoFormer-base [63] encoder configured with 12 transformer layers, 12 attention heads, a hidden size of 768, and rotary positional embeddings applied at each self-attention layer, with a maximum positional embedding length of 1536. The encoder operates over fixed-length event sequences of C = 1024 tokens and produces contextualized token representations, from which sequence-level representations are obtained using either CLS pooling or mean pooling. We set the chunk overlap to O = 12.5% of the original chunk length. Each event token is represented as the sum of multiple embedding components. In addition to the core event embedding, we incorporate categorical embeddings for care stage, visit order, and event type, all projected to the same hidden dimension (768) and learned jointly with the model. Numerical values associated with events are normalized and projected through a small MLP before being fused with the categorical embeddings. In the absence of a numeric value, we use a learnable null embedding of the same dimensionality. Temporal information is encoded using Time2Vec [71], applied to scaled time-delta features and added to the token representation. For retrieval, the pretrained backbone encoder is used to embed historical timeline chunks, which are stored in a vector database. During training, the query representation is used to retrieve the top-M most similar history chunks based on cosine similarity in the embedding space. Retrieval is non-parametric and does not involve gradient updates during the lookup stage; however, the encoder itself is fine-tuned end-to-end during downstream training. The prototype module consists of a learnable set of task-specific prototype vectors defined in the same embedding space as the query and history representations. Both query and retrieved history embeddings are projected into the prototype space, where relevance scores are computed via similarity between history embeddings and prototypes, conditioned on the query representation. These scores are used to weight or filter retrieved history chunks, enabling selective integration of long-range context. The fusion module combines the query representation with the filtered history representations using a shallow transformer with 2 layers and 4 attention heads. The fused representation is then passed to a task-specific prediction head implemented as a lightweight feedforward network. For all downstream tasks, the prediction head outputs a single logit optimized using binary cross-entropy (BCE) loss. All architectural components beyond retrieval are differentiable and optimized jointly during supervised fine-tuning. C.4

Pretraining Details

We adopt masked language modeling (MLM) objective applied to structured EHR event sequences for EHR-RAGp pretraining. During pretraining, patient timelines are tokenized into fixed-length sequences. We consider the full patient timeline during pretraining and pass it to the encoder backbone as overlapped chunks of 1024 tokens with 128 tokens of overlap. Hence, we produce dense sequences, whereas padding occurs only in the last chunk to fulfill the sequence length. During pretraining, a subset of 15% of the chunk events is randomly selected for masking following the RoBERTa [60] masking strategy. Specifically, masked positions are replaced with a special [MASK] token with probability 80%, substituted with a random event token with probability 10%, or left unchanged with probability 10%. The model is trained to predict the original event identity at masked positions using a cross-entropy loss over the entire vocabulary of 47,377 tokens. All embedding components, including event embeddings and categorical embeddings, are trained jointly during this stage. No temporal, or numeric value components are used during pretraining; our objective at this stage is solely to learn contextualized representations of structured EHR sequences. Pretraining is conducted on the full pretraining cohort described in Section B.1. We use AdamW [80] optimizer with an initial learning rate specified using cyclical learning rates following [81], a weight decay of 1 × 10−2 , a batch size of 16, and a cosine annealing learning scheduler. Training is performed for 100 epochs over the pretraining corpus. We monitor pretraining quality using accuracy on a held-out validation set, as described in Section C.1. Early stopping is implemented with a patience of three epochs if validation accuracy does not improve. The resulting pretrained backbone is subsequently reused for downstream fine-tuning and retrieval embedding generation. All downstream evaluation experiments are conducted using 4 NVIDIA A100 GPUs. 22

C.5

Vector Database Setup

To enable retrieval-augmented conditioning, we construct a patient-specific vector database that stores embeddings of historical visit segments derived from the pretrained EHR-RAGp backbone. We use Facebook AI Similarity Search (FAISS2 ) as the vector index for EHR-RAGp. For each patient in the downstream cohort, the history portion of the timeline (as defined in Section B.4) is segmented into fixed-length chunks, where each chunk corresponds according ti the various chunking strategies proposed in our work. Each chunk is independently encoded using the pretrained backbone encoder, and the resulting pooled representation is stored as a dense vector. We use mean pooling over the sequence representations and normalize the resulting embeddings prior to storage in the index. The vector index is constructed offline before downstream training and remains fixed during supervised fine-tuning. This design avoids repeated recomputation of historical embeddings and ensures scalable retrieval over long patient histories. During retrieval, we perform similarity search based on cosine similarity between the query embedding and historical chunk embeddings. It is worthy to note as well that we do not store chunked data along with their encoded vector, instead; we keep track of the chunk indices in the full chunked patient timeline. Upon retrieval, we perform online chunking during retrieve selected candidates by similarity search via their indices. This approach allow for faster retrieval and reduce the need for maintaining large database that accommodates the chunked data along with the vector index. C.6

Downstream Training Details

Table S4 summarizes the hyperparameter settings and search spaces used for downstream training of EHR-RAGp. Core backbone parameters are fixed across all experiments, including the RoFormer-base encoder with a hidden dimension of 768, 12 layers, and 12 attention heads. Retrieval-related hyperparameters including query chunk size, history chunk size, history chunk overlap, number of retrieved chunks, and similarity metric. all retrieval hyperparameters are held constant with values of 1024, 256, 32, 24, cosine similarity, respectively. For prototypes module, we use fixed prototype dimensionality of 768. For temperature, we experiment with temperature pairs , (Tq , Th ), for prototype assignment of [(0.025, 0.1), (0.05, 0.2), (0.02, 0.08)] and Weighing temperature Ts values of {0.1, 0.15, 0.2, 0.25}, whereas the number of prototypes is set to be optimized with discrete search space in {64 128, 256, 512}. For fusion, we set up a transformer encoder of 2 layers and 4 attention heads, and dropout of 0.1. We set the pooling strategy to be set during hyperparameter optimization with search space of {mean, query}. For training hyperparameters, we experiment with Stochastic Gradient Descent (SGD) optimizer, learning rate in [1e−5, 5e−4], Weight decay in [1e−3, 1e−2] and cosine annealing learning rate scheduler. For usage penalty λu , we experiment with fixed values in {0.004, 0.005, 0.006, 0.007}. Batch size is set to be optimized with discrete search space of {8, 12, 16, 20} and maximum epochs remain fixed with values of 75. For hyperparameters optimization, we run 25 trails for each prediction task using Bayesian search. All downstream evaluation experiments were conducted using a single NVIDIA H100 NVL 94GB GPUs. Maximum batch size with cane be used with such a GPU in our experiments is 20 query chunks of length 1024 events along with 24 history chunk of length 256, while average epoch time given the max settings is 38 minutes. It is worthy to note that lower capacity GPUs can be also used with reduction in the max settings mentioned earlier.

Baselines

D.1

Clinical Baselines

We consider encoder-based clinical baselines, as they match the structure of EHR-RAGp as an encoder-based model. 2 https://github.com/facebookresearch/faiss

Table S4: Hyperparameter settings and search spaces used in EHR-RAGp experiments. Category

Hyperparameter

Fixed

Search Space

Backbone

Encoder architecture Hidden dimension (d) Number of layers Number of attention heads pooling

RoFormer-base 768 12 12 –

– – – – {mean, CLS}

Retrieval

Query Chunk size History Chunk size History Chunk overlap∗ Number of retrieved chunks (M ) Similarity metric

1024 256 32 24 Cosine

– – – – –

Number of prototypes Prototype dimension Prototype temperature (Tq , Th ) Weighing temperature Ts

– 768 – –

{64, 128, 256, 512} – {(0.025, 0.1), (0.05, 0.2), (0.02, 0.08)} {0.1, 0.15, 0.2, 0.25}

Fusion layers Fusion attention heads pooling Dropout

2 4 – 0.1

– – {mean, query} –

Optimizer Usage penalty λu Learning rate Weight decay Batch size Max epochs

SGD – – – – 75

– {0.004, 0.005, 0.006, 0.007} [1e−5, 5e−4] [1e−3, 1e−2] {8,12,16,20} –

Prototypes

Fusion

Training

∗

overlap is used only with event-based chunking strategy

• DescEmb [58]: is a text-based framework that represents patient history as textual descriptions of medical events, covering the event type, numeric value (if any), along with its unit of measurement. These sequences are then processed via clinical or general-purpose pretrained language models. The main goal of DescEmb is to reduce reliance on predefined medical code vocabularies through a shared linguistic vocabulary that can be utilized across different EHR datasets. We consider two variants of DescEmb in our experiments, including BERT Fine-Tune and CLS Fine-Tune. • GenHPF [13]: is an extension of DescEmb that goes beyond the basic description of the medical event type, value, and unit to include all features associated with an event in the raw EHR tables, enhancing medical event representation. For instance, a medication event will cover type, name, frequency, dosage, route, etc. It shares the same objective as DescEmb of being agnostic to differing EHR schemas across institutions. • Med-BERT [32]: is a BERT-based foundation model that is pretrained on structured EHR records for disease prediction tasks. It relies solely on structured diagnosis ICD codes as pretraining and evaluation data. The pretraining approach typically follows masked language modeling as proposed in the BERT paper [27]. • BEHRT [29]: is another BERT-based foundation model that learns contextual embeddings of the patient’s structured records in a similar fashion to Med-BERT via masked language modeling using diagnosis codes. However, it differs from Med-BERT by incorporating patient age as a temporal signal. • Hi-BEHRT [30]: is a hierarchical extension of BEHRT that aims to handle longer, more comprehensive patient records. It captures short- and long-term patient patterns from hierarchically segmented medical sequences and is pretrained with the self-supervised BYOL framework [82]. • CEHR-BERT [31]: is another BERT-based model that incorporates temporal and visit information to improve disease prediction. It represents patient visits as sequences of medical codes, augmented with artificial time tokens, age embeddings, and time embeddings. D.2

Transformer-based Baselines • RoBERTa [60]: is a transformer-based language model that builds upon BERT by removing the next-sentence prediction objective and adopting dynamic masking and larger pretraining corpora. RoBERTa proves that careful optimization of the pretraining pipleines significantly improves contextual representation learning, which makes it a strong general-purpose encoder baseline. • LongFormer [61]: is a transformer variant designed to handle long input sequences efficiently through a sparse attention mechanism that combines local windowed attention with task-specific global attention. This design reduces the quadratic complexity of standard self-attention, which in turn enables modeling of substantially longer contexts compared to standard transformers. 24

• BigBird [62]: is another long-context transformer architecture that replaces full selfattention with a sparse attention pattern composed of random, local, and global attention. BigBird is theoretically proven to be a universal approximator while enabling scalable processing of long sequences, which making it suitable for tasks that involve extended contextual dependencies. • RoFormer [63]: is a transformer model that introduces rotary positional embeddings (RoPE) to encode relative positional information directly into the self-attention mechanism. This approach improves extrapolation to longer sequences and better captures relative token relationships, which makes it effective for sequence modeling tasks with variable-length inputs. D.3

Long Context Baselines • ModernBERT [59]: is a re-engineered BERT-style encoder that incorporates recent architectural and training improvements, including optimized attention implementations, improved normalization, and enhanced efficiency on modern hardware. ModernBERT aims to retain the strengths of encoder-based transformers while achieving better performance and scalability. • EHRMamba [10]: typically follows the embedding structure of CEHR-BERT but uses the state-space architecture MAMBA [83] instead of a BERT encoder. The use of a state-space model aims to enable efficient processing of relatively long EHR sequences compared to transformer models.

D.4

Retrieval-based Baselines • REMed [54]: is a retrieval-based foundation model that processes textual descriptions in a similar fashion to both DescEmb and GenHPF. It operates by randomly sampling events from patient history and ranking them to retrieve the most relevant ones for different downstream tasks. In contrast to our proposed model, REMed operates at the medical event level, whereas EHR-RAGp operates at historical segment level to cover broad history beyond limited set of events, 128 in case of REMed. • Vanilla EHR-RAGp: refers to the retrieval-augmented version of the proposed framework without the prototype-guided alignment module. In this setting, historical chunks are retrieved solely based on embedding similarity between the query and patient history, and the retrieved chunks are aggregated directly without prototype-aware refinement. This baseline isolates the contribution of the prototype-guided retrieval mechanism and allows evaluation of how latent prototype alignment improves retrieval quality and downstream prediction performance beyond standard semantic retrieval alone.

D.5

Large Language Model Baselines • Qwen2.5-7B-Instruct [65]: is a general-purpose instruction-tuned large language model designed for conversational reasoning and text understanding across diverse domains. • Mistral-7B-Instruct [66]: is an instruction-following variant of the Mistral family of open-source large language models optimized for efficient inference and general reasoning capabilities. • MedGemma-1.5-4B-it [67] is a medically adapted instruction-tuned language model developed for healthcare and biomedical applications. It is designed to support clinical reasoning and medical text understanding, and is evaluated here as a medical LLM baseline. • BioMistral-7B [68]: is a biomedical domain-adapted variant of Mistral trained on largescale biomedical corpora to improve performance on medical and scientific language tasks.

D.6

LLM Evaluation Protocol

We evaluated both general-domain and medical-domain large language models using a zero-shot next-token prediction framework for clinical risk prediction. For each downstream task, patient timelines were converted into sequential textual event representations and truncated according to predefined prediction windows. To ensure fair comparison with encoder-based baselines, the same 25

temporal slicing strategy and prediction horizons were used across all models. A fixed precomputed slices 1024 clinical events along with their numeric values events used. Each patient sequence was formatted as a chronological list of clinical events and passed directly to the language model without supervised fine-tuning. We used causal language modeling inference and extracted the logits corresponding to the “Yes” and “No” tokens from the next-token distribution. All other token logits were ignored, and probabilities were normalized over the binary “Yes/No” subset using a softmax operation. The normalized probability assigned to the “Yes” token was used as the prediction score. We present the prompt used for each task in The prompts used for zero-shot evaluation are shown in Box D.6. Performance was evaluated using AUROC and AUPRC. Confidence intervals were estimated using nonparametric bootstrap resampling with 1000 iterations. LLM Prompt Templates In-Hospital Mortality You are an expert clinical risk prediction model using electronic health records. --- PATIENT DATA --Electronic Health Records: {ehr_text} --- TASK --Will this patient die during this hospital admission? Answer only using one word: Yes or No. Answer: 1-Year Mortality You are an expert clinical risk prediction model using electronic health records. --- PATIENT DATA --Electronic Health Records: {ehr_text} --- TASK --Will this patient die within 1 year after hospital discharge? Answer only using one word: Yes or No. Answer: ICU Readmission within 30 Days You are an expert clinical risk prediction model using electronic health records. --- PATIENT DATA --Electronic Health Records: {ehr_text} --- TASK --Will this patient be readmitted to the ICU within 30 days after ICU discharge? Answer only using one word: Yes or No. Answer: Long Length of Stay > 7 Days You are an expert clinical risk prediction model using electronic health records. --- PATIENT DATA --Electronic Health Records: {ehr_text} --- TASK --Will this patient have a hospital length of stay longer than 7 days? Answer only using one word: Yes or No. Answer:

Table S5: Performance comparison between EHR-RAGp and LLM baselines Model

ICU-Readmit 30d AUROC (CI) AUPRC (CI)

In-hospital Mortality AUROC (CI) AUPRC (CI)

Long LOS 7d AUROC (CI) AUPRC (CI)

1YR Mortality AUROC (CI) AUPRC (CI)

General LLMs Qwen2.5-7B-Instruct Mistral-7B-Instruct

0.533 (0.508, 0.557) 0.515 (0.490, 0.540)

0.041 (0.036, 0.045) 0.039 (0.035, 0.045)

0.747 (0.733, 0.761) 0.723 (0.708, 0.740)

0.367 (0.342, 0.391) 0.313 (0.287, 0.340)

0.537 (0.522, 0.551) 0.610 (0.595, 0.624)

0.147 (0.138, 0.157) 0.186 (0.174, 0.200)

0.548 (0.534, 0.561) 0.511 (0.498, 0.526)

0.142 (0.134, 0.150) 0.130 (0.124, 0.138)

medgemma-1.5-4b-it BioMistral-7B

0.500 (0.472, 0.525) 0.509 (0.479, 0.535)

0.039 (0.035, 0.044) 0.041 (0.036, 0.047)

0.594 (0.578, 0.612) 0.594 (0.578, 0.610)

0.138 (0.127, 0.151) 0.147 (0.134, 0.161)

0.618 (0.604, 0.632) 0.523 (0.509, 0.538)

0.191 (0.179, 0.206) 0.136 (0.129, 0.145)

0.531 (0.517, 0.547) 0.516 (0.502, 0.530)

0.146 (0.137, 0.155) 0.136 (0.128, 0.144)

EHR-RAGp (Ours)

0.747 (0.724, 0.768)

0.156 (0.128, 0.189)

0.940 (0.933, 0.945)

0.716 (0.693, 0.738)

0.885 (0.876, 0.893)

0.628 (0.603, 0.652)

0.821 (0.811, 0.831)

0.396 (0.372, 0.422)

Clinical LLMs

Care stage OUTP ED HOSP ICU Prototype

12 10

Care stage OUTP ED HOSP ICU Prototype

UMAP-2

8 6 4

2 5

0 0.0

2.5

5.0

7.5

10.0 UMAP-1

12.5

15.0

17.5

0.0

(a) 1-Year Mortality Post-Discharge

2.5

5.0

7.5 10.0 UMAP-1

12.5

15.0

17.5

(b) ICU Readmission-30 Days

Figure S4: UMAP visualization of history chunk embeddings and prototypes for (a) long length of stay and (b) in-hospital mortality tasks.

Additional Results

E.1

LLM Baseline Results

Table S5 compares EHR-RAGp against both general-purpose and medically adapted LLMs evaluated in a zero-shot setting. Across all tasks, LLM baselines substantially underperform specialized EHR models, with particularly weak performance on ICU readmission and long-term mortality prediction. For example, the strongest LLM baseline, Qwen2.5-7B-Instruct, achieves an AUPRC of only 0.041 for ICU readmission compared to 0.156 for EHR-RAGp. Similarly, medical LLMs such as MedGemma and BioMistral do not demonstrate meaningful improvements over general LLMs, despite domain adaptation. Overall, these results suggest that current zero-shot LLMs struggle to effectively reason over structured longitudinal EHR trajectories, highlighting the importance of architectures specifically designed for temporal clinical data and patient-history retrieval. E.2

Latent Representation Structure

Figure S4 presents UMAP projections of history chunk embeddings for 1-Year Mortality PostDischarge and ICU Readmission-30 Days. E.3

Qualitative Analysis of Prototype-Guided Retrieval

This section provides a qualitative analysis of the prototype-guided retrieval mechanism in EHR-RAGp. To illustrate the behavior of the model, we present two representative case studies from the in-hospital mortality task, corresponding to positive and negative prediction examples. For each case, we analyze prototype usage, query-history alignment, chunk-level relevance weighting, and the contribution of visits and care stages to the final prediction. These examples provide insight into how EHR-RAGp organizes longitudinal patient history and selectively integrates clinically relevant context during inference. 27

E.3.1

Case of Positive Example

(a)

0.06

Query Info: Visit ID=V2 Label=1 Prob=0.998 Pred=1

0.04 0.02 0.00

Top query prototypes

Assignment Prob.

Figure S5 presents a representative positive example from the in-hospital mortality task, corresponding to a correctly predicted positive case (label=1, pred=1) with high confidence (prob=0.998) for a query segment originating from Visit V2 in the patient timeline. Panel (a) shows the query prototype distribution across the top activated prototypes, where prototype P 17 receives the highest assignment probability. Panel (b) illustrates the alignment between retrieved history chunks and the top query prototypes, revealing consistent prototype agreement across several retrieved chunks, particularly among the early ICU-related chunks. Panel (c) shows the final prototype-guided chunk weights, where the model assigns the highest importance to chunks originating from the second visit and predominantly associated with the ICU care stage. Chunks from lower-acuity stages such as ED and outpatient care receive comparatively lower weights. Panels (d) and (e) further aggregate these contributions at the visit and care-stage levels, showing that the prediction is primarily driven by the recent visit trajectory and high-acuity clinical context. Overall, the example illustrates how EHR-RAGp selectively emphasizes clinically severe and temporally relevant historical context when forming the final prediction.

P17 P0 P21 P2 P37 P29 P11 P48 P57 P3 P55 P56 P8 P27 P23 P33

P17

P21

P37

P29

P11

P48 P57 Prototype ID

P55

P56

P27

P23

P33

(b) 0.022 0.020 0.018 0.016 0.014 C1

C6 C7 C8 History chunk index

C10

C11

C12

C13

0.075

OUTP ED HOSP ICU

0.050 0.025 0.000

V2 C1

V2 C2

V2 C3

V2 C4

V2 C5

Importance

(d) 0.4 0.2 0.0

V2 V2 V2 C6 C7 C8 History chunk / visit

Visit ID

V1 C9

V1 C10

V2 C11

V1 C12

V2 C13

(e)

0.2 0.0

OUTP

HOSP Care Stage

ICU

Figure S5: Qualitative analysis of prototype-guided retrieval for an in-hospital mortality positive example. (a) Top prototype assignments for the query segment. (b) Alignment between retrieved history chunks and the query prototypes. (c) Prototype-guided relevance weights assigned to retrieved chunks. (d) Aggregated contribution of retrieved visits. (e) Aggregated importance across care stages. E.3.2

Case of Negative Example

Figure S6 presents a representative negative example from the in-hospital mortality task, corresponding to a correctly predicted negative case (label=0, pred=0) with low predicted probability (prob=0.006) for a query segment originating from Visit V5 in the patient timeline. Panel (a) shows 28

(a)

0.03

Query Info: Visit ID=V5 Label=0 Prob=0.006 Pred=0

0.02 0.01 0.00

Top query prototypes

Assignment Prob.

the query prototype distribution across the top activated prototypes, with a more diffuse assignment pattern, indicating less concentration around highly activated latent modes. Panel (b) illustrates the prototype agreement between the query and retrieved history chunks, where alignment is more uniformly distributed across chunks without a dominant subset receiving consistently strong activations. In panel (c), the final prototype-guided chunk weights are relatively balanced across visits and care stages, with no strong emphasis on ICU-related segments. Panels (d) and (e) further confirm this behavior, showing that importance is distributed across multiple visits and care stages rather than concentrated around a single high-acuity trajectory. Overall, this example illustrates how EHR-RAGp leverages broader longitudinal history across multiple visits and care stages to support a confident negative prediction, without relying on a single dominant high-acuity trajectory.

P45 P1 P15 P60 P46 P12 P2 P40 P53 P24 P3 P30 P52 P25 P29 P6

P45

P15

P60

P46

P12

P40 P53 Prototype ID

P24

P30

P52

P25

P29

(b)

0.024 0.022 0.020 0.018 0.016 0.014 0.012

C7 C8 C9 History chunk index

C10

C11

C12

C13

C14

C15

Importance

0.00

V1 C1

V4 C2

V3 C3

V4 C4

V2 C5

V1 C6

V3 V2 V5 C7 C8 C9 History chunk / visit

V4 C11

V2 C12

V3 Visit ID

V1 C13

V5 C14

V5 C15

OUTP ED HOSP ICU

(e)

0.1 0.0

V4 C10

(d)

0.2

Importance

Weight

0.06

0.2 0.1 0.0

OUTP

HOSP Care Stage

ICU

Figure S6: Qualitative analysis of prototype-guided retrieval for an in-hospital mortality negative example. (a) Top prototype assignments for the query segment. (b) Alignment between retrieved history chunks and the query prototypes. (c) Prototype-guided relevance weights assigned to retrieved chunks. (d) Aggregated contribution of retrieved visits. (e) Aggregated importance across care stages.

E.4

Analysis of Prototype Regularization

This section analyzes the effect of prototype regularization on the behavior of EHR-RAGp. Specifically, we compare the model with and without regularization from two perspectives: (1) global prototype utilization patterns, illustrating how regularization affects the distribution and diversity of prototype assignments, and (2) qualitative case studies, showing how regularization influences prototype-guided retrieval and chunk weighting for individual patient trajectories. Together, these analyses highlight the role of regularization in preventing prototype collapse and encouraging more balanced use of the latent prototype space. 29

Query Prototypes Usage Assignment prob.

Assignment prob.

Query Prototypes Usage 0.15 0.10 0.05 0.00

0.015 0.010 0.005 0.000

History Prototypes Usage Assignment prob.

Assignment prob.

History Prototypes Usage 0.03 0.02 0.01 0.00

30 Prototype ID

(a)

0.015 0.010 0.005 0.000

30 Prototype ID

(b)

Figure S7: Global prototype utilization patterns. (a) Without Regularization. (b) With regularization. E.4.1

Global Prototypes Usage

Figure S7 illustrates the effect of prototype regularization on global prototype utilization. Without regularization as in Figure S7a, query assignments exhibit highly imbalanced usage patterns, where a small subset of prototypes receive disproportionately large assignment probabilities while many prototypes remain minimally utilized. A similar, although less severe, imbalance is also observed for history chunk assignments. This behavior indicates partial prototype collapse, where the model concentrates representation learning around a limited number of latent prototypes. In turn, such a situation makes agreement between history and query rather difficult due to the sharper distribution of the query as compared to the history assignments. In contrast, enabling regularization, as in Figure S7b, produces a substantially more uniform assignment distribution across prototypes for both queries and history chunks. The resulting increase in prototype diversity suggests that regularization encourages broader utilization of the latent prototype space, leading to more balanced representation learning and improved retrieval behavior. E.4.2

Qualitative Analysis of Prototypes Usage

To further understand the role of usage regularization λu , we examine the model behavior with regularization, as in in Figure S8, and without regularization, as in Figure S9. In the settings where usage regularization is activated, prototype assignments are distributed across multiple latent prototypes, resulting in relatively balanced alignment scores and chunk weights across retrieved history segments. This behavior indicates that the model integrates information from several relevant historical chunks rather than relying on a single dominant prototype or retrieval path. On the other hand, without usage regularization, the query representation collapses strongly toward a single prototype (P0), which leads to highly concentrated retrieval behavior where few retrieved chunk receives the majority of the final weight while the remaining chunks contribute minimally. With respect to the model performance, it can be observed that despite the collapse the model can still provide correct prediction but less confident. Overall, these examples qualitatively demonstrate how prototype regularization stabilizes retrieval dynamics, encourages broader utilization of patient history, and prevents highly concentrated dependence on a narrow subset of latent prototypes, which in turn deteriorates the retrieval quality.

Query Info: Visit ID=V1 Label=1 Prob=0.907 Pred=1

0.02 0.00

Top query prototypes

Assignment Prob.

(a) Query top prototypes 0.04

P42 P19 P29 P53 P11 P46 P63 P16 P55 P47 P38 P28 P58 P27 P56 P32

P42

P19

P29

P53

P11

P46

P63

P16 P55 Prototype ID

P47

P38

P28

P58

P27

P56

P32

(b) History chunk alignment to query top prototypes 0.020 0.018 0.016 0.014 0.012 C1

C3 C4 History chunk index

0.15

OUTP ED HOSP ICU

0.10 0.05

V1 C1

0.00

V1 C2

V1 V1 C3 C4 History chunk / visit / care stage

V1 C5

V1 C6

Figure S8: Qualitative analysis of sample prediction for long length of stay (7 days) task with regularization (a) Top prototype assignments for the query segment. (b) Alignment between retrieved history chunks and the query prototypes. (c) Prototype-guided relevance weights assigned to retrieved chunks.

Top query prototypes

Assignment Prob.

(a) Query top prototypes 0.6

Query Info: Visit ID=V1 Label=1 Prob=0.571 Pred=1

0.4 0.2 0.0

P0 P37 P54 P11 P13 P49 P17 P36 P3 P23 P44 P2 P56 P35 P57 P55

P37

P54

P11

P13

P49

P17

P36 P3 Prototype ID

P23

P44

P56

P35

P57

P55

(b) History chunk alignment to query top prototypes 0.050 0.045 0.040 0.035 0.030 0.025 0.020 0.015 C1

C3 C4 History chunk index

0.6 0.4 0.2 0.0

V1 C1

V1 C2

V1 V1 C3 C4 History chunk / visit / care stage

V1 C5

V1 C6

OUTP ED HOSP ICU

Figure S9: Qualitative analysis of sample prediction for long length of stay (7 days) task without regularization (a) Top prototype assignments for the query segment. (b) Alignment between retrieved history chunks and the query prototypes. (c) Prototype-guided relevance weights assigned to retrieved chunks.

Record · ID 178900 · SHA-256 c8ce60f9d2418f1a
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.