Échantillonnage au Niveau de l'Étape Tronqué avec Récompenses de Processus pour le Raisonnement Augmenté par Récupération

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

February 26, 2026

Auteurs: Chris Samarinas, Haw-Shiuan Chang, Hamed Zamani

cs.AI

Résumé

L'entraînement des grands modèles de langage à raisonner avec des moteurs de recherche par apprentissage par renforcement est entravé par un problème fondamental de crédit des actions : les méthodes existantes, telles que Search-R1, ne fournissent qu'une récompense finale éparse après une trajectoire multi-étapes complète, ce qui rend impossible d'attribuer le succès ou l'échec à des décisions individuelles de raisonnement et de recherche. Les méthodes à récompense de processus comme StepSearch atténuent ce problème en introduisant une supervision au niveau de l'étape, mais elles reposent sur des récompenses heuristiques telles que le chevauchement TF-IDF avec des documents de référence, et échantillonnent toujours k trajectoires complètes par exemple, conservant une variance de gradient élevée. Nous proposons SLATE, un cadre fondé sur deux idées complémentaires : (1) un échantillonnage tronqué au niveau de l'étape, qui génère k trajectoires partageant un préfixe commun et ne différant qu'à l'étape suivante, et (2) des récompenses denses de type "LLM-comme-juge", qui remplacent le score heuristique par un évaluateur LLM capable d'évaluer la qualité de chaque étape de raisonnement, chaque requête de recherche et chaque réponse, fournissant ainsi une supervision plus riche et plus fiable. Nous prouvons théoriquement que, sous la même structure de récompense dense, l'échantillonnage tronqué réduit la variance des estimations de l'avantage jusqu'à un facteur T par rapport à l'échantillonnage de trajectoires complètes pour des trajectoires de T étapes, produisant des gradients de politique à variance plus faible et mieux ciblés. Les expériences sur sept benchmarks de question-réponse confirment que SLATE surpasse constamment les lignes de base à récompense éparse et à récompense de processus, avec les gains les plus importants sur les tâches multi-sauts plus difficiles et pour les modèles plus petits.

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.