PaperlayerView exact revision on arXiv

Source-linked AI summary

MedRAG: Enhancing Retrieval-augmented Generation with Knowledge Graph-Elicited Reasoning for Healthcare Copilot

Xuejiao Zhao, Siyan Liu, Su-Yin Yang, Chunyan Miao

arXiv:2502.04413v2cs.CLcs.AIcs.IR

TL;DR

Medical RAG systems can produce inaccurate or vague diagnoses for diseases with similar manifestations, limiting their usefulness for healthcare decision support. MedRAG constructs a four-tier diagnostic knowledge graph, integrates diagnostic differences with retrieved EHRs, and reasons over them with an LLM. Across public and private datasets, it outperforms existing RAG approaches in diagnostic accuracy and specificity, while generating treatment guidance and follow-up questions.

  • Problem

    Existing heuristic-based medical RAG models have inadequate diagnostic accuracy and specificity, especially when diseases share similar manifestations.

  • Method

    MedRAG constructs a four-tier diagnostic knowledge graph, dynamically integrates diagnostic differences with similar retrieved EHRs, and elicits reasoning through an LLM.

  • Results

    MedRAG outperforms existing RAG approaches in diagnostic accuracy and specificity across evaluations on DDXPlus and CPDD, with robust generalization across LLM backbones.

  • Takeaways & Limitations

    The framework provides more specific diagnostic insights and supports personalized treatment recommendations, medication guidance, and proactive reasoning-based follow-up questions.

  • Takeaways & Limitations

    Adding KG knowledge decreased L1 performance in some small-parameter closed-source models because granular disease differences complicated reasoning and blurred classification boundaries.

Abstract

from arXiv · show

Retrieval-augmented generation (RAG) is a well-suited technique for retrieving privacy-sensitive Electronic Health Records (EHR). It can serve as a key module of the healthcare copilot, helping reduce misdiagnosis for healthcare practitioners and patients. However, the diagnostic accuracy and specificity of existing heuristic-based RAG models used in the medical domain are inadequate, particularly for diseases with similar manifestations. This paper proposes MedRAG, a RAG model enhanced by knowledge graph (KG)-elicited reasoning for the medical domain that retrieves diagnosis and treatment recommendations based on manifestations. MedRAG systematically constructs a comprehensive four-tier hierarchical diagnostic KG encompassing critical diagnostic differences of various diseases. These differences are dynamically integrated with similar EHRs retrieved from an EHR database, and reasoned within a large language model. This process enables more accurate and specific decision support, while also proactively providing follow-up questions to enhance personalized medical decision-making. MedRAG is evaluated on both a public dataset DDXPlus and a private chronic pain diagnostic dataset (CPDD) collected from Tan Tock Seng Hospital, and its performance is compared against various existing RAG methods. Experimental results show that, leveraging the information integration and relational abilities of the KG, our MedRAG provides more specific diagnostic insights and outperforms state-of-the-art models in reducing misdiagnosis rates. Our code will be available at https://github.com/SNOWTEAM2023/MedRAG

1 Introduction

The paper motivates MedRAG by the harms of diagnostic error and the limitations of heuristic medical RAG when diseases share manifestations. It introduces a knowledge-graph-enhanced framework designed to improve diagnostic specificity, recommendations, and follow-up questioning.

  • Approximately 795,000 people in the United States each year suffer permanent disability or death due to misdiagnosis of dangerous diseases.
  • Healthcare copilots must provide accurate diagnoses, appropriate treatments and medications, and precise follow-up questions when patient information is insufficient or ambiguous.
  • Existing medical RAG and fine-tuned medical LLMs often use heuristic approaches that produce incorrect or vague outputs when diseases have similar manifestations.
  • MedRAG combines RAG with a four-tier diagnostic knowledge graph containing critical diagnostic differences to support more precise and specific decisions.
  • MedRAG is evaluated on DDXPlus and CPDD against state-of-the-art RAG models, with experiments reporting improved diagnostic accuracy and specificity and generalization across multiple LLMs.
  • The paper contributes diagnostic knowledge graphs for chronic pain and DDXPlus and a KG-elicited RAG approach that also generates treatment guidance and follow-up questions.

2 Related Works

Related work positions RAG as a way to ground healthcare generation in private clinical data, while highlighting limitations of retrieve-and-read approaches and the potential of structured knowledge graphs.

  • RAG uses domain-specific private datasets to answer queries without additional model training and is suited to retrieving privacy-sensitive EHR information.
  • Healthcare RAG grounds predictions in retrieved medical data to support more accurate, specific, and context-aware diagnostic and treatment outputs.
  • Typical RAG retrieves top-k relevant documents from a query and generates an answer from those documents, but this retrieve-and-read simplicity restricts its capabilities.
  • Knowledge graphs provide structured and inferable information that can augment LLMs and RAG more effectively than document repositories for some medical-response tasks.

3 Preliminaries

The preliminaries define MedRAG’s four-tier diagnostic knowledge graph, its diagnostic-difference search, and the conventional two-phase RAG process used for comparison.

  • A diagnostic knowledge graph contains four hierarchical levels linking broader categories, subcategories, diseases, and manifestations.
  • The graph uses an is_a relation for hierarchy and a has_manifestation_of relation linking diseases to manifestations, including augmented and EHR-derived features.
  • Diagnostic Differences KG Searching extracts the diagnostic-difference subgraph associated with a subcategory identified from the patient’s manifestations.
  • Conventional RAG retrieves top-k relevant documents with an embedding model and then uses a generative LLM and prompt to produce an answer.

4 Methods

MedRAG combines a four-tier diagnostic knowledge graph, manifestation matching, EHR retrieval, and KG-elicited LLM reasoning. The pipeline produces diagnoses and recommendations while asking targeted follow-up questions when crucial information is missing.

  • MedRAG accepts structured EHRs or unstructured descriptions of patient manifestations and outputs diagnoses, treatment recommendations, medication guidance, and follow-up questions when necessary.
  • Diagnostic Knowledge Graph Construction: The diagnostic KG is built by clustering disease descriptions, hierarchically aggregating diseases into categories, and decomposing manifestations into feature nodes.External knowledge augments database-derived features with diagnostic differences.
  • KG-elicited Reasoning RAG: The KG-elicited reasoning RAG retrieves relevant top-k EHRs and integrates them with diagnostic differences to trigger LLM reasoning for clinical recommendations.The framework is designed for distinguishing diseases with similar manifestations using both retrieved documents and structured KG information.
  • Diagnostic Differences KG Searching: Manifestation features are embedded, matched to diagnostic differences, and filtered by similarity thresholds before relevant KG nodes are selected.The system retrieves the top-m closest matches for each of n patient features, yielding n × m candidate nodes before threshold filtering.
  • Diagnostic Differences KG Searching: Upward traversal identifies the disease subcategory receiving the most votes from matched feature nodes, then downward traversal collects related diseases and feature triples.Votes are based on shortest-path distances, and the resulting triples form the diagnostic differences KG used for reasoning.
  • Proactive Diagnostic Questioning Mechanism: The Proactive Diagnostic Questioning Mechanism generates targeted follow-up questions when the initial patient description lacks information needed for precise decisions.

5 Experiments

MedRAG is evaluated on a public synthesized EHR benchmark and a private chronic-pain clinical dataset, with comparisons against six RAG baselines. The datasets support assessment of general applicability and real-world diagnostic performance.

  • MedRAG is evaluated on DDXPlus for general applicability and CPDD for chronic-pain diagnosis in real-world clinical settings.
  • Datasets: DDXPlus contains 49 diagnoses and over 1.3 million patients, while the balanced subset used in the study contains 13,230 patients’ EHRs.
  • Datasets: CPDD contains 551 patients with 33 distinct diagnoses collected from Tan Tock Seng Hospital.
  • Baselines: MedRAG is compared against six models: Naive RAG with COT, FL-RAG, FS-RAG, FLARE, DRAGIN, and SR-RAG.

6 Experimental Results

Experiments compare MedRAG with state-of-the-art RAG methods across datasets, backbone LLMs, proactive questioning, and component ablations. MedRAG generally achieves stronger diagnostic accuracy and specificity, while targeted questions and KG-elicited reasoning improve performance.

  • Quantitative Comparison (RQ1): MedRAG achieved the best or second-best performance across multiple metrics in both datasets, outperforming second-best scores by 11.32% on CPDD and 1.23% on DDXPlus.Accuracy at L3 is emphasized as the strongest indicator because higher specificity increases diagnostic difficulty.
  • Compatibility, Generalizability and Adaptability (RQ2): KG-elicited reasoning improved diagnostic accuracy across L1, L2, and L3 for all tested backbone LLMs.For Mixtral-8x7B, L3 accuracy increased from 22.34% to 63.46%.
  • Compatibility, Generalizability and Adaptability (RQ2): GPT-4o achieved the strongest performance among the evaluated backbone models, while MedRAG performed best among closed-source models.Token-level methods such as DRAGIN and FLARE were reported to face adaptation challenges with closed-source models.
  • Compatibility, Generalizability and Adaptability (RQ2): KG integration reduced L1 performance in small-parameter closed-source models, where granular disease differences created knowledge conflicts and blurred classification boundaries.The authors contrast this with GPT-4o, whose larger scale and stronger reasoning capacity produced higher L1 accuracy.
  • Proactive Diagnostic Questioning (RQ3): Targeted follow-up questions increased L3 accuracy from 52.83% with little manifestation information to 66.04% when 100% of key manifestations were covered.The reported L3 score also rose from 55.10% to 66.04% as coverage increased from 33.3% to 100%.
  • Ablation Study (RQ4): Combining retrieval with KG-elicited reasoning produced the best ablation outcomes, particularly for high-specificity diagnosis.Adding correct KG-augmented knowledge increased average accuracy by 18.88% for L1, 26.92% for L2, and 18.89% for L3 relative to random or absent KG reasoning.

7 Conclusion

MedRAG integrates KG-elicited reasoning with RAG to improve diagnostic accuracy and specificity, while supporting personalized recommendations and proactive questioning. Evaluation on public and private datasets indicates stronger performance against state-of-the-art RAG models, especially for diseases with similar manifestations.

  • 7 Conclusion: MedRAG integrates retrieved EHRs with critical diagnostic differences from a knowledge graph to improve diagnostic precision and specificity.The system also provides personalized treatment recommendations and medication guidance.
  • 7 Conclusion: MedRAG’s proactive diagnostic questioning mechanism provides impactful questions that may enhance diagnostic performance and consultation efficiency.
  • 7 Conclusion: Evaluation on public and private datasets shows MedRAG outperforms state-of-the-art RAG models, particularly in reducing misdiagnosis rates for diseases with similar manifestations.
  • 7 Conclusion: Future work will incorporate multimodal data and conduct real-world hospital testing to broaden applicability and assess clinical effectiveness.Planned extensions include medical imaging, physiological signals, blood tests, and deployment within healthcare copilot systems.

B Data Preprocessing and Experimental Setup

The experimental setup prepares balanced retrieval and test datasets, simulates missing manifestations for follow-up questioning, and evaluates MedRAG against multiple baseline models and prompt configurations.

  • B.1 Settings for Datasets: CPDD uses a 9:1 training-to-test split, while DDXPlus uses an 8:1:1 training, test, and validation split with balanced sampling.The DDXPlus retrieval subset contains 13,230 patient EHRs, with 240 samples selected per pathology.
  • B.2 Follow-up Question Generation: MedRAG masks selected patient manifestations to simulate missing information, then generates follow-up questions and repeats diagnostic reasoning after recovered features are added.
  • B.2 Follow-up Question Generation: The missing-feature simulation removes a proportion of manifestation nodes ranked by discriminability, along with corresponding sentences.
  • B.3 Prompt Configuration: Disease clustering groups diseases with similar manifestations and characteristics into categories for knowledge-graph construction.The clustering prompt uses manifestations, affected locations, and other characteristics as grouping criteria.
  • B.5 Generative Model Configuration: The generative prompt combines patient manifestations, retrieved documents, and diagnostic differences as inputs to the model.
  • B.4 Baseline Models: Experiments compare MedRAG with six baselines, including Naive RAG with COT, FL-RAG, FS-RAG, FLARE, DRAGIN, and SR-RAG.

D.1 Disease Clustering Result

Disease clustering groups different representations of the same disease in the CPDD EHR database, creating a consistent foundation for knowledge-graph construction and augmentation. Diagnostic differences can then distinguish diseases with overlapping features, such as lumbar canal stenosis and sciatica.

  • D.1 Disease Clustering Result: Disease clustering groups different forms and representations of the same disease and assigns each cluster a topic.This unifies disease representations for consistency and comparability.
  • D.1 Disease Clustering Result: The clustering output provides a unified foundation for subsequent disease knowledge-graph construction and augmentation.
  • D.1 Disease Clustering Result: For lumbar canal stenosis and sciatica, the response to sitting is a critical diagnostic difference: sitting alleviates stenosis symptoms but exacerbates sciatica discomfort.

E Report Generation Evaluation

MedRAG reports are evaluated objectively with standard text-generation metrics and subjectively using Mini-CEX criteria, with doctor validation. In the subjective evaluation, MedRAG scores higher than SRRAG.

  • E Report Generation Evaluation: Report-generation evaluation is conducted on CCPD because DDXPlus does not contain report data.
  • E Report Generation Evaluation: Report generation is evaluated with BERTScore, BLEU, ROUGE, and METEOR, while subjective assessment uses four Mini-CEX criteria.The subjective results were assessed by GPT-4o and validated by doctors.
  • E Report Generation Evaluation: MedRAG scores 290 out of 360 versus SRRAG’s 277 in subjective evaluation of reports from 10 randomly selected patients.

F Ablation Study on KG Components

The ablation study evaluates key diagnostic knowledge-graph components and shows that both component choices and hierarchical structure affect MedRAG’s effectiveness.

  • Component ablations: The ablation examines diagnostic key difference nodes, augmented feature nodes, patient clinical feature matching, and diagnostic-difference augmentation.These are the components assessed for their effects on MedRAG.
  • Component ablations: The four examined KG components significantly contribute to MedRAG’s overall effectiveness.Table A2 reports the contribution of diagnostic key difference nodes, augmented feature nodes, patient clinical feature matching, and diagnostic-difference augmentation.
  • Hierarchical structure: The hierarchical structure of the constructed diagnostic-differences KG directly impacts the experimental results.The ablation findings therefore cover both individual KG components and the organization of the KG itself.
Loading 2502.04413v2…