Import What You Need: Learning When and How to Augment EHR Graphs with External Knowledge

• Existing KG augmentation is often static, unbudgeted, and single-pass, risking weakly related concepts, excess cost, and knowledge mismatched to changing patient states.The cited limitations are associated with static expansion (Jiang et al., 2024; Zhu et al., 2024), missing budget control (Zhang et al., 2024), and single-pass augmentation.Source: §1 · p. 1
• ReTA builds an offline, quality-filtered pool of compact LLM-distilled templates grounded in PrimeKG, clustered for redundancy, and retrieved using code similarity and trajectory context.At each visit, it retrieves the top-K relevant templates.Source: §1 · p. 2
• A reinforcement-learning policy selects one action per visit: Soft Import adds semantic features, Hard Import grafts a KG subgraph, or Skip abstains when the encoder is confident.The policy balances predictive gain against augmentation cost across the patient trajectory.Source: §1 · p. 2
• ReTA encodes semantic and structural signals in separate channels and fuses them through an adaptive gate.This decoupling addresses the heterogeneous signals produced by Soft Import and Hard Import.Source: §1 · p. 2
• ReTA outperforms strong baselines across diagnosis prediction, mortality, and readmission at lower inference cost, with gains widening under sparse supervision.Its visit-level augmentation patterns are interpretable, and cross-dataset transfer outperforms the strongest baseline’s in-dataset results.Source: §1 · p. 2

2.1 Problem Setup

• The model builds each visit graph from observed ICD codes and CCS ancestors linked by hierarchy edges, and predicts next-visit CCS categories or binary mortality and readmission outcomes.Visit graphs are augmented with PrimeKG knowledge; binary outcomes use a sigmoid output and binary cross-entropy while retaining the other model components.Source: §2.1 · p. 2

2.2 Knowledge Pool

• An offline pool turns independently distilled, KG-grounded concept summaries and compact subgraphs into reusable templates clustered by semantic similarity.LLM-generated definitions and complication cascades vary in length with PrimeKG neighborhood density; mentions are mapped to biomedical identifiers and retained entities and links are verified against the KG. ClinicalBERT embeddings are clustered to form normalized template vectors.Source: §2.2 · p. 2, §2.2.1 · p. 3, §2.2.2 · p. 3

2.2.3 History-Aware Retrieval

• At each visit, ReTA retrieves the top-K templates using both current-code similarity and alignment with the patient’s accumulated clinical context.The additive score allows a template to match either a current diagnosis or the broader trajectory; the state is computed before augmentation.Source: §2.2.3 · p. 3

2.3 Augmentation Policy

• A state-conditioned policy selects one template-mode action per visit—Soft Import, Hard Import, or Skip—using patient history, the current base embedding, and prediction uncertainty.Soft Import enriches node features without changing topology; Hard Import grafts a compact subgraph while retaining base features; Skip leaves the visit unchanged. The policy observes uncertainty to favor Skip when the base representation is sufficient.Source: §2.3 · p. 4, §2.3.1 · p. 4

2.3.2 Reward

• The paired reward measures prediction improvement against the same encoder’s unaugmented input and penalizes Hard Import in proportion to added nodes.Skip receives zero reward by construction; running template utilities are appended to the policy input, and REINFORCE uses a running-mean baseline.Source: §2.3.2 · p. 4

2.4 Decoupled Encoding

• ReTA separates feature-based semantic encoding from graph-based structural message passing, then fuses the two views with a learned gate.The semantic channel uses observed visit codes, while the structural channel includes imported nodes and edges and incorporates prior-visit context; injected nodes contribute through message passing but do not directly enter the prediction.Source: §2.4 · p. 4, §2.4.1 · p. 4, §2.4.2 · p. 5

2.5 Training

• Training first warm-starts the encoder across all three action types, then optimizes the policy while continuing to refine the encoder on policy-augmented visits.The encoder snapshot after Stage 1 supplies the base prediction used for the policy’s uncertainty signal.Source: §2.5 · p. 5

3.1 Setup

• Experiments evaluate diagnosis prediction, in-hospital mortality, and 30-day readmission on MIMIC-III and MIMIC-IV against sequence, ontology, structure-learning, and KG-augmented baselines.Diagnosis metrics are AUPRC, Micro-F1, and Acc@20; mortality and readmission use AUROC, AUPRC, and F1. KG-augmented baselines share PrimeKG and the ICD-to-KG linker; ReTAUMLS tests a UMLS-derived knowledge pool.Source: §3.1 · p. 5

3.2 Main Results

• ReTA outperforms all baselines on every diagnosis dataset–metric pair, with AUPRC gains over KARE of +1.99 on MIMIC-III and +1.76 on MIMIC-IV, and beats strongest binary-task baselines by +2.5/+2.2 AUROC for mortality and +3.7/+3.4 for readmission.Source: §3.2 · pp. 5–6
• ReTA loses less AUPRC than baselines under cross-dataset transfer, while ReTAUMLS reaches 34.63% AUPRC on MIMIC-IV, 0.55 below PrimeKG.Baselines lose 8.6–18.6% under transfer, versus 4.2% to MIMIC-III and 7.9% to MIMIC-IV for ReTA; ReTA’s MIMIC-III transfer AUPRC of 33.08 exceeds KARE’s in-domain 32.53. GPT-4o scores below graph-based methods, and adding PrimeKG context narrows but does not close the gap.Source: §3.2 · pp. 5–6

3.3 Ablation

• Ablations identify the refined knowledge pool as the largest contributor, followed by the per-visit policy; dual-mode import complements the policy’s routing role.Removing the refined pool costs 3.12 AUPRC, rule-based expansions instead of LLM-distilled cascades cost 1.85, the policy costs 1.78, and dual-mode import costs 1.41. A rate-matched shuffle loses 1.16 and 1.42 AUPRC; the contextual bandit trails by 0.50 at similar latency. The template audit found 94.2% correct definitions and 93.8% correct relations; Qwen3-32B rebuilding shifts AUPRC by at most 0.40.Source: §3.3 · p. 7

3.4 Analysis

• ReTA achieves 35.18% AUPRC at 11.8 ms/visit, while Skip concentrates on confident visits and improves accuracy and latency together.Skip rates are 27% on MIMIC-III and 31% on MIMIC-IV, declining from 68% in the most-confident quartile to 9% in the least-confident; removing Skip costs 0.62 AUPRC and adds 1.0 ms/visit. The policy overrides code-only Top-1 retrieval on 53.4% and 59.2% of visits, and forcing Top-1 lowers AUPRC by 0.49 and 0.56, respectively.Source: §3.4 · pp. 7–8, Figure 4 · p. 7

3.5 Robustness

• ReTA’s advantage grows with data and label scarcity and persists under missing diagnoses and corrupted templates.At 10% of training data, its margin over KARE is 4.53; at 60% diagnosis masking it retains 87% of full-input AUPRC versus 78% for KARE. The gain over KARE rises from +1.54 on frequent diagnoses to +4.72 on the rarest, and at 30% template corruption ReTA loses 0.87 AUPRC versus 1.71 for KARE.Source: Figure 5 · p. 8, §3.5 · p. 8
• With prior medication and procedure codes, ReTA reaches 36.29 AUPRC versus 34.61 for KARE at lower latency, and uses Hard Import less often.Latency is 12.9 versus 16.3 ms; Hard Import falls from 34% to 25% of actions. The evaluation covers structured codes, while clinical notes remain future work.Source: §3.5 · p. 8

3.6 Case Study

• A case study shows the policy shifting from Soft Import to Skip to Hard Import as a patient moves from chronic to acute visits.At V1 the fusion gate is semantic-dominant (β=0.76), while at V5 it is structure-dominant (β=0.33); the policy overrides code-level retrieval to select Sepsis, and cascade nodes connect chronic context at V1 to V5. Head diagnoses trigger Skip most often (38%), while tail diagnoses trigger Hard Import most often (47%).Source: Figure 6 · p. 8, Figure 7 · p. 8, §3.6 · p. 9

EHR prediction

• Longitudinal EHR prediction has progressed from recurrent models to Transformer-style architectures, but purely data-driven models degrade under sparsity, motivating external knowledge as complementary priors.The cited progression includes Choi et al. (2016), Luo et al. (2020), and Rasmy et al. (2021).Source: §4 · p. 9

KG augmentation

• Topology-oriented methods use fixed hierarchies or corpus-level co-occurrence, while semantics-oriented methods enrich features without changing topology; neither controls per-visit whether, what, and how to augment.Examples include GRAM (Choi et al., 2017), G-BERT (Shang et al., 2019), SeqCare, GraphCare (Jiang et al., 2024), KARE (Jiang et al., 2025), and RAM-EHR (Xu et al., 2024).Source: §4 · p. 9

Selective augmentation

• ReTA selects topology augmentation, semantic augmentation, or abstention from a quality-filtered template pool according to each patient’s evolving state and a per-visit budget.Unlike RL-based KG reasoning over fixed graphs and learnable prompts that add soft tokens without topology changes or abstention, ReTA makes patient-specific augmentation decisions.Source: §4 · p. 9

5 Conclusion

• ReTA selects Soft Import, Hard Import, or Skip from a quality-filtered knowledge pool as a budgeted, per-visit decision based on each patient’s evolving clinical state.The three actions respectively support feature enrichment, topology augmentation, or leaving the visit unaugmented.Source: §5 · p. 9
• Experiments on MIMIC-III and MIMIC-IV show consistent gains across diagnosis prediction, mortality, and readmission, with advantages growing under sparse supervision and persisting across datasets and with UMLS.Source: §5 · p. 9
• Future work could extend ReTA to multimodal EHR inputs and uncertainty-aware retrieval.Source: §5 · p. 9

Limitations

• Rare-disease cascades had a higher major-error rate (2.5% vs. 0.8%), and KG grounding establishes structural support rather than clinical correctness.External KGs may also be incomplete for some clinical domains, and GPT-4o’s undisclosed pretraining corpus leaves indirect overlap with MIMIC-derived literature unresolved.Source: Limitations · p. 9
• The skip policy can fail when the base encoder is confidently wrong because it treats confidence as a proxy for augmentation need.The limitation passage notes that this assumption may not hold in such cases.Source: Limitations · p. 9
• ReTA’s evaluation covers structured codes but does not incorporate clinical notes.Source: Limitations · p. 9

Ethics Statement

• ReTA is a research prototype requiring extensive validation, regulatory review, and clinical integration before deployment; no fairness audit was conducted, and performance may vary across patient subgroups.LLM-generated templates are ontology-grounded and blindly audited, but residual errors may affect underrepresented conditions; the audit found higher error rates on rare diagnoses.Source: Ethics Statement · pp. 9–10
• Experiments use MIMIC-III and MIMIC-IV under the PhysioNet Credentialed Health Data License, with de-identified records and no attempted re-identification.Source: Ethics Statement · p. 9

Appendix Contents

• The listed appendix analyses cover policy behavior, training stability, paired-reward ablation, and qualitative analysis.Source: Appendix Contents · p. 12

A.1 Dataset and Preprocessing

• After 24-hour aggregation, the cohort retains patients with at least two visits; MIMIC-IV admissions are restricted to ICD-10, and each trajectory yields Ti − 1 next-visit samples.Table 3 presents dataset statistics after preprocessing; diagnoses within each visit are treated as a deduplicated set.Source: Table 3 · p. 12, §A.1 · p. 12
• All methods use fixed shared temporal splits: the earliest 70% of discharge times for training, the next 10% for validation, and the final 20% for testing.Model selection uses validation only.Source: §A.1 · p. 12
• Diagnosis targets use CCS categories; diagnoses alone form model inputs and visit graphs except in the richer-input experiment, and online message passing uses untyped, symmetrized adjacency.Medications and procedures are otherwise excluded from graph expansion.Source: §A.1 · p. 12
• Mortality is labeled by death during admission and readmission by return within 30 days; both tasks reuse diagnosis-prediction visit graphs and knowledge pools, changing only the prediction head.Each uses a single sigmoid output with binary cross-entropy.Source: §A.1 · p. 12
• Cross-dataset transfer trains on one dataset and evaluates on the other without retraining, using shared 283-category CCS labels, PrimeKG linker, visit graphs, templates, and encoder architectures.Only input ICD codes and patient populations differ; ICD-9 and ICD-10 codes use their respective AHRQ CCS mappings.Source: §A.1 · p. 12

A.2 Implementation and Hyperparameters

• The policy selects among 2K+1 actions using REINFORCE with a running-mean baseline; short patient trajectories and an MLP-only policy make a critic unnecessary.Ablations include CCS-based cascades, a frozen encoder, uncertainty-threshold and supervised selectors, and a test-time greedy oracle; bounded Hard Import adds O(Ks̄) edges and encoder cost O(L(|V|d² + (|E| + Ks̄)d)).Source: §A.2 · pp. 12–13
• Table 5 reports inference latency per visit, in milliseconds, on MIMIC-IV.Source: Table 5 · p. 13

B Knowledge Pool Construction

• The knowledge-pool pipeline runs offline using external resources and distilled artifacts, without patient trajectories, outcome labels, or split-dependent statistics.Source: B Knowledge Pool Construction · p. 13

B.1 CCS Hierarchy and Hard Import

• Hard Import grafts shared CCS ancestors as a small subgraph, creating short message-passing paths between weakly connected diagnoses with few added edges.ICD codes are leaf nodes linked to CCS categories, which are connected by is-a edges; online message passing treats links as untyped and symmetrized.Source: §B.1 · pp. 13–14

B.2 LLM Distillation

• GPT-4o produces a definition for Soft Import and a clinical cascade for Hard Import for each concept, with cascade length adapted to PrimeKG neighborhood density.Concepts with sparse verified neighborhoods receive up to five cascade items, while dense-neighborhood concepts may receive one.Source: §B.2 · p. 14

B.3 Grounding and Filtering

• Mentions are grounded by exact vocabulary matching or ClinicalBERT similarity above τmap = 0.90, and triples are retained only when both endpoints and their relation have external ontology support.Relation support comes from a direct PrimeKG edge or an ancestor/descendant relation within two CCS hierarchy levels; thresholds are fixed across datasets.Source: §B.3 · p. 14

B.4 Template Construction

• Grounded concepts are clustered by ClinicalBERT cosine distance; each template uses a normalized cluster-centroid vector and the cluster medoid’s representative subgraph.Template subgraphs deduplicate canonical nodes, drop self-loops, merge duplicate edges, and collapse multi-relational links into untyped edges.Source: §B.4 · pp. 14–15

B.5 Pool Diagnostics

• 11–13% of concepts are lost to low external support; retained Hard-Import relations have grounded endpoints and PrimeKG or CCS support, and clustered templates average 7–8 nodes and 17–18 edges.The 83.7% and 81.9% externally supported rates are pre-filter candidate pass rates over distilled links, not support rates for the final pool; unsupported candidates are removed before template construction.Source: Table 7 · p. 15, §B.5 · p. 15

B.6 Clinical Audit

• Physicians blinded to provenance and performance audited 240 retained templates, and the authors report higher error rates for rare diagnoses.The sample included 120 templates mapped to common diagnoses (>500 training occurrences) and 120 mapped to rare diagnoses (<5); disagreements were resolved by discussion.Source: Table 8 · p. 15, §B.6 · p. 15

B.7 Generator Dependence

• With Qwen3-32B, downstream AUPRC remains within 0.36 and 0.40 of the GPT-4o pool, with comparable candidate support and majority Jaccard overlap.The pool was rebuilt using the identical prompt, grounding, filtering, and clustering procedure.Source: §B.7 · p. 15, Table 9 · p. 16

C Sensitivity

• Pool cutoff τ changes AUPRC by less than 0.2, while α=0 drops it by 0.41; moderate α=0.2–0.3 performs best, and retrieval size K has an inverted-U pattern.K=10 misses relevant templates, whereas K=30 dilutes quality with noise; the α=0 drop is 0.41 for retrieval scoring alone, compared with −0.54 for the history-aware retrieval ablation.Source: C Sensitivity · p. 16
• Hard Import penalty λ2 degrades at both extremes, while Soft Import scale ξ is stable at moderate values and worsens when ξ>1.0 or ξ=0.The default λ2=0.10 is in the stable middle; cascade-depth thresholds and template utility decay γr are also stable within tested ranges.Source: C Sensitivity · p. 16
• Increasing m from 1 to 2 adds only +0.09 AUPRC for 4.5 ms/visit and 3.9 training hours; m=3 lowers AUPRC at greater cost.The authors identify m=1 as the best accuracy-efficiency trade-off; the discount factor performs best near γ=0.95 and degrades at γ=0.Source: Table 10 · p. 16, C Sensitivity · p. 16

D Extended Ablation Results

• Rate-matched shuffle preserves the Stage 1 checkpoint, Stage 2 updates, training budget, and action rates while shuffling visit assignments, and runs 1.4 ms below the full model.It omits the policy forward pass and runs at the same skip rate.Source: D Extended Ablation Results · p. 16
• Removing the uncertainty signal u_t cuts the skip rate from 31% to 12% in the reported MIMIC-IV deltas.The lower block of Table 11 reports MIMIC-IV deltas only; the supplied caption identifies the table as component ablations with AUPRC deltas but gives no values.Source: D Extended Ablation Results · p. 16, Table 11 · p. 16
• Always augment retrains Stage 2 without Skip and matches the Skip ablation in Figure 3.The passage does not provide the numerical values behind that match.Source: D Extended Ablation Results · p. 16
• Table 12 defines Labels as per-action supervision required during training and Eval./visit as policy evaluations per visit at test time.These are table interpretation notes, not reported comparative results.Source: D Extended Ablation Results · p. 16
• Table 13 measures template overrides on non-Skip visits, counts Skip among total departures, and treats positive-gain rates as post-hoc diagnostics rather than policy inputs.Source: D Extended Ablation Results · p. 16
• In Table 14, both methods receive prior medication and procedure codes under the same cutoff.Source: D Extended Ablation Results · p. 16

E.1 Label-Frequency Stratification

• +1.54 to +4.72 AUPRC points: ReTA’s gain over KARE rises monotonically from head diagnoses to the rarest frequency bin.The per-stratum results are reported in Table 15.Source: §E.1 · p. 16, Table 15 · p. 17
• At the 50–200 occurrence boundary, the KARE gap widens from +1.11 to +2.37; Transformer loses 52% of head-bin AUPRC in the rarest bin, versus 31% for ReTA.Non-KG baselines degrade more steeply across frequency strata than KG-augmented methods.Source: §E.1 · p. 17

E.2 Calibration

• ReTA has the lowest expected calibration error and Brier score among methods; its calibration gap over KARE is −0.71 ECE, versus an AUPRC gap of +1.76.Calibration is macro-averaged over CCS labels across five seeds; lower is better.Source: §E.2 · p. 17, Table 16 · p. 17

F Mechanism Analysis

• Mechanism diagnostics use the validation split with a fixed encoder snapshot, require no additional training, and are not used for model selection.These constraints apply to all diagnostics in the mechanism-analysis section.Source: F Mechanism Analysis · p. 17

F.1 Policy Behavior

• Below approximately 50 training occurrences, Hard Import dominates; above that threshold, Soft Import and Skip become increasingly preferred as the base encoder can learn meaningful representations.Figure 11 shows the frequency-dependent crossover in import probabilities.Source: §F.1 · pp. 17–18, Figure 11 · p. 18
• Template retrieval concentrates on a small set near cluster centers, while peripheral templates are rarely selected; median intra-cluster cosine distance is 0.14–0.15.This population-level pattern is shown in Figure 12(a) and is consistent with Table 7.Source: §F.1 · p. 18, Figure 12 · p. 18
• At the population level, Tail diagnoses flow more to Hard than Soft Import, and Hard routes mainly to the structural channel while Soft routes to the semantic channel.The distinct flows indicate complementary roles for the two augmentation modes.Source: §F.1 · p. 18, Figure 12 · p. 18

F.2 Training Stability

• Reward standard deviation falls roughly 40%, maximum action probability stays below 12%, and Hard Import share declines from 0.52 to 0.34 as entropy decreases, indicating diverse convergence.The declining Hard Import ratio reflects a shift toward lighter augmentation when structural enrichment is unnecessary.Source: Table 17 · p. 18, §F.2 · p. 18

F.3 Paired Reward Ablation

• Removing the paired reward drops AUPRC by 1.15 points and halves the skip rate.The paired comparison evaluates raw and augmented predictions on the same mini-batch with dropout disabled, rewarding augmentation only when it improves on the base encoder.Source: Table 18 · p. 18, §F.3 · p. 18

F.4 Qualitative Analysis

• At V4, Hard Import connects V1’s chronic context to cascade nodes, creating shorter message-passing paths to V5; the fusion gate shifts from β=0.74 at V3 to β=0.48 at V4 and β=0.33 at V5.The cascade nodes are Systemic Inflammation and Septic Shock; without them, messages would traverse distant shared ancestors in the CCS hierarchy.Source: §F.4 · p. 19
• Because V3 is skipped and triggers no retrieval, the V4 retrieval context edge originates at V1, illustrating how history-aware state can use earlier visits despite skipped intermediate visits.Source: §F.4 · p. 19
• For ambiguous symptoms, the Hard Import penalty can deter structural grafting, while frequent template matches may dominate retrieval for rare conditions and prune a clinically plausible dependency.Figure 13 depicts over-pruning under ambiguity; the example contrasts a viral gastroenteritis match with structurally informative adrenal insufficiency.Source: §F.4 · p. 19, Figure 13 · p. 19
• In the ambiguous-symptom example, Soft Import selects the semantically closest template, viral gastroenteritis, rather than the structurally informative adrenal insufficiency template.Source: §F.4 · p. 19
• The paired reward cannot fix this failure when raw and augmented predictions agree on the wrong answer, yielding near-zero reward; greater template diversity or uncertainty-aware retrieval may help.Source: §F.4 · p. 19