RARE: Retrieval-Augmented Reasoning Enhancement for Large Language Models (2412.02830v4)

Abstract: 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 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.

• The paper introduces a novel method that integrates retrieval-based actions into MCTS to generate structured and accurate reasoning paths.
• It replaces traditional discriminators with a factuality scorer that assesses evidence to ensure coherent, fact-supported answers.
• RARE demonstrates scalable improvements on benchmarks like MedQA and CommonsenseQA, outperforming established models such as GPT-4.

RARE: Enhancing Reasoning in LLMs through Retrieval Augmentation

The research paper introduces RARE (Retrieval-Augmented Reasoning Enhancement), a methodology designed to improve the performance of LLMs in complex reasoning tasks, specifically targeting domains such as commonsense and medical reasoning. The paper details how RARE builds upon the existing rStar framework and demonstrates notable improvements in reasoning accuracy and factual integrity, without the need for extensive model fine-tuning.

RARE leverages a Monte Carlo Tree Search (MCTS) framework augmented with retrieval-based actions, allowing the model to generate structured reasoning paths by integrating external information dynamically. This integration is particularly beneficial for tasks requiring extensive domain-specific knowledge, such as medical question answering (QA), where factual accuracy and context are paramount.

Key Contributions

1. Novel Retrieval-Augmented Actions: RARE introduces innovative actions within the MCTS, specifically designed to generate search queries and retrieve relevant documents that enrich the reasoning process. This is realized through two critical actions:
   • A6: For search query generation and information retrieval, which supports LLMs in forming contextually relevant answers.
   • A7: For refining and re-answering sub-questions using retrieved information, enhancing both the accuracy and coherence of the reasoning trajectory.
2. Retrieval-Augmented Factuality Scorer (RAFS): Replacing the traditional discriminator used in rStar, RAFS assesses the factual reliability of each reasoning path by analyzing individual statements against retrieved evidence. This factuality scorer assigns scores to ensure the selected reasoning path is logically coherent and factually supported.
3. Scalable Framework: RARE operates effectively with open-source LLMs such as LLaMA, demonstrating competitive performance against top-tier models like GPT-4. This scalability underscores RARE's potential as a viable solution to enhance reasoning capabilities across diverse domains where accuracy is critical.

Experimental Results

RARE was tested on medical QA tasks such as MedQA, MedMCQA, and MMLU-Medical, and on commonsense reasoning benchmarks including StrategyQA and CommonsenseQA, using multiple model sizes (e.g., LLaMA 3.2 3B and LLaMA 3.1 70B). The results were substantial; RARE consistently enhanced performance over baseline methods including Chain of Thought (CoT) and Self-Consistency.

For example, RARE-enabled LLaMA models achieved superior accuracy on MedQA and MMLU-Medical benchmarks surpassing even well-established models like GPT-4. These improvements illustrate RARE's robust capability to address complex, domain-specific reasoning tasks effectively.

Implications and Future Directions

The development of RARE signifies a crucial step forward in augmenting LLMs with retrieval-based reasoning capabilities, which is especially valuable in domains like healthcare. Practically, the integration of fact-checked reasoning could significantly aid clinical decision support systems, educational tools, and patient care optimization by providing precise, evidence-based insights.

Theoretically, RARE highlights the significance of retrieval-augmented methodologies in expanding the cognitive bandwidth of LLMs. Future research may explore refining reward mechanisms within the MCTS framework to optimize reasoning paths or enhancing the model's adaptability to various linguistic and multi-modal contexts.

Conclusion

RARE represents a significant advancement in the field of Natural Language Processing by coupling reasoning with retrieval dynamics to position LLMs as factual and coherent problem solvers across complex domains. This innovation not only paves the way for enhanced model performance in specialized areas like medical QA but also inspires further development of retrieval-augmented reasoning techniques in AI research. As a model-agnostic framework, RARE holds the promise of broad applicability and scalability, potentially transforming AI applications in knowledge-intensive industries.