Transformers Can Navigate Mazes With Multi-Step Prediction

Authors: Niklas Nolte, Ouail Kitouni, Adina Williams, Michael Rabbat, Mark Ibrahim
Primary area: learning theory
Keywords: transformers, planning, learning objectives

Abstract:
Despite their remarkable success in language modeling, transformers trained to predict the next token in a sequence struggle with long-term planning. This limitation is particularly evident in tasks requiring foresight to plan multiple steps ahead such as maze navigation. The standard next single token prediction objective, however, offers no explicit mechanism to predict multiple steps ahead---or revisit the path taken so far. Consequently, in this work we study whether explicitly predicting multiple steps ahead (and backwards) can improve transformers' maze navigation. We train under identical settings, parameter-matched transformers from scratch to navigate mazes of varying types and sizes with standard next token prediction and MLM-U: an objective explicitly predicting multiple steps ahead and backwards. We find MLM-U considerably improves transformers’ ability to navigate mazes compared to standard next token prediction across maze types and complexities. We also find MLM-U training is 4x more sample efficient and converges 2x faster in terms of GPU training hours relative to next token training. Finally, for more complex mazes we find MLM-U benefits from scaling to larger transformers. Remarkably, we find transformers trained with MLM-U outperform larger transformers trained with next token prediction using additional supervision from A* search traces. We hope these findings underscore the promise of learning objectives to advance transformers' capacity for long-term planning.

Decision: reject
Reviewer scores: 3, 5, 3, 5

Reviews and discussion:
<thread>
creation date:
1730466773693

thread_type:
official_review

<thread_object>
creator: reviewer_vuwk

<content>
summary:
The paper "Transformers Can Navigate Mazes with Multi-Step Prediction" investigates whether an alternative training objective, termed MLM-U, can enhance transformers' maze navigation abilities compared to the standard next token prediction objective. The authors note that traditional transformers struggle with long-term planning tasks, such as navigating mazes, as they are typically trained to predict the next token based on the previous tokens. The MLM-U objective, inspired by masked language modeling, aims to explicitly predict multiple steps ahead and backward, encouraging better long-term planning.

The authors train transformers using both the standard next token prediction and MLM-U objectives to navigate various types of mazes, comparing performance across different maze complexities and transformer model sizes. They find that MLM-U significantly outperforms the next token objective in maze navigation, particularly for more complex tasks. Additionally, MLM-U demonstrates improved sample efficiency, reduced training time, and better generalization. For example, transformers trained with MLM-U achieve higher accuracy on mazes up to 30x30 in size and outperform larger transformers trained with next token prediction, even those with additional A* search trace supervision.

Key contributions of the paper include:
1. Demonstration of MLM-U's superior performance in long-term planning tasks, achieving better maze navigation accuracy across all tested maze complexities.
2. MLM-U is found to be more sample-efficient and computationally efficient than next token training.
3. Detailed exploration of the impact of model size, showing that scaling up transformers trained with MLM-U leads to significant gains in performance.

These findings suggest that modifying the learning objective to include explicit multi-step planning can significantly improve the capabilities of transformers in navigation tasks, highlighting a promising direction for further research on long-horizon planning with transformers.

strengths:
### Originality
The paper introduces **MLM-U**, a novel training objective that encourages explicit multi-step prediction both forward and backward, enhancing transformers’ long-term planning capabilities. This approach is original in combining masked prediction with maze navigation, showcasing transformers in a novel, decision-making context.

### Quality
The study is methodologically robust, providing a controlled comparison between MLM-U and traditional next token prediction. It employs diverse maze environments and measures both **accuracy** and **training efficiency**, ensuring the findings are well-supported. The visualizations and detailed experimental analysis add credibility to the work.

### Clarity
The paper is **clearly written**, with effective motivation for MLM-U and well-organized results. Visual aids enhance understanding. Slightly more explicit detail on the training process could further improve accessibility, but overall, the flow is logical and easy to follow.

### Significance
The **MLM-U objective** significantly improves the planning capabilities of transformers, demonstrating better sample efficiency and faster convergence. This work paves the way for **more efficient models** that require fewer computational resources, potentially impacting areas like robotics and game-playing. Its contributions to extending transformers into new sequential decision-making tasks are noteworthy.

weaknesses:
### Limited Novelty in Training Objective
MLM-U's novelty may be limited as similar multi-step prediction methods have been explored in reinforcement learning and trajectory modeling. Including a **comparative analysis** with existing methods would help strengthen the originality.

### Limited Experimental Complexity
The experiments focus only on maze navigation, which may not fully demonstrate MLM-U's broader applicability. Extending experiments to tasks requiring **complex long-term planning** (e.g., robotics or control tasks) would provide stronger evidence of general utility. Testing in **dynamic or uncertain environments** would further demonstrate robustness.

### Limited Baseline Comparisons
The paper mainly compares MLM-U against next token prediction, missing comparisons to stronger baselines like **reinforcement learning approaches** or **graph-based transformers**. Including such baselines would provide a more rigorous benchmark of MLM-U's performance.

### Scalability Concerns
The scalability analysis is limited. More discussion on **scaling challenges** and **computational trade-offs** when using larger models is needed. An ablation study on **hyperparameters** like masking ratios and model depth would help understand how to best tune MLM-U for different settings.

### Insufficient Analysis of Failure Cases
The paper lacks a detailed analysis of **failure cases**. A more in-depth examination of **common failure patterns** and scenarios where MLM-U struggles would provide valuable insights for improving the model and identifying its limitations.

questions:
### Limited Novelty in Training Objective
- Could you provide a more detailed comparison between MLM-U and existing multi-step prediction methods? Specifically, how does MLM-U address limitations in prior approaches?

### Limited Experimental Complexity
- Can you provide results or discuss potential outcomes for applying MLM-U to tasks beyond maze navigation?
- Do you foresee any challenges or necessary modifications when extending MLM-U to dynamic or uncertain environments?

### Limited Baseline Comparisons
- Have you considered comparing MLM-U to reinforcement learning-based methods? What would be the expected advantages or drawbacks?
- Could additional baselines be included in future work to provide a clearer benchmark?

### Scalability Concerns
- Could you expand on the potential challenges when scaling MLM-U to larger models or more complex environments?
- Would it be possible to include an ablation study on hyperparameters like masking ratios to better understand their impact?

### Insufficient Analysis of Failure Cases
- Could you provide more detailed examples of MLM-U's failure cases and potential reasons behind these failures?
- What changes could be made to the model or training process to address these failure modes?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730483269335

thread_type:
official_review

<thread_object>
creator: reviewer_zjxn

<content>
summary:
This paper explores the applicability of MLM-U (Masked Language Modeling + Uniformly Sampled Masking Ratios) to training transformers to perform simple maze navigation tasks. It finds that compared to next token prediction, MLM-U based transformers can more accurately navigate mazes, completing 93.8% of 30x30 mazes, while next-token prediction models complete only 18.8% (with training on 3K epochs of 100K mazes). Further, MLM-U based transformers are more data and training efficient, saturating navigation accuracy on 5x5 mazes 2x faster than next-token prediction (training efficiency), and with almost 4x fewer training mazes.

strengths:
- The paper tackles an interesting and challenging issue—transformers' limitations in long-term planning.
- The application of MLM-U to the task is well-motivated, and it seems logical and reasonable that using MLM-U-based approaches would improve the performance on maze navigation tasks.
- The paper clearly demonstrates that MLM-U is more efficient as a training procedure than next token prediction given the setup, and the gains reported in the paper are quite impressive.
- The scaling results are impressive, suggesting that this is a promising direction of research for larger models, and more complex tasks.
- As far as I am aware, no existing work has clearly demonstrated the effects of MLM-U, or MLM-based training approaches on maze prediction tasks.
- The paper is generally well-structured, and the visualizations (e.g., maze complexity graphs, accuracy charts) are helpful.

weaknesses:
While the paper does have some notable strengths, there are also several weaknesses:
- Some of the key setup and explanation details are missing from the paper, for example: What is the input representation of the mazes (does it match the Lehnert paper? Or are there modifications)? What is the task definition (how much information is given about the maze - is the full representation given, or only limited information, requiring the model to remember across the trajectory)? What is the definition of "navigation accuracy" (is this the same as A* solution accuracy in the Leonhart paper)? It seems like the paper overlooks some of these basics, without making reference to an identical setup, or clearly explaining in the text/appendix. \
- Given that the setup is identical to the Lehnert paper, I'm not sure that this paper actually contributes much to the conversation on long-horizon planning with transformers. While this is solid evidence that for a toy task such as A* navigation matching, adding MLM-based objectives can help improve the performance of the model, maze navigation itself does not really require foresight (especially in the full-information version of the task), and it's not clear that these results will generalize to other long-horizon planning problems. This paper would be significantly bolstered by adding additional domains (for example, In Lehnert, they solve Sokoban puzzles), or potentially other puzzle setups such as Sudoku, which might have slightly more complex mechanics and still require planning. As it stands, I think that primarily focusing on maze navigation is interesting, but perhaps too limited.
- While the motivation for the paper is that MLM-U allows the model to predict multiple steps both forward and backward, I'm not convinced that the setup actually does this. While the details are a bit thin, it seems like the the MLM approach is applied to full trajectories - which while it does encourage predicting multiple steps in the past, seems to be significantly harder to apply during the inference phase. Indeed, I'm not sure that this encourages "foresight" so much as "hindsight" or "justification", and maybe helps the model with memory instead of planning. 
- The paper has no baselines beyond next-token prediction. While the main motivation of the paper is to compare MLM-U with vanilla next token prediction, it would be good to include both the Bachmann & Nagarajan approach (predicting the entire path before gradient) and the Gloeckle approach (multiple future prediction heads) as baselines - given that both of these methods claim to be quite effective for this task, and represent comparable methods. 
- There is no statistical analysis of the data, and it seems like many of the results are drawn from a single training run on each of the methods. It would be good to see the variance of the data, particularly in Figures 3 and 4, to understand exactly how variable the convergence is (it seems like Figure 4 in particular could be susceptible to fairly high variance).

questions:
- How is navigation accuracy defined? 
- What are the setup details (input representations, task definition, etc.)? 
- What precisely is the training algorithm for the approach? Is MLM applied to full trajectories? Or is MLM applied at each next-token prediction step? What are the sets of uniform masks used for MLM-U
- During decoding, how is the length of the initialized X determined? It seems like you need to already know the length of the final trajectory T in order to predict the masked tokens. 
- How many of the failure cases are equivalent shortest paths that are not equivalent to the deterministic approach (as discussed in Appendix. C)? Are these a significant portion of the results? 
- Are there any explicit failure modes recognized in the MLM-U examples (particularly for larger mazes) - or does MLM-U struggle with any categories of maze structure in particular (i.e. those with higher branching factors)?
- Are there any limits on data efficiency gains with MLM-U? How well does it perform on even larger mazes (perhaps 50x50 mazes)? 
- Are there data effects from using a smaller number of mazes (i.e. 100K) over a larger number of epochs (3000) rather than a larger number of mazes/more diverse set of mazes in general? How does this compare if you use more dynamic sets of mazes such as in the Lehnert paper (which uses up to 1M unique mazes)? It would be good to replicate Fig 11/12 from that paper in order to get a better idea of the training data scaling. 

Some minor comments:
- There's a couple of un-supported claims in the paper that would benefit from citation (L35: "transformers encounter challenges when tasked with planning and decision-making over extended horizon", L37: "particularly evident in tasks requiring foresight")
- Saying that masking rates are "uniform" is somewhat under-specified, and upon reading the appendix, it seems like "uniform" means that the masking rate is _drawn from a uniform distribution_. It would be good to clarify this use of terminology, and clearly specify the exact distribution used for the masking rate in MLM-U.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730631337503

thread_type:
official_review

<thread_object>
creator: reviewer_f3iu

<content>
summary:
In this work, the authors investigated whether Transformers can perform well in long-term planning tasks, and conduct experiments on the maze navigation task. As the conventional next-token prediction (NTP) objective is notorious about its shortcomings with navigation and planning, the authors propose to leverage MLM-U, an objective that encourages the model to learn with the entire traces in a non-autoregressive manner to better grasp the ability of multi-step planning. In fair experimental settings, the authors have demonstrated that the MLM-U objective outperforms NTP in terms of both performance and efficiency.

strengths:
- A timely work in the field of the planning ability of Transformers. The study with the MLM-U objective is nice to read.
- The superiority of MLM-U over NTP with additional supervision from A* search traces, as well as the efficiency in the use of training data, demonstrates the effectiveness of MLM-U.

weaknesses:
- A general weakness point is that this work is mostly about the application of the MLM-U objective, proposed in prior work, to the maze navigation domain. This limits the novelty of the work, and further demands the empirical experiments to be more sound and sufficient.

For other points, please refer to the Questions listed below.

questions:
- Are all of the experiments covered in this work equipped with the setting that the test data falls in distribution with the training data in terms of the size of the maze? In other words, Did the experiments on 30x30 mazes leverage the traces of 30x30 mazes for training, and further test on 30x30 mazes as well? It would be great to study the generalization of MLM-U in comparison with NTP in terms of the task difficulty. More specifically, two lines of experiments could be added:
  1. Train on traces of N_1 x N_1 mazes, then test on mazes of N_1 x N_1, N_2 x N_2, N_3 x N_3 ... (we can set N_1 < N_2 < N_3, or N_1 > N_2 > N_3; different meanings are implied).
  2. Train on traces of both N_1 x N_1 and N_2 x N_2 mazes, then test on mazes of N_1 x N_1 and N_2 x N_2. This studies whether MLM-U could exploit positive transfer or handle negative interference within data; And compare with NTP to get the understanding of which one is better.

- In Section 5.4, why would the precision of the positional encoding especially matter in this MLM-U training paradigm? Is it related the difficulty of training with the MLM-U objective? More analysis on the failed training dynamics with half-precision positional encodings will further contribute to the soundness of the paper.

- In this work, the Transformer models are trained from scratch using NTP or MLM-U objectives. However, I believe the contributions of the work would be much greater if the authors conduct experiments with LLMs in the finetuning stage: Instead of leveraging the NTP objective in SFT, try using MLM-U and discover whether the most capable open-sourced LLMs nowadays like LLaMa 3.1 could be empowered with the multi-step prediction ability. Then an interesting further question would be whether using MLM-U instead of NTP in the SFT stage would be harmful to other SFT data that is to be trained with the NTP objective. I encourage the authors to conduct additional experiments to respond with at least the first question (the utility of MLM-U in the SFT stage for OSS LLMs like LLaMa) to further improve this work.

I would consider raising my scores if the authors address my concerns properly.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730693562304

thread_type:
official_review

<thread_object>
creator: reviewer_c6d6

<content>
summary:
This paper proposes MLM-U, a training objective for improving multi-step prediction for transformers. Specifically, they propose to use MLM-U, an objective which predicts a span based on surrounding information, similar to the early MLM objective in fine-tuning BERT. Differently, they only use masks and don't use token substitution. Though the approach is very simple, they show that fine-tuning the transformer on two synthetic datasets (created with DFS and A*) with MLM-U objective achieves significantly better performance compared with the next-token-prediction loss usually used in decoder-only models. Besides, they find that fine-tuning with MLM-U tends to converge much faster and demonstrates much better generalization capability to unseen mazes. Their findings suggest that MLM-U objective can be effectively used to improve transformer based agents' capability in multi-step prediction.

strengths:
1. The proposed objective MLM-U, though simple, is very effective in improving the performance of multi-step planning for maze navigation. 
2. This paper provides a detailed analysis of the convergence speed, and generalization capability of the fine-tuned model, both suggesting that the transformer fine-tuned with MLM-U tends to achieve much better performance. 
3. This paper is well-written and easy to follow.

weaknesses:
1. Some implementation details are not clear to me. For example, when using different objectives, is it the case that MLM-U is applied to encoder-decoder architecture while next-token prediction is used for decoder-only architecture? Does the base model architecture influence the final performance (e.g., next token prediction on encoder-decoder architecture)?
2. The paper emphasizes the importance of RoPE positional encoding. However, the experiments related seem to be about training hyperparameters (i.e., fp16 vs. fp32). I feel it's more reasonable to compare with other positional encodings if the author wants to claim the importance of RoPE positional encoding, otherwise it might be better to call it training hyperparameter ablation. 
3. The novelty of the MLM-U objective might be small, as it has been widely explored in many encoder-decoder architecture works.

questions:
Please refer to weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732225444114

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
Thanks to all reviewers for the detailed feedback and suggestions.  

Based on reviewers’ feedback we’ve revised the manuscript to address several of the reviewers' concerns with changes highlighted in blue.  

**We would like to state explicitly the goal of this paper to set expectations:**  
The goal is to **demonstrate superior planning capabilities** of transformers when trained with a different objective. We believe this is a timely, valuable contribution to encourage the research community addressing one of the main limitations of today’s best transformer models (as noted by reviewers ZjXN and C6d6): multi-step planning. We show via controlled experiments with identically sized transformers trained on the same data, by only modifying the learning objective, transformers can learn to efficiently navigate mazes.

We choose maze navigation because it is a synthetic task that requires non-trivial planning in the context and has a simple knob to scale complexity via the maze size. **A model has to reason about the context, perform planning of the path by repeatedly referring to this context and then decode/execute** the shortest path. While it is true that this is a full-information environment, the requirements needed to solve the task are otherwise very general. In that respect, more complex problems like Sudoku and Sokoban are interesting and future work, but we believe that the presented evidence is worthy of publication, given that it shows such strong hints for shortcomings of next token prediction.


We appreciate reviewers found our work “timely” (f3iu) in   
**Tackling a key limitation of transformers**  
-  “an interesting and challenging issue—transformers' limitations in long-term planning.”—ZjxN

**Prosposing a novel application of a well-motivated method:**  
- “As far as I am aware, no existing work has clearly demonstrated the effects of MLM-U, or MLM-based training approaches on maze prediction tasks.”—ZjxN
- “The proposed objective MLM-U, though simple, is very effective in improving the performance of multi-step planning for maze navigation.”—C6d6

**Providing clear experimental evidence**  
- “The paper clearly demonstrates that MLM-U is more efficient as a training procedure than next token prediction given the setup, and the gains reported in the paper are quite impressive.”—ZjxN
- “This paper provides a detailed analysis of the convergence speed, and generalization capability of the fine-tuned model, both suggesting that the transformer fine-tuned with MLM-U tends to achieve much better performance.”—C6d6
- “Demonstration of MLM-U's superior performance in long-term planning tasks, achieving better maze navigation accuracy across all tested maze complexities.”—vUWK

**Exploring the benefits of scaling this approach**  
- “The scaling results are impressive, suggesting that this is a promising direction of research for larger models, and more complex tasks.”—ZjxN
- “Detailed exploration of the impact of model size, showing that scaling up transformers trained with MLM-U leads to significant gains in performance.”—vUWK
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper addresses a limitation of transformers trained for next-token prediction (NTP) in planning tasks that involve predicting multiple steps ahead. As a use case, it considers maze navigation. The authors propose an adaptation of MLM-U (Masked Language Modeling with a uniform sampling ratio) for this navigation task. Through experiments on various maze settings, they demonstrate that their approach outperforms the next-token prediction criterion, converges faster, and generalizes better to unseen mazes.

The reviewers appreciate the paper for its clarity and well-motivated approach. They consider it a timely contribution that implements a simple but effective idea, leading to significant improvements compared to NTP training.
However, the reviewers highlight several weaknesses, including limited novelty, restricted applicability that does not extend to other planning problems, and insufficient comparison with baselines. The authors are encouraged to further extend the application of their approach to a broader class of problems.