s3: You Don't Need That Much Data to Train a Search Agent via RL

TL;DR

s3 introduces a modular RL framework that trains only the search component of a retrieval-augmented system, using a novel Gain Beyond RAG (GBR) reward to maximize downstream generation quality while keeping the generator frozen. By optimizing over a multi-turn search-select-serve loop with a fixed top-k baseline, s3 achieves substantial gains with as few as 2.4k training examples, surpassing baselines trained on orders of magnitude more data. The GenAcc-based reward aligns well with human judgments and transfers across general and medical QA benchmarks, demonstrating strong data and compute efficiency and robust generalization. This approach decouples retrieval optimization from generation, enabling practical, model-agnostic improvements to RAG systems and reducing data requirements for training effective search agents.

Abstract

Retrieval-augmented generation (RAG) systems empower large language models (LLMs) to access external knowledge during inference. Recent advances have enabled LLMs to act as search agents via reinforcement learning (RL), improving information acquisition through multi-turn interactions with retrieval engines. However, existing approaches either optimize retrieval using search-only metrics (e.g., NDCG) that ignore downstream utility or fine-tune the entire LLM to jointly reason and retrieve-entangling retrieval with generation and limiting the real search utility and compatibility with frozen or proprietary models. In this work, we propose s3, a lightweight, model-agnostic framework that decouples the searcher from the generator and trains the searcher using a Gain Beyond RAG reward: the improvement in generation accuracy over naive RAG. s3 requires only 2.4k training samples to outperform baselines trained on over 70x more data, consistently delivering stronger downstream performance across six general QA and five medical QA benchmarks.

Figures (13)

• Figure 1: Training Data vs Averaged Performance across six general and five medical QA Datasets (tested with Claude-3-Haiku as the generator LLM).
• Figure 2: RAG has progressed from fixed or supervised retrieval to RL-based agentic methods. While prior work trains retrieval or generation jointly, s3 focuses solely on the searcher, improving generation without tuning the generator LLM.
• Figure 3: Decomposition of Agentic RAG. End-to-end approaches fine-tune the entire model using the entire generation accuracy, making it difficult to isolate the contribution of search. In contrast, s3 freezes the generator and trains only the searcher with Gain Beyond RAG (GBR), a novel reward that quantifies the added value of retrieved context over naïve RAG, enabling modular, efficient optimization.
• Figure 4: Overview of the s3 framework. The search agent iteratively retrieves documents, selects useful documents, and generates queries until completion. The final context $D_{s3}$ is then passed to a frozen generator LLM. The search agent is trained using Gain Beyond RAG (GBR), which quantifies improvement over naïve top-$k$ retrieval from the original question.
• Figure 5: Reward Curves for top $k=\{3,5,8\}$ and $\#\text{turns}=\{3,4\}$. The maximum selection is kept as 3.