[Paper Review] Improving Medical Reasoning through Retrieval and Self-Reflection with Retrieval-Augmented Large Language Models

Minbyul Jeong, Jiwoong Sohn+2 more|arXiv (Cornell University)|Jan 27, 2024

Topic Modeling4 citations

TL;DR

Self-BioRAG is a retrieval-augmented, self-reflective large language model framework specifically designed for biomedical and clinical question-answering tasks. It improves medical reasoning by dynamically deciding when to retrieve domain-specific evidence, integrating it with parametric knowledge, and self-evaluating its reasoning; it achieves a 7.2% absolute accuracy gain over state-of-the-art 7B open-source models on three major medical benchmarks.

ABSTRACT

Recent proprietary large language models (LLMs), such as GPT-4, have achieved a milestone in tackling diverse challenges in the biomedical domain, ranging from multiple-choice questions to long-form generations. To address challenges that still cannot be handled with the encoded knowledge of LLMs, various retrieval-augmented generation (RAG) methods have been developed by searching documents from the knowledge corpus and appending them unconditionally or selectively to the input of LLMs for generation. However, when applying existing methods to different domain-specific problems, poor generalization becomes apparent, leading to fetching incorrect documents or making inaccurate judgments. In this paper, we introduce Self-BioRAG, a framework reliable for biomedical text that specializes in generating explanations, retrieving domain-specific documents, and self-reflecting generated responses. We utilize 84k filtered biomedical instruction sets to train Self-BioRAG that can assess its generated explanations with customized reflective tokens. Our work proves that domain-specific components, such as a retriever, domain-related document corpus, and instruction sets are necessary for adhering to domain-related instructions. Using three major medical question-answering benchmark datasets, experimental results of Self-BioRAG demonstrate significant performance gains by achieving a 7.2% absolute improvement on average over the state-of-the-art open-foundation model with a parameter size of 7B or less. Overall, we analyze that Self-BioRAG finds the clues in the question, retrieves relevant documents if needed, and understands how to answer with information from retrieved documents and encoded knowledge as a medical expert does. We release our data and code for training our framework components and model weights (7B and 13B) to enhance capabilities in biomedical and clinical domains.

Motivation & Objective

• To address the limitations of general-purpose LLMs in biomedical reasoning, particularly hallucination and poor factual grounding.
• To improve performance on domain-specific medical question-answering by integrating on-demand retrieval with self-reflection capabilities.
• To demonstrate that domain-specific components—such as a biomedical retriever, curated corpus, and instruction-tuned model—are essential for high-fidelity medical reasoning.
• To develop a framework that mimics expert medical reasoning by combining retrieved evidence and parametric knowledge with self-assessment.
• To release a fully trained 7B and 13B model, along with data and code, to advance research in clinical and biomedical LLM applications.

Proposed method

• The framework uses a domain-specific instruction-tuned LLM (7B/13B) fine-tuned on 84k filtered biomedical instruction sets to generate and reflect on answers.
• It employs a retriever (MedCPT) to fetch relevant biomedical documents from curated corpora: PubMed Abstract, PMC Full Text, Clinical Guidelines, and Medical Textbook.
• Self-BioRAG uses reflective tokens to self-assess whether retrieval is needed, whether retrieved evidence supports the answer, and whether the final response is useful.
• The model dynamically decides between direct generation using parametric knowledge or retrieval-augmented generation based on question complexity and evidence need.
• It integrates both retrieved evidence (highlighted in blue) and model-internal knowledge (in red) to produce rationales and answers.
• The framework is trained end-to-end with a focus on multi-choice question answering, summarization, and clinical reasoning tasks.

Figure 1 : Comparison between three frameworks: generation using language model (LM), retrieval-augmented generation (RAG) using LM, and our Self-BioRAG. (A) depicts the process of sequence-to-sequence generation of LM. (B) The RAG framework first finds relevant documents from large-scale corpus suc

Experimental results

Research questions

• RQ1 Can a retrieval-augmented, self-reflective LLM framework outperform general-purpose models in biomedical question-answering?
• RQ2 Does the use of domain-specific components—retriever, corpus, instruction data, and reflective mechanism—significantly improve reasoning accuracy?
• RQ3 Can self-reflection tokens effectively guide retrieval decisions and answer quality in medical contexts?
• RQ4 How does the integration of retrieved evidence and parametric knowledge compare to relying solely on internal knowledge in medical reasoning?
• RQ5 Does the framework generalize across different types of medical questions, including those requiring clinical guidelines or textbook knowledge?

Key findings

• Self-BioRAG achieves a 7.2% absolute improvement in average accuracy over state-of-the-art open-source 7B models on three major medical question-answering benchmarks.
• On the PubMedQA dataset, Self-BioRAG achieves 54.6% accuracy, slightly below models like LLaMA2 (56.0%) and MEDITRON (58.1%), attributed to noisy evidence in retrieval.
• The model retrieves a significantly higher proportion of relevant evidence from the Medical Textbook corpus when solving USMLE-style questions.
• Self-BioRAG correctly identifies when retrieval is unnecessary, relying on parametric knowledge for simpler questions, and uses retrieval for complex, evidence-dependent queries.
• The integration of retrieved evidence and parametric knowledge leads to more accurate and explainable diagnoses, as demonstrated in case studies involving PCOS and impaired glucose tolerance.
• The framework demonstrates that domain-specific components are essential, as general-purpose RAG methods fail to generalize effectively in clinical and biomedical settings.

Figure 2 : Overview of our Self-BioRAG process: data construction, training, and inference of Self-Reflection Language Model (critic LM $C$ ) and Domain-specific Instruction-tuned Language Model (generator LM $M$ ). We construct 120k biomedical instruction sets using two off-the-shelf instruction se