# Knowledge Graph Augmented Large Language Models for Disease Prediction

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

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 14 of its 26 references that are in the map.
- First posted on arXiv 1 Dec 2025.

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

- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [s1: Simple test-time scaling](https://mlsanity.com/paper/2501.19393/) (2025)
- [Gemma: Open Models Based on Gemini Research and Technology](https://mlsanity.com/paper/2403.08295/) (2024)
- [Towards Reasoning in Large Language Models: A Survey](https://mlsanity.com/paper/2212.10403/) (2022)
- [Do NOT Think That Much for 2+3=? On the Overthinking of o1-Like LLMs](https://mlsanity.com/paper/2412.21187/) (2024)
- [LIMO: Less is More for Reasoning](https://mlsanity.com/paper/2502.03387/) (2025)
- [RETAIN: An Interpretable Predictive Model for Healthcare using Reverse Time Attention Mechanism](https://mlsanity.com/paper/1608.05745/) (2016)
- [Multitask learning and benchmarking with clinical time series data](https://mlsanity.com/paper/1703.07771/) (2017)
- [O1-Pruner: Length-Harmonizing Fine-Tuning for O1-Like Reasoning Pruning](https://mlsanity.com/paper/2501.12570/) (2025)
- [OpenR: An Open Source Framework for Advanced Reasoning with Large Language Models](https://mlsanity.com/paper/2410.09671/) (2024)
- [Deep Representation Learning of Electronic Health Records to Unlock Patient Stratification at Scale](https://mlsanity.com/paper/2003.06516/) (2020)
- [A Preliminary Study of o1 in Medicine: Are We Closer to an AI Doctor?](https://mlsanity.com/paper/2409.15277/) (2024)
- [RAM-EHR: Retrieval Augmentation Meets Clinical Predictions on Electronic Health Records](https://mlsanity.com/paper/2403.00815/) (2024)

## Built on by (most cited first, 4 of 4)

- [Closing Reasoning Gaps in Clinical Agents with Differential Reasoning Learning](https://mlsanity.com/paper/2602.09945/) (2026)
- [Structure-Aware RAG: Structured Retrieval Augmented Generation from Noisy Data for Conversational Agents](https://mlsanity.com/paper/2605.24366/) (2026)
- [Lung-R1: A Knowledge Graph-Guided LLM for Pulmonary Diagnostic Reasoning](https://mlsanity.com/paper/2606.11675/) (2026)
- [NSIDDx: A Design Framework for Neuro-Symbolic, Practitioner-First Differential Diagnosis in Low-Resource Settings](https://mlsanity.com/paper/2609.00256/) (2026)

## Abstract (arXiv)

Electronic health records (EHRs) enable strong clinical prediction, but explanations are often coarse and hard to use for patient-level decisions. We propose a knowledge graph (KG)-guided chain-of-thought (CoT) framework for visit-level disease prediction on MIMIC-III. We map ICD-9 codes to PrimeKG, mine disease-relevant nodes and paths, and use these paths to scaffold temporally consistent CoT rationales, retaining only samples whose conclusions match observed outcomes. We fine-tune lightweight instruction-tuned LLMs (LLaMA-3.1-Instruct-8B and Gemma-7B) on two small cohorts (400 and 1,000 index visits) across ten PrimeKG-mapped diseases. Our models outperform strong classical baselines, reaching AUROC 0.66-0.70 and macro-AUPR 0.40-0.47. Without additional training, the models transfer zero-shot to the CRADLE cohort, improving accuracy from 0.40-0.51 to 0.72-0.77. In a blinded clinician study, KG-guided CoT rationales are consistently preferred for clarity, relevance, and correctness. Code is available at: https://github.com/JonathanWry/KG-guided-LLM-pipeline