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

1. Problem Statement

Agentic Search involves training Large Language Models (LLMs) to iteratively reason, issue search queries, and synthesize retrieved information to answer complex questions. While Reinforcement Learning (RL) has successfully trained agents to decide when and what to search, existing approaches (e.g., Search-R1) treat the retrieval system as a fixed tool. The agent learns to use the tool, but the tool (the retriever/ranker) does not adapt to the agent's specific reasoning patterns.

The authors identify a critical bottleneck: retrieval quality.

• The Gap: A preliminary "Oracle Retrieval" experiment revealed that if the retrieval system were perfect (promoting gold-answer documents to the top), performance could improve by +14.7% for a 7B-parameter agent and +26.8% for a 3B-parameter agent across seven QA benchmarks.
• The Issue: Most errors in agentic search stem not from the reasoning process but from the retrieval system failing to return relevant documents, forcing the agent to retry or draw incorrect conclusions from noise.
• The Goal: Can we jointly optimize the retrieval system alongside the reasoning agent so they learn to complement each other?

2. Methodology: COSEARCH

COSEARCH is a framework that jointly trains a multi-step reasoning agent (Main Agent) and a generative document ranker via Reinforcement Learning, specifically using Group Relative Policy Optimization (GRPO).

System Architecture

The retrieval system is decomposed into two stages:

1. Fixed Dense Retriever: A standard retriever (e.g., E5) that fetches a large candidate set ($N=50$) of documents.
2. Trainable Generative Ranker: An LLM that takes the candidate set and the agent's sub-query, reasons about relevance, and selects/reranks a subset ($K=5$) for the Main Agent to observe.
3. Main Agent: Operates in a ReAct-style loop, generating thoughts and search queries based on the ranked documents provided by the ranker.

Both the Main Agent and the Ranker are optimized simultaneously using GRPO, with rewards derived from the correctness of the final answer.

Key Technical Challenges & Solutions

Challenge 1: Applying GRPO to the Ranker

• Problem: GRPO requires grouping multiple responses generated from the same input prompt to compute advantage estimates. For the Main Agent, all trajectories start with the same user query. However, the Ranker receives sub-queries generated by the Main Agent, which vary across different reasoning trajectories (rollouts).
• Solution: Semantic Grouping Strategy.
  • Instead of grouping by the initial user query, the authors cluster sub-queries by token-level F1 similarity.
  • Sub-queries that are semantically equivalent (e.g., "Who played X?" vs. "X's actor") are grouped together to form valid GRPO optimization groups.
  • This allows the ranker to be trained using the trajectories already generated for the main agent, requiring zero additional rollouts.
  • Small groups (below a minimum size threshold) are discarded to ensure stable advantage estimation.

Challenge 2: Designing a Reward for the Ranker

• Problem: The ranker's impact on the final answer is indirect. A good retrieval can still lead to a wrong answer if reasoning fails, and vice versa. A simple trajectory-level reward (final answer correctness) is too noisy for the ranker. Conversely, a relevance reward (matching gold documents) fails to supervise steps where retrieved docs are useful for reasoning but don't contain the final answer.
• Solution: Composite Reward.
  The ranker reward ($r_{rank}$) combines two signals:
  1. Relevance Reward ($r_{rel}$): Based on Hit@k metrics using pseudo-relevance labels (documents containing the gold answer). This provides immediate feedback on retrieval precision.
  2. Trajectory Reward ($r_{main}$): The token-level F1 score of the final answer.
  • Logic:
    • If a relevant document exists in the candidate set but is poorly ranked ($r_{rel}$ is low), the ranker is optimized solely on $r_{rel}$.
    • If the ranking is strong ($r_{rel}$ is high), the reward includes $r_{main}$ to capture long-term utility.
    • If no relevant document exists in the candidate set, the ranker relies on $r_{main}$ as the only signal.

3. Key Contributions

1. Identification of the Retrieval Bottleneck: Quantified the performance ceiling imposed by fixed retrieval systems, showing potential gains of up to +26.8% with oracle retrieval.
2. COSEARCH Framework: A novel RL framework for the joint end-to-end training of a reasoning agent and a generative ranker.
3. Semantic Grouping for GRPO: A method to enable GRPO training for rankers with varying inputs (sub-queries) by clustering semantically similar queries, avoiding the exponential cost of dedicated rollouts.
4. Composite Reward Design: A hybrid reward signal that balances immediate retrieval quality (Hit@k) with downstream task success (Answer Correctness).

4. Experimental Results

The framework was evaluated on seven QA benchmarks (3 single-hop: PopQA, NQ, TriviaQA; 4 multi-hop: HotpotQA, 2Wiki, Musique, Bamboogle).

• Performance Gains:
  • 7B Agent: COSEARCH achieved an average F1 of 0.568, outperforming the strong baseline Search-R1 by +6.6% relative gain.
  • 3B Agent: When paired with a 7B ranker, COSEARCH achieved a +10.8% relative gain over Search-R1.
  • Comparison: It consistently outperformed "Retrieval Only" (no ranker), "Fixed Ranker" (frozen ranker), and other baselines like ZeroSearch and Search-o1.
• Closing the Gap: COSEARCH closed approximately 27% of the retrieval gap for the 7B agent and 26% for the 3B agent relative to the Oracle upper bound.
• Efficiency: COSEARCH required fewer search turns on average (1.67 turns) compared to Fixed Ranker (2.62 turns), indicating higher-quality document selection reduces the need for iterative retries.

Ablation Studies:

• Removing the semantic grouping strategy caused a 3.3% drop in F1, confirming its necessity for stable training.
• Removing the relevance reward ($r_{rel}$) caused a larger drop (2.8%) than removing the trajectory reward, highlighting the importance of direct retrieval supervision.
• Using a smaller 3B ranker caused training divergence, suggesting the ranker needs sufficient capacity to handle permutation outputs over 50 candidates.

5. Significance

This paper challenges the prevailing paradigm in agentic search where the retrieval system is treated as a static tool. By demonstrating that jointly training the reasoning agent and the retrieval system is not only feasible but significantly more effective, COSEARCH points to a new direction for future search agents.

The work proves that:

1. Retrieval quality is a primary bottleneck in RL-based agentic search.
2. End-to-end optimization of the retrieval component via RL yields substantial performance gains, especially for smaller reasoning models.
3. Efficient training strategies (like semantic grouping) can overcome the technical hurdles of applying group-based RL to components with dynamic inputs.

This approach suggests that future "Deep Research" agents will likely require co-adaptation between their reasoning logic and their information retrieval mechanisms to achieve human-level or super-human performance on complex knowledge tasks.