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

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.

Rating

9

Confidence

4

Ethics flag

1

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: - 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? - 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? - 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?

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*. This 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. The paper is also written quite clearly and is easy to follow.

Rating

7

Confidence

4

Ethics flag

1

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*? How 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? One 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?

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.

Rating

6

Confidence

5

Ethics flag

1

Reasons to accept

- The paper is well-written. - 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. Section 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. Section 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.

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

Rating

6

Confidence

3

Ethics flag

1

Reasons to accept

Pros: 1. The paper presents a new method, search dynamics bootstrapping, for training Transformer models to solve complex planning tasks, which opens up new possibilities for applying deep learning to planning and reasoning problems. 2. The Searchformer model can outperform A* algorithm in efficiency and problem-solving ability, which highlights the potential of this approach to planning processes and tackling more complex problems. 3. This paper presents a well-designed set of experiments that systematically test their approach under different conditions. The use of synthetic datasets and the Sokoban domain demonstrate the approach's effectiveness, while also suggesting its potential for generalization to other planning problems.

Reasons to reject

Cons: 1. While the paper demonstrates improvements over the A* algorithm, it does not compare the Searchformer to other recent advances in neural network-based planning or symbolic models, making it difficult to assess its relative performance. 2. The paper does not provide theoretical guarantees on the optimality or completeness of the Searchformer's planning process. Further analysis is needed to understand the conditions under which the model's performance can be guaranteed. 3. The search dynamics bootstrapping approach relies on the A* algorithm to generate the initial training data. If the underlying algorithm is suboptimal or fails to find solutions for certain problems, it could limit the performance of the Searchformer model.

Questions to authors

1. Could this method be applied to more advanced symbolic planners or other types of heuristic search algorithms? 2. Is there a theoretical bound on the efficiency gains that can be achieved through this bootstrapping process?

Official Comment by Reviewer Loks

I appreciate the authors' responses; however, I do not feel as if the response alleviated all of my concerns. It remains difficult to assess the efficacy of the approach, relative to alternatives, since direct comparisons with strongly motivated LLM-based and/or symbolic planning approaches are missing. Without an analysis on optimality/completeness (requested by Reviewer KTvv) or on the relationship between some characterization of the task and the number of fine-tuning steps, it remains difficult to extract useful insights from the proposed approach that can guide an investigation in a different problem or domain. I will retain my current score.