Improving Retrieval-Augmented Generation through Multi-Agent Reinforcement Learning

TL;DR

This work addresses misalignment in retrieval-augmented generation (RAG) by modeling the RAG pipeline as a cooperative multi-agent reinforcement learning problem. It introduces MMOA-RAG, which jointly optimizes three trainable modules—Query Rewriter, Document Selector, and Answer Generator—via Multi-Agent PPO with a fixed retriever, guided by a unified final-answer reward (F1) and module-specific penalties. Through warm-started supervision and extensive ablations, the approach demonstrates consistent improvements across HotpotQA, 2WikiMultihopQA, and AmbigQA, and shows generalization across different RAG configurations and retrievers. The results highlight the benefits of coordinated optimization over independent module tuning and suggest a scalable framework for complex RAG systems with multiple interacting components.

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.

Figures (4)

• Figure 1: The overall framework of MMOA-RAG.
• Figure 2: Ablation about Optimizing Different Agents. In this figure, MMOA-RAG w/o $i$ ($i \in \{\text{QR, S, G}\}$) denote the variant where agent $i$ is excluded from the complete optimization process of multi-agent joint optimization.
• Figure 3: Ablation experiments on AmbigQA dataset. The horizontal axis represents the number of training samples, while the vertical axis denotes the shared reward $R_{\text{shared}}$ (F1 score) during the training process.
• Figure 4: The convenient approach to construct the SFT data for Selector.