Stream of Search (SoS): Learning to Search in Language

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

Rating

7

Confidence

4

Ethics flag

1

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. 2) 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. 3) The authors have constructed a dataset for training the countdown game, which includes a vast amount of search trajectory data. 4) 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. 2) 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. 3) 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'.

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." This 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: "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). " This 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).

Rating

5

Confidence

3

Ethics flag

1

Reasons to accept

* Detailed discussions of the problem which is useful and clear motivation of the problem * 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 * 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. * 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. * 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". * 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).

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.

Rating

8

Confidence

4

Ethics flag

1

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. - Significant improvement in search accuracy and problem-solving capability demonstrated in the Countdown game. - Good empirical analysis, including baselines and alternate approaches, as well as measuring search strategy similarities and differences. - Thoughtful analysis -- the insight that making some steps implicit can force the model to learn strategies is interesting. - Policy refinement approaches show additional benefit -- this feels particularly promising. - 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. - No exploration of tasks beyond Countdown games -- not clear that it'll generalize. - 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. - 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. - 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? - Can you use more accurate methods for arithmetic operations? Or are there other games you could pick? - 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? - "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?

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. The 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.

Rating

7

Confidence

4

Ethics flag

1

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. * The experiments and the analysis on the `Countdown` dataset are comprehensive. The results indicate the effectiveness of the proposed approach on the `Countdown` task. * 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. * 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. Still, I appreciate that the authors honestly discuss the limitations of this work in the Discussion section. [1] Beyond a*: Better planning with transformers via search dynamics bootstrapping. [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? * Do you observe any performance degradation on longer trajectories?

Thank you

Thank you very much for the additional results and the response. I raised my score to 7.

Reviewer, please respond to rebuttal

Reviewer, please respond to rebuttal even if your score hasn't changed. The discussion period ends soon (Thursday)