Knowledge Graph Augmented Large Language Models for Disease Prediction

TL;DR

This paper introduces a knowledge graph-guided chain-of-thought framework for disease prediction using large language models, enhancing interpretability and transferability in clinical settings with promising results on MIMIC-III and CRADLE datasets.

Contribution

It presents a novel KG-guided CoT approach for disease prediction that improves interpretability and zero-shot transferability of lightweight instruction-tuned LLMs in healthcare.

Findings

Models outperform classical baselines with AUROC 0.66-0.70

Zero-shot transfer improves accuracy to 0.72-0.77 on CRADLE

Clinicians prefer KG-guided rationales for clarity and relevance

Abstract

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