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

arXiv: https://arxiv.org/abs/2609.01839 · HTML: https://mlsanity.com/paper/2609.01839/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 19 of its 33 references that are in the map.
- First posted on arXiv 1 Sep 2026.

## Builds on (most cited first, 19 of 19)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling](https://mlsanity.com/paper/1412.3555/) (2014)
- [Publicly Available Clinical BERT Embeddings](https://mlsanity.com/paper/1904.03323/) (2019)
- [Med-BERT: pre-trained contextualized embeddings on large-scale structured electronic health records for disease prediction](https://mlsanity.com/paper/2005.12833/) (2020)
- [RETAIN: An Interpretable Predictive Model for Healthcare using Reverse Time Attention Mechanism](https://mlsanity.com/paper/1608.05745/) (2016)
- [K-BERT: Enabling Language Representation with Knowledge Graph](https://mlsanity.com/paper/1909.07606/) (2019)
- [DeepPath: A Reinforcement Learning Method for Knowledge Graph Reasoning](https://mlsanity.com/paper/1707.06690/) (2017)
- [ScispaCy: Fast and Robust Models for Biomedical Natural Language Processing](https://mlsanity.com/paper/1902.07669/) (2019)
- [Self-Alignment Pretraining for Biomedical Entity Representations](https://mlsanity.com/paper/2010.11784/) (2020)
- [Go for a Walk and Arrive at the Answer: Reasoning Over Paths in Knowledge Bases using Reinforcement Learning](https://mlsanity.com/paper/1711.05851/) (2017)
- [GRAM: Graph-based Attention Model for Healthcare Representation Learning](https://mlsanity.com/paper/1611.07012/) (2016)
- [Pre-training of Graph Augmented Transformers for Medication Recommendation](https://mlsanity.com/paper/1906.00346/) (2019)
- [Dipole: Diagnosis Prediction in Healthcare via Attention-based Bidirectional Recurrent Neural Networks](https://mlsanity.com/paper/1706.05764/) (2017)
- [Sample Efficient Reinforcement Learning with REINFORCE](https://mlsanity.com/paper/2010.11364/) (2020)
- [Biomedical Entity Representations with Synonym Marginalization](https://mlsanity.com/paper/2005.00239/) (2020)
- [GraphCare: Enhancing Healthcare Predictions with Personalized Knowledge Graphs](https://mlsanity.com/paper/2305.12788/) (2023)
- [Reasoning-Enhanced Healthcare Predictions with Knowledge Graph Community Retrieval](https://mlsanity.com/paper/2410.04585/) (2024)
- [Learnable Prompt as Pseudo-Imputation: Rethinking the Necessity of Traditional EHR Data Imputation in Downstream Clinical Prediction](https://mlsanity.com/paper/2401.16796/) (2024)

## Abstract (arXiv)

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.