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

TL;DR

This paper introduces Guided-ReST, a fine-tuning algorithm that improves language models' search abilities by using optimal solutions as landmarks, leading to better reasoning and code repair performance.

Contribution

The paper proposes Guided-ReST, a novel data generation and fine-tuning method that enhances language models' search strategies using optimal solutions as guidance.

Findings

Improved search performance on arithmetic reasoning tasks.

Enhanced code self-repair capabilities.

Significant gains on benchmarks like Countdown, CodeContests, and CodeForces.

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.