Learning to Better Search with Language Models via Guided Reinforced Self-Training

TL;DR

This work addresses the challenge of efficient, reliable multi-step reasoning in language models by proposing Guided-ReST, a self-training algorithm that uses optimal subgoals as landmarks to guide search traces. It combines subgoal augmentation with supervised fine-tuning and an operation-level PPO to distill improved search strategies, achieving substantial accuracy gains on Countdown and code self-repair benchmarks while reducing test-time compute. The approach demonstrates that guiding search with high-quality local corrections and aligning RL objective with an operation-level view yields better generalization, particularly under constrained token budgets. While effective, the method assumes access to optimal solutions and highlights future directions for teacher-student guidance and reasoning-backtracking in less ideal settings.

Abstract

While language models have shown remarkable performance across diverse tasks, they still encounter challenges in complex reasoning scenarios. Recent research suggests that language models trained on linearized search traces toward solutions, rather than solely on the final solutions, exhibit improved generalization, despite the search traces being potentially noisy or suboptimal. However, relying on such imperfect traces can result in inefficient use of test-time compute. To address this, we propose guided reinforced self-training (Guided-ReST), a fine-tuning algorithm designed to improve the model's capability for effective search during inference. The key insight behind Guided-ReST is that optimal solutions can serve as valuable step-by-step landmarks to guide the model's search process. Based on this insight, we introduce a novel data generation method that seamlessly incorporates optimal solutions into the model's search procedure, enabling the generation of high-quality search traces. By fine-tuning the model on these search traces, we effectively distill improved search strategies into the model. Our method significantly enhances the search capabilities of language models on arithmetic reasoning and code self-repair tasks, including Countdown, CodeContests, and CodeForces. We release the source code at https://github.com/snu-mllab/guided-rest.

Figures (13)

• Figure 1: Overview of guided reinforced self-training, illustrating how search traces are generated through the progressive integration of optimal solutions during self-generation. The numbers indicate the order in which nodes are explored.
• Figure 2: Overview of the subgoal augmentation algorithm. The numbers indicate the order in which nodes are explored. The red box highlights the modifications through node replacement.
• Figure 3: Accuracy of our method and baselines on Countdown. Solid lines represent seen targets, and dotted lines represent unseen targets.
• Figure 4: Accuracy of PPO initialized with ReST on Countdown. PPO does not benefit from ReST in the high-budget regime.
• Figure 5: Accuracy of our method when trained under the token-level MDP on Countdown.