Beyond A*: Better Planning with Transformers via Search Dynamics Bootstrapping

1 Introduction

• Transformers and LLMs still struggle with multi-step planning and higher-order reasoning, while effective approaches may not transfer across reasoning types.Source: §1 · p. 1

Our work

• Search dynamics bootstrapping fine-tunes the search-augmented model to generate shorter token sequences while retaining an optimal plan.Source: Our work · p. 2
• The method represents A* computation as an execution-trace token sequence and trains a Transformer to generate search dynamics together with an optimal plan.Source: Our work · p. 2
• 93.7% of Sokoban test tasks are solved, with 26.8% fewer search steps on average than the A* reference implementation.Source: Our work · p. 2
• Search-augmented models improve performance on independent test tasks despite generating sequences 10× to 100× longer than solution-only sequences.Source: Our work · p. 2

• Unlike methods that learn policies from synthetic data or treat symbolic planners as black boxes, search dynamics bootstrapping trains a Transformer on the planner’s internal search dynamics [3, 7].Source: §2 · p. 2
• Where prior work predicts reasoning traces or optimal actions, this approach trains a Transformer to generate the entire A* search process for an optimal multi-step plan [1] [2, 7].Source: §2 · p. 2
• Unlike neuro-symbolic systems with dedicated memory or recursion, Searchformer relies on next-token prediction, long contexts, and position embeddings [1, 2].Source: §2 · p. 2
• Prior RL approaches model trial-and-error trajectories and predict actions, states, or rewards, while MCTSNet hard-codes search and is limited to 500-step rollouts [1, 2].Source: §2 · pp. 2–3
• Figure 1 illustrates how a maze task and its optimal plan are tokenized, and how A* search operations are recorded in an execution trace [Russell & Norvig (2021, Chapter 3)].Source: Figure 1 · p. 3

3 Problem Setup

• Sokoban is more computationally difficult than maze navigation because each state combines box and worker positions, and incorrect moves can lead to dead ends.Source: §3 · p. 3

3.1 Generating execution traces of A∗search.

• A* traces record frontier and closed-set insertions, node coordinates, heuristic and start-cost values, and the resulting plan.Source: §3.1 · p. 4
• Search-augmented sequences encode a task prompt, A* execution trace, and optimal plan; solution-only sequences contain only the prompt and plan.Source: §3.1 · p. 4

3.2 Training a Transformer model

• The training and test sets contain no duplicate tasks, and Transformers learn to predict the entire search process, including attempted but failed paths.Source: §3.2 · p. 4

3.3 Moving past algorithm imitation via search dynamics bootstrapping

• Randomized A* tie-breaking and child expansion order provide varied traces; the trained model samples candidate sequences, retaining shorter ones only when they end in an optimal plan.Source: §3.3 · pp. 4–5
• The model is fine-tuned on shorter sampled sequences, and repeated bootstrapping aims to reduce inference steps while improving performance beyond A* imitation.Source: §3.3 · p. 5

4 Experiments

• Experiments used deterministic and non-deterministic A* datasets, with matched train-test difficulty and five repetitions unless otherwise indicated.Deterministic data has unique traces and plans; non-deterministic data randomizes tie-breaking and allows multiple correct sequences.Source: §4 · p. 5, §4.1 · p. 5, Figure 2 · p. 6

Deterministic A∗

• In deterministic maze evaluation, search-augmented models outperformed most solution-only models, especially with 100,000 or fewer training sequences.The solution-only model matched only the weakest search-augmented model with sufficiently large training datasets; search-augmented models handled traces thousands of tokens long without an incorrect token on more than 90% of test mazes.Source: Deterministic A∗ · p. 6

Non-deterministic A∗

• With non-deterministic A* data, even the smallest search-augmented models found optimal plans on more test tasks than solution-only models, especially with small training sets.Each task was evaluated over 64 generated sequences, and counted as correct if any sequence contained an optimal plan; model size affected search-augmented performance only with 50,000 training sequences.Source: Non-deterministic A∗ · p. 6

Performance across different task difficulty levels

• As maze size and state space grew, solution-only performance dropped rapidly, while search-augmented models maintained comparatively high accuracy, including the smallest model.The comparison counts correctly solved test tasks using non-deterministic A* data.Source: Performance across different task… · pp. 6–7

4.2 Solving Sokoban puzzles

• On Sokoban, search-augmented models outperformed solution-only models; even a 747M-parameter solution-only model was slightly outperformed by a 175M-parameter search-augmented model.The comparison uses a different tokenization pattern and transition structure from maze navigation.Source: §4.2 · p. 7

4.3 Searchformer: Improving search dynamics via bootstrapping

• Bootstrapping repeatedly sampled 32 sequences per Sokoban puzzle, retained the shortest sequence containing an optimal plan, and fine-tuned for another 10,000 steps.The process generated a new training set and repeated fine-tuning to shorten traces while retaining an optimal solution.Source: §4.3 · p. 7
• 26.8% shorter than A* search dynamics, final Searchformer traces average fewer search steps when producing optimal Sokoban plans.After three fine-tuning steps, A* search dynamics are reported as approximately 34.3% longer than Searchformer sequences.Source: §4.3 · pp. 7–8, Figure 3 · p. 8
• Fine-tuning reduced incorrect solutions by 40% and non-optimal solutions by 30%, while the small Searchformer matched the largest solution-only model on SWC.SWC reflects whether test tasks are solved and how close predicted plans are to optimal length.Source: Table 1 · p. 7, §4.3 · p. 8

5 Discussion

• Searchformer robustly follows symbolic-planner execution traces, shortens them beyond the trained rule-based strategy, and scales better than solution-only models using fewer training sequences.The authors present the work as evidence that Transformers can solve complex decision-making tasks with appropriate training data.Source: §5 · p. 8

5.1 Limitations

• A* trace lengths may grow exponentially with optimal-plan length, making the resulting token-sequence training data computationally expensive.The experiments use sequences significantly longer than those used to train Llama 2.Source: §5.1 · p. 9

5.2 Future work

• Future work could use curriculum learning, alternative planners or heuristics, and hierarchical planning to reduce trace lengths or search depth.The authors suggest starting with simpler tasks and adapting the improved model to more complex tasks.Source: §5.2 · p. 9
• The authors propose combining their approach with Plansformer to fine-tune LLMs for more robust complex planning.They also aim to inform further research on Transformers’ multi-step planning and LLM reasoning capabilities.Source: §5.2 · p. 9

6 Broader Impact

• The work is a proof of concept for using Transformers to robustly solve complex symbolic planning tasks, trained on synthetic datasets.The authors note that the tasks can also be solved with simple symbolic solvers.Source: §6 · p. 9

A Training encoder-decoder Transformers to predict optimal plans

• The encoder processes each task prompt into an embedding, while the decoder generates either a <trace><plan> sequence or a <plan> sequence.The former is the search-augmented format; the latter is used by solution-only models.Source: A Training encoder-decoder… · p. 13

A.1 Encoder-decoder architecture

• A prompt is represented as a sequence of integer token indices, x1:n.The encoder maps this sequence to a sequence of real-valued vectors.Source: §A.1 · p. 13
• The decoder generates responses autoregressively, predicting next-token probability vectors conditioned on the encoder output and preceding response tokens.Generation begins with a beginning-of-sequence token and continues until an end-of-sequence token is sampled.Source: §A.1 · p. 13
• The encoder and decoder each use stacked causal attention layers, with their parameters concatenated into θ.Source: §A.1 · p. 13

• Training minimizes cross-entropy between the training distribution over prompt-response pairs and the model’s response distribution.Source: §A.2 · p. 13
• Samples are reweighted to emphasize shorter response sequences in the loss objective.Source: §A.2 · p. 13
• The empirical objective averages over the D training pairs and is optimized using gradient descent.The average approximates the expectation for a finite i.i.d. sample.Source: §A.2 · p. 13

• Tokens are mapped to one-hot vectors and projected through a linear embedding layer into vector representations.Source: B Network architecture and… · p. 14
• Encoder and decoder layers contain multiple attention blocks; the decoder processes encoder output and uses RoPE position embeddings.The architecture uses no dropout.Source: B Network architecture and… · p. 14
• The same architecture hyperparameters are used for encoder and decoder, while optimization uses AdamW with linear learning-rate increase through step 2000 followed by a cosine schedule.The listed table passages describe the hyperparameter categories and optimization schedule but do not provide all table values.Source: Table 2 · p. 15, Table 3 · p. 15

C Dataset generation

• Sokoban tasks used 7 × 7 grids with two added interior walls and randomly placed docks, boxes, and worker locations; only tasks solvable by A* were retained.Source: C Dataset generation · p. 15
• Maze tasks used 30–50% wall cells and were retained when their optimal path was at least the maze width or height.Source: C Dataset generation · p. 15
• Duplicate tasks were rejected; Sokoban sequences approaching 100,000 tokens were capped at 10,000 for GPU memory, and test sequence lengths were comparable to training lengths.Source: C Dataset generation · p. 15
• Sokoban used a non-deterministic A* implementation, and Figure 6 presents the token sequence for the level shown in Figure 5.Source: Figure 5 · p. 15, Figure 6 · p. 16
• Training sequences were multiple orders of magnitude longer than usual LLM sequences, especially when A* traces were included; training and test sequence lengths were designed to match closely.Source: Figure 7 · p. 16, Figure 8 · p. 16

D.1 Measuring plan optimality

• Exact-match evaluation requires matching deterministic A* output; Any-optimal-64 marks a task correct if any of 64 responses contains a feasible optimal plan.Source: §D.1 · p. 17
• SWC measures the best predicted correct-plan cost relative to optimal cost across test tasks, weighting each task by whether a correct plan was found; it equals one when all generated sequences end optimally.Source: §D.1 · p. 17

D.2 Measuring search dynamics length

• Average-on-optimal length averages dynamics lengths for responses ending optimally; ILR compares the shortest generated dynamics with A* and includes tasks with an optimal or correct plan.Source: §D.2 · p. 17
• ILR values above one indicate that the model generates shorter search dynamics than the A* reference.Source: §D.2 · p. 17

E Searchformer performance analysis

• Each model generated 64 sequences per Sokoban test task, which could end in an optimal, incorrect, or correct but sub-optimal plan.Source: E Searchformer performance analysis · p. 17
• At each improvement step, generated sequences move below the A* length diagonal, and Searchformer searches in fewer steps to compute an optimal plan.Source: E Searchformer performance analysis · p. 17
• Correct but sub-optimal plans occurred with less than 5% chance, and bootstrapping brought their average sequence lengths closer to A* lengths.Source: E Searchformer performance analysis · p. 18

F Supplemental figures for experiments

• Solution-only model performance is primarily influenced by training-sequence count, while increasing model size does not always improve performance.Source: Figure 9 · p. 18
• Figure 10 plots A* average token-sequence length on the x-axis against model average length on the y-axis, with separate panels for optimal, correct, and correct sub-optimal plans.Source: Figure 10 · p. 19
• Figure 11 reports the percentage of correctly answered prompts for solution-only models, averaged across five seeds.Source: Figure 11 · p. 20
• Search-augmented models significantly boost optimal-plan prediction performance; the figure separately plots correctly generated responses and prompts with any sampled optimal plan.Source: Figure 12 · p. 21
• Figure 13 compares average-on-optimal sequence lengths for A* and search-augmented and Searchformer models across tasks where models yield optimal plans, using 64 trials per task.Source: Figure 13 · p. 21