PaperlayerView exact revision on arXiv

Source-linked AI summary

Improving Medical Reasoning through Retrieval and Self-Reflection with Retrieval-Augmented Large Language Models

Minbyul Jeong, Jiwoong Sohn, Mujeen Sung, Jaewoo Kang

arXiv:2401.15269v3cs.CLcs.AIcs.IR

TL;DR

Biomedical question answering can require knowledge beyond what LLMs encode, and existing retrieval methods may generalize poorly to domain-specific questions. Self-BioRAG combines biomedical instruction training with on-demand retrieval and self-reflection; it reports a 7.2% absolute improvement over a state-of-the-art open-foundation model with 7B or fewer parameters.

  • Problem

    Existing retrieval-augmented methods can generalize poorly to biomedical and clinical questions, retrieving incorrect documents or making inaccurate judgments.

    Source: §1 · p. 2

  • Method

    Self-BioRAG uses biomedical instructions, domain-specific retrieval, and reflective assessment to decide when to retrieve and answer using evidence and encoded knowledge.

    Source: §1 · pp. 2–3

  • Results

    7.2% absolute improvement over the state-of-the-art open-foundation model with 7B or fewer parameters, with significant gains over open-foundation LLMs and RAG approaches on multi-choice QA.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    Domain-specific components contribute to performance gains, with training on domain-specific instructions showing the highest improvement.

    Source: §1 · p. 2

Abstract

from arXiv · show

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.

1 Introduction

General retrieval-augmented methods can generalize poorly to biomedical and clinical questions, while Self-BioRAG combines domain-focused training with on-demand retrieval and self-reflection. Across biomedical question-answering benchmarks, it reports a 7.2% absolute improvement over a state-of-the-art open-foundation model with 7B or fewer parameters.

1 Introduction

  • Its workflow predicts whether retrieval is needed, then selects relevant evidence and generates an answer using that evidence and encoded knowledge.The framework uses a domain-specific retriever when a query requires factual content.Source: Figure 1 · p. 2
  • Existing methods, including Self-RAG, can generalize poorly to biomedical and clinical questions, fetching incorrect documents or making inaccurate judgments.RAG supplies supporting facts to LLMs, while Self-RAG uses reflective tokens to decide when to retrieve and assess evidence and answers.Source: §1 · p. 2
  • Self-BioRAG is trained on biomedical and clinical instructions and combines on-demand retrieval with self-reflection while preserving generation quality and reasoning ability.Here, reasoning means providing explanations for answers.Source: §1 · p. 2
  • 7.2% absolute improvement over the state-of-the-art model with 7B or fewer parameters; Self-BioRAG significantly outperforms open-foundation LLMs and RAG approaches on multi-choice QA.The evaluation also covers long-form QA, where answers were similar despite substantial terminology differences.Source: §1 · pp. 2–3
  • The authors report constructing 120k biomedical instruction sets and retaining 84k filtered sets for generator training; they also release data, code, and 7B and 13B weights.The 120k sets combine existing instruction sets with synthetically generated instructions.Source: §1 · pp. 2–3, Figure 2 · p. 3

2 Background

The background motivates open, domain-specific language models and reviews reflective reward strategies and retrieval-augmented generation. It identifies biomedical retrieval and documents as domain-specific needs.

2.1 Proprietary & Open Language Models

  • Open foundation and domain-specific models offer research-friendly alternatives to proprietary systems, with Self-BioRAG targeting tasks from multiple-choice questions to long-form generation.The authors frame this as a labor-inexpensive approach for vertical domains including biomedical and clinical fields.Source: §2.1 · p. 3

2.2 Learning with Reward Strategy

  • Self-RAG uses a critic model and reflective tokens to assess retrieval need, context appropriateness, rationale support, and output utility; Self-BioRAG adapts this approach to biomedical text.The passage contrasts this with RLHF reward strategies.Source: §2.2 · p. 3

2.3 Retrieval-Augmented Generation

  • Self-RAG retrieves on demand for cost efficiency, but Self-BioRAG uses biomedical- and clinical-specific retrieval methods and documents to provide field-aligned context.The passage contrasts on-demand retrieval with approaches that fetch evidence in advance or iteratively.Source: §2.3 · p. 4

3 Self-BioRAG

Self-BioRAG is built from biomedical instructions, a domain-specific retriever, a critic that annotates reflective criteria, and a generator trained on those annotations. Its components are designed to support biomedical and clinical question answering.

3 Self-BioRAG

  • Self-BioRAG construction combines biomedical instructions, MedCPT retrieval, critic-generated reflective annotations, and generator training.The critic annotates instruction sets; the generator is trained using the resulting data.Source: §3 · p. 4

List of Instruction Datasets for Biomedical and Clinical Domains.

  • The instruction collection totals 120k biomedical sets, assembled from Mol-Instructions, MedInstruct, and 18k synthetically generated instructions.The sets address tasks including information extraction, question answering, and summarization.Source: List of Instruction Datasets for… · p. 4

3.2 Biomedical Retriever

  • MedCPT retrieves evidence from PubMed abstracts, PMC full text, clinical guidelines, and English medical textbooks, then reranks the top-k results.Documents are chunked into 128-word segments with 32-word overlaps; the system retrieves k=10 from each source before reranking.Source: §3.2 · p. 4

Data Construction of Critic LM C.

  • The authors sample 5k of the 120k instructions and use GPT-4 API calls to generate reflective tokens for critic-model training.Exploring other reflective tokens suited to specific domains is left for future work.Source: Data Construction of Critic LM C · p. 4

Process of Training Critic LM C.

  • The released Self-RAG critic was unsuitable for predicting reflective tokens in biomedical domains, motivating development of a domain-specific critic.Source: Process of Training Critic LM C · p. 4

7 See Appendix 2 for a comparison of critic LMs’ performances

The supplied passages describe reflective-token annotations and critic training, but do not provide comparative critic-performance results.

7 See Appendix 2 for a comparison of critic LMs’ performances

  • Reflective tokens encode criteria used in Self-BioRAG data construction; the table caption identifies x, y, and e as input, output, and evidence.The caption says highlighted tokens are desirable because they help preserve existing instruction data.Source: Table 3 · p. 5
  • The critic is trained to predict predefined reflective tokens from sampled instruction sets, with its likelihood objective specified in the training passage.Source: §7 · p. 5
  • The trained critic annotates instruction sets, and instances with mispredicted tokens outside the predefined set are filtered out.The passage gives [Continue Generation] as an example of a mispredicted token.Source: §7 · p. 5

Annotating Biomedical Instruction Sets Using Critic LM C.

Critic LM C annotates biomedical instructions with reflective tokens that assess retrieval needs, evidence relevance, answer support, and response usefulness.

Annotating Biomedical Instruction Sets Using Critic LM C.

Data Construction Using Critic LM C & Training Generator LM M.

MedCPT retrieves biomedical evidence for instructions that require context, and critic LM C annotates examples used to train generator LM M.

Data Construction Using Critic LM C & Training Generator LM M.

Inference Process of Self-BioRAG.

At inference, Self-BioRAG uses reflective-token scores to select retrieved evidence and generate an answer using that evidence and encoded knowledge.

Inference Process of Self-BioRAG.

  • In a MedQA example, generator M decides whether retrieval is needed and selects the best of the top-k documents using a reflective-token weighted score.Source: Inference Process of Self-BioRAG · p. 5
  • For the example patient with obesity, acne, and type 2 diabetes history, the answer connects these clues to PCOS and mentions a glucose tolerance test.Source: Inference Process of Self-BioRAG · p. 5
  • Increasing the REL weight can prioritize documents relevant to the question, and inference-time weights adjust behavior without additional training.The example's selected evidence includes family history of type 2 diabetes and a diagnosis of PCOS.Source: Inference Process of Self-BioRAG · p. 5
  • The authors characterize the process as identifying question clues, retrieving factual content when needed, and answering with retrieved evidence and encoded knowledge.Source: Inference Process of Self-BioRAG · p. 5

Benchmark Evaluations.

The evaluation measures downstream performance of Self-BioRAG fine-tuned on Self-RAG across multiple-choice and long-form biomedical question-answering benchmarks.

Benchmark Evaluations.

Baselines.

The baseline comparison spans proprietary models, general open-foundation models, domain-fine-tuned models, and retrieval-augmented approaches. Table 4 also specifies how benchmark examples and evidence were selected.

  • Proprietary model scores are reported from Med-PaLM and Nori et al. to establish benchmark upper bounds.Source: Baselines · p. 5
  • Table 4 compares proprietary, general open-foundation, domain-fine-tuned, retrieval-augmented, and Self-BioRAG systems on biomedical benchmarks.Source: Baselines · pp. 5–6, Table 4 · p. 6
  • Table 4 uses three-shot examples selected by k-nearest neighbors, except for MMLU, which uses examples from Med-PaLM’s appendix.Source: Table 4 · p. 6
  • LLaMA2 is reported as the state-of-the-art open-foundation 7B model in the authors’ experiment and serves as the RAG base.Source: Baselines · p. 6
  • For LLaMA2 RAG, the authors provide top-10 biomedical-corpus evidence but input only the top item because of context-length limits.Source: Baselines · p. 6

Training and Inference Settings.

The settings combine biomedical instruction tuning, a specialized retriever, and adaptive retrieval at inference. Long-form QA evaluation reports both lexical-overlap and embedding-similarity metrics.

  • Self-BioRAG uses 84k critic-filtered biomedical instructions; its critic is fine-tuned on 5k GPT-4-annotated samples, and generator fine-tuning outperformed training directly on LLaMA2 or Meditron.Source: Training and Inference Settings · p. 6
  • MedCPT, specialized for biomedical queries, retrieves up to ten evidence items for each input.Source: Training and Inference Settings · p. 6
  • Long-form QA is evaluated with ROUGE-1/2/L for n-gram recall and BERTScore for token-embedding similarity.Source: Table 5 · p. 6
  • At inference, adaptive retrieval predicts [Retrieval] to decide dynamically whether to fetch top-ten biomedical-corpus evidence; vLLM speeds inference.Source: Training and Inference Settings · p. 6

5 Results and Analysis

Self-BioRAG outperforms open foundation models and RAG on the three biomedical multiple-choice benchmarks, while analyses examine its domain-adaptation components and evidence use. Results also cover adaptive retrieval, case behavior, instruction-set statistics, and critic-model performance.

5.2 Analysis

  • Reflective-token generation, four biomedical corpora, and MedCPT each contribute; biomedical instruction-set collection and processing is identified as most effective.The component analysis compares biomedical corpora with Wikipedia evidence and MedCPT with Contriever.Source: §5.2 · p. 7, Table 6 · p. 7
  • Retrieved evidence is distributed across four biomedical corpora despite smaller CPG and textbook indexes; textbook evidence has a higher share, while only 12%, 8%, and 11% is retrieved for MedQA, MedMCQA, and MMLU-Med.The figure reports evidence-use ratios for PubMed, PMC, CPG, and Medical Textbook.Source: §5.2 · p. 7, Figure 3 · p. 7
  • Adaptive retrieval performs comparably on average to forced top-10 retrieval, whereas forced retrieval improves MedMCQA but lowers performance on MMLU-Med versus no retrieval.The adaptive criterion uses δ = 0.2; the authors recommend adaptive retrieval.Source: §5.2 · p. 7, Table 7 · p. 7

Distinguishing When to Retrieve Documents in Self-BioRAG.

  • On MedQA, Self-BioRAG outperforms RAG and LLaMA2 both when it predicts retrieval is unnecessary and when it predicts retrieval is needed.The figure separates examples into [No Retrieval] and [Retrieval] according to Self-BioRAG's predictions.Source: Distinguishing When to Retrieve… · p. 8, Figure 4 · p. 8

Case Report of Using Retrieved Evidence.

What Contributes to the Performance Improvements in Self-BioRAG?

1 Statistics of our generated Biomedical Instruction Set

  • Generated biomedical instructions span varied tasks and lengths; a 100-instance review found text generation at 42%, classification at 15%, and symptom identification at 7%.The figures summarize instruction verbs and objects, instruction similarity to seed tasks, and instruction, input, and output lengths.Source: §1 · p. 10, Figure 5 · p. 10, Figure 6 · p. 10

2 Comparison of Critic LM’s Performance

  • The critic language model of Self-BioRAG, trained on Self-RAG, predicts utility and relevance reflective tokens well, but retrieval and supportiveness tokens far less well than expected.The authors split 5k instruction sets into train and dev to train and evaluate critic performance; Table 10 uses GPT-4 predictions as ground-truth predictions. On uncommon-domain questions, Self-RAG performs lower than Self-BioRAG.Source: §2 · p. 10, Table 10 · p. 11
Loading 2401.15269v3…