PaperlayerView exact revision on arXiv

Source-linked AI summary

CoSearch: Joint Training of Reasoning and Document Ranking via Reinforcement Learning for Agentic Search

Hansi Zeng, Liam Collins, Bhuvesh Kumar, Neil Shah, Hamed Zamani

arXiv:2604.17555v2cs.AIcs.CLcs.IR

TL;DR

Existing agentic search approaches leave retrieval fixed, and oracle retrieval shows a gap of up to +26.8% relative F1 across seven QA benchmarks. COSEARCH jointly trains the reasoning agent and a generative ranker, consistently improving over strong baselines across seven benchmarks.

  • Problem

    Existing RL-trained search agents keep retrieval fixed, while oracle retrieval improves average F1 by up to +26.8% across seven QA benchmarks.

    Source: §1 · p. 2

  • Method

    COSEARCH jointly trains a reasoning agent and a generative ranker with GRPO, using semantic grouping and a composite ranking-quality and trajectory-outcome reward.

    Source: §1 · p. 2

  • Results

    COSEARCH consistently improves across seven QA benchmarks, achieving +6.6% relative F1 over Search-R1 with a 7B agent.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The results indicate joint training of reasoning and retrieval is feasible and strongly performant, pointing to a key ingredient for future search agents.

    Source: Conclusion · p. 9

    Naively sampling H ranker responses at each of m steps would require G · H^m trajectories, an intractable rollout blowup.

    Source: §3.3 · p. 5

Abstract

from arXiv · show

Agentic search -- the task of training agents that iteratively reason, issue queries, and synthesize retrieved information to answer complex questions -- has achieved remarkable progress through reinforcement learning (RL). However, existing approaches such as Search-R1, treat the retrieval system as a fixed tool, optimizing only the reasoning agent while the retrieval component remains unchanged. A preliminary experiment reveals that the gap between an oracle and a fixed retrieval system reaches up to +26.8% relative F1 improvement across seven QA benchmarks, suggesting that the retrieval system is a key bottleneck in scaling agentic search performance. Motivated by this finding, we propose CoSearch, a framework that jointly trains a multi-step reasoning agent and a generative document ranking model via Group Relative Policy Optimization (GRPO). To enable effective GRPO training for the ranker -- whose inputs vary across reasoning trajectories -- we introduce a semantic grouping strategy that clusters sub-queries by token-level similarity, forming valid optimization groups without additional rollouts. We further design a composite reward combining ranking quality signals with trajectory-level outcome feedback, providing the ranker with both immediate and long-term learning signals. Experiments on seven single-hop and multi-hop QA benchmarks demonstrate consistent improvements over strong baselines, with ablation studies validating each design choice. Our results show that joint training of the reasoning agent and retrieval system is both feasible and strongly performant, pointing to a key ingredient for future search agents.

1 Introduction

The introduction identifies fixed retrieval as a bottleneck in agentic search and motivates jointly training the reasoning agent and retrieval system. COSEARCH combines a trainable generative ranker with a reasoning agent and reports consistent improvements across QA benchmarks.

1 Introduction

  • Agentic search iteratively combines reasoning, queries, and retrieved information, while common RL approaches leave the retrieval system unchanged.The agent learns when and what to search from final-answer correctness rewards.Source: §1 · p. 1
  • +14.7% and +26.8% average F1 gains from oracle retrieval for 7B and 3B agents, respectively, indicate a retrieval bottleneck across reasoning capacities.The experiment promotes gold-answer-matching documents at each retrieval step and leaves rankings unchanged when none match.Source: §1 · p. 2
  • COSEARCH jointly trains a main reasoning agent and a generative ranker with GRPO, while a fixed dense retriever supplies candidate documents.The ranker selects and reranks a subset of candidates for the agent to observe.Source: §1 · p. 2
  • Semantic grouping clusters ranker sub-queries by token-level F1 similarity for GRPO, while a composite reward combines Hit@k relevance with final-answer correctness.Both reward components use rule-based metrics and require no expensive LLM annotations.Source: §1 · p. 2
  • +6.6% relative F1 over Search-R1 with a 7B agent, with COSEARCH improving consistently across seven QA benchmarks.The evaluation spans single-hop and multi-hop settings.Source: §1 · pp. 2–3

2 Related Work

Related work covers agentic search and reinforcement learning for search agents, alongside approaches that adapt retrieval for language-model generation. COSEARCH differs by jointly applying reinforcement learning to a generative ranker and reasoning agent.

Agentic search and deep research

  • Agentic search interleaves reasoning and retrieval, and RLVR trains search agents from answer correctness (Jin et al., 2025; Li et al., 2025a; Song et al., 2025; Li et al., 2025b; Sun et al., 2025b; Guan et al., 2025; Zeng et al., 2026).The section also describes deep research agents coupling large reasoning models with web search.Source: §2 · p. 3

Training retrievers for LLMs

  • Agentic-R trains a dense retriever for agentic search with LLM-annotated labels and iterative contrastive learning; COSEARCH instead uses end-to-end RL with rule-based rewards.Prior retriever-training work also includes language-model supervision and LLM listwise ranking.Source: §2 · p. 3

RL for tool-augmented LLMs

3 Methodology

COSEARCH jointly optimizes a multi-step reasoning agent and a generative ranker while keeping first-stage dense retrieval fixed. It adapts GRPO to the ranker’s varying sub-query inputs and combines ranking and answer signals in its reward.

3.1 Problem Formulation

  • COSEARCH keeps the first-stage dense retriever fixed but trains a generative ranker to select and rank K of its N candidate documents as the agent’s observation.The framework jointly optimizes the ranker and main agent, making the observation depend on the trainable ranker.Source: §3.1 · p. 4

3.2 Main Agent Optimization

  • The main agent is trained on complete reasoning trajectories with GRPO; its token loss masks retrieved-document tokens, and its reward combines answer F1 with a format penalty.The format penalty is α = 0.2 when the trajectory violates the required ReAct format.Source: §3.2 · pp. 4–5

3.3 Ranker Optimization

  • Token-level F1 clusters semantically equivalent sub-queries into valid GRPO groups, filters groups below kmin, and reuses existing rollouts instead of dedicated ranker sampling.The grouping addresses varying ranker inputs across trajectories without additional sampling.Source: §3.3 · p. 5

3.4 Ranker Reward Design

  • The ranker reward uses Hit@k relevance and answer-level token F1: relevance alone for weak rankings, both signals for strong rankings, and answer F1 when candidates lack a gold-answer match.Trajectories that violate the required format are excluded from ranker reward computation and optimization.Source: §3.4 · p. 6

4 Experiments

The experiments evaluate COSEARCH across seven QA benchmarks and compare joint training with strong baselines and component variants. COSEARCH outperforms the baselines, while ablations and training analyses examine its design and behavior.

4.1 Experimental Setup

  • The evaluation covers seven single-hop and multi-hop QA benchmarks, using Search-R1 test splits and token-level F1 for consistent comparison.Baselines include Direct Inference, CoT, RAG, Search-o1, Search-R1, ZeroSearch, Retrieval Only, and Fixed Ranker.Source: §4.1 · p. 6

4.2 Main Results

  • COSEARCH achieves average F1 of 0.568 with the 7B agent, outperforming Search-R1 by a relative 6.6%, Retrieval Only by 4.0%, and Fixed Ranker by 2.7%.The improvement holds across all seven benchmarks.Source: §4.2 · p. 7, Table 2 · p. 7
  • COSEARCH closes approximately 27% of the retrieval gap for the 7B agent and 26% for the 3B agent compared with the oracle upper bound.The 3B main agent paired with a 7B ranker gains 7.0% over Retrieval Only.Source: §4.2 · pp. 7–8

4.3 Ablation Study

  • Removing trajectory-level reward reduces average F1 by 1.4%, while removing ranking-quality reward reduces it by 2.8%.Replacing Hit@k with nDCG@k reduces performance by 2.1%; substituting an LLM-as-judge for the composite reward yields a 1.8% drop.Source: §4.3 · p. 8
  • Removing semantic grouping and filtering causes the largest single-component degradation: 3.3% relative, from 0.568 to 0.549.The passage attributes this to noisy advantage estimates from grouping semantically heterogeneous sub-queries.Source: §4.3 · p. 8
  • A shared backbone degrades F1 by 1.6%, a cross-encoder ranker by 2.9%, and scaling the ranker to 3B causes training to diverge.The passage attributes the shared-backbone drop to conflicting gradients and the 3B failure to invalid permutation outputs over 50 candidates.Source: §4.3 · p. 8

4.4 Analysis

  • Figure 2 shows ranking quality rising alongside downstream F1, while COSEARCH and Oracle require fewer search turns.COSEARCH’s ranker starts below the Fixed Ranker but steadily surpasses it during training.Source: §4.4 · p. 8, Figure 2 · p. 9
  • When K falls from 5 to 1, COSEARCH degrades by 5.3%, compared with 8.8% for Retrieval Only and 7.2% for Fixed Ranker.COSEARCH at K=3 (0.561) surpasses Fixed Ranker at K=10 (0.559); increasing K from 5 to 10 yields diminishing returns across methods.Source: §4.4 · p. 9, Table 4 · p. 9

Conclusion

COSEARCH jointly trains a reasoning agent and generative document ranker, addressing retrieval quality as a bottleneck in agentic search. Across seven QA benchmarks, it improves on Search-R1, while oracle retrieval gaps reach +26.8% relative F1.

Conclusion

  • +6.6% relative F1 over Search-R1 with a 7B agent and +10.0% with a 3B agent; oracle retrieval gaps reach +26.8% relative F1.The conclusion reports consistent improvements across seven QA benchmarks.Source: Conclusion · p. 9
  • COSEARCH jointly trains the main reasoning agent and a generative document ranker via reinforcement learning.Semantic grouping enables ranker GRPO without additional rollouts, while a composite reward combines immediate ranking-quality and long-term trajectory-level signals.Source: Conclusion · p. 9

A.1 Main Agent Prompt

  • The main agent uses a ReAct-style loop, searching when more evidence is needed and otherwise returning a concise final answer.Its prompt specifies a reason block followed by either a tool call or an answer.Source: §A.1 · pp. 15–16

A.2 Ranker Prompt

  • The ranker considers the original question, current sub-query, and candidate documents, then selects and orders the K most relevant documents.Its output separates relevance analysis from the ranked document indices.Source: §A.2 · pp. 16–17

A.3 Ranker Input/Output Examples

  • The examples contrast an answer-containing candidate set with composite reward 2.00 and an answer-absent set where the ranker relies on rmain and receives reward 1.00.The first example reports rrel = 1.00 and rmain = 1.00; the second reports rrel = 0.00 and rmain = 1.00.Source: §A.3 · pp. 17–18

B Two-Stage Retrieval System

At each search step, a fixed dense retriever supplies candidate documents and a generative ranker selects the ranked documents returned to the reasoning agent.

B Two-Stage Retrieval System

  • The two-stage system takes the original query, the agent’s sub-query, and corpus, then returns an observation of K ranked documents.The fixed dense retriever recalls N candidates before the generative ranker reranks them.Source: B Two-Stage Retrieval System · p. 18
  • The fixed retriever obtains N candidates for the sub-query, and the algorithm records whether the gold answer is among them for the composite reward.The ranker prompt is built from the original query, sub-query, candidate documents, and ranking size K.Source: B Two-Stage Retrieval System · p. 18
  • The system parses the ranker output to extract the top-K documents and returns that set as the observation.The ranker produces an explicit ranked list.Source: B Two-Stage Retrieval System · p. 18

C Semantic Grouping Algorithm

Semantic grouping forms ranker GRPO groups from sub-queries that vary across trajectories, without requiring additional rollouts.

C Semantic Grouping Algorithm

  • Ranker calls are first split according to whether their candidate sets contain the gold answer, then clustered within each split by token-level F1 similarity.This two-level grouping addresses the varying sub-queries across rollouts.Source: C Semantic Grouping Algorithm · p. 19
  • The experiments use G=8 rollouts with up to six search steps per trajectory.Figure 4 instead shows a simplified example with G=5 rollouts and up to five steps.Source: C Semantic Grouping Algorithm · p. 19
  • Figure 4 depicts sub-queries grouped by token-level F1 similarity, with gray singleton groups smaller than kmin discarded before optimization.The illustration shows G=5 trajectories with up to five search steps each.Source: Figure 4 · p. 19, C Semantic Grouping Algorithm · p. 19

D Additional Implementation Details

The implementation uses a 51,200-question training set drawn from four QA datasets and compares against a separately trained fixed-ranker baseline.

D Additional Implementation Details

  • The RL training set contains 51,200 questions from Natural Questions, HotpotQA, Musique, and 2WikiMultiHopQA official training splits.The respective counts are 20,480, 14,220, 9,000, and 7,500 questions.Source: D Additional Implementation Details · p. 19
  • Evaluation uses each benchmark’s official test split, or its evaluation split when no test split is available.Source: D Additional Implementation Details · p. 19
  • The Fixed Ranker baseline uses 92,160 queries after filtering 50K MS MARCO and 50K Natural Questions queries for those with a relevant document in the top 50.It is trained with RL using the same Hit@{1, 3, 5} composite reward and generative listwise output format as COSEARCH’s jointly trained ranker.Source: D Additional Implementation Details · p. 19

E Semantic Grouping Examples

Semantic grouping clusters rollout sub-queries by token-level F1 similarity; the examples show how multi-hop questions can form distinct clusters while simpler questions yield near-paraphrases in one cluster.

E Semantic Grouping Examples

  • Multi-hop questions can produce multiple semantic clusters, while simpler questions tend to yield one large cluster of near-paraphrases.The examples show clusters corresponding to distinct sub-questions in a two-hop chain.Source: E Semantic Grouping Examples · p. 19
  • At training step 100, sub-queries from G=8 rollout trajectories are clustered by token-level F1 similarity using threshold δ=0.8.Tables 5 and 6 show the resulting clusters; column C gives each sub-query’s cluster index.Source: Table 5 · p. 20, Table 6 · p. 21
  • The procedure splits ranker calls by answer availability, greedily clusters normalized sub-queries when token-level F1 meets δ, then removes groups smaller than kmin and assigns UIDs.The two splits distinguish calls with and without answer availability.Source: E Semantic Grouping Examples · p. 22

F Oracle Retrieval Construction

The oracle retrieval procedure modifies observations online during rollout when retrieved candidates contain the gold answer, and otherwise leaves the default ranking unchanged.

F Oracle Retrieval Construction

  • At each rollout step, the agent generates a thought and sub-query, and the dense retriever returns candidate documents for that sub-query.The oracle trajectory is constructed online, so each modified observation conditions subsequent reasoning and sub-queries.Source: F Oracle Retrieval Construction · p. 22
  • When candidate documents contain the gold answer, the oracle promotes those documents and fills remaining observation slots with the highest-ranked nonmatching documents.The modified observation is then passed to the agent, which continues reasoning until it answers.Source: F Oracle Retrieval Construction · p. 22
  • When no retrieved candidate contains the gold answer, the oracle uses the default retriever ranking unchanged.This applies to intermediate sub-queries without a gold-answer match.Source: F Oracle Retrieval Construction · p. 22

G Search Turn Distribution

CoSearch and Oracle concentrate questions in fewer search turns than the Fixed Ranker, which sends more questions to three turns and requires more retrieval rounds on average.

G Search Turn Distribution

  • 91% of questions for COSEARCH and 87% for Oracle take 1–2 turns, while Fixed Ranker sends 55% to 3 turns; average turns are 2.62 for Fixed Ranker versus 1.67 for COSEARCH.Table 7 reports the validation-step-100 distribution of questions by search turns and the mean turns per question.Source: G Search Turn Distribution · p. 22, Table 7 · p. 23
  • Higher-quality ranking reduces the number of retrieval steps needed to resolve a question.The reported turn distribution contrasts COSEARCH and Oracle with Fixed Ranker.Source: G Search Turn Distribution · pp. 22–23
Loading 2604.17555v2…