[Paper Review] MKRAG: Medical Knowledge Retrieval Augmented Generation for Medical Question Answering

Yucheng Shi, Shaochen Xu+6 more|arXiv (Cornell University)|Sep 27, 2023

Topic ModelingComputer Science8 citations

TL;DR

The paper presents MKRAG, a retrieval-augmented generation approach that uses in-context medical knowledge editing to improve medical QA without fine-tuning, showing improved Vicuna performance on MedQA-USMLE.

ABSTRACT

Large Language Models (LLMs), although powerful in general domains, often perform poorly on domain-specific tasks such as medical question answering (QA). In addition, LLMs tend to function as "black-boxes", making it challenging to modify their behavior. To address the problem, our work employs a transparent process of retrieval augmented generation (RAG), aiming to improve LLM responses without the need for fine-tuning or retraining. Specifically, we propose a comprehensive retrieval strategy to extract medical facts from an external knowledge base, and then inject them into the LLM's query prompt. Focusing on medical QA, we evaluate the impact of different retrieval models and the number of facts on LLM performance using the MedQA-SMILE dataset. Notably, our retrieval-augmented Vicuna-7B model exhibited an accuracy improvement from 44.46% to 48.54%. This work underscores the potential of RAG to enhance LLM performance, offering a practical approach to mitigate the challenges posed by black-box LLMs.

Motivation & Objective

• Motivate improving medical QA by addressing gaps in LLM medical knowledge without fine-tuning.
• Propose a two-step medical fact retrieval and in-context editing to guide LLMs with external facts.
• Demonstrate the approach on the MedQA-USMILE/MedQA-USMLE dataset and compare retrieval models.

Proposed method

• Convert external medical facts to embeddings with a chosen embedding model.
• Broad search to collect top-K facts related to each answer candidate.
• Refined search to select top-k facts most similar to the question, forming F_R.
• Insert retrieved facts into the prompt to perform in-context editing of the LLM.
• Evaluate edits using a Vicuna-13B model on the MedQA-USMLE dataset.

Experimental results

Research questions

• RQ1 RQ1: Can model editing via in-context facts improve medical QA performance?
• RQ2 RQ2: Which retrieval embedding model (Contriever vs SapBERT) yields better factual retrieval for medical QA?
• RQ3 RQ3: How does varying the number of editing facts affect QA accuracy?

Key findings

• Edited Vicuna achieves 48.54% accuracy on MedQA-USMLE test, outperforming pre-edited Vicuna at 44.46%.
• Contriever slightly outperforms SapBert as the retrieval model in this setup (48.54% vs 48.07%), indicating retrieval quality matters.
• Increasing the number of editing facts from 4/8/16 shows a positive correlation with accuracy, up to model prompt size constraints.
• The method yields competitive gains without fine-tuning or retraining, compared to baselines like BioBERT, SapBERT, and QA-GNN.