PaperlayerView exact revision on arXiv

Source-linked AI summary

Truncated Step-Level Sampling with Process Rewards for Retrieval-Augmented Reasoning

Chris Samarinas, Haw-Shiuan Chang, Hamed Zamani

arXiv:2602.23440v4cs.CLcs.IR

TL;DR

Existing retrieval-augmented RL methods struggle to assign credit to individual decisions because outcome rewards are sparse and independently sampled process-reward trajectories mix variation across steps. SLATE combines shared-prefix truncated sampling with decomposed ternary process rewards, proving up to T-fold lower advantage variance and outperforming prior baselines across seven QA benchmarks, including 7.0% and 30.7% relative gains over SEARCH-R1 on 7B and 3B models.

  • Problem

    Outcome rewards provide no step-specific credit, while independently sampled process-reward trajectories leave each step’s advantage contaminated by randomness from other steps.

  • Method

    SLATE samples k continuations from a shared prefix and separately evaluates reasoning quality, query quality, and answer correctness on a ternary scale.

  • Results

    SLATE consistently outperforms sparse- and process-reward baselines across seven QA benchmarks, with a 7.0% relative improvement over SEARCH-R1 on 7B and 30.7% on 3B models.

  • Takeaways & Limitations

    Truncated sampling and dense rewards provide complementary gains, with truncated sampling improving results beyond step-level rewards alone.

  • Takeaways & Limitations

    The theoretical variance guarantee assumes non-negative future covariance, conditional step-reward independence, and comparable reward variances across steps.

Abstract

from arXiv · show

Reinforcement learning has emerged as an effective paradigm for training large language models to interleave reasoning with search engine calls. However, existing approaches face a fundamental credit assignment problem: methods like Search-R1 assign a single outcome reward to the entire multi-step trajectory, providing no signal about which reasoning or retrieval decisions were responsible for success or failure. Process-reward methods such as StepSearch introduce step-level supervision but still sample complete trajectories independently, so advantage estimates at any given step are contaminated by the randomness of all other steps. We propose SLATE (Step-Level Advantage estimation for Truncated Exploration), which addresses both problems through two complementary ideas. First, truncated step-level sampling generates k continuations from a shared prefix, isolating all variation to a single decision point. We prove this reduces the variance of advantage estimates by up to a factor of T compared to full-trajectory sampling for T-step trajectories, the first formal variance guarantee for step-level RL in retrieval-augmented reasoning. Second, dense, decomposed process rewards separately evaluate reasoning quality, query quality, and answer correctness on a ternary scale via an LLM judge, providing richer supervision than binary outcome signals or heuristic step-level scores. Experiments on seven QA benchmarks show that SLATE consistently outperforms both sparse-reward and process-reward baselines, achieving a 7.0% relative improvement over Search-R1 on the 7B model and 30.7% on the 3B model. Gains are largest on challenging multi-hop tasks, and ablations confirm that truncated sampling and dense rewards provide complementary benefits.

1 Introduction

SLATE targets credit-assignment and variance problems in retrieval-augmented reasoning by combining truncated step-level sampling with dense, decomposed process rewards. Across seven QA benchmarks, it consistently outperforms sparse- and process-reward baselines, with especially large gains for smaller models and challenging multi-hop tasks.

  • Motivation: Existing outcome-based and process-reward methods cannot reliably isolate the contribution of an individual reasoning or retrieval decision.SEARCH-R1 uses a single trajectory-level binary reward, while StepSearch and SWiRL independently sample complete trajectories, leaving step advantages affected by randomness elsewhere.
  • Motivation: High-variance advantage estimates aggregate reward variation across all T trajectory steps, slowing convergence and risking unstable training.The paper identifies variance reduction while preserving individual-action credit assignment as a central challenge.
  • SLATE: SLATE samples k truncated continuations from a shared prefix, varying only the action at step t to directly attribute rewards to that decision.The method provides a formal T-fold advantage-variance reduction over full-trajectory sampling.
  • SLATE: SLATE separately scores reasoning quality, query quality, and answer correctness on the ternary scale {−1, 0, +1}.An LLM judge uses a reason-then-score protocol to operationalize the decomposed process reward.
  • Results: 7.0% relative improvement over SEARCH-R1 was achieved on the 7B model, while the 3B model achieved a 30.7% relative improvement across seven QA benchmarks.The largest gains occurred on challenging multi-hop tasks, and ablations found complementary benefits from truncated sampling and dense rewards.

2 Related Work

Retrieval-augmented reasoning extends retrieval-augmented generation by interleaving reasoning with iterative retrieval, while reinforcement learning enables training beyond prompting alone. Prior RL methods improve credit assignment with step-level rewards but retain complete-trajectory sampling, and finer-grained alternatives can be computationally expensive or limited to short interactions.

  • Retrieval-Augmented Reasoning: Single-turn retrieval struggles with complex questions requiring iterative information gathering and multi-step reasoning.Retrieval-augmented reasoning addresses this by interleaving chain-of-thought reasoning with retrieval calls across multiple turns.
  • Retrieval-Augmented Reasoning: Zero-shot methods such as IRCoT, ReAct, and Search-o1 enable iterative retrieval through prompting but cannot improve through training.Reinforcement learning has become the primary optimization paradigm for training search-using language models.
  • RL for Retrieval-Augmented Reasoning: SEARCH-R1 pioneered reinforcement-learning training for language models that invoke search engines during multi-turn reasoning.It established outcome-based optimization for retrieval-augmented reasoning.
  • Process Rewards: StepSearch uses information-gain rewards and redundancy penalties to address reward sparsity but still samples complete trajectories.This preserves the cross-step randomness that complicates fine-grained credit assignment.
  • Turn-Level Credit Assignment: MT-GRPO combines turn-level and outcome advantages, but its tree-structured rollout cost grows as GK−1 with horizon K and its tractable GRPO setup is restricted to two-turn interactions.Longer horizons are handled with a critic-based PPO variant instead.

3 Methodology

SLATE combines truncated step-level sampling with dense, decomposed process rewards to improve credit assignment and supervision in retrieval-augmented reasoning.

  • Truncated Step-Level Sampling: SLATE samples k continuations from a shared prefix, isolating variation to the next reasoning step instead of varying entire trajectories.At each step, one candidate is selected to extend the prefix, after which retrieval and sampling continue.
  • Truncated Step-Level Sampling: Truncated sampling reduces advantage-estimate variance by up to T-fold compared with full-trajectory sampling, yielding lower-variance policy gradients and more stable training.The method samples single truncated steps of roughly L/T tokens rather than k complete rollouts.
  • Truncated Step-Level Sampling: The method computes step-level group-relative advantages for candidates and uses the selected action to construct the trajectory iteratively.Candidate actions may contain reasoning and search blocks or terminate with a final answer.
  • Dense, Decomposed Process Rewards: The decomposed reward evaluates reasoning quality, query quality, and answer correctness separately on the ternary scale {−1, 0, +1}.The design uses an LLM judge and requires trajectory context plus the gold final answer, without intermediate annotations.
  • Dense, Decomposed Process Rewards: Reasoning, query, and answer components use distinct evaluation criteria, with query quality judged before retrieval results are observed.This makes query rewards reflect intrinsic query quality rather than retrieval nondeterminism.
  • Dense, Decomposed Process Rewards: SLATE retains held-out exact-match evaluation and uses terminal rewards grounded in the gold answer to limit reward-hacking risks.Separate rubrics and a reason-then-score protocol further constrain how the learned reward can be optimized.

4 Theoretical Analysis

The analysis formalizes when truncated step-level sampling lowers advantage-estimation variance relative to full-trajectory sampling. Under stated covariance, conditional-independence, and variance-symmetry assumptions, the reduction can reach a factor of T, with corresponding efficiency and convergence benefits, while introducing a bias-variance trade-off.

  • Estimator comparison: Truncated sampling fixes a prefix and varies only the action at step t, isolating the decision whose step-level advantage is estimated.This contrasts with full-trajectory sampling, where rewards from all trajectory steps contribute variation to the estimate.
  • Theorem setting: The result applies to additive trajectory rewards R(τ) = ∑t rt, where rt is the reward at step t.The theorem explicitly compares scalar advantages formed from full returns with step-level advantages under this reward decomposition.
  • Limitations: The step-level and trajectory-level estimators target different quantities, creating a bias-variance trade-off because the step-level estimator discards future-reward information.The substitution is therefore justified only under conditions discussed in the paper’s accompanying remark.
  • Variance bounds: Under non-negative future covariance, the expected conditional variance of the truncated estimator is no greater than the variance of the full-trajectory estimator.The proof uses the law of total variance and decomposes the return into past, current, and future rewards.
  • Variance bounds: Under conditional independence and variance symmetry, truncated sampling achieves a 1/T variance factor, or up to a T-fold reduction, with equal group sizes k = G.The assumptions require comparable conditional step-reward variances across trajectory positions.
  • Implications: Truncated sampling also yields a T-fold reduction in total token-generation cost for achieving the same advantage variance as standard GRPO.Because policy gradients are linear in advantages, lower advantage variance directly yields lower-variance gradient estimates, enabling faster convergence and better final solutions.

5 Experiments

Experiments evaluate SLATE across seven QA benchmarks and show consistent gains over sparse- and process-reward baselines, especially on difficult multi-hop tasks and smaller models.

  • Main Results: SLATE achieves the best results across seven QA benchmarks, outperforming both SEARCH-R1 and StepSearch.The evaluation covers single-hop and multi-hop question answering, with Exact Match as the primary metric.
  • Main Results: 0.461 average EM on the 7B model is a 3.0% absolute, or 7.0% relative, improvement over SEARCH-R1’s 0.431.SLATE also outperforms the best prior result on every individual dataset.
  • Model Scale: 30.7% relative improvement on the 3B model corresponds to EM of 0.396 versus SEARCH-R1’s 0.303.The larger gain indicates greater benefit from dense step-level supervision for the smaller model.
  • Multi-Hop QA: SLATE’s largest gains occur on challenging multi-hop benchmarks, including +5.1% on Musique and +6.2% on Bamboogle versus SEARCH-R1.It is the only method consistently outperforming both SEARCH-R1 and StepSearch across all four multi-hop benchmarks.
  • Generalization: On single-hop benchmarks, SLATE improves over SEARCH-R1 by 1.3–1.9% absolute EM, while the largest gains occur on unseen multi-hop tasks.Five of seven evaluation benchmarks are out-of-domain relative to training on NQ and HotpotQA.
  • Ablations: Ablations show that truncated sampling adds gains beyond the decomposed reward model, while removing LLM-judge rewards reduces average EM by 2.4%.Using only EM reward performs close to SEARCH-R1, supporting complementary benefits from the two components.

6 Conclusions

SLATE argues that step-level optimization strategy is as important as the reward signal, combining truncated sampling with decomposed rewards to improve retrieval-augmented reasoning.

  • Conclusion: SLATE’s shared-prefix sampling provides provable advantage-variance reduction, while decomposed ternary rewards complement it with richer step-level supervision.Ablations confirm that truncated sampling improves performance beyond step-level rewards alone.

A.1 The SLATE Algorithm

Algorithm 1 trains SLATE by sampling multiple step-level continuations, scoring them with dense LLM-judge rewards, computing GRPO advantages, and updating the policy.

  • Sampling: For each question, the algorithm initializes an empty prefix and repeatedly samples k continuations while the step budget remains.The policy, search engine, judge, dataset, group size, temperature, and maximum steps are algorithm inputs.
  • Optimization: The method computes step-level GRPO advantages for the sampled actions and uses them to update policy parameters.The update is performed after trajectory completion, with gradients accumulated across trajectory steps.
  • Sampling: Each sampled continuation extends the shared prefix through reward-weighted sampling.This is the algorithm’s step-level exploration mechanism.
  • Termination: Generation stops when an answer is produced or the maximum step budget is reached, followed by a batched parameter update.The algorithm accumulates the update before applying it to the policy.

A.2 Experimental Setup

Experiments train Qwen2.5-7B-Base and 3B-Base with Wikipedia retrieval, GRPO, and Gemma3-27B judging, then evaluate Exact Match after parameter-efficient fine-tuning.

  • Models and Retrieval: Experiments use Qwen2.5-7B-Base and Qwen2.5-3B-Base with the 2018 Wikipedia dump and E5 retrieval of the top three passages per query.Training data combines NQ and HotpotQA, and Exact Match is the primary evaluation metric.
  • Training: Training uses GRPO with temperature η = 0.7, Gemma3-27B as the LLM judge, and a maximum action budget B = 4.The setup runs for 500 steps with a batch size of 32 and maximum sequence length of 4096 tokens.
  • Training: Parameter-efficient fine-tuning uses LoRA with rank 16 and α = 64 in bfloat16 precision on two NVIDIA A100 GPUs.Retrieved token loss masking follows SEARCH-R1.

A.3 Gradient Estimators for Full and Truncated Step-level Trajectories

The section contrasts full-trajectory and truncated step-level gradient estimators. Full-trajectory estimation uses complete returns, whereas SLATE fixes a shared prefix and varies only the current step.

  • Full-trajectory sampling computes the gradient for step t from rewards of complete trajectories.
  • Truncated step-level sampling computes the gradient at step t from continuations sharing prefix τ<t and differing only at that step.This isolates variation to the current decision point.

A.4 A Proof for Theorem 1: Variance Reduction via Truncated Sampling

The proof shows that fixing the prefix removes unrelated randomness from step-level advantage estimates. Under additional independence and variance-symmetry assumptions, this yields a T-fold variance reduction relative to full-trajectory sampling.

  • Part 1: General Variance Bound: Fixing prefix τ<t eliminates all randomness except the action at step t, reducing conditional advantage variance on average.
  • Part 1: General Variance Bound: Given a fixed prefix, total return decomposes into constant past rewards, the current reward r_t, and future rewards F_t.The resulting conditional variance includes current-reward variance, future-reward variance, and their covariance.
  • Part 1: General Variance Bound: Under nonnegative covariance between current and future rewards, per-step conditional variance is no greater than conditional full-return variance.
  • Part 1: General Variance Bound: The truncated estimator has no more variance in expectation than the full-trajectory estimator when group sizes are equal.
  • Part 2: T-fold Reduction Under Independence and Symmetry: Under variance symmetry, summing the per-step variances produces a T-fold reduction for truncated sampling relative to full-trajectory sampling.The argument relies on the stated independence and symmetry assumptions.
  • Part 2: T-fold Reduction Under Independence and Symmetry: Under conditional independence, per-step conditional variances represent distinct, nonoverlapping randomness sources across the trajectory.

A.5 Theoretical Remarks

The theoretical remarks characterize SLATE’s variance, cost, and bias trade-offs. They also identify assumptions and practical limitations surrounding step-level rewards, exploration, judge overhead, and task scope.

  • Variance comparison: T · Var[Â_i] summarizes the theorem’s comparison between trajectory-level and truncated step-level advantage variance.
  • Bias-variance trade-off: The truncated estimator is not unbiased for trajectory-level advantage because it targets step-level rewards instead of the full return.
  • Bias-variance trade-off: Short horizons T ≤4, externally grounded retrieval, and evaluative rewards support treating step-level rewards as low-bias proxies.
  • Sample efficiency: G/T truncated samples per step match the advantage variance of G full-trajectory samples under Theorem 1’s conditions.
  • Sample efficiency: A net T-fold reduction in token-generation cost results from fewer samples and shorter truncated samples, offset by repeating across T steps.
  • Optimization dynamics: Lower-variance advantages directly reduce gradient-estimate variance because policy gradients are linear in the advantages.
  • Limitations: SLATE limits exploration to one selected action per step and does not propagate retroactive credit from later rewards to earlier prefix actions.
  • Limitations: The method’s scope is constrained by LLM-judge overhead and evaluation on seven EM-based factoid QA benchmarks using a single search engine.
Loading 2602.23440v4…