PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Chen Chen, Mohsen Nayebi Kerdabadi, Dongjie Wang, Mei Liu, Zijun Yao

arXiv:2609.01839v1cs.LGcs.AI

TL;DR

Sparse, irregular EHR trajectories and static, unbudgeted KG augmentation motivate a more selective approach. ReTA learns per-visit, budget-aware choices among refined knowledge imports, and it consistently improves performance across clinical tasks while remaining efficient and transferable.

  • Problem

    Sparse, irregular EHR trajectories hinder modeling complex clinical patterns, while existing KG augmentation commonly uses fixed, context-agnostic, unbudgeted expansions.

  • Method

    ReTA builds a refined KG-grounded template pool and uses reinforcement learning to select Soft Import, Hard Import, or Skip per visit based on evolving patient state.

  • Results

    ReTA consistently outperforms strong baselines across diagnosis prediction, mortality, and readmission on MIMIC-III and MIMIC-IV, with gains also reported under transfer and sparse supervision.

  • Takeaways & Limitations

    Dynamic decisions about whether and how to augment each visit improve accuracy and latency relative to conventional augmentation strategies, with interpretable augmentation patterns.

  • Takeaways & Limitations

    The framework depends on external KGs and LLM-generated cascades, which may be incomplete or incorrect, especially for rare diseases, and it evaluates only structured codes rather than clinical notes.

Abstract

from arXiv · show

Longitudinal prediction from electronic health records (EHRs) is limited by the sparsity and irregularity in patient trajectories, and knowledge augmentation with external knowledge graphs (KGs) offers a promising way to alleviate these issues. However, most existing methods perform fixed, context-agnostic topology augmentation by adding the same KG nodes and edges regardless of a patient's evolving state. We propose ReTA, a Reinforcement learning-based dynamic Topology Augmentation framework that casts KG import as a per-visit, budget-aware policy. ReTA first constructs an offline refined pool of KG-grounded templates, then learns a policy to select one augment action per visit from three options: Soft Import, which enriches node features without modifying graph topology, Hard Import, which grafts a compact KG subgraph onto the visit graph to create message-passing shortcuts, and Skip, which leaves the visit unaugmented when the base encoder is already confident. To stabilize learning, ReTA employs a decoupled encoder that processes semantic and structural signals in separate channels and fuses them via adaptive gating. Experiments on MIMIC-III and MIMIC-IV across diagnosis prediction, mortality, and readmission show that ReTA consistently outperforms strong baselines while remaining efficient, transfers across datasets and knowledge graphs, and yields interpretable augmentation patterns. The robust gains under sparse supervision highlight the advantage of ReTA's dynamic decision to import knowledge, boosting accuracy while curbing costs.

1 Introduction

EHR trajectories are sparse and irregular, while existing KG augmentation is often static and unbudgeted. ReTA addresses this by learning budget-aware, visit-level augmentation decisions across complementary import modes.

  • Sparse and irregular EHR trajectories make complex clinical patterns difficult to capture, motivating the use of clinically validated KG concepts and relations.
  • Existing methods commonly apply fixed KG neighborhoods across patients and visits, risking weakly related concepts, diluted patient-specific signals, and higher message-passing costs.
  • ReTA formulates KG injection as a sequential, dynamic, budgeted decision problem based on each patient’s evolving history.
  • An offline refined pool retrieves compact KG-grounded templates using code-level similarity and trajectory-level context.
  • A reinforcement learning policy selects Soft Import, Hard Import, or Skip per visit, balancing predictive gain against augmentation cost.
  • Semantic and structural signals are encoded separately and combined through an adaptive gate to integrate the two augmentation modes.
  • Across MIMIC-III and MIMIC-IV, ReTA outperforms strong baselines across diagnosis prediction, mortality, and readmission while reducing inference cost.
  • Its gains widen under sparse supervision, and augmentation patterns remain interpretable at the visit level.

2 Methodology

ReTA models longitudinal EHR prediction with visit graphs augmented by a bounded, history-aware pool of KG-grounded templates. A reinforcement-learning policy selects one template-mode action per visit, while separate semantic and structural channels encode the resulting signals.

  • Problem Setup: The prediction setup represents each visit as diagnosis codes and predicts next-visit CCS categories or binary mortality and readmission outcomes.Visit graphs include ICD codes, CCS ancestors, and hierarchy edges; PrimeKG supplies additional biomedical relations.
  • Knowledge Pool: ReTA builds a reusable offline pool of compact KG-grounded templates to support bounded per-visit augmentation.Templates pair semantic summaries with compact subgraphs, are grounded against PrimeKG, and are compressed through clustering and de-duplication.
  • History-Aware Retrieval: At each visit, retrieval combines current-code similarity with trajectory context to select the top-K candidate templates.The trajectory signal uses the patient state formed from prior visit summaries and the current base embedding.
  • State and Actions: The policy chooses one of 2K+1 actions: Soft Import, Hard Import, or Skip, balancing predictive benefit against augmentation cost.Soft Import enriches features, Hard Import adds a compact KG subgraph and message-passing paths, and Skip leaves the visit unchanged; uncertainty informs the Skip decision.
  • Reward and Policy Learning: The reward measures augmentation-related prediction improvement while penalizing Hard Import according to its added nodes.The comparison uses augmented and raw inputs under the same dropout-disabled mini-batch; Skip receives zero reward by construction, and template utilities provide selection memory.
  • Decoupled Encoder: ReTA separates semantic and structural encoding channels before fusing their outputs with a learned gate.This prevents feature enrichment and topology changes from being conflated during representation learning.

3 Experiments

Across diagnosis prediction, mortality, and readmission on MIMIC-III and MIMIC-IV, ReTA outperforms strong baselines while improving efficiency, transfer, robustness, and interpretability. Ablations and analyses attribute these gains to refined knowledge, adaptive per-visit decisions, and selective augmentation.

  • Main results: ReTA outperforms all baselines on every diagnosis-prediction dataset-metric pair, with AUPRC gains of +1.99 on MIMIC-III and +1.76 on MIMIC-IV over KARE.
  • Main results: ReTA achieves higher diagnosis-prediction scores than graph-based baselines across reported MIMIC-III and MIMIC-IV results.The reported ReTA rows exceed GraphCare, RAM-EHR, and KARE across the displayed metrics.
  • Generalization: Cross-dataset transfer reduces ReTA’s AUPRC by 4.2% to MIMIC-III and 7.9% to MIMIC-IV, while transfer AUPRC of 33.08 on MIMIC-III exceeds KARE’s in-domain 32.53.
  • Main results: ReTA outperforms the strongest baselines by +2.5/+2.2 AUROC on mortality and +3.7/+3.4 on readmission across the two datasets.
  • Component ablation: The refined knowledge pool contributes the largest ablation gain, followed by the per-visit policy and dual-mode import, while rate-matched shuffling still loses 1.16 and 1.42 AUPRC.
  • Efficiency and robustness: ReTA reaches 35.18% AUPRC at 11.8 ms per visit, skips 27% of MIMIC-III visits and 31% of MIMIC-IV visits, and gains more under sparse or corrupted inputs.The skip rate falls from 68% in the most confident uncertainty quartile to 9% in the least confident; the margin over KARE grows from 1.76 at full data to 4.53 at 10%.

4 Related Work

Prior EHR prediction methods increasingly use external knowledge, but existing augmentation approaches are largely fixed, context-agnostic, or limited to feature-space injection. ReTA instead combines patient-specific retrieval with per-visit decisions over topology augmentation, semantic augmentation, and abstention.

  • Purely data-driven EHR models degrade under sparse trajectories, motivating external knowledge as complementary priors.
  • KG augmentation: Topology-oriented methods expand visit graphs using fixed ontology hierarchies or corpus-level co-occurrence rather than individual patient context.
  • KG augmentation: Semantics-oriented methods retrieve or generate patient-specific knowledge artifacts, but their augmentation remains distinct from topology modification.
  • Selective augmentation: ReTA coordinates topology augmentation, semantic augmentation, and abstention under a per-visit budget.
  • Selective augmentation: Unlike prior reinforcement-learning KG methods and learnable prompts, ReTA decides how to augment a patient-specific graph and can abstain.

5 Conclusion

ReTA frames knowledge augmentation as a budgeted, per-visit decision conditioned on each patient’s evolving clinical state. Across MIMIC-III and MIMIC-IV, it reports consistent gains across tasks, stronger benefits under sparse supervision, and transfer robustness across datasets and knowledge graphs.

  • ReTA selects Soft Import, Hard Import, or Skip from a quality-filtered knowledge pool using each patient’s evolving clinical state.
  • ReTA consistently improves diagnosis prediction, mortality, and readmission on MIMIC-III and MIMIC-IV.
  • Its advantage grows under sparse supervision and persists under cross-dataset transfer and replacement of PrimeKG with UMLS.

Limitations

ReTA depends on the completeness and correctness of external knowledge and generated cascades, with higher audited errors for rare diseases. Its confidence-based skipping can fail when the base encoder is confidently wrong, and evaluation excludes clinical notes.

  • External KGs may be incomplete, while LLM-generated cascades can produce plausible but incorrect complications for rare diseases.The audit reports a higher major-error rate for rare diseases: 2.5% versus 0.8%.
  • Grounding provides structural support rather than guaranteeing clinical correctness of generated knowledge.
  • The skip mechanism assumes base-encoder confidence reliably indicates augmentation need, which fails for confidently wrong predictions.
  • The evaluation covers structured codes but does not incorporate clinical notes.

Ethics Statement

The study uses de-identified MIMIC-III and MIMIC-IV records under a credentialed-data license without attempting re-identification. ReTA is a research prototype requiring further validation and oversight, and the work reports unassessed fairness and residual risks from generated knowledge.

  • All experiments use de-identified MIMIC-III and MIMIC-IV records under the PhysioNet Credentialed Health Data License, with no re-identification attempted.
  • ReTA is not intended for clinical deployment without extensive validation, regulatory review, and clinical oversight.
  • Performance may vary across age, sex, race, or insurance subgroups, and no fairness audit was conducted.
  • Residual errors in LLM-generated templates may propagate into predictions for underrepresented conditions, with higher audited error rates for rare diagnoses.

A.2 Implementation and Hyperparameters

ReTA combines offline knowledge-template construction with budget-aware policy learning and lightweight implementation choices. Its graph augmentation is structured around semantic enrichment, structural grafting, or skipping, with efficiency safeguards.

  • Policy learning: The policy selects among 2K+1 actions—two import modes for K templates plus Skip—using REINFORCE with a running-mean baseline.Short trajectories and a lightweight MLP make critic-based methods unnecessary.
  • Efficiency: Skip-eligible visits avoid graph assembly, yielding 0.0 ms for those visits and reducing average per-visit cost below GraphCare and KARE.Offline knowledge-pool construction is excluded from inference latency.
  • Efficiency: Hard Import adds O(Ks̄) edges, with total L-layer encoder cost O(L(|V|d^2 + (|E| + Ks̄)d)).The implementation uses bounded Top-K templates to control topology growth.
  • Knowledge construction: ReTA distills each concept into a definition for Soft Import and a clinical cascade for Hard Import, then grounds and filters the resulting relations.Grounding uses exact matching or ClinicalBERT similarity, while unsupported relations are discarded.
  • Knowledge construction: Concept-specific cascades adapt from one to five items according to PrimeKG neighborhood density, compensating for sparse relational context.Sparse neighborhoods receive longer cascades, whereas dense neighborhoods receive shorter ones.
  • Knowledge construction: Grounded outputs are clustered by ClinicalBERT cosine distance, with each template represented by a centroid vector and medoid subgraph.This clustering compresses redundant knowledge before online retrieval.

B.5 Pool Diagnostics

The refined knowledge pool is aggressively filtered, grounded, and compressed before deployment. Diagnostics indicate compact templates, generator robustness, and stable performance across moderate hyperparameter settings and budgets.

  • Pool quality: Only externally supported grounded links are retained before template construction; unsupported LLM-generated relations are removed by strict filtering.Support comes from PrimeKG edges or CCS ancestor/descendant relations within two hierarchy levels.
  • Pool quality: 11–13% of concepts are lost primarily because clinically plausible LLM relations lack external ontology support, while retained templates average 7–8 nodes and 17–18 edges.The compact templates keep Hard Import efficient at inference time.
  • Knowledge-graph transfer: Replacing PrimeKG with UMLS yields 874 MIMIC-III and 1,092 MIMIC-IV templates with comparable 7.2 / 15.3 average nodes/edges.The vocabulary bridge, prompting, grounding thresholds, and clustering procedure remain unchanged.
  • Generator robustness: Qwen3-32B produces downstream AUPRC within 0.36 and 0.40 of the GPT-4o pool under the identical construction pipeline.The comparison also reports comparable candidate support and majority Jaccard overlap.
  • Sensitivity: The pool cut changes AUPRC by less than 0.2 within the tested range, while retrieval size follows an inverted-U pattern and removing trajectory context drops AUPRC by 0.41.The history-aware ablation reports a −0.54 change in Table 48’s referenced comparison.
  • Budget sensitivity: Increasing the per-visit budget from 1 to 2 adds only +0.09 AUPRC while costing 4.5 ms/visit and 3.9 training hours; m=3 further increases cost and degrades AUPRC.These results support m=1 as the reported accuracy-efficiency trade-off.

D Extended Ablation Results

Extended ablations isolate the policy’s contribution under controlled comparisons and clarify how uncertainty affects Skip decisions. The appendix also documents evaluation controls and diagnostic conventions.

  • Controlled comparisons: Rate-matched shuffling preserves the Stage 1 checkpoint, Stage 2 updates, training budget, and action rates while changing only visit-action assignments.It runs 1.4 ms below the full model because it omits the policy forward pass.
  • Controlled comparisons: Always augment retrains Stage 2 without Skip and matches the Skip ablation in Figure 3.This comparison isolates the role of allowing unaugmented visits.
  • Policy signals: Removing the uncertainty signal reduces the skip rate from 31% to 12%.The appendix reports this as a MIMIC-IV component-ablation result.
  • Evaluation conventions: The appendix distinguishes label supervision, policy evaluations per visit, template overrides on non-Skip visits, and total departures including Skip.Positive-gain rates are post-hoc diagnostics rather than policy inputs.

E.1 Label-Frequency Stratification

ReTA’s augmentation behavior varies systematically with diagnosis frequency. Hard Import is favored for rare diagnoses, while Soft Import and Skip become more common when the base encoder has sufficient data.

  • Frequency-stratified performance: ReTA’s gain over KARE rises from +1.54 on head diagnoses to +4.72 on the rarest bin.The gap widens sharply at the 50–200 frequency boundary, from +1.11 to +2.37.
  • Calibration: ReTA achieves the lowest expected calibration error and Brier score among all methods, with a −0.71 ECE gap over KARE.The calibration gap is proportionally larger than the +1.76 AUPRC gap reported in the same comparison.
  • Policy behavior: Below approximately 50 training occurrences, Hard Import dominates; above that threshold, Soft Import and Skip become increasingly preferred.The reported explanation is that rare diagnoses have sparse co-occurrence statistics, whereas the base encoder can learn meaningful representations with more data.
  • Policy behavior: At the population level, Tail→Hard flow is wider than Tail→Soft, while Hard Import maps mainly to the structural channel and Soft Import to the semantic channel.The flow supports complementary roles for the two augmentation modes.

F.2 Training Stability

ReTA shows stable policy learning while adapting augmentation intensity, and its paired reward makes skipping beneficial when augmentation does not improve predictions. Qualitative analysis also illustrates structural shortcuts and a failure mode under ambiguous symptoms.

  • Training stability: Reward standard deviation drops by roughly 40%, while maximum single-action probability remains below 12% at convergence, indicating stable and diverse policy behavior.Entropy decreases during training, and the Hard Import ratio falls from 0.52 to 0.34 as the policy favors lighter augmentation when structural enrichment is unnecessary.
  • Reward ablation: AUPRC drops by 1.15 points and the skip rate halves when the paired reward is replaced by an unpaired −LCE reward.The paired reward credits augmentation only when it improves prediction beyond the base encoder on the same mini-batch.
  • Structural shortcut: At V4, Hard Import connects V1’s chronic context to cascade nodes, creating shorter message-passing paths to V5 Sepsis than the CCS hierarchy provides.The fusion gate shifts from β=0.74 at V3 to β=0.48 at V4 as structural enrichment is introduced.
  • Structural shortcut: History-aware retrieval at V4 can use V1 even when V3 is skipped, because skipped visits produce no retrieval.This case illustrates how earlier visits inform later augmentation decisions across the trajectory.
  • Failure mode: Under ambiguous symptoms, the policy may choose Soft Import’s frequent semantic match instead of Hard Import’s structurally informative template for a rare condition.The paired reward may remain near zero when raw and augmented predictions agree on the wrong answer; the paper suggests template diversity or uncertainty-aware retrieval as mitigations.
Loading 2609.01839v1…