プロセス報酬を用いた検索拡張推論のための段階的切り捨てサンプリング

要旨

大規模言語モデルに検索エンジンを活用した推論を強化学習で習得させる際、根本的な信用割り当て問題が障壁となる。既存手法であるSearch-R1などは、複数ステップからなる軌跡全体の完了後にのみ疎な結果報酬を提供するため、個々の推論や検索判断の成否を帰属させることが困難である。StepSearchのような過程報酬手法は、ステップ単位の監督を導入することでこの問題を緩和するが、正解文書とのTF-IDF重複度などヒューリスティックな報酬に依存し、かつ事例ごとにk本の完全な軌跡をサンプリングするため、勾配の分散が大きいという問題を残す。本研究では、二つの相補的なアイデアに基づくSLATEフレームワークを提案する。(1) 切り詰められたステップ単位サンプリング：共通の前置軌跡を持ち、次のステップのみが異なるk本の軌跡を生成する。(2) 密なLLM-as-judge報酬：ヒューリスティックなスコアリングを、各推論ステップ、検索クエリ、回答の質を評価する能力の高いLLM評価器に置き換え、より豊かで信頼性の高い監督を提供する。理論的に、同一の密報酬構造下では、Tステップの軌跡において、切り詰めサンプリングは完全軌跡サンプリングと比較してアドバンテージ推定値の分散を最大T分の1に低減し、分散が小さくより適切に焦点化された方策勾配をもたらすことを証明する。7つのQAベンチマークによる実験では、SLATEが疎報酬および過程報酬ベースライン手法を一貫して上回り、特に難易度の高いマルチホップタスクや規模の小さいモデルにおいて最大の性能向上を示すことを確認した。

English

Training large language models to reason with search engines via reinforcement learning is hindered by a fundamental credit assignment problem: existing methods such as Search-R1 provide only a sparse outcome reward after an entire multi-step trajectory, making it infeasible to attribute success or failure to individual reasoning and retrieval decisions. Process-reward methods like StepSearch alleviate this by introducing step-level supervision, but rely on heuristic rewards such as TF-IDF overlap with gold documents, and still sample k complete trajectories per example, retaining high gradient variance. We propose SLATE, a framework built on two complementary ideas: (1) truncated step-level sampling, which generates k trajectories that share a common prefix and differ only at the next step, and (2) dense LLM-as-judge rewards, which replace heuristic scoring with a capable LLM evaluator that assesses the quality of each reasoning step, search query, and answer, providing richer and more reliable supervision. We theoretically prove that under the same dense reward structure, truncated sampling reduces the variance of advantage estimates by up to a factor of T compared to full-trajectory sampling for T-step trajectories, yielding lower-variance, better-targeted policy gradients. Experiments on seven QA benchmarks confirm that SLATE consistently outperforms both sparse-reward and process-reward baselines, with the largest gains on harder multi-hop tasks and smaller models.

プロセス報酬を用いた検索拡張推論のための段階的切り捨てサンプリング

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

要旨

Support