Skip to content

Improving Retrieval-Augmented Generation through Multi-Agent Reinforcement Learning

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

MMOA-RAG frames a multi-component RAG pipeline (query rewriter, selector, generator) as a Cooperative Multi-Agent Reinforcement Learning (Co-MARL) problem, jointly optimizing all modules toward a single shared reward (F1 score) via MAPPO. This eliminates the objective misalignment that arises when RAG components are fine-tuned independently. The approach consistently outperforms single-module RL baselines across three QA benchmarks.

Problem

Standard RAG pipelines optimize each component (query rewriting, retrieval, filtering, generation) separately via SFT, creating misalignment between per-module objectives and end-to-end answer quality. Prior RL-based RAG work either targets only one module or handles at most two-component pipelines (retriever + generator), failing to model collaborative interdependencies in richer architectures.

Method

MMOA-RAG models the RAG pipeline as a fully cooperative multi-agent system under the tuple ⟨G, O, A, R⟩ formalism. Three LLM-backed agents share parameters: Query Rewriter (decomposes the question into sub-questions), Selector (filters K=10 retrieved documents to a useful subset), and Generator (produces the final answer). The Retriever (Contriever) is frozen and treated as part of the environment. All agents receive a shared reward R_shared = F1 score of the predicted answer, augmented by lightweight per-agent penalty terms (P_QR penalizes >4 sub-questions; P_S penalizes duplicate/malformed doc IDs; P_G penalizes excessively long outputs). Optimization uses MAPPO with a clipped Actor loss (PPO-style importance sampling + GAE advantage), a Critic loss over shared value estimates, and a KL penalty against the SFT reference model. Training proceeds in two stages: (1) SFT warm-start on task-specific data for each agent; (2) MAPPO joint optimization with parameter sharing across agents.

Key Contributions

  • First formulation of a multi-component RAG pipeline (≥3 modules) as a Co-MARL problem, enabling end-to-end joint optimization.
  • MAPPO-based training with a unified shared reward (F1) and per-agent format/length penalty terms that stabilize multi-agent exploration.
  • Parameter-sharing mechanism across Query Rewriter, Selector, and Generator to reduce compute.
  • Demonstrated generalizability: MAPPO improves over SFT across all tested pipeline configurations (QR+S+G, S+G, QR+G) and retrievers (Contriever, BGE, E5).

Results

  • On HotpotQA (multi-hop): MMOA-RAG achieves Acc/EM/F1 = 39.15/36.15/48.29, vs. best baseline RetRobust at 37.69/34.60/46.49 (+1.12/+1.55/+1.80 absolute).
  • On 2WikiMultihopQA: 42.73/41.52/46.40 vs. best baseline RAG-DDR at 40.49/39.45/44.18 (+1.71/+1.79/+1.89).
  • On AmbigQA (single-hop): 38.85/34.75/48.59 vs. best baseline Rewrite-Retrieve-Read at 35.94/31.90/45.92 (+2.60/+2.20/+2.67).
  • MAPPO over SFT warm-start yields ~+3 pp F1 on multi-hop datasets and ~+2 pp on single-hop across all pipeline configurations.
  • Ablation: removing any single agent degrades performance; removing the Generator (w/o G) causes the slowest convergence and worst results on AmbigQA, indicating Generator is the most critical module for single-hop QA.

Limitations

  • The retriever is frozen (dense Contriever); end-to-end optimization including a trainable retriever is left as future work.
  • Evaluated only on Llama-3-8B-Instruct; scalability to larger or instruction-tuned-only models is unverified.
  • Reward is purely answer-quality (F1); cost, latency, and efficiency trade-offs are not modeled in the current reward design.
  • Experiments use Wikipedia as the sole corpus; applicability to domain-specific or real-time retrieval contexts is unclear.
  • The penalty values (−0.5, −1) are hand-tuned hyperparameters with limited ablation.

Relevance to Agentic AI / LLM Agents

This paper directly instantiates a multi-agent cooperative paradigm for a structured LLM workflow, treating each pipeline stage as a distinct RL agent with its own observation, action space, and reward signal — a blueprint transferable to any modular agentic pipeline beyond RAG. The use of MAPPO with parameter sharing and a shared terminal reward demonstrates how to align heterogeneous sub-agents toward a global objective without per-agent reward engineering. For agentic RAG systems (e.g., iterative or DAG-structured retrieval loops), the Co-MARL framing offers a principled alternative to prompt chaining or isolated RLHF. The future work explicitly targets directed-acyclic-graph and dynamic agentic workflows, positioning MMOA-RAG as a stepping stone toward fully autonomous, self-optimizing LLM agent pipelines.