Source-linked AI summary
RARE: Retrieval-Augmented Reasoning Enhancement for Large Language Models
Hieu Tran, Zonghai Yao, Junda Wang, Yifan Zhang, Zhichao Yang, Hong Yu
TL;DR
Medical question answering must handle specialized, evolving knowledge and multi-step reasoning, while composite questions can require evidence at each step. RARE extends rStar with retrieval actions and factuality scoring; experiments report improved accuracy and competitive open-source model performance against proprietary systems.
Problem
Medical QA depends on specialized knowledge that may be absent from pretrained models and can change with updated treatments or guidelines; composite questions may need evidence at each reasoning step.
Source: §1 · p. 1Method
RARE extends rStar with question- and sub-question-based retrieval actions, uses MCTS to select action paths, and scores reasoning paths for factual accuracy.
Source: §1 · p. 2Results
Takeaways & Limitations
RARE operates without additional training or fine-tuning and is presented as a scalable solution for knowledge-intensive reasoning tasks.
Source: §6 · p. 9Evaluation was limited to open-source models such as LLaMA 3.1 because the API calls required for iterative retrieval and reasoning made closed-source model evaluations prohibitively costly.
Source: §7 · p. 10
Abstract
from arXiv · showhide
This work introduces RARE (Retrieval-Augmented Reasoning Enhancement), a versatile extension to the mutual reasoning framework (rStar), aimed at enhancing reasoning accuracy and factual integrity across large language models (LLMs) for complex, knowledge-intensive tasks such as commonsense and medical reasoning. RARE incorporates two innovative actions within the Monte Carlo Tree Search (MCTS) framework: A6, which generates search queries based on the initial problem statement, performs information retrieval using those queries, and augments reasoning with the retrieved data to formulate the final answer; and A7, which leverages information retrieval specifically for generated sub-questions and re-answers these sub-questions with the relevant contextual information. Additionally, a Retrieval-Augmented Factuality Scorer is proposed to replace the original discriminator, prioritizing reasoning paths that meet high standards of factuality. Experimental results with LLaMA 3.1 show that RARE enables open-source LLMs to achieve competitive performance with top open-source models like GPT-4 and GPT-4o. This research establishes RARE as a scalable solution for improving LLMs in domains where logical coherence and factual integrity are critical.
1 Introduction
RARE extends rStar with retrieval-augmented reasoning actions and factuality scoring to address the knowledge and multi-step demands of medical and other complex question answering.
1 Introduction
- Medical QA requires reasoning from patient conditions through diagnosis to evidence-based interventions, making structured multi-step reasoning important for accurate, contextually relevant answers.The paper also notes that medical QA involves intricate knowledge, clinical scenarios, and context-appropriate choices.Source: §1 · p. 1
- Medical QA depends on external, evolving domain knowledge, while composite questions may require relevant evidence retrieval at each reasoning step.The paper notes that newer treatments or guidelines may be absent from pretrained models.Source: §1 · pp. 1–2
- RARE builds on rStar: A6 retrieves documents from queries based on the question, while A7 retrieves targeted information for sub-questions and updates the next step.MCTS selects the action path leading to the final answer.Source: §1 · p. 2
- Across medical and general-domain QA tasks, RARE significantly enhances accuracy across LLMs and enables LLaMA 3 to perform competitively with GPT-4o.The paper reports applying RARE and baselines to three medical and three general-domain QA tasks.Source: §1 · p. 2
- The Retrieval-Augmented Factuality Scorer ranks reasoning paths for factual accuracy, supporting logical coherence and factual reliability throughout reasoning.This scorer is presented as a contribution alongside the retrieval actions.Source: §1 · pp. 2–3
2 Preliminaries
The section establishes MCTS and rStar as foundations for RARE: MCTS searches by iteratively refining action choices, while rStar adds diverse reasoning actions.
2 Preliminaries
- The preliminaries identify MCTS and rStar as the foundations for RARE’s proposed method.rStar augments MCTS with a diverse set of reasoning actions.Source: §2 · p. 3, §2.2 · p. 3
2.1 Monte Carlo Tree Search (MCTS)
- MCTS incrementally builds a search tree by simulating outcomes to estimate action values, using selection, expansion, simulation, and backpropagation to balance exploration and exploitation.The section identifies UCT as the mechanism for balancing exploration and exploitation.Source: §2.1 · p. 3
2.2 Mutual Reasoning Makes Smaller LLMs Stronger Problem-Solvers
- rStar extends MCTS with five actions: one-step or remaining thoughts, sub-question generation or re-answering, and question rephrasing.These actions support more adaptive exploration than approaches relying on a single action type.Source: §2.2 · p. 3
3 Methodology
RARE extends rStar with retrieval-augmented reasoning actions and a factuality scorer, selecting reasoning paths using retrieved evidence. Its generator retrieves for both original questions and generated sub-questions.
3.1 Overview of RARE Framework
- RARE combines an MCTS-based retrieval-augmented generator with RAFS, which replaces rStar’s discriminator and selects the trajectory with the highest factuality score.The framework integrates external evidence into intermediate reasoning and prioritizes paths aligned with that evidence.Source: §3.1 · p. 4
3.2 Retrieval-Augmented Generator
- A6 retrieves documents from queries based on the original question, while A7 retrieves for A3-generated sub-questions and re-answers them to update subsequent reasoning.The generator draws on PubMed, StatPearls, medical textbooks, and Wikipedia, using ColBERT for passage-level retrieval.Source: Figure 4 · p. 4, §3.2 · p. 5
3.3 Retrieval-Augmented Factuality Scorer (RAFS)
- RAFS splits each reasoning path into statements, retrieves evidence for generated queries, labels statements Supported or Not Supported, and scores paths by the proportion supported.It adapts SAFE using LLaMA 3.1 and a corpus index containing Wikipedia and medical resources.Source: §3.3 · p. 6
4 Results
RARE is evaluated on medical and commonsense reasoning using LLaMA 3.2 3B, LLaMA 3.1 8B, and LLaMA 3.1 70B Instruct models. Across both benchmark groups, it outperforms baselines, with further analyses of components, reasoning paths, scorer alignment, and token costs.
4 Results
- The evaluation uses three Instruct models: LLaMA 3.2 3B, LLaMA 3.1 8B, and LLaMA 3.1 70B.The paper may omit “Instruct” when referring to these models.Source: §4 · p. 6
4.1 Performance on Medical Reasoning tasks
- RARE consistently outperforms baselines on medical and commonsense reasoning; at 70B, it surpasses GPT-4o on two medical benchmarks and all three commonsense benchmarks.On LLaMA3.1 70B, RARE scores 87.43% vs. 85.55% on MedQA and 90.91% vs. 90.45% on MMLU-Medical; it reaches 85.74% on StrategyQA, 86.98% on CommonsenseQA, and 92.66% on PIQA, surpassing GPT-4o.Source: §4.1 · pp. 6–7, §4.2 · p. 7, Table 2 · p. 7
- RARE-enhanced LLaMA3.1 8B beats rStar by 5.17% on MedQA, 2.19% on MedMCQA, and 2.47% on MMLU-Medical.RARE consistently outperforms baselines across the three model sizes reported.Source: §4.1 · pp. 6–7
4.2 Performance on Commonsense Reasoning
- On LLaMA3.1 8B, RARE gains over rStar are 6.45% on StrategyQA, 4.26% on CommonsenseQA, and 2.87% on PIQA.Its largest reported improvement over CoT is 10.19% on StrategyQA, compared with 7.22% on CommonsenseQA.Source: §4.2 · p. 7, Table 2 · p. 7
4.3 Ablation Study
- In a 250-sample MedQA ablation, full RARE reaches 74.8% accuracy, above rStar’s 70.0% and configurations adding RAFS, A6, A7, or both actions.RAFS alone reaches 71.4%; combining A6 and A7 reaches 73.2%.Source: §4.3 · p. 7, Table 3 · p. 7
4.4 Common Reasoning Paths Patterns
- Successful MedQA and StrategyQA paths include A1 → A2, A3 → A2, and A1 → A6, while more complex sequences also appear.A3 → A7 → A3 appears among successful MedQA paths; StrategyQA also shows noticeable use of A6 and A2.Source: §4.4 · pp. 7–8
4.5 Human Evaluation of the Retrieval-Augmented Factuality Scorer
- RAFS agrees with medical experts at an average 85.14%, with 87.84% agreement for Annotator 1 and 82.43% for Annotator 2.The two experts’ inter-annotator agreement is 86.49%.Source: §4.5 · p. 8
4.6 Computational Analysis
- RARE averages 119.9k tokens per example versus rStar’s 92.6k and CoT’s approximately 400, while improving accuracy by 9.05% over CoT on average.The reported average accuracy improvement over rStar is 4.52%.Source: §4.6 · p. 8, Figure 6 · p. 8
5 Related Work
Prior work explores prompting operations and MCTS for LLM reasoning, while medical RAG develops retrieval strategies for knowledge-intensive tasks. RARE combines retrieval and non-retrieval actions in a flexible planning framework.
5.1 Prompting LLMs to reason
- Prompting methods include planning, decomposition, abstraction, and program-based reasoning, while rStar extends MCTS with five reasoning operations and an LLM discriminator.Other LLM applications of MCTS include word- or sentence-level search, hierarchical planning, and unified reasoning and action spaces.Source: §5.1 · pp. 8–9
5.2 Medical RAG
- Medical RAG work includes query rewriting and iterative retrieval; unlike SeRTS’s depth-first, one-query-at-a-time strategy, RARE supports breadth-first multi-query search and retrieval-free reasoning.RARE integrates RAG and non-RAG operations, adapting between retrieval-driven and retrieval-free reasoning according to task complexity.Source: §5.2 · p. 9
6 Conclusion
RARE aims to improve LLM reasoning accuracy and factual reliability through retrieval-augmented actions and factuality scoring. Experiments on medical and commonsense benchmarks show effectiveness and competitive performance with proprietary systems.
6 Conclusion
- RARE is an autonomous agent that requires no additional training or fine-tuning and uses real-time retrieval and reasoning across tasks and datasets.The conclusion describes it as a scalable solution for knowledge-intensive reasoning tasks.Source: §6 · p. 9
- Experiments on medical and commonsense benchmarks demonstrate RARE’s effectiveness and its potential to bridge open-source models with state-of-the-art proprietary systems.The conclusion frames its target scope as tasks requiring logical coherence and factual integrity.Source: §6 · p. 9
7 Limitations
RARE’s limitations concern evaluation scope, reasoning-path selection, and search guidance. The authors identify broader model testing and more robust path optimization as future directions.
7 Limitations
- RARE has been evaluated only with open-source models such as LLaMA 3.1, not larger proprietary models such as GPT-4, because iterative retrieval and reasoning require costly API calls.The framework is designed to be model-agnostic and could be applied to proprietary models if resources permit.Source: §7 · p. 10
- RARE selects a single reasoning path leading to a correct answer but does not necessarily optimize for the shortest or most robust path, such as one with the highest model confidence.Future work could explore reward functions to prevent reward hacking and improve selection of reliable paths.Source: §7 · p. 10
- RARE currently uses Monte Carlo Tree Search to explore action paths without a trained reward model dynamically guiding the search.The authors suggest reward models or alternative optimization strategies as possible extensions to improve reasoning quality and efficiency.Source: §7 · p. 10
8 Ethics Statement
The ethics statement emphasizes supervised clinical use, attention to bias and generalizability, and human expertise in healthcare deployment. Human evaluators were medical experts who participated voluntarily and were fairly compensated.
8 Ethics Statement
- RARE is not intended as a standalone diagnostic or treatment tool, and clinical workflow integration requires qualified-practitioner supervision for patient safety and ethical use.The authors describe its outputs as supplementary to human expertise and caution against overreliance on AI in healthcare.Source: §8 · p. 10
- Models used with RARE may reflect training-data biases that could affect fairness and reliability in sensitive medical contexts.The authors call for rigorous bias and fairness audits to minimize potential harm.Source: §8 · p. 10
- RARE has primarily been evaluated on English-language, text-based general and medical QA, leaving non-English and multimodal applicability untested.The authors recommend extending the framework to diverse linguistic and cultural contexts.Source: §8 · p. 10
- The authors frame ethical deployment as requiring guidelines that keep RARE supplementary to human expertise and preserve human decision-making in healthcare.The stated aim is to ensure AI enhances rather than replaces healthcare decision-making.Source: §8 · p. 10
- The study’s human evaluation involved two medical experts who passed the USMLE; they were paid $40/hour and participated voluntarily after being informed about their role.The authors state that annotators were compensated fairly.Source: §8 · p. 10
A.1 Implementation Details
Implementation settings vary the number of MCTS rollouts by model size and use a same-backbone factuality scorer. Inference, including factuality scoring, is parallelized.
A.1 Implementation Details
- MCTS uses 4 rollouts for LLaMA 3.2 3B and LLaMA 3.1 8B, and 2 for LLaMA 3.1 70B, balancing inference effectiveness and computational efficiency.The lower rollout count for the larger model is particularly motivated by its higher inference costs.Source: §A.1 · p. 14
- RAFS uses the generator’s backbone model for factuality evaluation, and all inference processes, including scoring, are parallelized.The passage gives LLaMA 3.1 3B generation paired with LLaMA 3.2 3B RAFS as an example and says the setup maintains efficient inference.Source: §A.1 · p. 14
A.2 Evaluation tasks
RARE is evaluated on medical and commonsense reasoning tasks. The medical set emphasizes domain-specific clinical knowledge, while the commonsense set includes implicit, physical, and general knowledge reasoning.
A.2 Evaluation tasks
- Medical reasoning evaluation uses MedQA-USMLE, MedMCQA, and MMLU-Medical, covering USMLE-based questions, multiple-choice clinical knowledge, and diverse medical topics.The passage characterizes these datasets as requiring complex, domain-specific reasoning.Source: §A.2 · p. 14
- Commonsense evaluation uses StrategyQA, CommonsenseQA, and PIQA for implicit reasoning strategies, commonsense knowledge, and physical interactions, respectively.The passage notes that StrategyQA requires complex implicit reasoning, while the other tasks may benefit from advanced reasoning to a lesser extent.Source: §A.2 · p. 14
A.3 Baselines
The appendices describe RARE’s baselines and retrieval resources, its MCTS and factuality-scoring components, and prompts for its reasoning actions.
A.3 Baselines
- Compared methods include step-by-step CoT, self-consistency voting, rStar, and medical retrieval systems MedRAG and i-MedRAG.The baselines are described as approaches for coherent reasoning, voting across paths, MCTS exploration, or retrieval-based medical reasoning.Source: §A.3 · p. 14
A.4 Retrieval Model and Corpus
- MedCorp retrieval draws on PubMed, StatPearls, medical textbooks, and Wikipedia; ColBERT ranks passages through fine-grained token-level matching.The corpus includes a 23.9-million-article PubMed subset and 9,330 StatPearls articles.Source: §A.4 · pp. 14–15
A.5 Monte Carlo Tree Search (MCTS)
- MCTS iteratively selects, expands, simulates rollouts, and backpropagates outcomes; UCT balances exploration and exploitation during selection.UCT uses average reward, parent and child visit counts, and a constant C that controls the exploration–exploitation balance.Source: §A.5 · p. 15
A.6 Evaluation of Reasoning Candidates with RAFS
- RAFS scores reasoning paths by alignment with retrieved evidence: the correct ketotifen path scores 1.0, versus 0.625 and 0.6 for alternatives.The figure illustrates how factuality scores distinguish candidate paths for a medical question.Source: §A.6 · p. 16, Figure 7 · p. 17
A.7 Prompt of Different Actions in MCTS
- A6 generates search queries from the question and retrieves documents; A7 re-answers a question or sub-question using relevant retrieved information.The prompts illustrate query generation and a response format that includes relevant documents.Source: §A.7 · pp. 25–26