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

✅ This paper is accepted by The Web Conference (WWW) 2025!

💻 This is the official implementation for our accepted paper MedRAG: Enhancing Retrieval-augmented Generation with Knowledge Graph-Elicited Reasoning for Healthcare Copilot.

Authors

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

Nanyang Technological University   |   Tan Tock Seng Hospital   |  Woodlands Health

* Both authors contributed equally to the paper

** Corresponding author

---

🔥 News

• [2025.03.24] MedRAG has recently drawn some interest from media outlets and bloggers, such as Medium, AI Era, and CSDN. Tks for all the support — we're continuing to improve! 🙏
• [2025.03.12] We release the officially generated diagnostic knowledge graph for the DDXPlus dataset. Have a try!
• [2025.02.25] The Chinese Demo of MedRAG is now available on Bilibili, and the English Demo is available on YouTube.
• [2025.02.10] MedRAG reported by the 'Quantum Heart' on RedNote.
• [2025.02.04] We release the official implementation of MedRAG.
• [2025.01.20] MedRAG has been accepted by WWW'25. Please check the latest paper version on ArXiv.

Overview

Figure 1: The overall framework of MedRAG.

MedRAG is to designed to enhance Retrieval-Augmented Generation (RAG) models by integrating Knowledge Graph (KG)-elicited reasoning, specifically for the medical domain. This model helps healthcare professionals generate diagnosis and treatment recommendations based on patient manifestations, improving diagnostic accuracy and reducing the risk of misdiagnosis, particularly for diseases with similar manifestations.

Key features of MedRAG include:

• Knowledge Graph-Enhanced Reasoning: Integrates a diagnostic knowledge graph to improve the reasoning ability of the RAG model.
• Accurate Diagnostic Support: Provides specific diagnostic insights and personalized treatment recommendations, even for complex or similar diseases.
• Follow-Up Question Generation: Proactively generates relevant follow-up questions to clarify ambiguous patient information and enhance decision-making.
• Evaluated on Real-World and Public Datasets: Demonstrated superior performance on the public DDXPlus dataset and a private chronic pain diagnostic dataset (CPDD) compared to existing RAG models.

---

Core Design of MedRAG: Knowledge Graph-Elicited Reasoning

The MedRAG approach addresses the following key challenges:

1. Knowledge Graph Construction: Using hierarchical aggregation to build disease knowledge graph, capturing complex relationships between diseases, categories, and their manifestations.
2. RAG-Based Reasoning: Combines EHR retrieval with diagnostic knowledge graph reasoning to enhance diagnostic accuracy.
3. Personalized Diagnostic Suggestions: Integrates multi-level information to provide personalized treatment and follow-up questions for doctors.

---

Dataset

The full MedRAG test set, including raw image data and annotations, can be downloaded from the links below. Due to the large size of the dataset, a lighter version is also available for quicker testing.

• Download Full DDXPlus: A large-scale, synthesized EHR dataset widely recognized for offering complex, diverse medical diagnosis cases. It includes comprehensive patient data such as socio-demographic information, underlying diseases, symptoms, and antecedents.
• CPDD: A private EHR dataset for chronic pain management from our partner hospital, Tan Tock Seng Hospital in Singapore.

Usage

To use MedRAG, follow these steps:

1. Get ready for repository and dependencies
   Clone this repository to your local machine and install requirements in requirements.txt

   git clone https://github.com/SNOWTEAM2023/MedRAG.git
   
   cd MedRAG
   pip install -r requirements.txt

2. Modify Tokens
   To use your own OpenAI and Hugging Face API tokens, replace the placeholders in the authentication.py with your actual tokens. The relevant sections in the code have been left blank for this purpose.

   # Replace with your OpenAI API token
   api_key = "your_openai_api_token"   
   
   # Replace with your Hugging Face API token
   hf_token = "your_huggingface_api_token"

3. Run the main.py script
   Once the paths and tokens have been updated, run the main.py file to start the program:

   python main.py

Experimental Results

Main results

Figure 1: Results of quantitative performance comparison.

Our proposed MedRAG achieved the best or second-best~(with only one exception) performance across multiple metrics in all datasets. Accuracy on the $L3$ metric is the best indicator of MedRAG's performance, as higher specificity increases diagnostic difficulty. MedRAG outperformed the second-best scores on the CPDD and DDXPlus datasets.

Figure 2: Performance of MedRAG on different LLM backbones with and without KG-elicited reasoning.

We evaluate KG-elicited reasoning on different LLM backbones, including both open-source and closed-source models. The results demonstrate that the inclusion of KG-elicited reasoning significantly enhances diagnostic accuracy across $L1$, $L2$, and $L3$ for all backbone LLMs, compared to models without its use.

Additional Visualizations

Clustering result

The result of disease clustering in CPDD.

Diseases knowledge graph

The result of hierarchical aggregation in DDXPlus.

The result of hierarchical aggregation in CPDD.

Diagnostic differences augmentation

Diagnosic difference example.

While lumbar canal stenosis and sciatica share some similar features, the critical distinguishing factor lies in the response to sitting. In lumbar canal stenosis, features are typically alleviated when sitting, whereas in sciatica, sitting tends to exacerbate the discomfort.

📖 Citation

If you find our work useful, please consider citing our paper:

@inproceedings{zhao2025medrag,
  title={MedRAG: Enhancing Retrieval-augmented Generation with Knowledge Graph-Elicited Reasoning for Healthcare Copilot},
  author={Zhao, Xuejiao and Liu, Siyan and Yang, Su-Yin and Miao, Chunyan},
  booktitle={THE WEB CONFERENCE 2025}
}