Knowledge Graph Augmented Large Language Models for Disease Prediction

Electronic health records (EHRs) support strong clinical prediction but often provide coarse, post hoc explanations that are 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 explanations, retaining only samples whose conclusions match observed outcomes. We then fine-tune lightweight LLaMA-3.1-Instruct-8B and Gemma-7B models on two small cohorts (400 and 1,000 index visits) across ten PrimeKG-mapped diseases. Our models outper-form strong classical baselines, reaching AUROC of 0.66-0.70 and macro-AUPR of 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. Blinded clinicians consistently prefer KG-guided CoT for clarity, relevance, and correctness. Code is available at https://github.com/JonathanWry/KG-guided-LLM-pipeline.