PaperlayerView exact revision on arXiv

Source-linked AI summary

CAS: Conformalized Agentic Search via Adaptive Retrieval and Policy Weighting

Zixi Zhu, Jiayuan Su, Jian Zhang, Yu Lin, Hongwei Wang

arXiv:2608.20771v1cs.AI

TL;DR

RL-trained Search Agents face unreliable fixed retrieval, hallucinations from overconfidence, and redundant searches. CAS applies APS to adapt retrieval sets and ACI to regulate GRPO training, with experiments showing improved reasoning accuracy and fewer redundant tool invocations.

  • Problem

    RL-trained Search Agents face critical-evidence loss or noise inclusion from heuristic Top-K retrieval, while progressive RL overconfidence can cause hallucinations and redundant searches.

  • Method

    CAS applies APS for adaptive retrieval prediction sets and ACI to quantify answer confidence and penalize low-confidence trajectories in GRPO.

  • Results

    CAS significantly improves reasoning accuracy and reduces redundant tool invocations across single-hop and multi-hop QA datasets.

  • Takeaways & Limitations

    CAS provides a reliability-oriented framework for Search Agents by combining statistical retrieval coverage with confidence-aware policy optimization.

  • Takeaways & Limitations

    Validation focuses primarily on general open-domain QA, while specialized professional domains remain unexplored; the framework also relies on a strong external teacher model.

Abstract

from arXiv · show

Search Agents face a severe reliability crisis during reinforcement learning (RL) fine-tuning. Heuristic Top-K retrieval often causes critical evidence loss or noise inclusion, while over-confidence induced by progressive RL leads to hallucinated answers and redundant searches. To build highly reliable agents, we introduce Conformal Prediction (CP) and propose Conformalized Agentic Search (CAS). This framework establishes reliability guarantees on both the retrieval and training sides: on the retrieval side, an Adaptive Prediction Set (APS), a specific CP realization, translates statistical coverage into dynamic document truncation to construct prediction sets that are adaptive in size; on the training side, Adaptive Conformal Inference (ACI), a dynamic CP algorithm, dynamically constructs prediction sets with controllable coverage to quantify answer confidence, which is then used to penalize low-confidence trajectories within the Group Relative Policy Optimization (GRPO) objective, ensuring the model learns only from reliable ones. Experiments across single-hop and multi-hop QA datasets demonstrate that our framework significantly improves reasoning accuracy while drastically reducing redundant tool invocations, establishing a highly reliable and efficient agent paradigm. Our code is available at https://github.com/S1llyBird/CAS.

1 Introduction

CAS addresses reliability risks in RL-trained Search Agents by applying conformal prediction to both retrieval and training. Its APS and ACI components improve reasoning accuracy and reduce redundant tool use.

  • RL fine-tuning risks omit critical evidence or include distracting noise under fixed Top-K retrieval, while overconfidence can produce hallucinations and redundant searches.
  • CAS applies Conformal Prediction to both retrieval and training, providing statistical guarantees for Search Agent reliability.
  • Training: ACI dynamically constructs prediction sets to quantify answer reliability and penalizes low-confidence trajectories during GRPO optimization.
  • Retrieval: APS adaptively sizes retrieved prediction sets to preserve correct evidence across queries with different difficulty.
  • Results: Experiments on single-hop and multi-hop QA significantly improve reasoning accuracy and training stability while drastically reducing redundant tool invocations.

2 Conformal Prediction

Conformal Prediction constructs prediction sets with finite-sample coverage guarantees, while APS adapts set size to predictive uncertainty. ACI extends conformal inference dynamically to accommodate changing data distributions through feedback on miscoverage.

  • Conformal Prediction uses a non-conformity score and calibrated quantile to construct prediction sets with marginal coverage of at least 1−α under i.i.d. sampling.
  • Adaptive Prediction Set: APS ranks predictive probabilities and accumulates them until a calibrated threshold is exceeded, forming the prediction set at the minimum qualifying index.
  • Adaptive Prediction Set: APS produces compact sets for confident inputs and expanded sets for ambiguous inputs while maintaining the 1−α coverage guarantee.
  • Adaptive Conformal Inference: ACI dynamically varies αt to handle data streams whose underlying distribution may change over time, unlike standard CP's exchangeability assumption.
  • Adaptive Conformal Inference: ACI decreases αt after miscoverage to expand subsequent sets and increases αt after success to tighten them, preserving dynamic uncertainty quantification.

3 Methodology

CAS combines adaptive conformal retrieval with confidence-aware policy optimization to improve reliability in agentic search. APS dynamically bounds retrieved evidence, while ACI uses answer confidence to weight RL trajectories and reduce reinforcement of unreliable behavior.

  • CAS bridges LLM reasoning and external retrieval through APS for dynamic retrieval uncertainty and ACI for penalizing low-confidence trajectories.The framework interleaves internal reasoning, search actions, and final answer generation within one policy trajectory.
  • Fixed top-k retrieval can add redundant noise or truncate critical information, motivating conformal control of the retrieved evidence.Standard CP provides marginal coverage, but does not guarantee conditional coverage adapted to individual input difficulty.
  • 3.2 Retrieval Side via APS: APS converts normalized retrieval scores into a dynamically truncated document subset using a calibrated threshold and cumulative probability mass.The resulting subset is intended to adapt retrieval size to query uncertainty rather than use a fixed number of documents.
  • 3.3 Reward Design: The RL policy uses exact-match and format rewards, with γ = 0.2 penalizing format violations on correct answers while providing intermediate structural feedback on incorrect ones.The format indicators encode grammatical validity and successful generation of the answer boundary.
  • 3.4 Training Side via ACI: ACI defines trajectory non-conformity as answer-token NLL and updates confidence thresholds from smoothed recent miscoverage over a rolling calibration window.The updated threshold partitions trajectories into low-confidence and higher-confidence sets during RL training.
  • 3.4 Training Side via ACI: ACI-guided GRPO discounts positive reinforcement for unconfident lucky guesses while retaining full penalties for confidently incorrect trajectories.The weighting mechanism uses η ∈ (0, 1) for low-confidence samples and combines confidence weighting with GRPO advantages and KL regularization.

4 Experiments

CAS is evaluated across seven QA datasets, model scales, task types, ablations, and conformal-parameter settings. It consistently improves accuracy while showing that APS and ACI jointly support adaptive retrieval, confidence control, and reduced unnecessary searching.

  • Experimental Setup: Experiments cover seven open-domain QA datasets spanning single-hop and multi-hop reasoning, with NQ and HotpotQA used to construct the mixed RL training corpus.The evaluation includes NQ, TriviaQA, PopQA, HotpotQA, 2Wiki, MuSiQue, and Bamboogle.
  • Experimental Setup: CAS is compared with non-retrieval, retrieval, supervised fine-tuning, RL fine-tuning, rejection-sampling, Search-R1, and Search-R2 baselines.The baselines are organized into four categories covering inference, retrieval, fine-tuning, and reference search agents.
  • Main Results: 0.464 is CAS’s average score on Qwen3-8B, exceeding Search-R2 at 0.446; on Qwen2.5-3B-Instruct, CAS reaches 0.401, +0.065 over Search-R1.The reported gains hold across both evaluated backbone scales.
  • Main Results: +0.079 and +0.078 are CAS’s Multi-Hop QA gains over Search-R1 on the 3B and 8B backbones, while the 8B model gains +0.045 over Search-R1 on General QA.On 8B General QA, CAS also slightly exceeds Search-R2 by +0.001; Search-R2 is unavailable for Qwen2.5-3B-Instruct because it is closed-source.
  • Ablation Study: Removing either APS or ACI causes notable degradation on both General QA and Multi-Hop QA tasks.The -APS configuration uses fixed top-k retrieval with k = 3.
  • Ablation Study: 0.384 is the overall score after removing ACI, down from 0.401, while disabling APS reduces accuracy to 0.389.Without ACI, tool calls become higher and more variable; without APS, fixed retrieval can add noise or miss critical multi-hop facts, prompting additional searches.
  • Sensitivity Analysis: ρ = 0.4 penalizes nearly 40% of trajectories and dilutes RL reward signals, whereas ρ = 0.1 affects only 10% and provides insufficient confidence guidance.The sensitivity analysis identifies a trade-off between overly strict and overly lenient confidence filtering.
  • Sensitivity Analysis: αAPS = 0.05 retrieves 4.8 documents on average with 95% marginal coverage, whereas αAPS = 0.35 retrieves 2.4; the default αAPS = 0.20 balances coverage and noise reduction.The high-guarantee setting benefits single-hop QA but introduces noise in multi-hop QA, while aggressive truncation leaves insufficient supporting facts.

5 Related Works

Related work positions CAS at the intersection of retrieval-augmented generation, RL-based search agents, and conformal prediction. These lines of work motivate adaptive retrieval and confidence control for autonomous agents.

  • Retrieval and RL Agents: RAG-based systems increasingly address when retrieval should be triggered, while retrieval-integrated RL methods connect external knowledge with agent training.The discussion cites Adaptive RAG, Search-o1, SAKI-RAG, and RL-based retrieval approaches.
  • Retrieval and RL Agents: RL has turned LLMs into autonomous agents, but sparse binary rewards are associated with overconfidence and hallucinated behavior in Search Agents.The related-work discussion includes SWE-agents, Web Agents, and Search Agents as examples of this broader agent paradigm.
  • Conformal Prediction: Conformal Prediction provides reliable marginal coverage without distributional assumptions, but existing CP work does not fully guarantee conditional coverage.The paper places CAS among efforts applying CP to classification, detection, LLMs, and RL fine-tuning.

6 Conclusion

CAS integrates Conformal Prediction into RL-trained search agents through APS and ACI, targeting reliable retrieval and calibrated training. The framework improves reasoning accuracy while reducing redundant tool invocations.

  • CAS integrates Conformal Prediction into RL fine-tuning to address reliability in search agents.
  • APS and ACI jointly support reliable document retrieval and mitigation of model overconfidence.
  • CAS significantly enhances reasoning accuracy and reduces redundant tool invocations.
  • The framework establishes a principled foundation for future reliable autonomous agents.

7 Limitations

The paper reports limitations in CAS’s evaluation scope, calibration dependence, and reliability guarantees, while the surrounding material specifies its implementation inputs and training loop.

  • Limitations: CAS is primarily validated on general open-domain QA, leaving highly specialized professional domains unexplored.
  • Limitations: The APS calibration set relies heavily on a strong external teacher model for query decomposition and relevance judging.
  • Limitations: CAS focuses on outcome reliability and does not extend statistical guarantees to the intermediate reasoning process.
  • Implementation: Algorithm 1 initializes calibration scores from an untrained policy before iterating reinforcement-learning updates.
  • Implementation: Each trajectory proceeds until termination, triggers retrieval when a search query is generated, and computes NLL scores and group advantages.
  • Theoretical supplement: The theoretical supplement begins by proving the finite-sample marginal coverage guarantee of standard Split Conformal Prediction.

B.1 Marginal Coverage Guarantee of Conformal Prediction

This section establishes finite-sample marginal coverage for conformal prediction and APS under exchangeability, then describes APS’s asymptotic conditional-coverage behavior and ACI’s long-run coverage under distribution shift.

  • Standard conformal prediction: Under the i.i.d. assumption, standard conformal prediction guarantees P(Ytest ∈ C1−α(Xtest)) ≥ 1−α.
  • Standard conformal prediction: Exchangeability makes the test non-conformity score uniformly rankable among calibration scores, yielding the finite-sample coverage result.
  • Adaptive Prediction Sets: APS retains rigorous marginal coverage while adapting prediction-set size to local input uncertainty and approximating conditional coverage.
  • Adaptive Prediction Sets: APS coverage is at least 1−α under exchangeable calibration and test samples with a model trained on a disjoint split.
  • Adaptive Prediction Sets: With distinct scores, APS coverage is bounded above by 1−α + 1/(|I2| + 1).
  • Adaptive Conformal Inference: ACI addresses distribution shifts that violate i.i.d. assumptions and targets long-run miscoverage through dynamic error-level updates.
  • Adaptive Conformal Inference: As T →∞, ACI’s average miscoverage converges almost surely to the target rate α without assumptions on the data-distribution shift.

C Supplementary Implementation Details

CAS uses separate training and retrieval services in a distributed environment, with dense retrieval infrastructure and configured rollout, context, and calibration workflows.

  • System architecture: The training service runs distributed GRPO across four GPUs using VeRL, Ray, sglang, flashinfer, and flash-attn.
  • System architecture: The retrieval service is an independent FastAPI endpoint using faiss-gpu and e5-base-v2 for dense retrieval.
  • System architecture: The retrieval backend is configured for a peak rate of 120 queries per second with a 30-second timeout.
  • Configurations: Training uses a unified search-integrated reasoning dataset with maximum prompt, response, and context lengths of 4096, 3000, and 15,000 tokens.
  • Configurations: GRPO samples five trajectories per prompt and allows up to four assistant turns.
  • Prompt templates: Prompt templates maintain consistency with Search-R1 and support APS calibration-set construction through query decomposition and retrieval relevance judging.

E Additional Experimental Analysis on Qwen3-8B

Additional Qwen3-8B analysis links model capacity and search behavior to training outcomes, and shows that enabling thinking mode undermines search-integrated training. The model’s behavior changes from aggressive searching to abandoning information gathering on many single-hop problems.

  • Performance and search behavior: Qwen3-8B consistently outperforms Qwen2.5-3B-Instruct, which the authors primarily attribute to greater model capacity.Figure 3a illustrates the capacity advantage.
  • Performance and search behavior: During approximately the first 80 training steps, Qwen3-8B makes fewer average search calls than Qwen2.5-3B-Instruct.The comparison is reported for the search behavior shown in Figure 3b.
  • Thinking mode: Enabling thinking mode drastically reduces training effectiveness because the model first searches aggressively, then stops invoking the tool after solving many single-hop problems internally.This complete cessation of active information gathering renders training ineffective for search-integrated tasks.

F Detailed Results for Ablation and Sensitivity Analysis

The complete individual-dataset results corresponding to Table 3 are provided in Table 8.

  • Detailed results: Table 8 provides the complete results for all individual datasets.These results correspond to Table 3 in the main text.
  • Detailed results: The detailed results are presented across individual datasets rather than only as the aggregate presentation in Table 3.The passage identifies Table 8 as the complete dataset-level counterpart to Table 3.
  • Detailed results: Table 8 serves as the detailed-results reference corresponding to Table 3 in the main text.The passage explicitly establishes this correspondence.

G Case Study

The case studies illustrate CAS behavior across simple and noisy single-hop queries and multi-hop reasoning. They show direct evidence extraction, distractor filtering, and iterative decomposition with targeted searches.

  • Case Studies: In a straightforward single-hop case, clean relevant documents let the agent answer quickly after one precise search action.The agent concludes reasoning and extracts the correct answer without unnecessary actions.
  • Case Studies: In a noisy single-hop case, the agent filters similarly named distractors by contextual relevance and grounds its answer in the correct source.The case demonstrates robust document-level relevance evaluation under retrieval noise.
  • Case Studies: For multi-hop queries, the agent decomposes the task, retrieves a missing bridge entity, and uses it to formulate a targeted second-hop query.The process interleaves internal reasoning with external tool interaction.
  • Prompt templates: The reasoning template defines the CAS reasoning process, while separate prompts decompose multi-hop queries and judge document relevance.These prompts are used across CAS stages, including construction of the retrieval-side APS calibration set.
  • Additional analysis: Figure 3 visualizes training accuracy, tool invocation frequency, thinking-mode effects on final EM scores, and search efficiency.Panels (a)–(d) cover these four training and behavior dimensions.
  • Additional analysis: Table 8 contains component-wise ablations and sensitivity analyses for ρ and αAPS on Qwen2.5-3B-Instruct.The supplied caption identifies the table’s experimental scope.
Loading 2608.20771v1…