# Transformers Can Navigate Mazes With Multi-Step Prediction

arXiv: https://arxiv.org/abs/2412.05117 · HTML: https://mlsanity.com/paper/2412.05117/

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 28 references that are in the map.
- First posted on arXiv 6 Dec 2024.

## Builds on (most cited first, 23 of 23)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [Diffusion-LM Improves Controllable Text Generation](https://mlsanity.com/paper/2205.14217/) (2022)
- [Gemma: Open Models Based on Gemini Research and Technology](https://mlsanity.com/paper/2403.08295/) (2024)
- [Simple and Effective Masked Diffusion Language Models](https://mlsanity.com/paper/2406.07524/) (2024)
- [Simplified and Generalized Masked Diffusion for Discrete Data](https://mlsanity.com/paper/2406.04329/) (2024)
- [The Reversal Curse: LLMs trained on "A is B" fail to learn "B is A"](https://mlsanity.com/paper/2309.12288/) (2023)
- [Better & Faster Large Language Models via Multi-token Prediction](https://mlsanity.com/paper/2404.19737/) (2024)
- [DiffusionBERT: Improving Generative Masked Language Models with Diffusion Models](https://mlsanity.com/paper/2211.15029/) (2022)
- [The pitfalls of next-token prediction](https://mlsanity.com/paper/2403.06963/) (2024)
- [Beyond A*: Better Planning with Transformers via Search Dynamics Bootstrapping](https://mlsanity.com/paper/2402.14083/) (2024)
- [A Configurable Library for Generating and Manipulating Maze Datasets](https://mlsanity.com/paper/2309.10498/) (2023)
- [The Factorization Curse: Which Tokens You Predict Underlie the Reversal Curse and More](https://mlsanity.com/paper/2406.05183/) (2024)
- [NavFormer: A Transformer Architecture for Robot Target-Driven Navigation in Unknown and Dynamic Environments](https://mlsanity.com/paper/2402.06838/) (2024)
- [ALPINE: Unveiling the Planning Capability of Autoregressive Learning in Language Models](https://mlsanity.com/paper/2405.09220/) (2024)
- [Structured World Representations in Maze-Solving Transformers](https://mlsanity.com/paper/2312.02566/) (2023)
- [DiSK: A Diffusion Model for Structured Knowledge](https://mlsanity.com/paper/2312.05253/) (2023)
- [Latent Plan Transformer for Trajectory Abstraction: Planning as Latent Space Inference](https://mlsanity.com/paper/2402.04647/) (2024)
- [Investigating Navigation Strategies in the Morris Water Maze through Deep Reinforcement Learning](https://mlsanity.com/paper/2306.01066/) (2023)
- [Scaling Value Iteration Networks to 5000 Layers for Extreme Long-Term Planning](https://mlsanity.com/paper/2406.08404/) (2024)

## Built on by (most cited first, 11 of 11)

- [Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking](https://mlsanity.com/paper/2505.24857/) (2025)
- [Roll the dice & look before you leap: Going beyond the creative limits of next-token prediction](https://mlsanity.com/paper/2504.15266/) (2025)
- [Empirical Analysis of Decoding Biases in Masked Diffusion Models](https://mlsanity.com/paper/2508.13021/) (2025)
- [Optimizing Decoding Paths in Masked Diffusion Models by Quantifying Uncertainty](https://mlsanity.com/paper/2512.21336/) (2025)
- [Looking beyond the next token](https://mlsanity.com/paper/2504.11336/) (2025)
- [Transformers in the Dark: Navigating Unknown Search Spaces via Bandit Feedback](https://mlsanity.com/paper/2603.24780/) (2026)
- [Do LLMs Build Spatial World Models? Evidence from Grid-World Maze Tasks](https://mlsanity.com/paper/2604.10690/) (2026)
- [Zipping the Thought: When and How Compressed Reasoning Data Works in LLM Post-Training](https://mlsanity.com/paper/2605.28008/) (2026)
- [Cognitive Maps in Language Models: A Mechanistic Analysis of Spatial Planning](https://mlsanity.com/paper/2511.13371/) (2025)
- [Lookahead Path Likelihood Optimization for Diffusion LLMs](https://mlsanity.com/paper/2602.03496/) (2026)
- [Constraints-Guided Diffusion Reasoner for Neuro-Symbolic Learning](https://mlsanity.com/paper/2508.16524/) (2025)

## Abstract (arXiv)

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 parameter-matched transformers from scratch, under identical settings, 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 that 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. The code can be found at https://github.com/facebookresearch/maze_navigation_MLMU