Beyond A*: Better Planning with Transformers via Search Dynamics Bootstrapping
Authors: Lucas Lehnert
Sainbayar Sukhbaatar
DiJia Su
Qinqing Zheng
Paul McVay
Michael Rabbat
Yuandong Tian
Venue: COLM (2024)
Keywords: Planning, Sequential Decision Making, Reasoning
TL;DR: We demonstrate how to train transformers to solve complex planning tasks and how they can be used to find more efficient search methods.

Abstract: While Transformers have enabled tremendous progress in various application settings, such architectures still struggle with solving planning and sequential decision-making tasks. In this work, we demonstrate how to train Transformers to solve complex planning tasks. This is accomplished by first designing a synthetic language that captures the computation performed by the $A^*$ search algorithm when solving a planning task. Then, an encoder-decoder Transformer model is trained to predict this language, resulting in a language model that can correctly solve novel planning tasks by generating $A^*$'s search dynamics. We fine tune this model to obtain a Searchformer, a Transformer model that optimally solves previously unseen Sokoban puzzles 93.7\% of the time, while using up to 26.8\% fewer search steps than our $A^*$ reference implementation. Searchformer significantly outperforms baselines that predict the optimal plan directly with a 5-10$\times$ smaller model size and a 10$\times$ smaller training dataset. Lastly, we demonstrate how Searchformer scales to larger and more complex decision making tasks with improved percentage of solved tasks and shortened search dynamics.

Published as a conference paper at COLM 2024
Beyond A": Better Planning with Transformers via Search Dynamics Bootstrapping
Lucas Lehnert*
Sainbayar Sukhbaatar FAIR team at Meta
Andy Su
University of Saskatchewan lucas. wensseeeeettuuakk..o
FAIR team at Meta
Qinqing Zheng FAIR team Meta
Paul McVay FAIR team Meta
Michael Rabbat FAIR team at Meta
Yuandong Tian FAIR team at Meta
yuandong@me ta. com
Abstract
While Transformers have enabled tremendous progress in various appli- cation settings, such architectures still struggle with solving planning and sequential decision-making tasks. In this work, we demonstrate how to train Transformers to solve complex planning tasks. This is accomplished by first designing a synthetic language that captures the computation per- formed by the A search algorithm when solving planning task. Then, an encoder-decoder Transformer model is trained to predict this language, resulting in language model that can correctly solve novel planning tasks by generating A 's search dynamics. We fine tune this model to obtain a Searchformer, a Transformer model that optimally solves previously un- seen Sokoban puzzles 93.7% of the time, while using up to 26.8% fewer search steps than our A reference implementation. Searchformer signifi- cantly outperforms baselines that predict the optimal plan directly with a 5-10x smaller model size and 10x smaller training dataset. Lastly, we demonstrate how Searchformer scales to larger and more complex decision making tasks with improved percentage of solved tasks and shortened search dynamics.
1 Introduction
Transformer-based architectures (Vaswani et al., 2017) have demonstrated impressive perfor- mance in different tasks, including holding conversations at the human level (Shuster et al., 2022; OpenAI, 2022; 2023; Touvron et al., 2023), high-quality image understanding (Caron et al., 2021; Oquab et al., 2024; Assran et al., 2023) and video generation (Singer et al., 2023), multi-modal generation (Girdhar et al., 2023; Radford et al., 2021), and code com comple- tion (Roziere et al., 2023; OpenAl, 2021). But despite these successes, Transformer based architectures and Large Language Models (LLMMs) still struggle when responding to tasks that involve multi-step planning (Valmeekam et al., 2023a;b) or higher-order reasoning (Mo- mennejad al., 2023; Fan et al., 2020).
Various methods have been proposed to move past these limitations. One popular approach to simulate thinking process and train a Transformer or language model generate intermediate "thoughts" before outputting response. These Chain-of- -fho-uhg-Thhggtt (CoT) prompting (Wei al., 2022) and Tree-of-thoughts (ToT) methods (Yao et al., 2023) encourage the model to "think" step by step. While these techniques are often effective, they may lead to worse performance in some cases, for example due self-enforcing (Huang et al., 2023). Techniques effective one dataset may not work well on others due to changes in the type of reasoning involved (e.g., spatial reasoning vs. mathematical reasoning). How to enable Transformers and LLMs to plan, solve complex sequential decision-making tasks, and perform reasoning still remains elusive and an active area of research.
" Jucas Lehnert performed this work as part of the FAIR team at Meta.

---

Published as a conference paper at COLM 2024
Our work. We demonstrate how to use Transformers and language modelling to robustly solve complex planning tasks. To accomplish this, we design our own synthetic language that expresses planning the computation involved in obtaining an optimal plan-in token sequence. In this synthetic language, we also include task specifications as well as plan solving the task optimally. Similar to LLMs,, we train Transformer neural networks to predict the next token given a sequence of tokens in this synthetic language. With this approach, we demonstrate how to build language models that imitate the computation performed by A search (Russell & Norvig, 2021, Chapter 3). Lastly, we present Searchformer, Transformer model that solves complex planning tasks in fewer search steps than our A* reference implementation. This model is obtained through search dynamics boolstrapping, method that first trains a Transformer to imitate A" search process and then fine-tunes the model to find an optimal plan within fewer search steps.
To train a Transformer to perform planning, we consider two dataset variants. The first variant uses token sequences only containing planning task specification and its optimal solution plan. We refer to this data as solution-only sequences and neural networks trained on this data solution-only models. The second variant incorporates A" search dynamics into token sequence. We refer to this data as search-augmented sequences and neural networks trained on this data as search-augmenled models. Lastly, pre-trained aarch-auggented model, that generates A" execution traces and therefore imitates A A'ss computation, is fine-tuned optimal plan. We refer to this final fine-tuned model as Searchformer. When solving Sokoban puzzles, our Searchformer models solve 93.7% of all test tasks while performing on average 26.8% fewer search steps than our own A' reference implementation.
execution trace encoding
generate shorter token sequences while still outputting an
Through sequence of experiments that control task complexity, dataset size, and model size, we demonstrate that including execution traces into the training data increases perfor- mance on an independent test task set del-ies a 10x to 100x increase in the length of the generated sequences. We find that search-augmented models (that include execution traces into their training data) generate an optimal plan more often on unseen tasks with ten times fewer training sequences than larger solution-only model (that is trained on sequences only including task description and task solution). This result highlights the power of including A's search dynamics into the training process of Transformer models.
2
Related work
While existing works (Trinh et al., 2024; Ruoss et al., 2024) leverage synthetic datasets to learn policies for reasoning, our study is fundamentally different in this regard. We demonstrate how to train Transformers that generate the computation performed by planning algorithm. The resulting system accepts a planning task as an input prompt and then reasons about this task via generation of an execution trace. Our approach stands in contrast to existing Reinforcement Learning (RL) architectures that focus on autonomous decision-mal maasionnmmaing agents and that a neural network predictions. RL. algorithms such as AlphaZero (Silver et al., 2018), MuZero (Schrittwieser et al., 2020), and AlphaGeometry (Trinh et al., 2024) optimize a neural network as part of a more complex planning or reasoning algorithm. For example, (Silver et al., 2017) use MCTS (Coulom, 2006) as policy improvement operator (Sutton & Barto, 2018, Chapter 4) to update the neural network's 's weights. In contrast, our approach first trains neural network that imitates the computation performed by a planning algorithm. In this work, the neural network is not part of a more complex planning routine. In similar vein, the presented search dynamics bootstrapping method driven by the Transformer model itself and not an improvement operator such as policy improvement. To generalize towards more efficient search patterns, this bootstrapping process uses one of our own pre-trained search- sarrc-au models to generate token sequences that encode novel and more efficient search patterns.
black box system to make
Hierarchical planning methods and temporal abstractions (Sutton et al., 2023; 1999; Diet- terich, 2000; Hafner et al., 2022) are methods for abstracting multi-step decision sequences, called skills or options, into subroutines. In contrast, this work does not focus on extracting
2

---

Published as a conference paper at COLM 2024
skills. Instead, we demonstrate how to generate entire execution traces and optimize their length. Unlike prior work (Yang et al., 2022; Pallagani et al., 2022; Ruoss et al., 2024), the trained Transformer models do not predict just single action and instead output an entire solution plan for a specific planning task.
The property of generating entire execution traces is perhaps most similar to the Scratchpad method (Nye et al., 2021). .) In contrast to Nye et al., our study focuses on discovering new and more efficient execution traces instead of solely imitating them. We also do not rely on fine-tuning pre-trained LLMs and instead train every model from scratch on synthetic language.
Our work bears some similarity with neuro-symbolic systems (Graves et al., 2014; Cai et al., 2017), which build differentiable architectures to mimic the functionality of existing sym- bolic systems. However, these methods use dedicated components (e.g., explicit memory components, built-in recursion), while Searchformer focuses on next-token prediction only. Here, Searchformer relies on generating long contexts and position embeddings (Chen et al., 2023; Peng et al., 2023) to predict in optimal plan. Ultimately, we aim to shed light on how to build more general language model-based architectures that automatically learn a planning mechanism.
Using Transformer architectures to solve complex sequential decision-making tasks has been studied in prior work in an RL setting (Chen et al., 2021; Janner et al., 2021; Laskin et al., 2023). These works present different methods for modelling trajectories of trial-and-erron interactions and focuses on predicting a next action, state, or reward or a combination of them. In contrast, we demonstrate how to use Transformer to model the search steps involved in computing an optimal plan. MCTSNet (Guez et al., 2018) also attempts to learn the search procedure itself, but still hard-codes the MCTS search procedure into neural network, which leads to quadratic backpropagation overhead and can only deal with up to 500 step rollouts, while our approach can deal with much longer search execution trace. We demonstrate that Transformers can not only imitate symbolic planning algorithm but can also be used to discover more efficient heuristics via fine-tuning.
3 Problem setup
Figure provides an overview of our synthetic dataset generation process. We consider two domains: maze navigation (Figure 1(a)) and solving Sokoban puzzles (Figure 5 in Appendix B). In maze navigation, the goal is to find the shortest path through an n-by-n maze. In Sokoban, worker can move up, down, left, right and has to push each box onto a dock to solve the puzzle. An incorrect move may immediately lead to a dead end and thus reasoning across many time steps is required to solve the puzzle. Each state in puzzle consists of a combination of box and worker positions, making Sokoban computationally more difficult to solve than maze navigation.
3.1 Generating execution traces of A* search.
The A algorithm computes an optimal plan by manipulating two sets of nodes: a frontier set containing the current search frontier, and closed set containing all searched nodes. In the maze example in Figure (a)), each node corresponds to an empty (non-wall) grid cell. For each node, the algorithm computes heuristic value and a cost from start value. Which node is searched next is determined by the frontier set, closed node set, heuristic values, and cost from start values (Figure 1(c), left panel). A execution trace is collected by tracking all insertion operations into the frontier and closed set along with heuristic and cost from start values (Figure 1(c), right panel). The right panel in Figure 1(c) illustrates the resulting trace for the maze example shown in Figure 1(b). Each row corresponds either to an insertion of a node into the frontier, indicated by a create token, or to moving node into the closed set, indicated by close token. The resulting plan is then appended to this trace. This trace is constructed such that given any prefix, the state of the A algorithm can be reconstructed to resume execution. In maze navigation, A uses the Manhattan distance
3

---

Published as a conference paper at COLM 2024
Prompt:
Response:
bus
2
wall cell
plan 0
start cell
gozl
plan
goal cell
Mall
2
plan
plan
1 2
plan
eos
egs
(a) Maze navigation task
Tokenization scheme for a planning task and its solution
planning algorithm
Tokenization algorithm execution
Require: Start node lubert and goal node Nigul-
bos
1: Scloeed
create
Add frontier
Syortee f- (Paaut)
close
0-2-c0
Add node closed set
while (Stmotic do
create
0 01 [ 42
Cost from start c(n)
4: Heur arg minnesponner cost(i))
Trace
close
1 cl c2
Heuristic value k(()
Seooed < Sooud L (eerr)
for Pdbila children( (Heur) do
create
1 c2
if pos()() = pos(Mans for any E Selsed U Smonter then
close
0 0 c2 el
c(t)) + h(x)) S. c(nana) k(Haaa) then
create
0 63
continue
10:
end if
plan
0 2
11:
end
plon
1
12:
Set parent(naaaa) 4-
Plan
plan
13:
Stontier +- Simmtir U (Raaa)
plan
14:
end for
eos
15:
while
16:
Compute and return plan by recursing on parents of Maur.
(c) A execution when solving planning task logged into execution trace
Figure Expressing a planning task in token sequences. (a): A x maze navigation task where the goal is to find the shortest path from start to goal without entering wall cell. (b): The 3 maze navigation task is expressed as a prompt token sequence (left panel) and the optimal plan is expressed as response token sequence (right panel). The start and end of a sequence is indicated by a
token, bos, and an end-of-sequence token, eos. Numbers indicate x, y coordinates. (c): The search dynamics of the A algorithm (left panel) is logged into an execution trace (right panel). The first two integers encode the node's 's x, y coordinates. The last two tokens in the trace encode the cost-since-start value c(n) and the heuristic value (n) (letter "c" distinguishes costs from coordinate numbers).
to the goal location as a heuristic. In Sokoban, every box matched to the closest dock and then A computes the sum of all Manhattan distances between each box and dock pair.
For each experiment, we generate two token sequence variants, as illustrated in Figure 1:
Solution-only sequences of the format <prompt><plan>, where the <prompt> part encodes task description and the <plan> part the optimal plan (Figure 1(b)).
Search-augmentted sequences of the format sprompt>ttaae>><plan. where the <trace> part encodes A"'s execution trace (Figure (c)).
Every model is trained from scratch and after training the model's output is parsed and evaluated it contains an optimal or feasible solution plan.
3.2 Training Transformer model
When generating token sequence dataset, each task unique and the test constructed such that it does not contain any duplicate of the training set. With this experiment design, we hope to gain insight into how Transformers can be used to solve planning tasks and generalize to previously unseen test tasks.
By including intermediate computation steps into the training data, the Transformer model is trained to imitate the computation performed by the A* algorithm. In contrast to Procedure Cloning (Yang et al., 2022), the -aacch-uugmented model learns to predict the entire thinking process, including paths searched by A that lead to dead end.

---

Published as a conference paper at COLM 2024
For each experiment an adaptation of the encoder-decoder T5 architecture (Raffel et al., 2020) is used that integrates Rotary Position Embeddings (RoPE) (Su et al., 2023). The encoder processes the <prompt> part of training sequence, and the decoder processes either a <trace><<lan'> aaaae>>ppl--ffommattee sequence (search-augmented model) or only a <plan>-formatie tted sequence (solution-only model). Depending the model variant, each network is trained to maximize the cross-entropy between the distribution of the decoder generations and the distribution of sampling a corresponding sequence from the training dataset. Appendix A describes our optimization setup in more detail.
3.3 Moving past algorithm imitation via search dynamics bootstrapping
To reduce the number of tokens generated by a search-augmented model during inference, we implement a method to shift the distribution with which the decoder generates execution traces. First, a search-augmented model is trained to imitate the search dynamics of A search. To discover new search dynamics with this search-augmented model and explore the different execution traces, the search-augmented model must generate different sequences for the same task prompt. We accomplish this by inducing non-determinism into the training data and use a non-determinsitic A implementation that breaks cost ties randomly and randomizes the order with which child nodes expanded. This approach does not decrease the efficiency of A search itself and merely changes the order with which different nodes are searched while still respecting A heuristic and cost calculations. The resulting search-augmented model will then approximate the probability distribution with which the training sequences were generated.
Once model is trained to imitate the search dynamics of non-deterministic A search, it used to generate a new training dataset consisting of shorter token sequences. This new dataset is constructed by using the trained search-augmented model to sample multiple different token sequences for each training prompt (test prompts are not used to ensure separation between training and test data). Each generated sequence is parsed and checked ends in an optimal plan. If this is the case and the sequence is also shorter than the corresponding sequence contained in the original training dataset, then this shortened sequence is included in the new short sequence training dataset. the generated sequence does not end in optimal plan or is longer than the original training sequence, then the sequence from the original training dataset is reused.
Subsequently, the search-augmenter model is fine-tuned on the new short sequence train- ing dataset. To distinguish from the search-augmented model that imitates A'ss search dynamics, call this new model Searcliformer This procedure can then be repeated by using the resulting fine-tuned model to generate the next even shorter sequence dataset and then fine-tuning the Searchformer model again. In Section 4.3 we demonstrate that this procedure does in fact reduce the number of steps performed during inference while further improving performance. The Searchformer model no longer imitates A search and has instead discovered new way of solving a planning problem using fewer steps.
4 Experiments
In our experiments, we use two different A implementations for sequence data generation:
1. Deterministic A' dataset: Sequences are generated by executing A' in a deterministic fashion (by ordering child nodes and breaking equal cost ties deterministically). Given task prompt, the optimal plan and A execution trace is unique, and the Transformer learns the deterministic tie breaking rules implicitly encoded in the data. Evaluating such a model is simple, because the generated sequences need to exactly match the sequences generated by A+.
2. Non-deterministic A dataset: Sequences are generated by executing A* in non- deterministic fashion (by randomly ordering child nodes and breaking equal cost ties randomly). Given a task prompt, the optimal plan and A* execution trace is no longer unique. Therefore there are multiple correct responses and the Transformer learns to imitate the random tie breaking rules implicitly encoded in the data. The generated
5

---

Published as a conference paper at COLM 2024
100-
100
100
15M Search Augmented
Correct Response
Sequences [in %]
Correctly Solved
Test Tasks(in %]
Correctly Solved
Test T Tasks|in %]
80
60
46M Search
Augmented
40
175M Search
20
20
Augmented
50k
100k
500k 1M
50k
100k
500k 1M
10c10
20c20
30030
175M Solution Only
Number
Number of TrainingSequnees
TrainingSequnoes
Maze Shape
(a) Deterministic case
(b) Non-deterministic case
(c) Performance across task difficulties
Predicting intermediate computation steps leads to robust performance in the low data regime. For each model, the number of free parameters (indicated in millions of parameters with "15M", "46M", and "175M") is varied. (a): Comparison of how many test tasks were answered with correct token sequence when training on the deterministic A* dataset (exact-match criterion in Appendix C). (b): Comparison for how many test task at least one optimal plan was found when training on the non-deterministic Á" dataset (any-optimal-64 criterion in Appendix C). (c): Performance degradation when increasing task difficulty (maze size). Here, the non-deterministic A dataset was used and models are evaluated with the any-optimal-64 criterion.
Figure 2:
sequences vary between different executions, but the resulting plans are still optimal and execution traces still respect A" cost and heuristic calculations as described in Section 3.3.
Figure 7 (Appendix B) presents an overview of the token sequence length for each dataset and shows that the generated A execution traces grow in length with task complexity. Figure& 8 (Appendix B) shows that training and test sets are matched in difficulty and have comparable trace lengths. For each task, one model may generate search sequence ending either in optimal plan, feasible plan (a plan that is correct but sub-optimal), an invalid plan. Appendix C outlines how each model's ability to predict feasible and optimal plan is scored and how the search dynamics of the augmeetd and Searchformer models are evaluated. Unless indicated otherwise, experiments are repeated five times and figures plot averages across all repeats. All reported errors indicate the Standard Error of Measurement (SEM).
4.1 Maze navigation
In the first experiment set, train set of encoder-decoder Transformer models to predict optimal plans for maze navigation tasks. We vary the training dataset size and model size (the number of optimized parameters) between different training runs and evaluate each model on the test tasks generated using the same hyper-parameters.
4.1.1 Deterministic A"
Figure 2(a) plots for how many test tasks correct response was generated. Both solution- only and search-augmented models are trained on the deterministic A dataset and are evaluated if they exactly re-produce the token sequences generated with A search (please refer to the exact-match criterion in Appendix C). Similar to results presented by Ivanitskiy et al. (2023), the solution-only models solve most test tasks for large enough training datasets. Nonetheless, the solution-only models are outperformed by most search-augmented models. Only for large enough training datasets, the solution-only model matches the performance of the worst search-augment sequences and less), performance of the solution-only model degrades significantly, while the performance of each search-augmented model stays relatively high.
model. In the low training data regime (100,000 training
This result is surprising, because for more than 90% of the test mazes, the search-augmented models generate Errac>><plan--fommale sequences that are thousands of tokens long

---

Published as a conference paper at COLM 2024
without predicting any single token incorrectly. Moreover, the solution-only models, that average predict ten times shorter sequences, are significantly outperformed by the search- a models. Even the smallest search-augmented model significantly outperforms the much larger solution-only model parameters. Appendix D.1 presents another ablation experiment testing how sensitive model performance is when the specific tokenization pattern is changed. Here, we find again that the search-augmented models are robust to changes in the specific tokenization pattern and significantly outperform the solution-only.
This result highlights the power of training Transformers to generate long algorithm exe- cution traces. We do not observe compounding prediction errors that usually limit deep model-based RL agents (Asadi et al., 2018), because the used backward-causal decoder network constructs for an n-step sequence an x и attention map. Here, this property of the Transformer architecture is used to boost performance when predicting an optimal plan.
4.1.2 Non-deterministic. A*
When trained on non-deterministic A" data, the model could output multiple different optimal paths for one task. Here, we use each model to generate 64 token sequences for each task. The test task is counted as correctly answered of any one of the 64 sequences contains an optimal plan (please refer to the any-optimal-64 criterion in Appendix C). Because we only test if at least one generated sequence contains an optimal plan, we obtain higher absolute numbers in Figure 2(b) than in Figure (a)).
Figure 2(b) plots for how many test tasks optimal plan was found when generating for each test task 64 token sequences. Here, we can observe pattern similar to the deter- ministic A" dataset: even the smallest search-augmented models outperform solution-only model, especially for a small training set. Moreover, we found that model size only im impacts the performance of each of the search-augmented models when using very small training datasets (50,000 training sequences). For larger training dataset sizes no significant differ- ence is found. Increasing the number of parameters of the solution-only models does not significantly improve their performance in the low data regime (Figure 9 in Appendix F).
4.1.3 Performance across different task difficulty levels and network architectures
Lastly, Figure 2() illustrates how task's s difficulty influences the performance of each model. Using the non-deterministic A data, we consider the number of correctly solved test tasks as a function of maze size. The larger the maze, the larger the task's state space and the more computation is required to find an optimal solution plan. While the solution-only model's performance drops rapidly as the tasks become more challenging, the search-augmented models maintain a comparably high accuracy, even for the smallest model size. Appendix presents full comparison across all maze sizes. Apart from overcoming computational resource constraints, we believe that this approach scales to even more complex tasks.
While the solution-only models learn to predict an optimal plan for large and diverse enough training datasets, aarchh-uggmened models perform significantly better in the low data regime for more difficult tasks. The search-augmented models reach higher performance because they perform additional computation during inference. These models imitate the search dynamics for grounded reasoning chain that leads to an optimal plan. In contrast, the solution-only models have to infer direct correlations between task description and optimal solution plan through supervised training where many of such correlations can be spurious and unreliable during evaluation on the test task set.
To test if the presented results depend on the use of an encoder-decoder architecture, we aatpettttrceeeplaannoooooo
trained decoder-only architecture in architecture had 46 million parameters and apart from network architecture no other hyper- parameters were changed. For this experiment we used 50,000 training sequences from the 20 x 20 maze datasets. The resulting decoder-only search augmented model had a test set performance of 94.8% while the corresponding encoder-decode search augmented models solved 91.8% of the test tasks. This result indicates that our approach can be implemented with different architectures to obtain even better test set performance.
sequences. This
7

---

Published as a conference paper at COLM 2024
Table 1: Test set performance in the Sokoban tasks. Over 200 unseen test Sokoban tasks, we report percentage of solved and optimally solved test tasks. For sequences ending in either an optimal and correct plan we report the SWC (Success Weighted by Cost) score, and ILR (Improved Length Ratio of Search Dynamics) scores. In every column, higher scores indicate better performance. Please check Appendix for detailed definitions of these scores.
Params.
Model
Solved (%)
Optimal (%)
SWC
ILR (solved)
ILR (optimal)
Solution only Search augmented
90.3 #1.0
86.8 =0.3
0.890 +0.009
92.5 11.0
90.8 1.6
0.924 40011
0.908
+00.20
0.919 0.019
45M
Searchformer, step 1
95.5 #1.0
93.5 =10
0,953 10.010
1.054
#0.055
1.062 =0.15
Searchformer, step
96.0 a0.55
93.4 10.6
0.957 40.005
1.158 +00025
1.181 0001
Scarchformer, step 3
95.5 =0.8
93.7 #16
0,953 40.009
1.292 #0.044
1.343 =0.067
Solution only
95.7 #0.2
90.0 =0.8
0.949 #0.008
175M
Search augmented
95.2 +0,9
93.2 110
0.949 +0.010
0.925 #0.010
0.933 =0011
757M
Solution only
96.5 =0.1
92.2 #1.2
0.958 +0.002
4.2 Solving Sokoban puzzles
To test similar results can be obtained on a different and more complex task with a different tokenization pattern and different transition structure, we repeat our experiments for Sokoban puzzles using our non-eeteriinistic A' implementation. Table lists how often each model generated correct optimal plan for each test task. As before, by training on execution traces, the search-augmented models outperform the solution-only models. Even increasing the parameterization of solution-only model to 747 million parameters only leads to a marginal performance improvement. On average, this 747 million parameter solution-only model is still outperformed slightly by smaller 175 million parameter search-augmented model. This experiment further confirms our findings on more complex planning tasks with a different transition structure and a different tokenization method.
4.3 Searchformer: Improving search dynamics via bootstrapping
In this last experiment, we investigate how the search-augmented model can be iteratively improved to compute an optimal plan while generating a shorter execution trace. Here, our goal is to shorten the length of the search trace while still producing optimal solution.
We start out with the smallest search-augmented model trained on the non-deterministic A Sokoban dataset and use it to generate new shorter sequence training dataset as outlined in Section 3.3. For each Sokoban puzzle in the training data, we generated 32 different <trace><plan>-forme formmtttt sequences by sampling tokens from the Transformer's output distribution and include the shortest generation (measured in tokens) if it contains optimal plan. Subsequently, we fine-tune the search-augmented model this newly created training data (by running an additional 10,000 training steps) to obtain the first Searchformer model. Using this Searchformer model, we subsequently generate another short sequence dataset and repeat the fine-tuning procedure to further improve the model.
Figure illustrates how the sequence lengths generated by the Searchformer model's are iteratively shortened. With every improvement step, the average length of the generated traces-the number of search seeps-decreases (Figure 3a)). The final Searchformer model generates search dynamics sequences that are on average 26.8% shorter than the sequences generated with A search. Consequently, the Searchformer model found a way find a plan in complex task that is more efficient in terms of search steps than the A implementation used to train the initial aurrch-d model. The search-augmented models match the response length distribution of our A* implementation on the test task set, while the Searchformer model's distribution is skewed towards shorter sequence lengths (Figure 3(b)).
As reported in Table 1, fine-tuning resulted in significant performance improvement, reducing the rate of incorrect and non-optimal solutions by 40% and 30% respectively. The Success Weighted by Cost (SWC) score (Wu et al., 2019) factors in how many test tasks were
8

---

Published as a conference paper at COLM 2024
TokenSqquentt Length
A'
A* Search Augmented
[in % relative x
1000
Search Augmented
Sequences [in %]
20
Number of
15
Searchformer
10
Searchformer
80
5
70
8k 10k Sequence Length
2k 4k
6k
2k 4k 6k 8k 10k
Iteration
Sequence Length Avg. per Task
Avg. per Test Task
(a) Search length improvement.
(b) Distribution of average-on-optimal length.
Improvement of search dynamics length via bootstrapping in Sokoban (a): For each test task answered with optimal plan, the average generated execution trace length is computed and averaged across experiment repeats. At each iteration, the subset of test tasks answered with an optimal plan by the search-augmented and Searchformer models are compared and the relative improvement over our A reference implementation plotted. Figure 12 provides more details about the sequence length distribution. (b): Dis- tribution of execution trace lengths for each optimally answered test task before fine-tuning (search-augmented model shown the left panel) and after three fine-tuning iterations (Searchformer shown in the right panel).
Figure 3:
solved correctly and how close the predicted plans are to the optimal length (Appendix C). perfect score would be one, and one can see in Table 1 that the comparably small Search- former matches the performance of the largest solution-only model (also note the small SEM values). The Improvzdd Length Ratio of Search Dynamicc (ILR) measures how much the length of each execution trace is shortened (Appendix C). With each improvement iteration the scores increase and climb above one.
The results reported in Figure 3 and in Table compare each model's performance on test tasks that were either correctly or optimally solved. To test a model overfits to easier test tasks with shorter execution traces, we plot in Figure 10 (in Appendix D) the execution trace length generated with A against the execution trace length generated by each model. Each point in this scatter plot corresponds to single test task. Here, the trend of shortening the execution trace via search dynamics bootstrapping is clear and one can also observe that neither model only specializes on solving easier test tasks with shorter execution traces.
5 Discussion
State-of-the-art LLMs exhibit remarkable capabilities responding to user prompts in seemingly intelligent manner. Despite these advances, LLMs often provide false or misleading responses when prompted with tasks involving goal-directed reasoning and sequential decision making. For example, state-of-the-art LLMs fail solve even small maze navigation tasks with in-context learning (Appendix E). To better understand how language models can be used for planning, this work aims build a connection between language modelling, planning, and sequential decision making. In this context, we demonstrate how to train specialized language models to robustly solve complex planning tasks.
In contrast to classic planning algorithms,, formal guarantees about test performance or the number of bootstrapping iterations needed to reach certain performance level are unknown for our approach. For some small number of test tasks, we found that all models fail to find an optimal or even feasible plan. This property a consequence using deep learning techniques and relying the empirical risk minimization framework (Vapnik, 1991). The aim of this work is to open new doors towards discovering new search heuristics and to develop better understanding of how language models can be applied to planning problems. The aim is not to replace traditional planning algorithms for which there exist formal guarantees with far more efficient implementations.

---

Published as a conference paper at COLM 2024
5.1 Limitations and future work
In this work, a backward-causal decoder is used to generate an entire execution trace. This makes our method computationally very expensive for complex tasks with very long execution traces. In fact, the presented experiments use token sequences that are significantly longer than the sequences used to train LLMs such as Llama 2 (Touvron et al., 2023). One way to mitigate this limitation and improve the efficiency of the presented method is to use curriculum learning: Starting from simple tasks with reasonably long execution traces, train and fine-tune a Searchformer to reduce the trace length, and then use this model to generate shorter execution traces for slightly more complex tasks. Another possibility is explore other planning algorithms or integrate value functions into A search, similar to MCTS, to cap the maximal depth the search algorithm explores. Integrating hierarchical planning methods and temporal abstractions (Sutton et al., 2023; 1999; Dietterich, 2000; Hafner et al., 2022) are another avenue to equip the resulting model with the ability to abstract over decision sequences and find an optimal plan using fewer computation steps.
key aspect of designing a synthetic language that captures algorithmic computation is ensuring that with any token sequence prefix, the algorithm execution state can be reconstructed and execution can be resumed. Although we focus on A" search as an example, our approach could be applied to various algorithms and encourages further research in that direction. For example, our data assumes that only optimal plans are encoded. If training data includes sub-optimal but feasible plans, the resulting solution- only and search-augmented models would produce feasible yet sub-optimal plans. This motivates further experimentation with different LLM fine-tuning techniques (Ouyang et al., 2022; Kaufmann et al., 2024) to enhance solution quality.
As illustrated in Figure 10, each bootstrapping iteration improves the Searchformer in terms of the generated execution trace lengths. One direction of future work would be to investigate how many iterations are optimal and to what degree the model can be improved. Here, variance in execution trace length of the initial training data may play an important role. For example, using an inefficient search heuristic will lead A to explore larger portions of the search tree until optimal solution is encountered by random chance, leading to a higher variance in the resulting trace dataset. In this case, improving the Searchformer model is easier. While there might be theoretical limits to how much particular algorithm can be improved, a Searchformer model may also exploit certain properties of the specific task distribution it is trained on.
In comparison to Plansformer (Pallagani et al., 2022), our approach focuses on training Transformers from scratch to solve complex planning tasks synthetic datasets. We believe that our results and methods could be combined with methods such as Plansformer to fine-tune LLMs and enable them to solve complex planning tasks more robustly. In this context, one important direction for future work is how our approach transfers to in-context learning in pre-trained LLMs that communicate with human users in natural language.
Conclusion
We demonstrate how to effectively train Transformers to solve complex planning tasks by designing synthetic language that mirrors computations performed by the A" planning algorithm. We also introduce a bootstrapping method to optimize this synthetic language, resulting in model that solves novel evaluation tasks in fewer search steps than our reference A implementation. With this study, we hope to shed light on how language models can be applied to sequential decision making and hope to inform further research about better understanding the reasoning capabilities of LLMss.
Acknowledgments
We thank Amy Zhang for helpful discussions on earlier versions of this work. We also thank the anonymous reviewers for valuable feedback that strengthened this paper.
10

---

Published as a conference paper at COLM 2024
References
Kavosh Asadi, Dipendra Misra, and Michael L. Littman. Lipschitz continuity in model-based reinforcement learning. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 264-273 PMLR, 10-15 Jul 2018. URL https://procceedingg war.smmlr.perrr p asadi 18a. html.
Mahmoud Assran, Quentin Duval, Ishan Misra, Piotr Bojanowski, Pascal Vincent, Michael Rabbat, Yann LeCun, and Nicolas Ballas. Self-supervised learning from images with joint-embedding predictive architecture. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 15619-15629, 2023.
Jonathon Cai, Richard Shin, and Dawn Song. Making neural programming architectures generalize via recursion. arXiv preprint or xxXoo:1004.0611, 2017.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 1650-9660, 2021.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling, 2021.
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending con- text window of large language models via positional interpolation. arXiv preprint arXiv:2306. 15599,,,,.,,5,,,, 2023.
Rémi Coulom.
Efficient selectivity and backup operators in monte-carlo tree search. In International conference on computers and games, pp. 72-83. Springer, 2006.
Thomas G Dietterich. Hierarchical reinforcement learning with the maxq value function decomposition, Journal of artificial intelligence research, 13:227-303, 2000.
Angela Fan, Thibaut Lavril, Edouard Grave, Armand Joulin, and Sainbayar Sukhbaatar. Accessing higher-level representations in sequential transformers with feedback memory. CoRR, abs/2002.0942 2020. URL https: http///axxvvvv////eee wwt.:/ogg/////....ooo..oo
Rohit Girdhar, Mannat Singh, Andrew Brown, Quentin Duval, Samaneh Azadi, Sai Saketh Rambhatla, Akbar Shah, Xi Yin, Devi Parikh, and Ishan Misra. Emu video: Factorizing text-to-video generation by explicit image conditioning. arv preprint ar arXin:2311.10709, 2023.
Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep learning. MIT press, 2016.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arxiv preprint WwXi:1440.5441,2211
Arthur Guez, Théophane Weber, Ioannis Antonoglou, Karen Simonyan, Oriol Vinyals, Daan Wierstra, Rémi Munos, and David Silver. Learning to search with mctsnets. In International conference on machine learning, pp. 1822-1831. PMLR, 2018.
Danijar Hafner, Kuang- Huei Lee, Ian Fischer, and Pieter Abbeel. Deep hierarchical planning from pixels. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https://openrevvie.. net/forum2id==Z2kkkk3
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, 2023.
Michael Ivanitskiy, Alexander F Spies, Tilman Räuker, Guillaume Corlouer, Christopher Mathwin, Lucia Quirke, Can Rager, Rusheb Shah, Dan Valentine, Cecilia Diniz Behn, Katsumi Inoue, and Samy Wu Fung. Linearly structured world representations in maze- solving transformers. In UniReps: the First Workshop on Unifying Representations in Neural Models, 2023. URL https: wwwprrvvveww.eeteforeeeee mett/formmidd-pZakRk
11

---

Published as a conference paper at COLM 2024
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning one big sequence modeling problem. Advances in Neural Information Processing Systems, 2021.
Timo Kaufmann, Paul Weng, Viktor Bengs, and Eyke Hüllermeier. A survey of reinforcement learning from human feedback, 2024. URL http:://aaxxi..org/aa///331 14925.
Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steiger- wald, DJ Strouse, Steven Stenberg Hansen, Angelos Filos, Ethan Brooks, maxime gazeau, Himanshu Sahni, Satinder Singh, and Volodymyr Mnih. In-context reinforcement learning with algorithm distillation. In The Eleventh International Conference Learning Representa- lions, 2023. URL https://opeereview..o www.perreeeeeeeeeeoounnnniiaiaee wtt.://www..nttoooomooooooooo
Ilya Loshchilov and Frank Hutter. SGDR: stochastic gradient descent with restarts. CoRR, 03983. Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, 2019.
00ps/1608.03983, 2016. URL http:://axxi..rrg//////
Ida Momennejad, Hosein Hasanbeig, Felipe Vieira, Hiteshi Sharma, Robert Osazuwa Ness, Nebojsa Jojic, Hamid Palangi, and Jonathan Larson. Evaluating cognitive maps and planning in large language models with cogeval, 2023.
MongoDB Inc. MongoDB. https://ww..nnggggg com/. Accessed: 2024-01-23.
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, Charles Sut- ton, and Augustus Odena. Show your work: Scratchpads for intermediate computation with language models, 2021.
OpenAL. Openai codex, 2021. URL https:://peeaai.co///ll/////////
OpenAl. Openai: Introducing chatgpt, 2022. URL http::///opnaal.co///lo/////////
OpenAL. Gpt-4 technical report, 2023.
Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy V. Vo, Mare Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel HAZIZA, Francisco Massa, Alaaeldin El-Nouby, Mido Assran, Nicolas Ballas, Wojciech Galuba, Russell Howes, Po-Yao Huang, Shang-Wen Li, Ishan Misra, Michael Rabbat, Vasu Sharma, Gabriel Synnaeve, Hu Xu, Herve Jegou, Julien Mairal, Patrick Labatut, Armand Joulin, and Piotr Bojanowski. DINOv2: Learning robust visual features without supervision. Transactions Machine Learning Research, 2024. ISSN 2835-8856. URL
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul F Chris- tiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 27730-27744 Cur- ran Associates, Inc., 2022. URL https::/prooceeeee neurs cc/paper.aap. www.fflllppppeee files/paper//////// pdf.
Vishal Pallagani, Bharath Muppasani, Keerthiram Murugesan, Francesca Rossi, Lior Horesh, Biplav Srivastava, Francesco Fabiano, and Andrea Loreggia. Plansformer: Generating symbolic plans using transformers, 2022.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-ppeffrrmacce deep learning library, 2019.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. Xiv preprint ar arXio:2309.0007 2023.
12

---

Published as a conference paper at COLM 2024
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agar- wal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 8748 8763. URL https: / pttp///doocceeensseeeee/eoc wwaiinnss..rr.....r. wwt..peeess/////eeeeeee roooo// html.
PMLR, 18-24 Jul 2021.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1-67, 2020. URL http: http://mrrrrrrrrr wwtp://orppppp/////eeeeeeee/se/e 201/1/22000111111 20-074. 774.hte html.
Baptiste Roziere, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Tal Remez, Jérémy Rapin, et al. Code llama: Open foundation models for code. arxiv preprint ar ArXi::2308 12950, 2023.
Anian Ruoss, Grégoire Delétang, Sourabh Medapati, Jordi Grau-Moya, Li Kevin Wenliang, Elliot Catt, John Reid, and Tim Genewein. Grandmaster-level chess without search, 2024.
Stuart J. Russell and Peter Norvig. Artificial Intelligence: A Modern Approuch. Pearson Education, edition, 2021.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839): 604-609, 2020.
Kurt Shuster, Jing Xu, Mojtaba Komeili, Da Ju, Eric Michael Smith, Stephen Roller, Megan Ung, Moya Chen, Kushal Arora, Joshua Lane, Morteza Behrooz, W.K.F. Ngan, Spencer Poff, Naman Goyal, Arthur Szlam, Y-Lan Boureau, Melanie Kambadur, and Jason Weston. Blenderbot 3: deployed conversational agent that continually learns to responsibly engage. ArXiv, abs/ abbs/2008.00m 2208.03188, 2022. URL https://api. wtw.ppis.mmaattccolllssee 251371589.
David Silver, Julian Schrittwieser, Karen Simonyan, loannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, 550(7666) 2017.
David Silver, Thomas Hubert, Julian Schrittwieser, Toannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, Timothy Lillicrap, Karen Simonyan, and Demis Hassabis. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, 362(6419) 1144, 2018. doi: 10.1126 wewence..arre///. URL https: wwtp://science.or/ddooooo 1126/ wttp:/ccc.cc.aaaaaa..
Uriel Singer, Adam Polyak, Thomas Hayes, Xi Yin, Jie An, Songyang Zhang, Qiyuan Hu, Harry Yang, Oron Ashual, Oran Gafni, Devi Parikh, Sonal Gupta, and Yaniv Taigman. Make-a-video: Text-to-video generation without text-video data. In The Eleventh Interna- fional Conference Learning Representations, 2023. URL https: h/tpps:/openreviieee wtw.eerrriieeeeeneoorrnn id-nJfylDvgzlq.
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, 2023.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Richard S Sutton, Doina Precup, and Satinder Singh. Between mdps and semi-mdps: A framework for temporal abstraction in reinforcement learning, Artificial intelligence, 112 (1-2):181-211, 1999.
Richard S Sutton, Marlos C Machado, G Zacharias Holland, David Szepesvari, Finbarr Timbers, Brian Tanner, and Adam White. Reward-respecting subtasks model-based reinforcement learning. Artificial Intelligence, 324:104001 2023.
13

---

Published as a conference paper at COLM 2024
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023.
Trieu H. Trinh, Yuhuai Wu, Quoc V. Le, He He, and Thang Luong. Solving olympiad geometry without human demonstrations. Nature, 625(7995):476-492 2024. doi: 10.1038/ s41586-023-00747-55 URL https://doi. https/////oo..///..... org/10. 1038/s
Karthik Valmeekam, Matthew Marquez, Sarath Sreedharan, and Subbarao Kambhampati. On the planning abilities of large language models critical investigation, 2023a.
Karthik Valmeekam, Alberto Olmo, Sarath Sreedharan, and Subbarao Kambhampati. Large language models still can't plan (a benchmark for ilms on planning and reasoning about change), 2023b.
V. Vapnik. Principles of risk minimization for learning theory. In J. Moody, S. Hanson, and R.P. Lippmann (eds.), Advances in Neural Information Processing Systems, volume 4. Morgan-Kaufmann, 1991. URL httpss//prooceeddngg meerridiiinnnnnnnii...se.. fwtpppppprflllsppppapapsseeepeep 1991 999//iieff44fhhaarrfffff666e22b2 hatpeeddddddaaae...ddp
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, LlionJones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. CoRR, abs/1706 2ess/770.03762, 2017. URL httt:://arkie htttpp///arii..orr/b////666 03762.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, brian ichter, Fei Xia, Ed Chi, Quoc V Le, and Denny Zhou. Chain-of-thought prompting elicits reasoning in large language models. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 24824-24837. Curran Associates, Inc., 2022.
Yi Wu, Yuxin Wu, Aviv Tamar, Stuart Russell, Georgia Gkioxari, and Yuandong Tian. Bayesian relational memory for semantic visual navigation, 2019.
Mengjiao (Sherry) Yang, Dale Schuurmans, Pieter Abbeel, and Ofir Nachum. Chain of thought imitation with procedure cloning. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 36366-36381. Curran Associates, Inc., 2022.
Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Thomas L. Griffiths, Yuan Cao, and Karthik Narasimhan. Tree thoughts: Deliberate problem solving with large language models, 2023.
14

---

Published as a conference paper at COLM 2024
A Training encoder-decoder Transformers to predict optimal plans
All program code used for the experiments presented in this paper as well as all model check- points can be found on Github at https://github. comm/facbbookeesercchh wtwprereeearrrressccrrrrroooo
In this work, we consider encoder-decoder Transformers (Raffel et al., 2020) and process each prompt- the task specification-
encoder network to obtain an embedding of planning task. Using the decoder, the resulting embedding is then decoded into a response of the format <trace><plan> or the format <plan>. We refer to a model predicting sequences of the format Strace><plan>> as a search-augmented model and a model predicting sequences of the format <plan> as solution-only models.
A.1 Encoder-decoder architecture
A token sequence is processed by neural network by first indexing the set of all tokens contained in a specific dataset. This indexing is used to map any token sequence to set of integers. Formally, we denote prompt as sequence of integers x1xx (x,,..,, xn)- An encoder Transformer neural network with network weights dene is a function feanc mapping tokenized prompt Xim of arbitrary length n to sequence of real-valued vectors:
fe Xim ト Zim
(1)
Each vector z; in this sequence of vectors 1 has the same dimension. The decoder network is then used to generate response auto-regressively: Starting with specific beginning-of- sequence token bos to cue the decoder, a sequence is recursively built by first mapping the start token bos to a probability distribution over next-tokens. This probability distribution is stored as vector P1 whose dimension is equal to the size of the vocabulary. The next token then generated by sampling from this distribution and the sampled token is appended to the response sequence. Subsequently the two-token response sequence is fed into the decoder again to compute the next probability vector P2 and sample the next token. This procedure is repeated until an end-of-sequence token eos is sampled. While only the last computed probability vector is needed to sample the next token, the decoder network simultaneously predicts a sequence of next token probability vectors Piam given an input sequence y1.m.. Furthermore, this prediction is conditioned on the encoder output 1.m. The decoder Transformer neural network with weight parameters d is therefore function
80 Zine V1:m
Pi.n-
(2)
The encoder network fore and decoder network 80 internally both use a number of stacked causal attention layers to form an encoder-decoder Transformer (Vaswani et al., 2017) as outlined in more detail in Appendix A.3. We denote concatenation of all encoder parameters enc and decoder parameters ddec with the vector 0.
A.2 Training with teacher forcing
An encoder-decoder architecture is optimized to generate responses that follow the dis- tribution of training dataset by minimizing the cross-entropy between the training data distribution PD over prompt-response pairs (Xn,ymi ym) and the distribution Pe with which the encoder-decoder model is generating responses. This cross-entropy loss objective
H(pp. pe)
EE [- log po (y [1,]]
(3)
EE
Σ
log Po 1am Vy
where line (4) follows from the auto-regressive generation procedure described before. Within the same planning dataset, different prompt-response pairs can have different prompt and response lengths. To emphasize shorter response sequences during training, we re-weigh each sample resulting in the loss objective
(( =
1 D
1
m 11
log Po (v 1mn Iv A.mng))
(5)
m4
15

---

Published as a conference paper at COLM 2024
where the first summation averages over all D prompt-response pairs of the training dataset. In Equation (5) the super-script d indexes individual prompt-response pairs in the training dataset. This average is an empirical approximation of the expectation in Equation (3) for finite i.i.d. sample of size D. This loss objective is optimized using gradient descent (Good- fellow et al., 2016, Chapter 10).
Hidden activations
Linear Layer
Decoder Block
Attention
Hidden activations
Encoder Output
Layer Norm
Encoder Block
Causal Self-Attention
Causal Self-Attention
Apply RoPE to and K only
Apply RoPE
Q and K only
Layer Norm
Layer Norm
Hidden activations
Hidden activations
Figure Attention blocks architecture used in the encoder-decoder Transformer archi- tecture. The encoder network consists of set of feed-forware layers where each layer processes hidden activations parallel with a set of encoder attention blocks (left diagram). Similarly, the decoder network consists of a set of feed-forward layers composed of a number of decoder attention blocks (right diagram). The number of blocks each layer is referred to the number of heads. Token sequences are first mapped into integer sequences using look-up dictionary. Then, these sequences are fed through a PyTorch (Paszke et al., 2019) torch.m nn. Embedding layer map the integer sequence into sequence of hidden activation vectors. After the last decoder layer, hidden activations are mapped through linear layer into sequences of logits vectors to compute the next token probability vectors Pi.n.
A.3 Network architecture and hyper-parameters
The encoder-decoder Transformer first maps every token to one-hot vector of the same dimension as the token vocabulary space. These one-hot vectors are then projected through linear embedding layer into set of vectors.
The encoder then consists of multiple feed-forwan layers and each layer consists multiple encoder blocks (left part of Figure 4). The output of these layers then mapped through another linear layer to project the hidden activations into tensor of the correct shape. The decoder also consists of multiple feed-forward layers and each layer also consists of multiple decoder blocks (right part of Figure 4) processing the hidden activations in parallel. As illustrated in Figure 4, the decoder network conditioned on the encoder by processing the encoder's 's output tensor directly in the second attention map. Furthermore, each tokens position is encoded by applying RoPE embeddings (Su et al., 2023) as indicated in Figure We did not use dropout in our architecture.
Table lists the used architecture hyper-parameter and Table lists the prper-parameters used for optimizing each model. All experiments et al., 2019) and default parameters were used unless reported here otherwise.
were implemented in PyTorch 2.0 (Paszke
16

---

Published as a conference paper at COLM 2024
Table 2: Architecture Hyper-parameters. The same parameters were used in both the encoder and decoder network. The number of heads indicates how many attention blocks are used in one layer. Layer dimension indicates the dimension of the feature vectors processed through each attention block (dimension of K, V, and Q in Figure 4). All models
used a RoPE frequency of 10000.
Parameter
15M model
46M model
175M model
747M model
Layers
8
9
16
Heads
Layer dim.
64
96
192
96
Table 3:
Optimization hyper-parameters. Every model was optimized using AdamW (Loshchilov Hutter, 2019) with setting 6o 0.9 and ß1 = 0.99, Initially, the learn- ing rate was linearly interpolated: Starting at zero and then increasing linearly to the value listed below until step 2000. Then a cosine learning rate schedule was followed (Loshchilov & Hutter, 2016). to the between zero at the first training step and the listed value below for the
Parameter
Model
Maze lasks
Sokoban Puzzels
Learning rate
15M
2.5 10
2.5 10 4
46M
7.5-10 10-5
7.5
175M
10-
5.0110-3
747M
10-5
Batch size
all
16
64
Training steps
all
400000
80000
B Dataset generation
All datasets including token dataset generated our models are publicly available through
inference with our Github repository at
through
https://githh: thub.
All datasets were generated by first randomly sampling a task and then executing A to obtain an optimal plan. Maze tasks were generated first by randomly selecting 30-50% of all cells to be wall cells. Then start and goal location was randomly selected and A was executed to obtain an optimal plan. If the plan had length of at least the mazes width or height (e.g. for 10 x 10 mazes the optimal plan needs to contain at least 10 steps), then the task was added into the dataset. For Sokoban, a 7 x 7 grid was sampled and two additional wall cells were added as obstacles to the interior of the map. Then two docks, boxes, and the worker locations were randomly placed. If the sampled task is solvable by A", then the task was admitted to the dataset.
Figure 5: Example Sokoban puzzle?. The prompt, A execution trace, and optimal plan for this task is illustrated in Figure in Appendix B. For Sokoban, we only use non-deterministic A implementation.
Due to the large volume of generated data, all data was stored in and transformed with a MongoDB (MongoDB Inc.) instance and map-reduce techniques. Furthermore, when sampling tasks, the dataset was constructed to reject duplicate tasks ensuring that training and test data and each prompt is distinct. Once each task and trace dataset was generated, each execution trace is converted into prompt and response token sequences, as illustrated
2This level image composed using image icons from https: //gi thub. comm//ooeenod///ooooann (accessed 2023-11-21).
17

---

Published as a conference paper at COLM 2024
in Figure 1(c) and Figure 8. Because the Sokoban tasks are very complex and difficult to solve for A", the resulting token sequences were partially very long and reached almost 100000 tokens. Due to GPU memory requirements, the Sokoban dataset was further sliced to only include sequences of with at most 10000 tokens. Figure 8 compares the sequence length distribution for each dataset. During training, each dataset was also sorted and sliced to only contains the reported number of training sequences. Furthermore, each model was evaluated on the same test tasks within each task dataset. The test dataset contains plans and traces that are of comparable length to the training data (Figure 8).
bos
bos
plan
worker
create close
plan 1
box
4
Trace (2583 tokens)
worker bex 4 box c0
plan 1 4
box
3 4
plan 1
dock
3
create worker box 3
c10
plan 2
Prompt dook
close
c12
plan 2
plan 3
wall
0
Plan
plan 2
wall
2
plan
plan 2
wall
plaa 1
GoG
plan 1
plan
Figure 6: Token sequence example for Sokoban This figure lists the token sequence for the Sokoban level depicted in Figure
Sokoban
30x30 Maze
20x20 Maze
10x10 Maze
100 200 300 400 500
2k 4k
6k
8k
10k
Solution-only Sequence Length
Search-Augmenter Sequence Length
Figure 7: Training Sequence Length Comparison. The left panel plots the length of the solution-only sequences and the right panel plots the length of the search-augmented sequences, excluding the start and end of sequence tokens bos and eos. The whiskers indicate the range of all sequence lengths and the box plots indicate the 25%, 50% and 75% quantiles. Because we focus on planning the token sequences are multiple orders of magnitude longer than usual token sequences used to train L.LMss-sspccalle when A* execution traces are included in the responses. For example, fine-tuning of the Llama 2 model on human preference data is performed with sequences that are on average 600 tokens long (Touvron et al., 2023).
complex sequential decision making tasks,
18

---

Published as a conference paper at COLM 2024
Plan Sequence Length
A Execution Trace Length
non-det.
Trann
10x10 Maze
det.
Test
non-det,
20x20 Maze
non-det,
30x30 Maze
det.
Sokoban
nen-det,
100
200
300
400
2k
6k
8k
10k
Number Tokens
Number of Tokens
Figure 8: Sequence length distribution for each dataset. The training and test sets are designed such that their sequence lengths match closely.
C Evaluation criteria
In the presented experiments evaluate if each model outputs an optimal feasible plan and how long the generated search sequences are for the search-augmented and Searchformer models.
C.1 Measuring plan optimality
We evaluate whether the output plan is optimal using one of three criteria:
Exact-match criterion. For each task, if the generated sequence from a trained model matches the output of deterministic A* exactly, is labelled as correct, otherwise labelled as incorrect. This is only used to evaluate supervised cloning of deterministic A" Any-optimal-64 criterion. For each task, sample 64 responses from trained model. Each response is parsed and evaluated if contains a feasible or optimal plan, regardless of the generated <trace> part. If any of the 64 plans is feasible and optimal, then the task is labelled as correct.
SWC score. To further measure the sub-optimalness of the resulting plans, we also report the Success Weighted by Cost (SWC) (Wu et al., 2019), a statistic that factors in how close the cost I the best predicted correct plan (over 64 trials) is the optimal plan cost 1i, averaged across all n test tasks and weighted by a binary variable cj E (0, 1):
1 n
Σ SWC == 11 1=1
max
14,17
When computing the SWC score, the binary variable ej is set to one if correct plan was found and zero otherwise. This SWC score lies between zero and one. If all generated sequences end in an optimal plan, then this value one.
C.2 Measuring search dynamics length
For sequences ending in optimal or feasible plan, evaluate the length of sequence dynamics in terms of number of tokens using one of the two criteria:
Average-on-optimal length. For each task, we sample 64 responses from a trained model and compute averaged search dynamics length for sequences that lead to an optimal plan.
ILR score. To further measure how much improvement of the moneeraeee search dynamics against A planning, we report the Improved Length Ratio of Search Dynamics (ILR) score. Specifically, for each test task compute the ratio between the length
19

---

Published as a conference paper at COLM 2024
of the shortest generated search dynamics sequence and the A token sequence length f. We then average across all test tasks while only including ratios for tasks for which either an optimal or correct (and potentially sub-optimal) plan was found, specified by the binary variable ej. The corresponding measures are thus called on--onnmpl and oon-ood- The ILR is defined
ILR: = G
The ILR measure can take non-negative values and values above one indicate that the model generates shorter search dynamics than the A* reference. Consequently, if the numbers lie significantly above one, then the model has found more efficient way to search a task's state space to compute an optimal plan.
D Searchformer performance analysis
In Section 3.3, each search-augmented and Searchformer model is evaluated by generating 64 token sequences for each Sokoban test task. For the same test task the same model can generate sequences that end in an optimal plan, incorrect plan, sub-optimal plan. In Figure 10 we compare the length of the generated sequences with the length of sequences generated when using A search for each case. The percentages in each panel's caption list how many out of all generated sequences end in an optimal plan, correct plan (that can be optimal or sub-optimal), and a correct but sub-optimal plan.
correct but
In the left panel of Figure 10(a) and Figure 10(b), points are centered around the diagonal axis indicating that the search- aarchhauggented models do approximately match the A' search algorithm in terms of token sequence lengths. Figure 10(a) and Figure 10(b) further confirm the results presented Sec. 3.3: With every improvement step, the points move down and below the diagonal line. This highlights that the improved Searchformer models generate token sequences that are shorter than sequences generated with A search. The Searchformer has found method of searching task to compute an optimal plan in fewer search steps than A search uses.
Figure 10(c) illustrates what happens when each model generates a correct but sub-optimal plan. Here, the search-augmented model, that is trained to imitate A', generates trace sequences that are significantly longer than the sequences generated with A*. This suggests that the model struggles in computing a correct plan and generates too many search steps, ultimately leading in finding a correct but sub-optimal plan. Similarly, the Searchformer models also generate sequences that are longer than the sequences generated with A search. Despite these inefficiencies, our bootstrapping method is still able to improve the model and bring the average sequence length closer to the length of sequences generated with A search (right most panel in Figure 10(c)). While we would desire the trace length to be low in either case, we found that each model generates a correct but sub-optimal plan with less than 5% chance. Consequently, Figure 10(b) shows that the final Searchformer model still generates correct plan with on average fewer search steps than A search. Statistically, the differences between Figure 10(a) and Figure 10(b) are marginal.
D.1 Ablation experiment comparing plan tokenization pattern
This ablation experiment test how sensitive our performance results are to changes in how plans are encoded in token sequences. For this experiment we use the deterministic * dataset for 20 mazes and encode solution plan as a decision sequence instead of encoding the plan as a position sequence. For example, a token sequence of the form plan 2 plan 1 plan 0 plan
a token sequence of the form plan down down 1eft. We refer to this tokenization pattern as the action token style and emphasize that the tokens plan, up down, left, and right are only used to encode the optimal plans. The do not occur anywhere else in the used training sequences.
Table 4 compares the test set performance of the solution-only and search-augmento models for each tokenization style. One can observe that the solution-only models, which
20

---

Published as a conference paper at COLM 2024
Table Comparison between different plan tokenization styles This table compares how changing the plan tokenization pattern impacts performance. This experiment uses the deterministic A dataset for 20 x 20 mazes and compares performance numbers for a 100,000 training sequence dataset. Each number lists the number of correct response sequences (in %) for the test task set. This table reports averages and SEMs across five repeats.
Params.
Model
Position token style
Action token style
Solution only
59.9 =4.2
31.8 +:.11
15M
Search augmented
83.0 2.0
84.2 +1.3
Solution only
67.7 4:2
48.2 1.3.8
45M
Search augmented
=18
90.8 +:07
Solution only
68.6 =4.1
38.1 +7.0
175M
Search augmented
=1.7
93.0 +0.3
are trained on <erompttxplan--forrttte sequences, degrade in performance significantly. In comparison, the performance of the search-augmented models, which are trained on cetpxtttrcexxpllln-omttttt
sequences, increases by few percent points. This result underlines the benefits of our approach of including execution traces into the training data, leading to datasets with more tightly correlated token sequences that are easier to learn.
E State-of-the-art LLMs on Maze Task
Below we list examples that even state-of-the-art LLMs still fail to solve a simple maze navigation problem.
Zero-Shot-Example
Prompt:
Here's the generated 10x10 maze along with Locations of Walls: (0, (0, 8) (1, 6) (3, 0) (3, 1)
(3, 9) (4, 6) (4, 8) (4, 9) (5, 7) (6, 7) (8, 2) (8, 8) (9, 2) (9, 5) (9, 6) (9, 7) (9, 8). Start Location: (0, 0).
End Location: (9, 9). Could you solve this? Give me the optimal path in (x.y)
Gemini Answer: (0, (1, 0000000------------00-000 (4,2) (5,2
(5, (6, (6,4) 4)
-> 5) + -60 - (8.8) - (9,8) (9,9)
Note: The path incorrectly includes (3, 0) and (3,1), (8, 8), (9, which are wall locations.
Few-Shot-Example
Prompt:
Here's the generated 10x10 maze along with Locations of Walls: (0, (0,
(3, 9) (4, 6) (4, 8) (4, 9) (5, 7) (6, 7) (8, 2) (8, 8) (9, (9, 5) (9, 6) (9, 7) (9, 8). Start Location: (0, 0).
End Location: (9, 9). Could you solve this? Give me the optimal path in (x.y)
Example 1:
Locations of Walls: (0,4), (0, (0, 8)
0) (1, 3), (1, (3, (4, (4, (4, (4,5)
(5,
6), (5, 8) (6, 1), (6, 2), (6, 4), (6, 9) (7, 9) (8, 4), (8, 5) (8, 9) (9, 1). Start Location: (0, 0). End Location: (9, 9). Optimal Path: (1)
(2,3) (2,5), (2, 6), (2, (3, 7), (4, 7), (5, 7), (6, 7), (6, 8), (6
Example 2:
Locations of Walls: (1100000000000000000000000404410 (10TG//0000000000000000044444440
(5)
(6, 8) (6, ) Start Location: Loctt 0). End Location: 99,9). Optimal Path: (0, 0) (0, 1) (0, 2) (1, (1,2)
(7,
(2 PPlk(000000000223333232333333330
(4,5) (5, (6, 5) (7,5) (8, (8,
7) (8, (8, 9) (9,9)
21

---

Published as a conference paper at COLM 2024
Gemini Response: (0, 0), (1, 0), (2, 0), (2, 1), (2, 2), (3, 2), 3), (3, 4), (3, 5) (3, 6), (4, 6) (5, 6), (5, 5), (6, 5), (6, 4),
(7,4), (7, (7, 6), (7, 7), 7), (8, 6), (8, 5) (8, 4), (8, 3), (8.2), (8, 0) (9, 0) (9, 1), (9,2), 2),
(9,3), (9, 4) (9, (9, (9, (9,
(9, 6)
Note: The path breaks into the walls (4,6) (8,2) (9.2) (9.6) (9.7),
Supplemental figures for experiments
100
Solution 15L
Correctly Solved
Test Tasksfin %.
80
Solution SolutionOnly.44e Only,46M
60
Solution Only,175M
50k 100k
500k 1M
Number of TrainingSequencss
Figure 9: Solution-only model performance. Performance of the solution-only models is primarily influenced by the number of training sequences. Increasing model's size does not always improve performance.
22

---

Published as a conference paper at COLM 2024
Search Augmented (56.3% of gen. seqs.)
Searchformer, Step (62. 2% gen. seqs.)
Searchformer, Step 4% gen. seqs.)
Searchformer, Step 1% all gen. seqs.)
Generated Trace Length
10k
5k
0
10k
5k A Trace Length
5k
10k
A Trace Length
A Trace Length
Trace Length
(a) Comparison of all sequences ending with an optimal plan
Search Augmented
Searchformer, Step (64. 8% of all gen. sequ.)
Searchformer, Step (65. 7% of all gen. seqs.)
Searchformer, Step (64.5% of all gen. seqs.)
Generated Trace Length
(58.5% of all gen. seqs.)
10k
5k
10k
5k
10k
0
5k
10k
5k
10k
A TTacce Length
A TTace Length
A Trace Length
A Trace Length
(b) Comparison of all sequences ending with correct plan
Search Augmented
Searchformer Step of all gen. seqp.)
Searchformer, Step (3.3% of all gen. seys.)
Searchformer, Step (4.4% all gen. seqs.)
Generated Trace Length
(2.2% all gen. sequs)
10k
5k
10k
10k
10k
5k
10k
A Trace Length Lenghh
A Trace Length
Trace Length
Trace Length
(c) Comparison of all sequences ending with correct but sub-optimal plan
Figure 10: Sequence length comparison between sequences generated with A search and sequences generated with each model. Each dot in each scatter plot corresponds to one specific test task. On the x-axis, we plot the average token sequence length when A search is used. On the y-axis, we plot the average token sequence length when each model is used. Error bars indicate standard deviations. Percentages indicate the fraction of the generated sequences that are included in each scatter plot. (a): Sequence length comparison for all test prompts for which an optimal plan was generated. (b): Sequence length comparison for all test prompts for which a correct plan was generated. This plot aggregates across sequences ending in an optimal plan and sequences ending in a correct but sub-optimal plan. (c): Sequence length comparison for all test prompts for which correct but sub-optimal plan This plot only aggregates across sequences
was generated using the corresponding ending in correct but sub-optimal plan.
23

---

Published as a conference paper at COLM 2024
ColutionOOnll15500
Solution Only, 46M 100k Seqs.
Solution SolutionOnlly17773 Only,175M 500k Seqs.
Correct Responses
50k Train Seqs.
TrainSeqs.
100
80
[in %d
10x10
20x20
30x30
10x10
20x20 Maze Shape
30x30
10x10
20x20 30x30
10x10
20x20
30x30
Maze Shape
Maze Shape
Maze Shape
(a) Deterministic plan prediction with solution-only models
Search Augmented, 15M Search Augmented, 46M - Search Augmented, 175M
Solution Only,175M
Correct Responses
50k Train Segs.
100k Seqs.
500k TrainSeqs.
TrainSeqs.
100
80
[in %6]
60
40
10x10
20x20
30x30
10x10
20x20
30x30
10x10
20x20 Maze Shape
30x30
10x10
20x20
30x30
Maze Shape
Maze Shape
Maze Shape
(b) Deterministic prediction with search-augmented models
Solution 15M
Solution Only,46M
Solution 175M
Correctly Answered
50k Train Seqs.
100k Train
500k Seqs.
IM TrainSeqs.
100
Prompts lin %6]
80
60
40
20
10x10
20x20
30x30
10x10
20x20 Maze Shape
30x30
10x10
20x20
30x30
10x10
20x20 30x30
Maze Shape
Maze Shape
Maze Shape
(c) Non-deterministic plan prediction with search-augmented models
Search Augmented,
Search Augmented, 46M
Search Augmented, 175M
Solution OolutionOnll,1733.
Correctly Answered
50k Train Seqs.
100k Srqin
500k Seqss
IM TrainSeqs.
Prompts jin %]]
1100
40
20
10x10
20x20
30x30
10x10
20x20
30x30
10x10
20x20
30x30
10x10
20x20 Maze Shape
30x30
Maze Shape
Maze Shape
Maze Shape
(d) Non-deterministic prediction
predictnn with search-augmented models
Figure 11: Optimal plan prediction performance for solution-only and search-augmented models on maze datasets. (b), Percentage of correctly generated responses. Note that the solution-only models only need to generate few hundred token long plan to answer prompt correctly. The search-augmented model must generate a much longer A execution trace correctly to correctly answer prompt. This requires the trace plan models to generate response sequences that are hundreds or thousands of tokens long (cf. Figure 7). If one token is incorrectly predicted, the response is counted as incorrect. (c): Percentage of correctly generated plans. Each model was used to generate multiple responses for each prompt. If one of the generated responses contains optimal plan, the test prompt is counted as correctly answered.
24

---

Published as a conference paper at COLM 2024
Search Augmented
Step
Searchformer
Step 2
Searchformer
Step
Searchformer
5000
10000
Sequence Length (Avg. per TestTask
Figure 12: Comparison of search dynamics length (in terms of number of tokens) between A search and our models (search-augmented in blue, Searchformer step 1-3 in orange), on the test subset in which our models yield optimal plans. Here, for each test task, we average the lengths of sequences that ends in an optimal plan, out of 64 trials (i.e., average-on- optimal length (Appendix C)). The box plots show the distribution of average-on-optimal lengths, in which the left boundary, mid solid line, right boundary represents the 25%, 50% and 75% quantiles. Dotted lines are means and whiskers show the range.
25


Reviews:
[{"review_id": "OXlSQ3Vs10", "reviewer": "Reviewer_KTvv", "summary": "This paper presents a novel and effective approach to train Transformer models for complex planning tasks and demonstrates improvements over traditional algorithms like A*. The search dynamics bootstrapping methodology, which trains the model to imitate A*'s\nsearch process and then fine-tunes it to discover shorter solution sequences, is a creative and promising technique. And it could be extended to other algorithms and problem domains, paving the way for more efficient and generalizable neural network-based planners.", "questions": "", "limitations": "", "rating": 6, "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": "**Re Q1:**\nOur method could be applied to different planning algorithms by designing a new synthetic language. For A*, we found it is sufficient to express all of its steps in a token sequence. Given any prefix of the token sequence, the state of A* can be reconstructed to resume execution. This recipe of logging all operations could be applied to different algorithms to construct the corresponding synthetic language and train Searchformer.\n\n**Re Q2 and Re Reason to Reject Pt. 2:**\nA limitation of most deep learning methods is that guaranteeing test performance is often very difficult, and the theoretical properties of transformers still remain wide open. This applies to both (1) theoretical guarantees of planning like what A* gives, and (2) whether the bootstrapping process has guarantees. The motivation of this work is not for theoretical guarantees, but to open a new door to discover more efficient algorithms (in terms of “number planning of steps”, not wall-clock time) than traditional ones, verified by extensive empirical analysis.\n\nFor bootstrapping, we choose to use non-deterministic sequences to diversify the generations, which makes it easier to pick up better planning samples for the next iteration. The fine-tuning only needs a reward function (shortening the search trace while maintaining optimality), and does not require hand-crafted heuristics, as in traditional methods. \n\n**Re Reason to Reject Pt. 1:** Sec 2 provides extensive literature review. In summary, our work provides a novel training paradigm to improve traditional search methods by first doing imitation learning on search paths, and then further improving them by leveraging differentiability of Transformers. In comparison, previous works either start from sequence models (e.g., pre-trained LLMs) and explore heuristics for better reasoning in the natural language space (e.g., Chain-Of-Thoughts), or learn heuristics for existing search procedures (e.g., AlphaGo). \n\n**Re Reason to Reject Pt. 3:** Our approach should still work in this case with minor changes. If the training data have suboptimal solutions, we can generate multiple sequences per input, and pick the one that is closer to the optimal plan in length (and have shorter search trace), and use that data for the next iteration. This will lead to better plans that is close to optimal. In the paper, we focused only on the search trace length as we are interested in optimizing the internal search and planning process of the model.", "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": "Pros:\n1. The paper presents a new method, search dynamics bootstrapping, for\ntraining Transformer models to solve complex planning tasks, which opens\nup new possibilities for applying deep learning to planning and reasoning\nproblems.\n2. The Searchformer model can outperform A* algorithm in efficiency and\nproblem-solving ability, which highlights the potential of this approach to\nplanning processes and tackling more complex problems.\n3. This paper presents a well-designed set of experiments that systematically\ntest their approach under different conditions. The use of synthetic\ndatasets and the Sokoban domain demonstrate the approach's\neffectiveness, while also suggesting its potential for generalization to other\nplanning problems.", "reasons_to_reject": "Cons:\n1. While the paper demonstrates improvements over the A* algorithm, it does\nnot compare the Searchformer to other recent advances in neural network-based\nplanning or symbolic models, making it difficult to assess its relative\nperformance.\n2. The paper does not provide theoretical guarantees on the optimality or\ncompleteness of the Searchformer's planning process. Further analysis is\nneeded to understand the conditions under which the model's performance\ncan be guaranteed.\n3. The search dynamics bootstrapping approach relies on the A* algorithm to\ngenerate the initial training data. If the underlying algorithm is suboptimal or\nfails to find solutions for certain problems, it could limit the performance of\nthe Searchformer model.", "questions_to_authors": "1. Could this method be applied to more advanced symbolic planners or other\ntypes of heuristic search algorithms?\n2. Is there a theoretical bound on the efficiency gains that can be achieved\nthrough this bootstrapping process?", "extra_scores": {}}, {"review_id": "axebvFt3y2", "reviewer": "Reviewer_Loks", "summary": "The manuscript proposes a framework for optimising Transformer-based models for search-based planning tasks, through iteratively fine-tuning the models on traces generated from effective (albeit heuristic) symbolic planners like A*, in order to surpass those heuristic approaches in terms of number of search steps.", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "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": "**Re Sec. 1:** We aim to improve the performance of Transformers in complicated reasoning/planning tasks, which serve as LLM's core components. We train from scratch to make clean-cut comparison, since many pre-trained models are trained on proprietary datasets to excel in major benchmarks, which may contain search dynamics of puzzles. We do not assert that Searchformer can be integrated into other frameworks, but suggest the possibility (Sec. 5.2 future work).  Our initial tests show that in-context learning with SoTA LLM cannot solve our tasks (will put in Appendix). \n\n**Re Sec. 2:** Our approach is different from hierarchical planning, in which it typically defines temporal abstractions (eg., options) and relies on end2end training to learn such abstraction from scratch via RL. In contrast, our approach starts with imitation learning on well-defined steps of A*’s search process, and improve via fine-tuning. Therefore, unlike hierarchical planning that may be very tricky to tune and learn, and may require a lot of data, our approach may be more data efficient. We will include the discussion in the revision. \n\n**Re Sec. 4.3:** Our bootstrapping method finds a shorter execution trace by repeatedly sampling sequences with the same prompt. These shorter execution traces will still lie in the support of the distribution modeled by the trained Searchformer. By repeating the bootstrapping process, we aim to shift the execution trace length distribution to discover new execution traces that lie outside the initial A* training dataset. This shift can be observed in Fig 10a, where most points lie below the main diagonal indicating that the execution trace length distribution has been shifted significantly after multiple iterations.\n\nAt this point, it is not clear how many bootstrapping iterations are optimal. Variance in execution trace length of the initial training data may play an important role. For example, using an inefficient search heuristic will lead A* to explore larger portions of the search tree until an optimal solution is encountered by random chance, leading to a higher variance in the resulting trace dataset. In this case, improving the Searchformer model should be easier. If the training data already outlines an optimal search process, it may be much more difficult for an improvement.\n\nOne iteration currently takes more than 24 hours to complete while using 512 V100 GPUs. So we leave a more detailed study of additional iterations for future work.", "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": "- The paper is well-written.\n\n- The manuscript considers an interesting subject — improving the inductive bias of Transformer-based models. This is potentially impactful, as they remain a component of larger-capacity pre-trained models, like foundation models.", "reasons_to_reject": "Section 1 — On the subject of potential impact: the manuscript presents a framework for improving the effectiveness of Transformer-based models and later asserts that the method can be integrated into other frameworks that similarly optimise LLMs. However, the manuscript never showed whether this is even necessary (e.g., for the tasks it considers). The manuscript should also provide comparisons with LLMs that learns A* search dynamics, in context, for reference.\n\nSection 2 — Rather than highlighting only the works that leverage Transformers to generate single actions, the manuscript may also discuss and compare against hierarchical approaches as well as works that leverage Transformer-based models for generating plans and trajectories. Whereas the manuscript makes reference to some of these works in Section 5, it is attempting to derive its novelty by drawing a comparison between its multi-step planning capability and single action generation. Therefore, it does not seem appropriate that the manuscript should omit the other methods from this discussion.\n\nSection 4.3 — The manuscript states, \"Using this Searchformer model, we subsequently generate another short sequence dataset and repeat the fine-tuning procedure to further improve the model.\" How many times, and — more importantly — why? What general insight can the manuscript provide about, e.g., the relationship between the nature of the task and the number of fine-tuning steps that would be required to obtain the best Searchformer model?", "questions_to_authors": "N/A — see above.", "extra_scores": {}}, {"review_id": "1e2JVx8ziG", "reviewer": "Reviewer_6qVT", "summary": "This paper presents SearchFormer, a Transformer architecture trained to solve multistep planning and reasoning tasks. SearchFormer is retrained on token sequences representing the dynamics of A* search and then fine-tuned via expert feedback. SearchFormer successfully solves Sokoban puzzles in fewer steps than standard A*.\n\nThis is a novel and interesting use of Transformers, and overall I found the results to be convincing. Some shortcomings that I see concern the lack of comparison to other pathfinding algorithms besides A*, and a lack of guarantees similar to those provided by A*. The significance of the work is undercut by these absences, because without them it comes across as only slightly more than a toy demonstration of Transformers' ability to generate A* sequences and thereby solve A* problems, similar to ChessGPT's ability to generate chess moves and thereby passably play chess.\n\nThe paper is also written quite clearly and is easy to follow.", "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": "We thank the reviewer for the helpful comments and feedback.\n\n**Re. guarantees for finding a shortest path:** Indeed the property that every plan is optimal is captured implicitly in the training data, and imitated by the Searchformer. Note that even trained with optimal paths, Searchformer is still not guaranteed to output optimal plans for unseen tasks. For more discussion on theoretical guarantees, see *“Re Q2 and Re Reason to Reject Pt. 2”* part in our response to reviewer KTvv (sorry, can't fit here due limited space)\n\n**Re. other planning algorithms:** Our method is a novel training paradigm, and is orthogonal to other search/planning algorithms. See *“Re Q1”* part in our response to KTvv about applying our method to different algorithms. We will incorporate this argument into Sec 3.1 and add a comment in the future work section that studying other algorithms is left for future work as it is beyond the scope of the paper.\n\n**Re. relevance to the Conference on Language Modelling:** Our motivation is that existing LLMs perform poorly on most planning tasks, and we believe that Searchformer is an important step to improve planning capabilities in the future. What we demonstrate is how to express planning (the computation involved in obtaining an optimal plan as well as the optimal plan itself) in a synthetic language. Then, we show how to train Language Models on this synthetic language and optimize the language itself by using the trained LM to reduce the number of search steps needed to arrive at an optimal plan. Therefore, we believe that our manuscript contributes to research on language models as it builds a connection between language modeling, planning, and sequential decision making. As a future work, it may be possible to apply our method to language-based search and planning traces, or perhaps to tasks where the input and output sequences are in natural language while the search trace is still in synthetic language. \n\nThis argument is also highlighted in your comment about reasons to accept.\n\nWe will revise the introduction and future work sections to highlight this important point.", "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": "I find this to be an interesting application of Transformers. Much of the success of Transformers in non-NLP domains has been in representing in the input in as much of a \"language-like\" fashion as possible vis-à-vis token sequences. This provides more evidence in that direction. If there is a \"language\" in pathfinding sequences generated by algorithms like A*, then it appears to be discernible by Transformers.", "reasons_to_reject": "I have some concerns regarding the lack of comparison to other pathfinding algorithms, even those that baseline A* itself surpasses, and questions about whether SearchFormer can make the same guarantees that baseline A* can, and how that can be demonstrated.  I also have some questions regarding relevance to the Conference on \"Language Modeling.\"", "questions_to_authors": "A* can guarantee a shortest path, and this property is presumably implicitly captured in the training data. Can SearchFormer similarly guarantee a shortest path, in fewer steps than A*?\n\nHow does SearchFormer compare to other pathfinding algorithms? Presumably if it outperforms A* it would also outperform Dijkstra's algorithm, but what about other algorithms like D* where constraints vary?\n\nOne question regards relevance to the Conference on Language Modeling: not all Transformers are language models; they just happen to have arisen from NLP and are adept at modeling sequences. Does this research tell us anything new about language or language models specifically?", "extra_scores": {}}, {"review_id": "kr5izOftK1", "reviewer": "Reviewer_pZre", "summary": "This paper presents Searchformer, a fine-tuned T5 model for planning. The key findings are that (1) training a transformer to predict not only the plan actions but also the A* search dynamics (called \"search-augmented sequence\"), offers better (data and model size) efficiency in training and results in more generalizable planners; and that (2) further fine-tuning the transformer on shorter plans that are non-deterministically synthesized yields faster planners.", "questions": "", "limitations": "", "rating": 9, "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": "We thank the reviewer for the positive review and insightful comments.\n\n**Re. more complex tasks:** In Fig 12, we already present a study where maze size and complexity is scaled, resulting in longer execution traces (Figs 7 and 8). As tasks become more complex, more training data is needed to reach a certain performance level. Apart from computational resource constraints, our results do not indicate that there is any upper limit on task complexity. Searchformer should still work in more complex environments. We will highlight this point in Sec 4.1.3. \n\n**Re. specific output representations:** The exact tokenization scheme does influence the model’s performance. We ran an experiment with the deterministic A* dataset for 20x20 mazes and changed the plan text representation as suggested in the comment (e.g. \"plan down down left\"). We refer to this as \"action token style\". \n\nWith the \"action token style\", the 45M search-augmented model's performance increases slightly (from 88.1%±1.8% to 90.8%±0.7%), while the solution-only model's performance drops significantly (from 67.7%±4.2% down to 48.2%±3.8%). This shows that Searchformer is insensitive to the exact tokenization scheme, due to incorporated execution traces that lead to more tightly correlated tokens sequence that are easier to learn, while search-only models may be very sensitive to it. Furthermore, 15M and 175M models show similar trends. We will put full results in the Appendix. \n\n**Re. decoder-only architectures:** We trained one 46M decoder-only search-augmented model on <prompt><trace><plan>-formatted sequences using our 20x20 maze dataset (50k training sequences generated with non-det. A*). This model solved 94.8% of all test tasks with an optimal plan, while the corresponding 45M encoder-decoder search-augmented model solved 91.8% of the test set. This result suggests that our approach also works with a decoder-only architecture. We will include this result in the discussion of our experiments.\n\n**Re. in-context learning:** While the focus of the submitted paper is on training LMs from scratch on synthetic data, we believe that understanding how our approach transfers to in-context learning in pretrained LLMs is an important direction of future research to improve the planning capabilities of LLMs. Our tentative results show SoTA LLMs fail to solve our tasks with in-context learning, even using our text representation. We will revise the future work section to incorporate this point.", "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": "The paper presents many interesting findings. The idea of training a transformer to predict the search dynamics is novel yet intuitive. It is also surprising to see that a transformer can precisely predict the dynamics consisting of thousands of tokens. The non-deterministic experiment is also interesting; it shows the potential for transformers to discover better plan solutions themselves.", "reasons_to_reject": "I don't see obvious reasons to reject the paper, but a few thoughts are: \n- Generalizing the idea to more complicated navigation tasks. For example, it can be helpful to analyze Searchformer on larger maze environments with obstacles following geometric shapes, which could potentially make the desired trajectories longer. In such complicated environments, does Searchformer still work?\n- How sensitive is Searchformer to the specific output representation in Figure 1? For example, instead of using (x, y) coordinates, the plan can also be described as \"down, down, right\", which is even shorter than the sequence of coordinate tokens. Did the authors see any impact from following different representations?\n- Does the idea of Searchformer apply to transformer architectures beyond T5 (e.g., decoder-only LMs)? And does the same idea apply to in-context learning of LMs?", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The manuscript presents a recipe for training Transformers to perform planning and shows that the resulting models can effectively perform planning for task requiring complex action sequences. Fine-tuned, the model can sometimes outperform A*, the oracle algorithm that was used to seed training.
The manuscript offers a new take on how to use Transformer models in service of planning algorithms. Traditional approaches such as MCTS and A* use handcrafted rules for building and refining search trees. Using network models to learn such rules is an interesting and underdeveloped direction of research. While search is initially learned from A* data, they develop a mechanism for bootstrapping search dynamics that creates better planners that take fewer steps to reach solution on Sokoban. The paper is written clearly with well executed analysis and figures.The manuscript could be even more compelling if comparisons to other network-based planning approaches.More difficult benchmarks would always be a plus, but the paper develops an interesting idea clearly and Sokoban is a difficult RL/control problem on its own even though it is situated in a gridworld environment.
On the whole, the simple and compelling results make the paper easy to read and will inspire more research into the underexplored direction of using large networks to perform planning algorithms.
Decision: Accept