Afgekapte stapgewijze bemonstering met procesbeloningen voor retrieval-augmented redenering

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

February 26, 2026

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

cs.AI

Samenvatting

Het trainen van grote taalmodellen om te redeneren met zoekmachines via reinforcement learning wordt belemmerd door een fundamenteel credit assignment-probleem: bestaande methoden zoals Search-R1 bieden slechts een schaarse uitkomstbeloning na een volledige multi-stap traject, waardoor het onmogelijk is om succes of falen toe te schrijven aan individuele redeneer- en retrievalbeslissingen. Procesbeloningsmethoden zoals StepSearch verlichten dit door supervision op stapniveau in te voeren, maar steunen op heuristische beloningen zoals TF-IDF-overlap met gouden documenten, en nemen nog steeds k complete trajecten per voorbeeld, waardoor de variantie van de gradiënt hoog blijft. Wij stellen SLATE voor, een raamwerk gebaseerd op twee complementaire ideeën: (1) afgekapte steekproeven op stapniveau, die k trajecten genereren die een gemeenschappelijk voorvoegsel delen en alleen verschillen bij de volgende stap, en (2) dichte LLM-als-rechter beloningen, die heuristische scoring vervangen door een capabele LLM-evaluator die de kwaliteit van elke redeneerstap, zoekopdracht en antwoord beoordeelt, wat rijkere en betrouwbaardere supervision biedt. Wij bewijzen theoretisch dat onder dezelfde dichte beloningsstructuur, afgekapte steekproeven de variantie van advantage-schattingen met een factor T verminderen in vergelijking met steekproeven van volledige trajecten voor T-staps trajecten, wat resulteert in policy gradients met een lagere variantie en een beter gericht bereik. Experimenten op zeven QA-benchmarks bevestigen dat SLATE consistent beter presteert dan zowel baseline-methoden met schaarse beloning als procesbeloning, met de grootste winst op moeilijkere multi-hop taken en bij kleinere modellen.

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.