Source-linked AI summary
Stream of Search (SoS): Learning to Search in Language
Kanishk Gandhi, Denise Lee, Gabriel Grand, Muxin Liu, Winson Cheng, Archit Sharma, Noah D. Goodman
TL;DR
Language models often lack training examples of productive mistakes and therefore struggle with error compounding and multi-step lookahead. This paper represents search trajectories as Streams of Search, trains models on diverse heuristic-generated processes, and iteratively improves them with policy methods. SoS training improves search over optimal-trajectory training, while finetuned models solve previously unsolved problems and may discover new strategies.
Problem
Language models are rarely trained on mistakes or decision processes, leaving their ability to search, backtrack, and handle error compounding and lookahead limited.
Method
The paper serializes exploration, backtracking, and pruning as Streams of Search, trains models on diverse search trajectories, and applies APA and STaR for policy improvement.
Results
SoS training improves performance over models trained solely on optimal trajectories, while policy-improved models solve previously unsolved problems and use different search strategies.
Takeaways & Limitations
Language models can learn to solve problems through search, self-improve with diverse strategies, and potentially discover new search strategies.
Takeaways & Limitations
The empirical evaluation is restricted to Countdown, and externally structured search methods may remain more efficient for challenging real-world tasks in the short term.
Abstract
from arXiv · showhide
Language models are rarely shown fruitful mistakes while training. They then struggle to look beyond the next token, suffering from a snowballing of errors and struggling to predict the consequence of their actions several steps ahead. In this paper, we show how language models can be taught to search by representing the process of search in language, as a flattened string -- a stream of search (SoS). We propose a unified language for search that captures an array of different symbolic search strategies. We demonstrate our approach using the simple yet difficult game of Countdown, where the goal is to combine input numbers with arithmetic operations to reach a target number. We pretrain a transformer-based language model from scratch on a dataset of streams of search generated by heuristic solvers. We find that SoS pretraining increases search accuracy by 25% over models trained to predict only the optimal search trajectory. We further finetune this model with two policy improvement methods: Advantage-Induced Policy Alignment (APA) and Self-Taught Reasoner (STaR). The finetuned SoS models solve 36% of previously unsolved problems, including problems that cannot be solved by any of the heuristic solvers. Our results indicate that language models can learn to solve problems via search, self-improve to flexibly use different search strategies, and potentially discover new ones.
1 Introduction
Language models are usually trained on correct outcomes rather than mistakes, limiting their ability to search, plan, and backtrack. Stream of Search represents search trajectories in language and trains models on diverse search processes before policy improvement.
- Training data usually shows outcomes rather than decision-making processes, so models do not learn mistakes, search, planning, or backtracking.
- Autoregressive models face snowballing errors and difficulty predicting action consequences several steps ahead.
- Stream of Search serializes exploration, backtracking, and pruning in a unified language for training language models to search.
- The framework uses Countdown, where input numbers and arithmetic operations must be combined to reach a target in a high-branching search space.
- SoS models are trained on diverse search trajectories and then improved with APA and STaR to use different strategies and solve previously unsolved problems.
2 Related Works
Prior approaches use language models within externally specified search systems, demonstrations, process supervision, or fixed search traces. These approaches constrain search procedures or require substantial supervision, motivating autonomous use of diverse strategies.
- Search systems commonly use language models to generate or evaluate actions while symbolic algorithms determine exploration strategy.
- In-context demonstrations let models perform tree search, but restrict them to the demonstrated exploration strategy or heuristic.
- Process supervision trains models from verifier feedback on intermediate steps, but requires large datasets of human annotations.
- Prior trajectory-learning methods mimic fixed procedures such as Monte Carlo Tree Search, BFS, or A* rather than autonomously using or discovering different procedures.
3 A Language for Search
The paper models problem solving as movement through a search tree and introduces a language of primitive operations for representing that process. Streams can explicitly encode states and search actions, including exploration choices and backtracking.
- The problem is modeled as an MDP with states, actions, transitions, and rewards for reaching a goal state.
- A search tree contains explorations from an initial state through possible child states until leaf states are reached.
- A correct path is a sequence of valid state-action transitions ending at the goal state.
- The proposed vocabulary represents search using primitive operations such as current and goal states, state queues, expansion, exploration choice, pruning, backtracking, goal checks, and heuristics.
- The trajectory explicitly represents current and goal states, backtracking, goal checks, and exploration choices, while heuristic functions, state values, and pruning remain implicit.
4 Problem Setup
The experiments use Countdown as a challenging planning task and train on synthetic streams generated by diverse suboptimal search strategies. Accuracy depends on whether generated trajectories contain a correct solution path, while additional metrics compare strategy alignment.
- Countdown requires combining input numbers with arithmetic operations to reach a target and has a high branching factor requiring planning, search, and backtracking.
- The study uses four-input problems to obtain long search traces without exceeding a standard language-model context window.
- The synthetic dataset contains 500,000 trajectories from 12 BFS- and DFS-based strategies, with about 57% leading to solutions.
- The evaluation holds out seen targets with new inputs and new targets with new inputs for generalization tests.
- Accuracy is the percentage of problems whose generated trajectory contains the correct solution path.
- Strategy alignment measures either shared solved and unsolved problems or overlap between visited states, normalized by the longer trajectory.
5 Learning from Suboptimal Search Strategies
Training on suboptimal search trajectories substantially improves Countdown solving over training only on optimal paths, while the resulting model does not simply reproduce one symbolic search strategy.
- Experimental comparison: The study compares language models trained on clean optimal paths with models trained on diverse, sometimes unsuccessful search trajectories.The experiment asks whether messy search experience is more useful than clean solutions alone.
- Results: 51.27% accuracy on held-out inputs was achieved by the SoS model, versus 25.73% for the optimal-path model.This advantage occurred despite the SoS dataset containing fewer correct solutions.
- Results: The SoS model must learn both state transitions and arithmetic, unlike symbolic search methods that receive an environment model exposing transitions.Its state-exploration error was 0.8%, with about two arithmetic errors per trajectory on average.
- Search strategy analysis: The SoS model’s visited-state alignment is not highly correlated with any single symbolic strategy.Its highest reported correlation is 0.57 with DFS using the sum heuristic, and its lowest is 0.27 with BFS using breadth size 5 and the sum heuristic.
- Search strategy analysis: SoS shows generally higher alignment with strategies using the sum heuristic without predominantly adopting one strategy from its training data.Correctness alignment exhibits a similar pattern across strategies.
6 Policy Improvement with Stream of Search
The paper improves Stream of Search models with STaR and APA, using correct generated trajectories or advantage-based policy optimization. Both methods improve performance and alter the model’s search behavior.
- Advantage-Induced Policy Alignment: APA uses an actor-critic setup with a value network, a correctness-and-length reward, and a reference policy that limits policy drift.The method was selected for its reported stability and robustness to hyperparameter changes.
- Advantage-Induced Policy Alignment: Updating APA’s reference policy when validation reward converges produces further improvement and is more stable than scheduling reductions in the APA coefficient.The reference distribution is reset to the current policy during training.
- Results: About 6% improvement over the base SoS model is obtained with APA, while STaR converges after three iterations and adds 5% on held-out-input tests.APA validation accuracy stops improving after about 4,000 training steps, with three reference-policy resets.
- Search behavior: After finetuning, the models prefer states associated with certain symbolic strategies, while APA is less aligned with those strategies than the base SoS model.These visitation changes suggest flexible use of search strategies and possible discovery of novel heuristics.
- Self-Taught Reasoner: STaR generates trajectories, filters them for correctness, and finetunes the SoS model on the resulting correct solutions.The procedure repeats until validation performance converges.
- Results: Approximately 36% of previously unsolved problems and about 4% of difficult problems are solved by the improved SoS models.Difficult problems are those unsolvable by any symbolic strategy used to train the models.
7 Discussion
The SoS framework represents structured search, including backtracking and heuristic reasoning, within language-model sequence modeling. Results support self-improvement and improved search, but empirical evidence remains restricted to Countdown and leaves several extensions open.
- SoS enables language models to learn structured search with backtracking, heuristic state evaluation, and world modeling inside a sequence-modeling paradigm.The framework represents search operations in a unified language and exposes productive mistakes rather than only optimal solutions.
- Policy improvement produces fewer arithmetic errors and more efficient correct solutions.
- Training on exploration and backtracking outperforms training solely on optimal solution steps.
- SoS supports internal search that can address error snowballing and lookahead challenges by backtracking and exploring alternative paths.
- Empirical results are restricted to Countdown, while state evaluation, additional search operations, and cross-domain transfer remain open questions.The authors also note that generating initial SoS datasets may be difficult when symbolic search algorithms are unavailable.
A Reproducibility Statement
The reproducibility statement provides code for the experiments, APA, and an optimized training fork.
- The authors release code for replicating the experiments and provide links to the APA implementation and optimized training fork.The training fork includes optimizations for parallelism, memory, mixed precision, and flash-attention.
B Creation of the SoS Dataset
The SoS dataset is generated from BFS and DFS symbolic search strategies guided by simple heuristics. It contains correct and optimal-path trajectories for Countdown problems.
- BFS and DFS generate SoS trajectories using heuristic functions based on target distance and factor distance.
- DFS uses a target-based heuristic threshold, while BFS retains the top b branches with b ranging from 1 to 5.
- 285,501 trajectories, or 57.1%, in the SoS dataset are correct.
- The optimal-path dataset contains only correct steps and provides 500,000 correct solutions.
- The search procedures track current and goal states, expand states, and append explored states and goal checks to the SoS trajectory.
C.1 Model Architecture
The models use a GPT-Neo transformer architecture and are trained with specified optimization settings. STaR samples and filters trajectories, while APA uses its own rollout and optimization configuration.
- The base model is GPT-Neo with 16 layers, 16 heads, a 4096-token context, and hidden size 1024.
- Base training uses 50,000 steps, batch size 24, cosine learning-rate scheduling, and learning rate 1 × 10^-5.
- STaR samples trajectories at temperature 0.8, retains only correct trajectories, and resets to the base SoS model each improvement iteration.
- APA uses 32 rollouts, temperature 1.0, batch size 8, two online epochs, and learning rate 1 × 10^-6.
D Analysis of generated trajectories
The evaluation tests trained models on seen targets with unseen input-number sets and on held-out targets. Error analysis categorizes trajectory failures, while policy improvement methods reduce exploration in correct trajectories.
- Models were evaluated on 10,000 tests with seen targets and unseen input-number sets, plus 10,000 tests with unseen or held-out targets.
- Both policy improvement methods reduced the average number of states or nodes explored before reaching the goal in correct trajectories compared with the SoS LM.
- Trajectory errors were grouped into arithmetic, formatting, exploration, and other errors involving incorrect associated number sets.
- APA training significantly suppressed arithmetic errors while preserving a low rate of exploration errors.
- Table 1 reports trained-model statistics averaged over 10,000 tests with 95% confidence intervals.
- Table 2 reports average errors per trajectory over 10,000 tests with 95% confidence intervals.