[Paper Review] Reasoning-Enhanced Healthcare Predictions with Knowledge Graph Community Retrieval

Pengcheng Jiang, Cao Xiao+5 more|arXiv (Cornell University)|Oct 6, 2024

Biomedical Text Mining and OntologiesBiochemistry, Genetics and Molecular Biology3 citations

TL;DR

KARE is a novel framework that enhances clinical prediction by integrating knowledge graph (KG) community retrieval with large language model (LLM) reasoning. It constructs a multi-source medical KG, organizes it into hierarchical communities for precise retrieval, and dynamically enriches patient contexts with focused, multi-faceted medical insights, achieving up to 15.0% improvement in mortality and readmission prediction over state-of-the-art models on MIMIC-III and MIMIC-IV.

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.

Motivation & Objective

• To address hallucinations and lack of fine-grained medical knowledge in LLMs for clinical decision support.
• To overcome the limitations of traditional retrieval-augmented generation (RAG) by retrieving semantically relevant, contextually precise medical knowledge.
• To enhance prediction accuracy and interpretability in high-stakes clinical tasks like in-hospital mortality and readmission prediction.
• To leverage hierarchical graph community detection and summarization for structured, retrieval-optimized medical knowledge graphs.
• To integrate reasoning-enhanced LLMs with dynamically retrieved KG communities for trustworthy, evidence-based predictions.

Proposed method

• Constructs a comprehensive, multi-source knowledge graph by integrating biomedical databases, clinical literature, and LLM-generated insights.
• Applies hierarchical graph community detection and summarization to organize the KG into semantically meaningful, contextually relevant subgraphs.
• Employs a dynamic knowledge retrieval mechanism that selects and retrieves pre-summarized KG communities tailored to the patient's clinical context.
• Augments patient input with retrieved KG community insights to enrich context for LLM reasoning.
• Utilizes a reasoning-enhanced prediction framework that generates step-by-step reasoning chains grounded in retrieved medical knowledge.
• Employs a retrieval-augmented generation pipeline where LLMs generate predictions based on both patient EHRs and contextually relevant KG communities.

Experimental results

Research questions

• RQ1 Can KG community-level retrieval improve the relevance and precision of retrieved medical knowledge compared to traditional RAG in clinical prediction?
• RQ2 How does integrating hierarchical graph community detection and summarization enhance the quality of retrieved knowledge for LLMs?
• RQ3 To what extent does dynamic, context-aware retrieval of KG communities improve prediction accuracy in mortality and readmission tasks?
• RQ4 Can reasoning-enhanced LLMs leveraging retrieved KG insights produce more interpretable and trustworthy clinical predictions?
• RQ5 How does KARE compare to state-of-the-art models in terms of performance and robustness across diverse clinical prediction benchmarks?

Key findings

• KARE achieves up to 15.0% relative improvement in in-hospital mortality prediction on the MIMIC-III dataset compared to the best baseline models.
• On the MIMIC-IV dataset, KARE improves readmission prediction performance by up to 12.7% over state-of-the-art methods.
• The framework demonstrates consistent gains across both MIMIC-III and MIMIC-IV, with improvements ranging from 10.8% to 15.0% on mortality and 12.6% to 12.7% on readmission tasks.
• The integration of reasoning chains with retrieved KG insights significantly enhances model interpretability, allowing clinicians to trace predictions to specific medical evidence.
• Case studies show that KARE correctly identifies high-risk patients with complex comorbidities, such as heart failure with shock and multiple cancers, by synthesizing multi-faceted knowledge from KG communities.
• The dynamic retrieval mechanism successfully avoids irrelevant or sparse knowledge by focusing on semantically coherent, contextually relevant subgraphs, reducing hallucination risks.