Stream of Search (SoS): Learning to Search in Language
Authors: Kanishk Gandhi
Denise H J Lee
Gabriel Grand
Muxin Liu
Winson Cheng
Archit Sharma
Noah Goodman
Venue: COLM (2024)
Keywords: search, planning, reasoning, backtracking, self-improvement, bootstrapping
TL;DR: We show that language models can learn to search in language, autonomously using and discovering new search strategies to solve problems.

Abstract: 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 --- 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.

Published as a conference paper at COLM 2024
Stream of Search (SoS): Learning to Search in Language
Kanishk Gandhi' Stanford University
Denise Lee Stanford University
Gabriel Grand MIT
Muxin Liu Harvey Mudd
Winson Cheng Stanford University
Archit Sharma Stanford University
Noah D. Goodman Stanford University
Abstract
Language models are rarely shown fruitful mistakes while training. They then struggle to look beyond the next token, suffering from snowballing errors and struggling to predict the consequence their actions several steps ahead. In this paper, we show how language models can be taught search by representing the process of search in language, as flattened string a stream of search (SoS). We propose unified language for search that captures an array of different symbolic search strategies. We demon- strate our approach using the simple yet difficult game of Countdown, where the goal is to combine input numbers with arithmetic operations to maansorreer-bsee language model
reach target number. We pretrain 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-Inducel Policy Alignment (APA) and Self-Taught Reasoner (STaR). The finetuned SoS models solve 36% of previously unsolved problems, including prob- lems 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.
"To err is human, to backtrack divine" -Apocryphal
1 Introduction
Imagine, only ever seeing the right solutions to problems, never mistake or recovery from it You might learn that problems must be solved in one clean pass, rather than through exploration and error. Most data used to train language models (LMs) only reflects the outcome of a decision making process, not the process itself. LMs never learn to make mistakes. They never learn to search, plan or backtrack. Complex decision-making and reasoning requires search. In this paper we explore the impact of training a LM on the search process, including mistakes, and then allowing them to self-improve.
Transformer-basel ruto-segeessive models have been shown to struggle with planning (Valmeekam et al., 2024; Pallagani et al., 2023; Momennejad et al., 2024). Recent work has highlighted this weakness in autoregressive models by identifying two main issues (LeCun, 2023; Bachmann & Nagarajan, 2024): 1) the snowballing of errors, where single mistake can compound and lead to increasingly poor performance in subsequent steps (Ross et al., 2011; Arora et al., 2022), and 2) difficulty in 'lookahead tasks', where the model must predict the consequences of its actions several steps ahead (credit assignment, Cf. Sutton & Barto, 2018). Both of these issues can be attributed to limited ability to search and backtrack. While recent efforts have combined language models with symbolic search algorithms (Ahn et al., 2022;
"Corresponding author: kanishk, Sanishh.gandhi@stt gandhi@stanforr edu !Code Available Here: https:/ httpss//gittub.oom htwp:/hww..ooo/aaaiikkk..gce /waiiskkggggaaaea wttp://stt-e-eeeaaaaeaa
1

---

Published as a conference paper at COLM 2024
(a) Search problem
(b) Streams Search
Target: 50
Solution: 39 13 52
Numberr::
auu Esoloring Opunation uprs Hooslin Numbors 10s
Start state
66,
renented hods on tot 20 Openaiinn 06-13 53
52-2
Ganeet Stote poitts 22l Openatinn 186-13 zz)
Exoloring Openaton: 30-33772 fleaiing Numbera 1gs
Operition 39/3377
Mexiing to hisde mu.op
Cument Stuwe 72l Openatore, teo-ges) 30n230721
Openttion
Reasttin Mumbert
56 60 umiqust Selution
Backtracking
Diverse Search
Gumeet Stute
sd. Openatins (56-1pe5p)
Exploring Openalion 50-30-40 Rasiling Numbura las 20)
Strategies
Moving hodis 40.2
Search Strategy
Gureet Stete 5m3M 13 sPl Opecationn hapettrn
Gn Geneested hoxtn trc.
Cpectton 66/12
Moxing to 40.25
Fansoring Opeettion go- so, Seoutting Numbore tpsl 5k equal: Reached
Goal state
(c) Pre- Tre-Training Policy Improvement
Improved Search Strategies
APA & STaR
Language Model
Figure 1: Overview of the Stream of Search (SoS) framework. (a) A search problem in Countdown is instantiated with input numbers and target number. The input numbers need to be combined with simple arithmetic operations to get the target. (b) The Stream of Search dataset contains search trajectories generated by diverse search strategies, including exploration and backtracking. (c) The language model is first trained on the SoS dataset and then iteratively improved using policy improvement techniques such as APA and STaR.
Yao et al., 2024) to mitigate some of these problems, they are limited - only supplementing language models during inference -they leave open the question of whether language models could effectively carry out search themselves. The most consequential outcome from learning to search could be during the training itself (Silver et al., 2018). If language models can learn to search during training, they may be able to discover more flexible search strategies through eef-impprovemen. This could lead to models that are better equipped to handle the challenges posed by error compounding and lookahead tasks.
In this paper, we demonstrate that language models can taught to search and backtrack in language, representing the process as serialized string, Stream of Search (SoS). We systematize the different components of search, such as exploration, backtracking, and pruning in unified language. We instantiate this unified language in the context of a search problem inspired by the game of Countdown, a generalized version of the game of 24 (Yao et al., 2024; Countdown, 2024). problem consists of input numbers and target number. The goal is to combine the input numbers with arithmetic operations to reach the target (see Fig. la). Countdown presents challenging search problem due to its high branching factor and the need to efficiently navigate the combinatorial search space towards the target number. We generate an initial training dataset of search trajectories or streams of search using different symbolic planners and simple heuristic functions, then train a language model on this diverse dataset. This can be contrasted with training on the optimal path once solution is found, without the process of search and backtracking, We find that the Stream of Search LM substantially outperforms models trained to predict the optimal steps. Moreover, we observe improvements in the search and planning abilities of the Stream of Search model when finetuned to optimize for correctness using Advantage-Indl Alignment (APA, Zhu et al., 2023) and expert iteration with STaR (Zelikman et al., 2022; Gulcehre et al., 2023).
Policy
Our results indicate that transformer-based language models, when shown how to recover from mistakes and search through different options, can learn to solve problems by searching. More importantly, our results indicate that these models can self-improve to autonomously
2

---

Published as a conference paper at COLM 2024
use different search strategies, solving previously unsolved problems. Finally, we see some evidence that they discover new search strategies when trained to optimize for accuracy.
2
Related Works
Language models as components of a search system. A line of recent work integrates language models as part of larger search and planning systems (Yao et al., 2024; Ahn et al., 2022). In these methods, LMs typically play two roles: (1) to generate candidate actions or successor states in the reasoning process, and (2) evaluate proposed actions or states, by determining validity and/or or assigning a heuristic value. A symbolic search algorithm such as BFS DFS dictates the strategy for exploration, and how the steps or evaluators are called (Yao et al., 2024; Besta et al., 2023). While these methods have been shown to improve search accuracy on certain problems, the LM components are typically used only for inference, so their reasoning ability is not improved. In contrast, our work focuses on training LMs that are capable of exploration, backtracking, and other critical components of reasoning. Relative to these "extrinsic" methods, which use fixed search strategies, our method learns an "intrinsic" policy that allows the LM to autonomously search the solution space. In doing so, we avoid the high inference costs (Sel et al., 2023) required by tree-of-thoughts style approaches.
In-context demonstrations of search. In contrast to using separate symbolic search algo- rithms to guide search, in-context demonstrations can also be used to demonstrate search procedures in language (Gandhi et al., 2023; Sel et al., 2023). These methods allow the language model to perform tree search based on the demonstrated search procedures. While improving efficiency in search, these methods are restricted by the choices of the search procedure demonstrated in the in-context examples, such as the demonstrated exploration strategy or the search heuristic.
Process supervision. Process supervision of LMs is another method where an external verifier model is trained to learn the rewards of each intermediate reasoning step and the LM is then trained on this detailed feedback. Process supervision was shown to outperform outcome supervision in mathematical reasoning tasks (Lightman al., 2023). However, the training of the verifier requires very large labelled dataset with human-generated annotations for each intermediate step, which may not scale well for other, more complex domains. Comparatively, our method directly improves the model's planning and search abilities and removes the need to train a verifier reward model.
Learning from search trajectories. Similar to work, Yang et al. (2022) train transformer models on search trajectories to mimic search strategies such as Monte Carlo Tree Search or BFS. While their approach focuses on mimicking single fixed search procedure, we are interested in autonomous usage of different search procedures and the discovery of new ones. In recent concurrent work, Lehnert al. (2024) employ A search traces train base transformer model, where each trace includes the state and fixed A heuristic value and search procedure. Their work aims to train transformers to closely imitate A*, and enhance its efficiency. In contrast, our method emphasizes discovery of new search strategies. We let the model learn its strategies, with the choice of what to keep implicit vs explicit the search process shaping what the model improves upon.
related line of research, starting with Daumé et al. (2009) (c.f. Ross et al. (2011); Ran- zato et al. (2015)), suggests supervising the process of finding solution into structured prediction problem. This is similar in spirit to SoS where there is supervision over the process of search. While Daumé et al. (2009) provides theoretical framework for learning policy over structured 'process' data, SoS demonstrates how the search process itself can be represented
flat language sequence, defining explicit vocabulary to represent elements of search algorithms, such as backtracking and heuristics, allowing LMs to directly learn diverse search strategies from data.
3

---

Published as a conference paper at COLM 2024
Stream of Search
Current State
5Rk,
State Exparsion
Moving to
Moving betovenn states
Optimal Path
Impl citl evalucte
Sala?
heuristic
Ser SEG Moving to Rox x. Rexum SEtkal Poct- Gool Reached!
Mosing to Aoi
Only cornect steps included
2. Aou
Impikeit erenndy
X
Fo GElsy)
eeploring
Peas
Gooll Check
Mosing to Av
Backtrackng
"ard!ng
A Aut
Ru SElsa)
Suns Gool Reached!
Goal Stare
Figure 2: A visualization of how a search process translated into stream of search with language for search. (left) The search process represented as a tree, with different states and operations. The colored states represent the search trajectory T, the green states represent the correct path to the goal P and the arrows represent transitions between the states. The black circles represent unexplored states. (center) The search process serialized as text to create stream of search. The labels specify the different components of the process. See Fig. lb for how this is realized in Countdown. (right) The optimal path, P, to the goal state. Backtracking, exploration and the messy process of search are excluded.
3
A Language for Search
The problem space can be modeled as a Markov Decision Process (MDP), with set of states, S, representing the steps in solving problem; a set of actions A, representing operations that the model can perform on states or to transition between them; a transition function T S A - S, defining the transition from one state to another based on an action; and reward function R: S > IR, assigning a reward for reaching the goal state.
The process of search can be modeled as follows. The query for search has an initial state so E S and a goal state Sg E S. The search tree is defined by the problem itself. It contains all possible explorations of all actions from the initial state so (containing the input) to each possi- ble child state of so, and so on until a leaf state (a state with no children) is reached. A correct path to the solution P is sequence of states and actions (50, ag, $1, a1, .0g.............. - 1.g1 1g in the search tree, where each successive state $;+1 is obtained by applying valid action aj E A to the previous state S, ie,, 1 = T(s;, 4;), and the final state in the sequence is the goal state 8g (see Fig. 2).
We are interested in the process of search, which corresponds to moving around the search tree until solution is found. To represent this process, we propose vocabulary of primitive operations that can be used define components of different search algorithms. This is in spirit similar to defining domain-specific language for planning (Fikes & Nilsson, 1971).
Current State: Se, The state s. that is being explored.
Goal State: $g, The state 5g that is the target.
State Queue: Sq The states at the 'frontier' of the the the trajectory that haven't been explored yet.
State Expansion Function: SE S , Explore a state adjacent to the current
state based on a transition function T.
Exploration Choice: Choosing the order of states to explore following the state expansion. Eg: Should the breadth be explored first (BFS) the depth (DFS) or any the frontier states (A*).
Pruning: An action discarding states subtrees that are unlikely to lead to solution.
Backtracking: An action move between nodes that have been explored. This allows the algorithm to choose which state should be expanded next.
Goal Check: An action to check if the current state is the goal state (s. == $g).
Heuristic: A function It E H. : S x S + R that approximates the distance of the current state s, from the goal 8g- This can be used to decide which states should be explored or pruned. The heuristic function serves as an approximation of the
2This is non-exhaustive list, choose to represent the most salient operations here.
4

---

Published as a conference paper at COLM 2024
value function, guiding the search or decision-making process by estimating the desirability of states in terms of their distance to the goal. goal.
Each of these operations can be left implicit, affecting how the trajectory unfolds, or made explicit in language as part of the search trajectory T. When operations are implicit, a model is more likely to internalize abstract representations for them that can be improved with training. Explicit operations will turn into explicit reasoning moves made by the LM. We choose to represent the current state, the goal state, the backtracking operations, the goal checks and the exploration choices explicitly in the trajectory (written using language, Fig. 2). We choose to keep the heuristic functions, values of states and the pruning strategy implicit.
4 Problem Setup
Task Description: Countdown To show the utiliiyy of using streams of search, we look at generalization of the 24 Game (Yang et al., 2022) called Countdown (Countdown, 2024). Countdown is a game where a set of input numbers need to be combined with simple arithmetic operations to reach target number (see Fig. la). We choose this task, since it has a high branching factor ((H) 4 for a depth that has N inputs) and thus requires planning, search and backtracking to be solved. We consider problems with 4 input numbers since these problems are challenging enough to have long search traces without the search traces exceeding a standard LM context window (for example, the search trajectories for games with 5 input numbers can be 60,000 tokens long). The range of target numbers in our problems is from 10 to 100. We randomly hold out 10% of the targets for 'out-of- distribution' evaluation.
Training data To train a model on streams of search for Countdown, we construct a syn- thetic dataset using set of diverse, and suboptimal symbolic search strategies. To construct our Stream of Search dataset, we define 12 search strategies based on breadth first search (BFS) and depth first search (DFS) (see App. Alg. 4, Alg. 3) that rely on two simple and interpretable heuristic functions. The heuristics we use to guide search are 1) the absolute difference between the sum of the remaining options and the target and 2) distance to the factors of the target (see App. B). We generate stream of search dataset of 500,000 search trajectories and corresponding trajectories with just the optimal solution. Of these 500,000 trajectories, only about 57% of the trajectories lead to the solution. In our dataset, a search trajectory stream of search is serialized
string that represents a list of tree nodes / states in the order of traversal (either generation or exploration). We hold out two kinds of generalization cases: 1) seen targets with new sets of inputs, and 2) new targets with new sets of inputs.
Metrics To measure accuracy in Countdown, evaluate the percentage of problems for which the model is able to generate a correct solution trajectory. More formally, we define correctness as a binary function, 1, if the correct path to the solution, P, is present in the generated trajectory, T.
To quantitatively understand the search strategies used by the trained models, we define two ways to measure alignment between different search strategies:
1. Alignment of Correctness: This measures whether two search strategies solve the same set of problems correctly and incorrectly. We calculate this as the Pearson correlation between the solved and unsolved problems for the two search strategies.
2. Alignment of States Visited: This measures the overlap in the states visited by two search strategies. To calculate this, we parse the search trajectories, Ti and Tz, into their constituent states and count the number of states that are common between them. We normalize this count by dividing it by the maximum number of states in ITina the two trajectories. More formally: State Alignment( Tı, T2) maxi (2ll) where |Ti and |T2] denote the number of states in trajectories Ti and T2, respectively. To measure alignment between two models, we calculate the state alignment score for each problem and then compute the mean of these scores.
5

---

Published as a conference paper at COLM 2024
Mediods
(b) Aligmmmt Correttres
Aligrment of Statee Vistted
Taguts New
New Taigrt New Seper
DM5 tam and
DIFE with Starphy
07-
BFS vith Sax.
IITS St (hredtrf1.
Nutaac Broodk+41.
II'S with Shulppy
02
505
Ustand
015 vit Sam BFF vith Mikippy SoS
D25
3P5 Sun
wTh McRply
Aseaga
7dth
Swith
Figure 3: (a) Average accuracy of the symbolic search strategies used to construct the Stream of Search dataset compared to the model trained with optimal paths and the Stream of Search model. Error bars represent 95% binomial confidence intervals.
(b) Alignment of
Correctness: The correlations of performance on different problems for the Stream of Search
model and different symbolic search strategies. (c) Alignment of States Visited: Number of times each search strategy visited the same states, normalized by the maximum number of states in the two trajectories. We find that the alignment scores for states visited for the SoS model are not highly correlated with any single symbolic strategy.
5 Learning from Suboptimal Search Strategies
Is it more useful to learn from clean, optimal solutions, from messy, and sometimes unsuccessful, search trajectories? In this section we train LMs from scratch on one or the other and evaluate there performance at solving held out Countdown problems.
Experiment Setup. For all experiments, we train a GPT-Neo model (Gao et al., 2020) with 250M parameters and context length of 4096 tokens (see appendix for more details). We train models with a causal language modeling objective on the 500,000 Countdown problems in two conditions: 1) Optimal Paths (OP): the model is trained to predict the correct, and optimal path P for all problems in the dataset. 2) Stream of Search (SoS): The model is trained on search trajectories T sampled from different search strategies. The SoS dataset has 285,000 correct solutions while the OP dataset has 500,000 correct solutions. Both models are trained up to the same number of gradient steps.
Results. We find that the model trained on streams of search outperforms the model trained on the optimal solutions (see Fig. 3a). The SoS model achieves an accuracy of 51.27% on held-out inputs compared to 15.73% of the OP model. A similar pattern is seen for the test set with held-out targets. This shows that despite having fewer 'correct' examples for training, the SoS model outperforms the OP model. The SoS model has a slightly lower accuracy compared to the average accuracy of the symbolic search methods used to construct the SoS training dataset. However, it is worth noting that SoS is solving a harder problem: symbolic search relies on an environment model to expose state transitions, while the SoS LM simulates those transitions itself, effectively learning to search with an internal 'world model'; the language model must also learn how to perform arithmetic operations. Despite these challenges, the SoS model generates valid trajectories with low error rate in state exploration (0.8%) and only about 2 arithmetic errors per trajectory on average (see App. Tab. 2).
To understand the strategies that the trained SoS model uses, we measure the alignment of the model generated search trajectories with symbolic strategies. We find that the alignment scores for states visited for the SoS model are not highly correlated with any single symbolic strategy (see Fig. 3c). The highest correlation of SoS is with DFS using the sum heuristic (0.57), and the lowest correlation is with BFS using breadth size of 5 and the sum heuristic (0.27). A similar pattern is observed when we measure the alignments of correctness (see Fig. 3b) for the SoS model and the symbolic search strategies. Overall, the SoS model has higher scores for alignment with strategies that use the sum heuristic, but does not seem to predominantly use any one strategy from its training data.
6

---

Published as a conference paper at COLM 2024
6 Policy Improvement with Stream of Search
The SoS LM has learned to use search to solve new problems. Can it learn to improve upon the symbolic strategies seen in its training data? In this section, we explore if the model can self-improve with feedback based on correctness and efficiency. We measure ability to solve previously unsolved problems (by the symbolic search strategies) from the training dataset and solve difficult problems from the training set that none of the symbolic search strategies can solve. To improve the model, we use two RL strategies, expert iteration using STaR (Zelikman et al., 2022), and Advantage-Induces Policy Alignment (Zhu et al., 2023).
Experiment Setup. To improve the models with STaR (Alg- 1), we use problems in the training dataset to generate 100,000 correct trajectories. We sample with a temperature of 0.8. These trajectories are then used to finetune the model. We repeat this process until we see convergence in performance on the validation set.
Alternatively, we can use advantage-induced policy alignment (APA; Alg. 2). APA is an Actor-Critic reinforcement learning technique that involves creating copy of the language model to serve as a value network that then used to enhance the policy, the original language model. We define straightforward reward function that takes into account the correctness and length of the generated trajectory. We try APA, as we wanted to see if using a separate value network would improve exploration and the path-stitching abilities of the language model. We chose APA over other methods like Proximal Policy Optimization (PPO) due to its stability and robustness changes in hyperparameters.
APA uses a reference policy, Te,, to prevent the policy from drifting from its initial state. We observe that updating the reference policy whenever the validation reward converges results in further policy improvement (see Fig. 4b). This shifting of the reference distribution can be interpreted as a means to reduce the weight assigned to staying close to the reference distribution (the À parameter in the APA objective). In practice, we found the stra strategy of shifting the reference distribution to be more stable for training when compared to designing a schedule for reducing over training.
Algorithm 1 Expert Iteration with STaR (Zelikman et al., 2022)
1: Input The SoS model trained on Stream of Search Dataset D = ((xx, Tm)
2: Mo < SoS
D Initialize the SoS model
3: for n in 1...N N do
4:
Ti + Mn-1 (x: Vi E [1, m]
Perform trajectory generation Filter trajectories based on correctness
5:
D, + (xx, li E [1, s.t. Pi T
6: Mn < train (M, DR)
Finetune the model on correct solutions
7: end for
Algorithm 2 Advantage Indued Policy-Alignment (APA) (Zhu et al., 2023)
1: Input: An initial policy parameter Tini a given reward function R, Advantage coeffi- cient A.
2: 7g + Tinit
3: neef t- Ninit
Copy the SoS model to create reference network. b Copy the SoS model to create value network.
4: TTaalue +- Tinit
5: for i .....o
(f) (()) 6: Roll out ne, , to produce dataset D; = (5 a1
(f)
(f) (s an (r
7: Update policy function according to
8:
0, arg maxy APAA (0; D:).
We omit the critic loss for simplicity
9: where
Adu 70 (s,a log Kye (als))
10:
CAPA (0; D) = [D] Eisa)eD (log (als)) 11: If validation reward converges, update Ref
12:
Tref & ne,
13:
end for
7

---

Published as a conference paper at COLM 2024
Policy Impmyement with Stream of Search
AA Validation Accuray: Resetting the Reference Policy
STaR Test coutary
Soon Targen New Ieguse
Targen lepets
Thegare tepare
New Targts- New Irpurs
0.36
0.54
8.50
Avcuac
0.52
0.43
StetEt- OE
5:5 45TAR
SOSHARL
1000
2000
3000
4000
3.40
Screzz DE
STAR 1
STER
STIR3
Search
Trmiinn Steps
Search
Figure 4: (a) Improvement in accuracy with different policy improvement methods. (b) Accuracy on the validation set of problems during APA training. We reset the reference policy after the validation accuracy converges. Different colors represent training after resetting the reference policy to the current model. We stop seeing improvements after 3 such resets. (c) Improvements in test accuracies with each STaR iteration. Error bars
represent 95% Binomial Confidence Intervals.
The SoS models converge after 3 iterations of STaR finetuning (Fig. 4). We see that after three iterations, the finetuned SoS+STaR model solves an additional 5% of the held-out inputs test set beyond the base SoS model. A similar pattern is seen for the held-out targets test set. When SoS models are finetuned with APA, we that the validation accuracy stops improving after about 4000 training steps. We reset the reference policy 3 times, when the validation accuracy stops improving see Fig. 4b; different colors represent training after resetting the reference policy. Overall, we see an improvement of about 6% over the base SoS model.
Results.
When we analyze the difference between the base and finetuned models in terms of the alignment of the states visited, we observe (Fig. 5a) that both the STaR and APA models visit more states associated with the 'multiply' heuristic, which measures distance to the factors of the target. Further, we note that the APA model is less aligned with the symbolic strategies compared to the base SoS model, indicating that it diverges more from the symbolic strategies and employs different strategies for searching. These state visitation metrics provide insights into how the SoS+STaR and SoS+APA models can flexibly utilize various search strategies, potentially discovering novel heuristics and search methods.
To further evaluate the performance of the improved models, we select 10,000 problems from the SoS training set that were unsolved by symbolic strategies when the dataset was generated, and 10,000 difficult problems that none of the symbolic strategies used to train the SoS models can solve. Remarkably, the models are able to solve approximately 36% of the previously unsolved problems (Fig. 5b) and about 4% of the difficult problems (Fig. 5c). Finally, SoS + APA and SoS STaR models also have better models of the environment, making fewer errors while searching (Fig. 6; App. Tab. 2), and finding the solution more quickly (Fig. 6).
7 Discussion
We have introduced the Stream of Search (SoS) framework enabling language models to learn solve problems by searching in language, without any external structure or components. By systematizing the elements of search into unified language, we are able to represent various search strategies common format to construct a dataset with diverse streams of search. Our experiments demonstrate that training language models to search leads to superior performance compared to models trained solely on optimal trajectories. This highlights the importance of exposing models to the messy process of problem solving, with exploration and backtracking, instead of only the ideal solution steps. SoS models can then self-improve by optimizing for correctness, using STaR and APA.
The SoS framework may address criticisms (LeCun, 2023; Bachmann & Nagarajan, 2024) of language models for planning and problem solving. The problem of snowballing errors is
8

---

Published as a conference paper at COLM 2024
(a) Difference in lignment Stazes visited aller Policy Ingeovement
(b) Solving Unsolved Problems
Solving Difficult Problems
(that bo symbolic strotegies solve)
0.36
D15 Maltale
6.33 bel
54 Dallhaa Dhideeenede Delved
0.24
0.40
-302
0.32
0-02
-206
Muttaty bel
0TS Mulupfy b-2
6.15
0.20
9,01
DPS Multaly be3
-A04
BPS Multaly b-d
4006
00.55
Multaly bed
4.00
0.23
Stremn of
9:5 4 -5768
9:00
Stwar
SuS STaal So5 ARR
565 STaR
565 IAR
Starch
Search
Figure 5: (a) Difference in the states visited compared after finetuning with APA and STaR. We see that the model prefers the states visited by certain symbolic strategies after improvement. (b) The percentage of previously unsolved problems from the training set that the trained models were able to solve. (c) The percentage of difficult problems, defined as those not solvable by any symbolic algorithms, that the trained SoS models successfully solved. Error bars are 95% Binomial Confidence Intervals.
the base Stream of Search Model
(a) A rithmetic Errors
(b) Length of Correct Solutions
Soen Targets + New Inputs
Soon Targets New Inputs
Figure (left) Aver- age number of arith- metic errors made per search trajectory by dif- ferent models. Policy improvement leads to fewer errors. (right) Average number of states visited per tra- jectory for correct solu- tions. Policy improve- ment leads more ef- ficient solutions.
Targets New Inputs
New Targets Inputs
Avarage of States Visited
Averuge Erruss per Tesjectayy
2.0
75
1.5
76
85
Streem of SoS+STaR Search
SoS+A PA
Stream of
SoS+STaR
SoS+A
Search
addressed by teaching model to backtrack. Search allows models to explore alternative paths, overcoming failures in lookahead tasks by considering multiple possible outcomes before committing to course of action. Crucially, SoS leads language models to learn an internal 'world model' for search. Unlike symbolic search that relies on an explicit environment model, SoS models simulate state transitions themselves. Using learned world models allows more adaptable and generalizable search (Cf. Schrittwieser et al., 2020) and addresses key criticism of pretrained L.Ms.
Our empirical results were restricted to the game of Countdown. Yet Countdown, with high branching factor and variable goal states, captures the characteristics of complex planning problems. We are optimistic that SoS extends to more challenging, real-world tasks. In the short-term, externally-structun et al., 2024) are likely to be more efficient for these tasks; in the longer run the increased flexibility and learnability of internally-structun plausible these approaches can be combined, for sinstance, with with ToT providing initial training data for SoS.
search methods such as Tree of Thought (Yao
search (i.e. SoS) may prevail. also
While we leave the evaluation of states to be done implicitly by the network in our current work, explicitly representing state evaluations (Gandhi al., 2023) and introducing other formalizable operations such as limits, summarization, cycle checks, and subgoal setting could enhance the SoS framework. Future research could explore integrating subgoals and hierarchical planning, well as incorporating reflection and self-evaluation to enable models to discover and improve novel search strategies (Huang et al., 2023; Shinn et al.,

---

Published as a conference paper at COLM 2024
2024; Stechly et al., 2024). Generating the initial SoS dataset can be challenging, it is not always feasible to create symbolic search algorithms to solve problems. An important question is how well search abilities transfer between domains and between formal and informal domains.
In conclusion, we have shown that hallmarks of symbolic reasoning structured search with backtracking, heuristic state evaluation and world modeling- can be achieved within sequence modeling paradigm. To do so requires showing language models examples of productive mistakes, not only the optimal final solutions. By embracing the diversity of search strategies and iteratively refining models, we can unlock the potential of language models to tackle complex problems and discover new ways to solve them.
Acknowledgements
We would like to thank Gabriel Poesia, Jacob Andreas, Joy He- He-Yueya, Dongwei Jiang, Joseph Feffer, Eric Zelikman, Jan-Philipp Fränken and Ced Zhang for their discussions and support. This worked was supported by the Stanford Human-Centere Artifical Intelligence (HAT)- Google grant, and the NSF Expeditions Grant, Award Number (FAIN) 1918771.
References
Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Chuyuan Fu, Keerthana Gopalakrishnan, Karol Hausman, et al. Do as i can, not as say: Grounding language in robotic affordances. arXiv preprint arXio:2204.01691, 2022.
Kushal Arora, Layla El Asri, Hareesh Bahuleyan, and Jackie Chi Kit Cheung. Why expo- sure bias matters: An imitation learning perspective of error accumulation in language generation. arXiv preprint ar arXiv:2204.01171, 2022.
Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls of next-token prediction. arviv preprint ar 4rXi::2433.06963 2024.
Maciej Besta, Nils Blach, Ales Kubicek, Robert Gerstenberger, Lukas Gianinazzi, Joanna Gajda, Tomasz Lehmann, Michal Podstawski, Hubert Niewiadomski, Piotr Nyczyk, et al. Graph of thoughts: Solving elaborate problems with large language models. arXiv preprint arXiv:2308.00 (66.:2308.09687, 2023.
Countdown. Countdown (game show). https://en httts://ee.wwkkieeiiaa wtw.pedda..rgg/iiik///nnnoooooo (game_show) 2024. [Online; accessed 29-March-2024],
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and mefore-efficiint exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344-16359, 2022.
Hal Daumé, John Langford, and Daniel Marcu. Search-based structured prediction. Machine learning, 75:297-325, 2009.
Richard E Fikes and Nils J Nilsson. Strips: A new approach to the application of theorem proving to problem solving. Artificial intelligence, 2(3-4):189-208, 1971.
Kanishk Gandhi, Dorsa Sadigh, and Noah D Goodman. Strategic reasoning with language models. arXiv preprint ar arXio:2305.19166 2023.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arxiv preprint arX:v:200100000 100000,007 2020.
Caglar Gulcehre, Tom Le Paine, Srivatsan Srinivasan, Ksenia Konyushkova, Lotte Weerts, Abhishek Sharma, Aditya Siddhant, Alex Ahern, Miaosen Wang, Chenjie Gu, et al. Reinforced self-training (rest) for language modeling. ar arviv preprint ar XrXi:2308.,88995, 2023.
10

---

Published as a conference paper at COLM 2024
Jie Huang, Xinyun Chen, Swaroop Mishra, Huaixiu Steven Zheng, Adams Wei Yu, Xinying Song, and Denny Zhou. Large language models cannot self-correct reasoning yet. arvio preprint ar Xii:23100001798 2023.
LeCun. Do large language models need sensory grounding for meaning and under- standing. In Workshop on Philosophy of Deep Learning, NYU Center for Mind, Brain, and Consciousness and the Columbia Center for Science and Suciety, 2023.
Lucas Lehnert, Sainbayar Sukhbaatar, Paul Movay, Michael Rabbat, and Yuandong Tian. Beyond a Better planning with transformers via search dynamics bootstrapping. Xiv preprint #rXio22402.14085 2024.
Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, llya Sutskever, and Karl Cobbe. Let's verify step by step. Xiv preprint ar Xii:23005.20033 2023.
Ida Momennejad, Hosein Hasanbeig, Felipe Vieira Frujeri, Hiteshi Sharma, Nebojsa Jojic, Hamid Palangi, Robert Ness, and Jonathan Larson Evaluating cognitive maps and planning in large language models with cogeval. Advances in Neural Information Processing Systems, 36, 2024.
Vishal Pallagani, Bharath Muppasani, Keerthiram Murugesan, Francesca Rossi, Biplav Sri- vastava, Lior Horesh, Francesco Fabiano, and Andrea Loreggia. Understanding the capa- bilities of large language models for automated planning, arXiv preprint ar 111123105.16151, 2023.
Mare' Aurelio Ranzato, Sumit Chopra, Michael Auli, and Wojciech Zaremba. Sequence level training with recurrent neural networks. arviv preprint arXiv:1511 (tto:151..063,, 2015.
Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pp. 627-635. JMLR Workshop and Conference Proceedings, 2011.
Julian Schrittwieser, loannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839): 604-609, 2020.
Bilgehan Sel, Ahmad Al-Tawaha, Vanshaj Khattar, Lu Wang, Ruoxi Jia, and Ming Jin. Algorithm of thoughts: Enhancing exploration of ideas in large language models. Xiv preprint ar Xir::2308.100000 10379, 2023.
Noah Shinn, Federico Cassano, Ashwin Gopinath, Karthik Narasimhan, and Shunyu Yao. Reflexion: Language agents with verbal reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al.
A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, 362(6419):1140-11444 2018
Kaya Stechly, Karthik Valmeekam, and Subbarao Kambhampati. On the self-verification limitations of large language models on reasoning and planning tasks. ar preprint arXiv:2402.08115, 2024.
Richard Sutton and Andrew G Barto. Reinforcementt learning: An introduction. MIT press, 2018.
Karthik Valmeekam, Matthew Marquez, Sarath Sreedharan, and Subbarao Kambhampati. On the planning abilities of large language models-a critical investigation. Advances in Neural Information Processing Systems, 36, 2024.
11

---

Published as a conference paper at COLM 2024
Mengjiao Sherry Yang, Dale Schuurmans, Pieter Abbeel, and Ofir Nachum. Chain of thought imitation with procedure cloning. Advances in Neural Information Processing Systems, 35: 36366-36381, 2022.
Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. Advances in Neural Information Processing Systems, 36, 2024.
Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah Goodman. Star: Bootstrapping reasoning with reasoning. Advances in Neural Information Processing Systems, 35:15476-15488, 2022.
Banghua Zhu, Hiteshi Sharma, Felipe Vieira Frujeri, Shi Dong, Chenguang Zhu, Michael I Jordan, and Jiantao Jiao. Fine-tuning language models with advantage-induced policy alignment. ar Xiv preprint arXi::2306.02231, 2023.
12

---

Published as a conference paper at COLM 2024
A Reproducibility Statement
We release the code to replicate our experiments here: We use the APA implementation available here: https:/ github..com www.b..com///om microsoot///////// SEHHH/PAHHHPPAA.AAA- tree/main examples. We opti- mize their code for parallelism, memory, mixed precision training and to use flash-attention (Dao et al., 2022). See our training fork here: https:/ /ttp//b.ttcboommmm/ htt.hhub..o//aaissskkk.... / /RHHF-APA.
https://githut httgtggt...coooo com catthhu.com/kaaissee /knnishkkk wotm///aiiik////rmroffssassn /trream-of--earcch
B Creation of the SoS Dataset
SoS Dataset.To generate the Stream of Search dataset, we use two symbolic strategies, BFS (Alg. 4) and DFS (Alg. 3). These strategies are guided by two heuristic functinos to generate trajectories to create the SoS training set. Let the input numbers for the Countdown problem be [11, .2.... nkk,, and the target number be T. The sum heuristic measures the distance between the sum of all input numbers to the target number: haum (I, T) = - Σ n, The multiply heuristic measures the minimum distance between the sum of inputs and the factors of the target. If T has factors [f, f2..2.... fm). So, the multiply heuristic can written as Hmudtiply = min minf E nj]lvj [1. m]). For DFS, the threshold for the heuristic is set n;
to the value of the target hig T. For BFS, breadth limit b E (,,2,3,4,) 4, is used to keep only the top b branches each state expansion. These hyperparameters provide a simple and intuitive way to design symbolic algorithms to generate the SoS dataset.
The SoS dataset has 285,501 (57.1%) correct trajectories.
Optimal Paths dataset. The optimal trajectories only include the correct steps to reach the goal state. There are 500,000 correct solutions shown to the optimal paths model.
Algorithm SoS-DFS
Require: Problem x, Heuristic State Expansion Function SE, threshold hig Goal State 5g- SoS ""
Sc x
if S. == Sg then
return SoS "Goal Reached"
end if
if G(..) is Null and se != Sg then
return SoS end if
Sg + -sss E SE(sc) s.t. ht(s) > hn
fors E S do
SoS +=
SoS DFS(s)
end for return SoS
Optimal Path Solution
Durrent State: 18: [74, 24, 443, Operations: []
Exploring Operation: 74424=98, Resulting Numbers: [36, 44, 98]
Generated Node #2: [36, 44, 98] from Operation: 74-24=98
turrent State: 18:[36, 44, 98], Operations: ["74-24-98
Exploring Operation: 36+44=80, Resulting Numbers: [98, 96]
Generated Node #3: [98, 88] from Operation: 36-44=80
Durrent State: 18: 801, Operations: ["74-24=98", '36-44=88"1
Exploring Operation: 98-88-18. Resulting Numbers: [18] 18,18 equal: Goal Reached
Figure 7: An example of an Optimal Path to the solution.
13

---

Published as a conference paper at COLM 2024
Algorithm 4 SoS-BFS
Require: Problem x, Heuristic h, breadth limit b, Next State Expander SE, Goal State $g
$4,0 f-
SoS
for i = .I...,Tdoo do
9c + s:-1 where si-1 E Sq.i1
SoS += (sc)
if Sg then
return SoS += "Goal Reached"
end if
S - (E(sc)
Hi h(,))
Sq + arg s H((s)
SoS (5q
end for return SoS
C Training Details
Models are trained with an Adam Optimizer on 4 x 80 GB A100 GPUs. For all methods, model selection is done based on accuracy on a validation set of 1,000 problems. All hyperparameters in this section are per GPU unless specified otherwise.
C.1
Model Architecture
We train a GPT-Neo model (Gao et al., 2020) with the following architecture:
Layers: 16 (all global attention)
Heads:
Context length: 4096
Hidden Size: 1024
Dropout: 0.1
Data Type: bf16
Attention: Flash-Attention-2 (Dao et al., 2022)
C.2
Stream of Search Model
Number of training examples: 5 x 105
Evaluation steps: 500
Batch size: 24
Gradient accumulation steps:
Number of training steps: 50,000
Weight decay: 0.01
Warmup steps: 1
Learning rate scheduler type: Cosine
Learning rate: 10
C.3
Optimal Paths Model
Number of training examples: 105
Batch size: 24
Gradient accumulation steps:
Number of training steps: 50,000
Weight decay: 0.01
Warmup steps: 1
14

---

Published as a conference paper at COLM 2024
Stream of Search Solution
Current State: 18:[74, 24, 36, 44], Ogerations: []
Exploring Operation: 74-44=30, Resulting Numbers: [24, 36, 30]
Generated Node #8 18:[24, 36,
Operation: 74-44=30
Exploring Operation: 44-36=8, Resul ting Numbers: [74, 24, 61
Generated Node #8.1: 18:[74, 24, 8] Operation: 44-36-8
Exploring Operation: 44-24-20, Resulting Numbers: [74, 36, 20]
Generated Node #8 18:[74,
Operation: 44-24=20
Exploring Operation: 74-24-50, Resulting Numbers: [36, 44, 50]
Generated Node #0 18:[36, 44, 50] Operation: 74-24-50
Exploring Operation: 74-36-38, Resulting Numbers: I24, 44, 38]
Generated Node #8.4: 18:[24, 44, Operation: 74-36-38
Moving Node #0,0
Current State: 18:124, 36. 381. Ogerations: ['74-44-38"]
Exploring Operation: 36-24-12. Resulting Numbers: [30, 12] Generated Node
18:[30, 12] Operation: 36-24:12
Exploring Operation: 36-30-6, Resuiting Numbers: [24.
Generated Node #8,0.1: 18: 5] Operation: 36-38-6
Exploring Operation: 24+38-54, Resulting Numbers: [36, 54]
88.8.2:
Operation: 24-30-54
Exploring Operation: 38-24-6,
Resulting Numbers: [36,
Generated Node v8,0,3: 18:136,
6] Operation: 36-24-6
Exploring Operation: 24+36-68.
Resulting Numbers: [38, 68]
Generated Node #0,0,4: 18:[30.
Operation:
Moving to Node v8.4
Current State: 18:[24, 44, 38). Ogerations: ['74-36-38")
Exploring Operation: 44-38-6,
Resulting Numbers: [24,
Generated
18:[24,
Operation: 44-38-6
Exploring Operation: 44-24-20.
Resulting Numbers: [38, 20]
Generated Nade v8,4, 18:138,
28] Operation: 44-24-28
Exploring Operation: 24+38-62,
Resulting Numbers: L44, 62]
#8.4.2:
Operation: 24+38-62
Exploring Operation: 24+44=68,
Resulting Numbers: [38, 68]
Generated Node #8.4.3: 18:138
68J Operation: 24+44-68
Exploring Operation: 38-24-14,
Resulting Numbers: [44, 14]
Generated v8.4,4: 18:[44,
Operation: 30-24=14
Moving to Node 1
Current State: 18:[74, 24. 8]. Operations: ['44-36-8"]
Exploring Operation: 74-8=66, Resulting Numbers: [24, 66]
Generated 18.1.0: 18:124, 663 Operation: 74-8-66
Exploring Operation: 24-8=16, Resulting ting Numbers: [74, 16]
Generated Nade 18.1,1: 18:[74, 16] Operation: 24-8=16
Exploring Operation: 24+8=12, Resulting ting Numbers: [74, 32]
Generated ue, 18:[74, 323 Operation:
Exploring Queration: 7448=82, Resulting Numbers: [24, 82]
Generated Node 1,.1,3: 18:124, 82) Operation: 74+8=82
Exploring Operation: 74-24=50, Resulting Numbers: [0, 58]
18.1,4: 18:[8, 50] Operation: 74-24-50
Moving to Nade #9,2
Current State: 18:[74, 36, 20], Operations: ['44-24-28']
Exploring Operation: 74+20=94,
Resulting Numbers: [36, 94]
Generated Node 88,2,8: 18:[36,
Operation: 74-28:94
Exploring Operation: 74-28=54,
Resulting Numbers: [36, 54]
Generated Node #8, 18:136,
54J Operation: 74-28-54
Exploring Operation: 36-20=16,
Resulting Numbers: [74, 16]
Generated u8,2,2: 18:174,
Operation:
Exploring Operation: 74+38-114,
Resulting Numbers: [28; 110]
Generated Node 28; 18:[20,
110] Operation: 74+36-118
Exploring Operation: 74-36=38,
Resulting Numbers: [28, 38]
Generated Nade #9,2,4: 18:[28,
38] Operation: 74-36=38
Moving to Nade 3 Current State: 18:[36, 44, 50], Ouerations: ['74-24-58"]
Exploring
44-38=8, Rasul Mimbers: [50,
Generated 80,3,8: 18:[58,
Operation: 44-36=8
Exploring Operation: 36+44=80,
Resulting Numbers: [50, 80]
Generated Node #8,3, 18:[50,
se] operation: 36-44=80
Exploring Operation: 58-36-14, Resulting Numbers: [44, 14]
Generated Node #0,3,2: 18:[44, 143 Operation: 52-36-14
Exploring Operation: 58-44=6, Resul ting Minbers: [38,
Generated Made #8,3,3: 18:[36, 5] Operation: 50-44=6
Exploring
36+50=86, Resulting Numbers: [44, 86]
Generated Nade 39,3,4: 18:[44, 86] Operation: 38-50=86
Moving to Node #0,0.0
current State: 18: 123, Operations: 1'74-44=38", "36-24=12"]
Exploring Operation: 38-12=18, Resulting Numbers: [18]
18, 18 equal: Goal Reached
Figure 8: An example of a Stream of Search trajectory.
Learning rate scheduler type: cosine
Learning rate: 1 10
15

---

Published as a conference paper at COLM 2024
C.4 Stream of Search + STaR
For STaR, we use a temperature of 0.8 to sample trajectories for problems in the training set. We filter search trajectories that only correct trajectories used finetune the next iteration. Similar to Zelikman et al. (2022) we reset the model to the base SoS model at each iteration of improvement.
Number of training examples: 1 109
Batch size: 24
Gradient accumulation steps:
Number training steps: 20,000
Weight decay: 0.01
Warmup steps: 100
Learning rate scheduler type: cosine
Learning rate: 1 x 10-5
C.5
Stream of Search + APA
Number rollouts: 32
Temperature for sampling: 1.0
Batch Size: 8
Online Epochs:
Gradient accumulation steps:
Scaling of reward: None
Value Loss Coefficient: 10
7: 1.0
A: 2.0
Learning rate scheduler type: None
Learningrate rate: x 10
D Analysis of generated trajectories
In each evaluation of the trained model, prompt the model to generate trajectories on 1) 10,000 tests with seen targets and unseen sets input numbers, and 2) 10,000 tests with unseen held-out targets. On top of improving accuracy, both policy improvement methods also reduced the average number of states or nodes explored before reaching the goal in correct trajectories, compared to the correct trajectories generated by the SoS LM. Under our framework, the number of states explored is correlated with the length of generation. This demonstrates that the policy improvement methods increase the efficiency of search.
To further study the factors contributing to the correctness of trajectories, error analysis was conducted for each one. Errors are grouped into four main types: 1) arithmetic (hallu- cinating invalid operations, dividing by zero, etc.), 2) formatting, 3) exploration (jumping to a node that has no parents), and 4) other (a child node has the wrong set of numbers associated with it). We can see that the APA training method suppressed arithmetic er- rors significantly and preserved a low rate of exploration errors. Meanwhile, the STaR method, while boosting accuracy, had an increased rate of exploration anomalies and other memory-based discrepancies.
We also analyze the correctness of solutions as a function of solution length in Tab. We find that accuracy scores remain similar across most lengths, suggesting no significant degradation. All 4096-token trajectories are incorrect.
16

---

Published as a conference paper at COLM 2024
Model
Target type
Accuracy
Num. states explored in correct trajectories
SoS
Seen
0.5127
75.043 + 0.640
SoS
Unseen
0,4900
81.747 = 0.643
STaR
Seen
0.5559
66.176 + 0.641
STaR
Unseen
0.5334
72.252 - 0.647
APA
Seen 0.5652
64.733 + 0.652
APA
Unseen 0.5423
70.849 0.655
Table 1: Details on statistics for each trained model (average over 10000 tests). Error margins represent 95% confidence intervals.
Model
Target type
Arithmetic
Formatting
Exploration
Other
SoS
Seen
2.247 = 0.097
0.014 + 0.003
0.008 + 0.003
0.018 + 0.005
SoS
Unseen
2.547 0.102
0.014 + 0.002
0.005 0.002
0.011 = 0.003
STaR
Seen
1.927 + 0.084
0.012 0.002
0
0.069 + 0.009
STaR
Unseen
2.179 + 0.084
0.013 + 0.002
0.022 + 0.004
0.074 1: 0.009
APA
Seen
0.668 + 0.036
0.014 1 0.003
0.008 1 0.003
0.022 1: 0.004
APA
Unseen
0.785 = 0.039
0.013 = 0.002
0.006 = 0.002
0.017 + 0.003
Table 2: Error analysis for trained models (average number of each error per trajectory, 10000 tests). Error margins are 95% Confidence Intervals.
Number of Tokens
SoS
STaRx3
APA
0-1000
2884/2994
3304/3373
3393/3445
1000-2000
1291/1318
1255/1304
1219/1273
2000-3000
534/545
551/588
523/661
3000-4000
389/429
419/467
388 /420
4000-5000
29/4764
30/4268
29/4201
Table 3: Correctness of SoS models function of solution length. Numbers represent correct solutions / total number of solutions. Accuracy scores remain similar across most lengths, suggesting no significant degradation.
17


Reviews:
[{"review_id": "OPONICzyuq", "reviewer": "Reviewer_7d71", "summary": "This work proposes SoS that trains a language model on trajectories that are imperfect/sub-optimal but encode the search and backtrack behaviors. To describe the search process, this work also proposes a language for search. \n\nThe results demonstrate the proposed approach outperforms a model trained on pure optimal trajectories (without search and backtracking). The proposed approach can be combined with approaches that optimize the correctness and further enhance the performance. Finally, the analysis indicates the trained model can go beyond the training data and solve challenging problems.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful review of our work. We are glad that you found our approach encoding search and backtracking behavior into trajectories to be well-motivated, insightful, and interesting.\n\n>  While this work has fairly extensive experiments and analysis on the Countdown dataset, … results on more datasets such as (1) math/coding datasets …\n\nIn this work, we wanted to perform a comprehensive study in a controlled domain to validate our key insights: If serialized search data that is suboptimal would be better than optimal data. Our goal was to establish: given search trajectory data, models can effectively learn to search & self-improve. Generating good search trajectories for more naturalistic datasets like GSM8K, HumanEval  is an important direction for future work building on SoS. Follow-up work will need to address the challenges in constructing data for such domains.\n>  discussions between SoS and the co-current work Beyond A* [1] and existing work esp [2] Procedure Cloning\n\nThe key distinction of our work from procedure cloning [2] is that rather than mimicking a single fixed search procedure like BFS, we aim to have the model learn to flexibly use & discover search strategies by training on a diverse set of search methods and heuristics. This diverse data sets up policy-improvement. In contrast to the recent Beyond A*, which trains models to imitate A* with a fixed heuristic, we let the model learn its strategies, with the choice of what to keep implicit vs explicit in the search process shaping what the model improves upon. The Beyond A* method relies on the A* heuristic being near-optimal. We will expand on these points in the related work.\n>  optimal trajectories with SoS\n\nIn our SoS dataset, 15% are optimal solutions. We have not experimented with combining additional optimal trajectories with the SoS data, this could be an interesting extension. \n>  … performance on longer traj?\n\nWe analyzed the correctness of our SoS models as a function of solution length. \n|# tokens|0-1000|1000-2000|2000-3000|3000-4000|4000-4096|\n|----|----|----|----|----|----|\n|SoS|2884/2994|1291/1318|534/545|389/429|29/4764|\n|STaRx3|3304/3373|1255/1304|551/588|419/467|30/4268|\n|APA|3393/3445|1219/1273|523/661|388/420|29/4201|\n\nAccuracy scores remain similar across most lengths, suggesting no significant degradation. All 4096-token trajectories are incorrect.\n\nWe appreciate your feedback and will incorporate it. Please let us know if you have any other questions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "* This work proposes a very interesting approach that encodes search and backtrack behavior to trajectories. Search and backtracking are very important topics for LLM, and the proposed approach is well-motivated and insightful. \n* The experiments and the analysis on the `Countdown` dataset are comprehensive. The results indicate the effectiveness of the proposed approach on the `Countdown` task.\n* The paper is very well written and easy to follow", "reasons_to_reject": "* The experiment setup is limited, making the conclusion somewhat less strong. While this work has fairly extensive experiments and analysis on the `Countdown` dataset, this is the *only* dataset this work experiments with. It will be interesting to see the results on more datasets such as (1) math/coding datasets (e.g., GSM8K, HumanEval) where getting trajectories with search and backtracking is less straightforward; (2) Interactive decision-making tasks (e.g., MiniWob++, Webshop, WebArena) where the tasks are more complex. \n* It will be interesting to have more in-depth discussions between SoS and the co-current work [1] and existing work [2] since all works train the models on trajectories that encode search algorithms. Although the paper claims that SoS discovers *new* search strategies and can tackle more challenging problems, side-by-side comparison with existing work (especially [2]) will help the community better understand what are the new takeaways. \n\nStill, I appreciate that the authors honestly discuss the limitations of this work in the Discussion section. \n\n[1] Beyond a*: Better planning with transformers via search dynamics bootstrapping.\n\n[2]  Chain of thought imitation with procedure cloning", "questions_to_authors": "* Have you experimented with combining optimal trajectories with SoS data? Or, in the SoS setting, what is the % of optimal trajectories?\n* Do you observe any performance degradation on longer trajectories?", "extra_scores": {}}, {"review_id": "5AvpEOWFi4", "reviewer": "Reviewer_wcQ5", "summary": "The authors describe innovative approaches to teaching search strategies to language models -- something they are already rather adept at in Natural Language. The authrors represent the search process as language, calling it Stream of Search (SoS), evaluating the efficacy of this approach using a game called Countdown . By pretraining a transformer-based model on streams of search generated by heuristic solvers and finetuning with policy improvement methods, the model shows a significant increase in search accuracy over just training on optimal policy traces. Furthermore, their approach can solve problems not solved by several heuristic solvers. The paper is well-written, provides a clear and original contribution, and addresses the interesting problem space of exploring complex search spaces. I enjoyed the paper.", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive reading of our work. We are glad that you enjoyed reading the paper. We appreciate your positive feedback on the clarity of the paper, the empirical analysis, and the insights around implicit strategy learning.\n\n>  SoS … outperform all symbolic search approaches.\n\nWe agree that outperforming well-tuned symbolic solvers would be a convincing demonstration  for SoS, but it is difficult to beat handcrafted symbolic solvers for a narrow domain like countdown. In our current setup, we think that the SoS models stop improving due to a lack of exploration rather than the policy improvement algorithm being a bottleneck. We will expand on how exploration of paths is a limiting factor for improvement in the discussion.\n\n>  The authors select a task that relies on arithmetic operations, then use a language model with known arithmetic errors to perform those operations -- an odd choice.\n\nYou raise a good point that arithmetic errors can be a confounding factor. Transformers trained on arithmetic can perform in-distribution calculations, but struggle with out-of-distribution calculations (larger length / numbers). So while arithmetic mistakes are present (see App. Table 2), they do not explain the performance differences of models. We will clarify the role of arithmetic errors in the main section of the paper.\n>  Not clear if training to the same number of gradient steps... it might not be well trained; … would like to see more analysis of how well-fitted the model is.\n\nThank you for raising this point. We will add curves showing validation performance over the course of training. The OT model converges more quickly than the SoS model and has a higher perplexity. We train both models up to a large number of gradient steps (perplexity stops decreasing) and then select the best model based on performance on a validation set.\n>  In addition to search trajectory similarity, I would love to some qualitative analysis of search paths -- that felt missing\n\nWe will add more qualitative examples of SoS trajectories to the appendix, highlighting specific behaviors such as backtracking, pruning, and different types of errors the model makes (arithmetic mistakes, failing to find a solution, etc).\n>  model size on SoS effectiveness?\n\nInvestigating the impact of model scale is an important question, we will address this in future work.\nLet us know if you have any other questions or concerns about the work that we could address.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- Innovative approach to teaching language models to search. The natural language search representation is flexible and can cover many different search strategies.\n- Significant improvement in search accuracy and problem-solving capability demonstrated in the Countdown game.\n- Good empirical analysis, including baselines and alternate approaches, as well as measuring search strategy similarities and differences.\n- Thoughtful analysis -- the insight that making some steps implicit can force the model to learn strategies is interesting.\n- Policy refinement approaches show additional benefit -- this feels particularly promising.\n- Well-written and clear presentation of the methodology, experiments, and results.", "reasons_to_reject": "- The results are encouraging but not entirely convincing yet -- it would be great to see SoS (or the STaR or APA variants) outperform symbolic search approaches. I realize this is a tough bar, but it would be more convincing.\n- No exploration of tasks beyond Countdown games -- not clear that it'll generalize.\n- The authors select a task that relies on arithmetic operations, then use a language model with known arithmetic errors to perform those operations -- an odd choice. \n- Not clear if  training to the same number of gradient steps is a fair comparison. It might not be well trained; there could be differences in learning difficulty across settings -- would like to see more analysis of how well-fitted the model is.\n- In addition to search trajectory similarity, I would love to some qualitative analysis of search paths -- that felt missing", "questions_to_authors": "- Can you provide insight into how SoS would generalize to more complex tasks beyond Countdown?\n- Can you use more accurate methods for arithmetic operations? Or are there other games you could pick?\n- Curious: Does the size of the model matter? Would a larger model or a smaller model act or learn differently? What's the impact of model size on SoS effectiveness?\n- \"Trained to same number of gradient steps\" -- is this a good measure of how well the model has fitted the data? Are some settings easier to learn? What does the learning curve look like?", "extra_scores": {}}, {"review_id": "3nDJvlWzR6", "reviewer": "Reviewer_oA5B", "summary": "During LLM pretraining, usually the de facto training admits only one gold-standard distribution which the model is trained to approximate/predict. This raises the issue of training and inference distribution mismatch. During search/inference, this mismatch could potentially lead to cascades of search errors, degrading the quality of search and hence the quality of the generated output distribution. This phenomenon is summarised in the Abstract (somewhat inaccurately) as \"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.\"\n\nThis paper proposes to address this by allowing the model to explore/backtrack (non-optimal) search paths/nodes with synthetic data. In the current setup, this data is generated using a \"template\" for the Countdown evaluation task -- where a set of numbers are given and the learner is required to find optimal way(s) by doing arithmetic to reach a pre-specified target number. To quote: \n\n\"To construct our Stream of Search dataset, we define 12 search strategies based on breadth first search\n(BFS) and depth first search (DFS) (see App. Alg. 4, Alg. 3) that rely on two simple and\ninterpretable heuristic functions. The heuristics we use to guide search are 1) the absolute\ndifference between the sum of the remaining options and the target and 2) distance to the\nfactors of the target (see App. B). \"\n\nThis training strategy is referred to as Stream of Search. Empirical results are demonstrated on solving \"difficult\" problems as well as solving \"unsolved\" problems, the metric is accuracy. Various comparisons/analysis of different search strategies are also done by comparing search strategies with the Pearson correlation metric and the Alignment of States Visited metric (towards the End of Sec 4).", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and feedback! We are glad that you found the problem motivation & discussion of the various search algorithms, policies, to be useful & clearly presented.\n>  related works … SEARN …relation to SoS, LLMs\n\nThank you for pointing out these relevant earlier works on the distribution mismatch problem in sequence models. Techniques for seq2seq models, structured prediction, such as SEARN, are highly relevant. Here is a draft we intend to include: A related line of research, starting with SEARN [1, see also 2,3], suggests supervising the process of finding a solution into a structured prediction problem. This is similar in spirit to SoS where there is supervision over the process of search. While SEARN provides a theoretical framework for learning a policy over structured ‘process’ data, SoS demonstrates how the search process itself can be represented as a flat language sequence, defining an explicit vocabulary to represent elements of search algorithms, such as backtracking & heuristics, allowing LMs to directly learn diverse search strategies from data. \n[1] Daumé+2009. [2] Ranzato, M+2015 [3] Ross+2011\n>  improvements small ... knobs to tune\n\nAlthough the improvements with policy improvement methods aren’t drastic, they are promising. The methods are relatively robust to the ‘knobs’ being tuned, for the STaR, we found that performance is relatively robust to these ‘knobs’.\n> single task…ad-hoc\n\nWe chose Countdown in this work as it is a challenging search problem with a large branching factor, where different heuristics can be naturally applied. This allowed us to systematically evaluate the key components - using search streams as training data, iterative improvement, different heuristics and solvers - in a controlled setting to understand their impact. We establish: 1) SoS significantly outperform those trained on optimal solutions, 2) Iterative finetuning leads to improvements. The core ideas behind SoS - training models to search via language - can be extended to other domains:  naturalistic datasets such as GSM8K, HumanEval. These will likely require additional insights, such as incorporating goal-setting & better exploration.\n\nWe will update the discussion section of our paper to more clearly articulate these points about the potential generalizability of SoS & lay out a roadmap for future work. We thank you for your feedback. Please let us know if you have any other questions or concerns.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "* Detailed discussions of the problem which is useful and clear motivation of the problem\n* Detailed presentation of various algorithms/policy/heuristics used especially those related to the search strategies and detailed analysis of the (mis) alignment of the different search strategies\n* Some discussions of related works", "reasons_to_reject": "* The distribution mismatch and the search problem has been a long standing problem not only in pretrained LLMs but dates back to earlier models such as seq2seq models (e.g., as discussed in this paper: https://arxiv.org/pdf/1511.06732) and even earlier works such as SEARN (https://link.springer.com/content/pdf/10.1007/s10994-009-5106-x.pdf). I found the current paper completely disregarded these highly related works, presenting the paper and the problem it's trying to tackle seemingly applicable only to LLMs. I believe the \"older\" techniques have much to say/illuminate the methods and issues the current paper is trying to address, but by reading the current version of this paper, this connection cannot be readily established, and this leaves a void that is hard to fill.\n* The paper focuses on one single task, the heuristics used to generate the synthetic training data seems to be largely ad-hoc to the evaluation task; there may be parts of meta-heuristics which could be generalized to other tasks, but  by and large it seems it's quite specific to the evaluation task. \n* The improvements over the single-path search strategy seems to be small (Fig. 4a). And it seems there a lot of knobs to tune to get such \"improvements\". \n* By reading this paper, I certainly learned something, but it feels mostly still a wip, especially in the experimental sections in terms of the range of tasks the proposal is tested on and the strength of it over the (missing?) baseline(s).", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "U7LmpsvVNT", "reviewer": "Reviewer_4vuw", "summary": "This paper introduces a novel approach called \"Stream of Search\" (SoS) that enables language models to learn how to search effectively within their own language space. The key idea is to train the model to generate a stream of search queries that incrementally refine an original query, rather than just producing a single output. This allows the model to actively explore and navigate the language space to find the most relevant information, similar to how humans find the correct solution from multiple erroneous explorations. Within supervised training and further RLHF, the SoS model solves 36% of previously unsolved problems, including problems that cannot be solved by any of the heuristic solvers. These 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", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and feedback on our work. We are glad that you find our approach to be novel and effective, with convincing experimental results demonstrating significant performance improvements over baselines.\n\n>  My main concern is whether this type of method is only applicable to problems with deterministic results, such as Countdown. … I'm unsure if a model trained on Countdown data can improve the accuracy of other related tasks.\n\nWe agree that demonstrating the generalizability of SoS to a wider range of problems is an important direction. In this initial study, we focused on the Countdown game as it represents a challenging search problem with a large branching factor, allowing us to test our key hypotheses in a controlled setting: 1) That linearized search strings, with mistakes and backtracking data, (SoS) help in training models that are better than models trained on optimal solutions 2) SoS models can be self-improved.\nHowever, we believe the core ideas behind SoS - training models to search via language- can be extended to other domains. In particular, we plan to explore applying SoS to more naturalistic datasets such as GSM8K for math word problems, and HumanEval for coding. These domains will likely require additional insights into generating effective search trajectories, such as incorporating goal-setting, exploration and self-correction into the trajectories. We will update the discussion section of our paper to more clearly articulate these points about the potential generalizability of SoS and lay out a roadmap for future work investigating its application to a wider range of problems.\n\nPlease let us know if you have any other questions or concerns about the work that we could address.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1) To the best of my knowledge, this is the first search model based on SFT and RL. All previous work has focused on zero-shot or few-shot learning approaches. Experiments also demonstrate significant performance improvement due to SFT and RL.\n\n2) A key distinction between this method and other common techniques like TOT (Tree of Trees) is that the model, during decision-making, has access to the global search history, whereas other tree-based search methods can only access the current state.\n\n3) The authors have constructed a dataset for training the countdown game, which includes a vast amount of search trajectory data.\n\n4) Experimental results convincingly demonstrate that the proposed method significantly outperforms other Optimal Paths methods by a large margin.", "reasons_to_reject": "1) My main concern is whether this type of method is only applicable to problems with deterministic results, such as Countdown and Math 24. Currently, I have no idea how to generalize SoS to other types of problems.\n\n2) The experiment did not demonstrate the generalizability of the SoS method. I suspect whether the trained model can only be used for the task of playing Countdown. I'm unsure if a model trained on Countdown data can improve the accuracy of other related tasks.\n\n3) The construction of the synthetic dataset is not generalizable. So, for other tasks, if I don't have the mentioned 12 search strategies, can I still obtain enough high-quality training data to train the SoS model?", "questions_to_authors": "Please refer to my concerns in 'reasons to reject'.", "extra_scores": {}}]

Author rebuttal:

Meta-review: This is an interesting paper that proposes a novel way to encode and flatten the search process in a textual way and then shows that transformers are capable of learning this representation. It's basically a structured alternative to chain of thought for problems that admit a structured search approach. The reviewers all found the contribution interesting and the experiments thorough and compelling. Nice work!
Decision: Accept