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

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.

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.

We 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. However, 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.

Please let us know if you have any other questions or concerns about the work that we could address.

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.

related works … SEARN …relation to SoS, LLMs

Thank 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. [1] Daumé+2009. [2] Ranzato, M+2015 [3] Ross+2011

improvements small ... knobs to tune

Although 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’.

single task…ad-hoc

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

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

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

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.

SoS … outperform all symbolic search approaches.

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

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.

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

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.

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

In addition to search trajectory similarity, I would love to some qualitative analysis of search paths -- that felt missing

We 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).

model size on SoS effectiveness?

Investigating the impact of model scale is an important question, we will address this in future work. Let us know if you have any other questions or concerns about the work that we could address.

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.

While this work has fairly extensive experiments and analysis on the Countdown dataset, … results on more datasets such as (1) math/coding datasets …

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

discussions between SoS and the co-current work Beyond A* [1] and existing work esp [2] Procedure Cloning

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

optimal trajectories with SoS

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

… performance on longer traj?

We analyzed the correctness of our SoS models as a function of solution length.

Accuracy scores remain similar across most lengths, suggesting no significant degradation. All 4096-token trajectories are incorrect.

We appreciate your feedback and will incorporate it. Please let us know if you have any other questions.