Knowledge Graph Augmented Large Language Models for Next-Visit Disease Prediction

2025-11-30

AI Summary

Existing EHR-based clinical prediction models suffer from coarse-grained post-hoc explanations and inadequate patient-level decision support. Method: We propose Knowledge Graph-guided Chain-of-Thought (KG-CoT), a framework that leverages multi-hop reasoning paths as structural scaffolds, integrating PrimeKG and ICD-9 code mappings to guide LLaMA-3.1-Instruct-8B/Gemma-7B in generating temporally coherent, clinically plausible explanations for disease predictions. Contribution/Results: KG-CoT significantly enhances interpretability and clinical utility: on MIMIC-III, it achieves AUROC of 0.66–0.70 and macro-AUPR of 0.40–0.47; under zero-shot transfer to CRADLE, prediction accuracy improves from 0.40–0.51 to 0.72–0.77. Clinical expert evaluation confirms superior explanation quality. Its core innovation lies in deeply embedding structured medical knowledge into large language model inference, enabling verifiable, traceable, and clinically grounded predictive explanations.

Technology Category

• Cognitive Modeling & Cognitive Systems: Conceptual Inference and Reasoning
• Reasoning under Uncertainty: Graphical Models
• Data Mining & Knowledge Management: Linked Open Data, Knowledge Graphs & KB Completion

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Electronic health records (EHRs) support powerful clinical prediction models, but existing methods typically provide coarse, post hoc explanations that offer limited value for patient-level decision making. We introduce a knowledge graph (KG)-guided chain-of-thought (CoT) framework that generates clinically grounded and temporally consistent reasoning for visit-level disease prediction in MIMIC-III. ICD-9 codes are mapped to PrimeKG, from which disease-relevant nodes and multi-hop reasoning paths are extracted and used as scaffolds for CoT generation; only explanations whose conclusions match observed outcomes are retained. Lightweight LLaMA-3.1-Instruct-8B and Gemma-7B models are then fine-tuned on this supervision corpus. Across ten PrimeKG-mapped diseases and limited training cohorts (400 and 1000 cases), KG-guided models outperform strong classical baselines, achieving AUROC values of 0.66 to 0.70 and macro-AUPR values of 0.40 to 0.47. The models also transfer zero-shot to the CRADLE cohort, improving accuracy from approximately 0.40 to 0.51 up to 0.72 to 0.77. A blinded clinician evaluation shows consistent preference for KG-guided CoT explanations in clarity, relevance, and clinical correctness.

Problem

• Predicts next-visit diseases using knowledge graphs for clinical reasoning
• Generates temporally consistent explanations for patient-level decision making
• Improves prediction accuracy on limited training data via KG-guided fine-tuning

Innovation

• Knowledge graph guides chain-of-thought reasoning for predictions
• Fine-tunes lightweight LLMs on clinically validated reasoning paths
• Transfers zero-shot to new cohorts with improved accuracy