Reasoning-Enhanced Healthcare Predictions with Knowledge Graph Community Retrieval

TL;DR

KARE introduces a knowledge-graph community retrieval framework to augment LLM reasoning for healthcare predictions. By constructing a multi-source medical concept KG, organizing it with hierarchical communities, and dynamically augmenting patient context with focused KG summaries, KARE enables step-by-step reasoning and accurate predictions. Evaluated on MIMIC-III and MIMIC-IV for mortality and readmission, KARE consistently outperforms baselines and provides interpretable reasoning chains. This approach demonstrates the practical promise of integrating structured medical knowledge with LLM reasoning to improve trustworthiness and decision support in clinical settings.

Abstract

Large language models (LLMs) have demonstrated significant potential in clinical decision support. Yet LLMs still suffer from hallucinations and lack fine-grained contextual medical knowledge, limiting their high-stake healthcare applications such as clinical diagnosis. Traditional retrieval-augmented generation (RAG) methods attempt to address these limitations but frequently retrieve sparse or irrelevant information, undermining prediction accuracy. We introduce KARE, a novel framework that integrates knowledge graph (KG) community-level retrieval with LLM reasoning to enhance healthcare predictions. KARE constructs a comprehensive multi-source KG by integrating biomedical databases, clinical literature, and LLM-generated insights, and organizes it using hierarchical graph community detection and summarization for precise and contextually relevant information retrieval. Our key innovations include: (1) a dense medical knowledge structuring approach enabling accurate retrieval of relevant information; (2) a dynamic knowledge retrieval mechanism that enriches patient contexts with focused, multi-faceted medical insights; and (3) a reasoning-enhanced prediction framework that leverages these enriched contexts to produce both accurate and interpretable clinical predictions. Extensive experiments demonstrate that KARE outperforms leading models by up to 10.8-15.0% on MIMIC-III and 12.6-12.7% on MIMIC-IV for mortality and readmission predictions. In addition to its impressive prediction accuracy, our framework leverages the reasoning capabilities of LLMs, enhancing the trustworthiness of clinical predictions.

Figures (21)

• Figure 1: A conceptual illustration of our KARE framework.Step 1 constructs a comprehensive medical concept knowledge graph by integrating information from multiple sources, organizing it into a hierarchical community structure. This structure allows for the generation of community summaries that facilitate precise knowledge retrieval. Step 2 dynamically augments the patient's EHR context with relevant summaries from the knowledge graph, offering the LLM focused and relevant medical insights. Step 3 generates training samples by employing an expert LLM to create reasoning chains based on the augmented patient context and ground truth labels. It then fine-tunes a local LLM using a multitask learning approach to produce interpretable reasoning chains and accurate predictions. By combining knowledge retrieval with LLM-driven reasoning, KARE significantly enhances the accuracy and reliability of clinical predictions.
• Figure 2: Validation loss of the label prediction during the fine-tuning with different settings. Loss is computed every 1/4 epoch. Task: mortality prediction on MIMIC-IV. "Base" and "Aug." denote base context and augmented context, respectively.
• Figure 3: Ablation study of (left) the metrics we proposed for patient context augmentation, and (right) the KG used as the knowledge source. N.H., Coh., Rec., and T.R. denote node hits, coherence, recency, and theme relevance, respectively. Tested task: MIMIC-IV-Readmission.
• Figure 4: Our pipeline to construct concept-specific KG $G^{KG}$ with bio KG (UMLS) and EHR.
• Figure 5: Our pipeline to construct concept-specific KG $G^{BC}$ from biomedical corpus with EHR.