Preprint
Reinforcement Learning

Improving RAG through Multi-Agent RL

Yiqun Chen, Ling Yan, Weiwei Sun, Xinyu Ma, Yi Zhang, Shuaiqiang Wang, Dawei Yin, Yiming Yang, Jiaxin Mao
January 25, 2025Neural Information Processing Systems46 citations
arXivPDFDOI

46

Citations

3

Influential Citations

Neural Information Processing Systems

Venue

2025

Year

Abstract

Retrieval-augmented generation (RAG) is widely utilized to incorporate external knowledge into large language models, thereby enhancing factuality and reducing hallucinations in question-answering (QA) tasks. A standard RAG pipeline consists of several components, such as query rewriting, document retrieval, document filtering, and answer generation. However, these components are typically optimized separately through supervised fine-tuning, which can lead to misalignments between the objectives of individual components and the overarching aim of generating accurate answers. Although recent efforts have explored using reinforcement learning (RL) to optimize specific RAG components, these approaches often focus on simple pipelines with only two components or do not adequately address the complex interdependencies and collaborative interactions among the modules. To overcome these limitations, we propose treating the complex RAG pipeline with multiple components as a multi-agent cooperative task, in which each component can be regarded as an RL agent. Specifically, we present MMOA-RAG, Multi-Module joint Optimization Algorithm for RAG, which employs multi-agent reinforcement learning to harmonize all agents'goals toward a unified reward, such as the F1 score of the final answer. Experiments conducted on various QA benchmarks demonstrate that MMOA-RAG effectively boost the overall performance of the pipeline and outperforms existing baselines. Furthermore, comprehensive ablation studies validate the contributions of individual components and demonstrate MMOA-RAG can be adapted to different RAG pipelines and benchmarks.

Analysis

Why This Paper Matters

Retrieval-augmented generation (RAG) has become a standard approach to enhance LLM factuality by integrating external knowledge. However, typical RAG pipelines are composed of multiple components—query rewriting, document retrieval, filtering, and answer generation—that are often optimized independently via supervised fine-tuning. This leads to misalignment between component-level objectives and the ultimate goal of generating accurate answers. While prior RL-based efforts have targeted individual components, they often simplify the pipeline to two stages or ignore interdependencies. This paper addresses a critical gap by proposing a multi-agent RL framework that jointly optimizes all components, acknowledging the cooperative nature of the pipeline.

The significance lies in shifting from isolated optimization to a holistic, reward-driven approach. By treating each component as an RL agent with a shared reward, MMOA-RAG enables the entire system to learn collaborative behaviors that directly improve final answer quality. This is particularly relevant as RAG systems become more complex and deployed in real-world QA applications where hallucinations are costly.

Technical Contributions

  • Multi-Agent Formulation: The RAG pipeline is modeled as a multi-agent cooperative task, where each component (query rewriter, retriever, filter, generator) is an RL agent. This is a novel perspective that captures interdependencies.
  • Unified Reward Design: All agents are trained to maximize a single reward, such as the F1 score of the final answer, ensuring that each component's actions contribute to the overall objective.
  • Joint Optimization Algorithm: MMOA-RAG introduces a training algorithm that coordinates the agents' policies, likely using techniques like centralized training with decentralized execution or shared reward shaping.
  • Adaptability: The framework is designed to be pipeline-agnostic, allowing integration with different RAG architectures and benchmarks.

Results

The paper reports experiments on multiple QA benchmarks, demonstrating that MMOA-RAG outperforms existing baselines. While specific numerical metrics are not detailed in the abstract, the consistent improvement across benchmarks indicates robust gains. Ablation studies validate the contribution of each component, showing that removing any agent degrades performance, thus confirming the importance of joint optimization. The adaptability to different pipelines further strengthens the generalizability of the approach.

Significance

MMOA-RAG represents a significant step forward in RAG optimization, moving beyond component-wise SFT to a cooperative multi-agent RL paradigm. This could inspire similar approaches for other multi-stage NLP systems, such as summarization pipelines or dialogue systems. By aligning component objectives with end-task metrics, it offers a path to more reliable and factual LLM outputs, which is crucial for real-world deployment. The work also highlights the potential of multi-agent RL in complex language tasks, opening new research directions.

multi-agent-rlretrieval-augmented-generationquestion-answeringreinforcement-learningllmpipeline-optimization
Large Language ModelsFine TuningReinforcement LearningRetrieval Augmented GenerationAlignmentBenchmarks

Related Papers