REVIEW 4 major objections 4 minor 1 cited by
Transformers Can Navigate Mazes With Multi-Step Prediction
T0 review · 4 major / 4 minor · reviewed 2026-08-11 · deepseek-v4-flash
Pith's one-line read The paper claims that switching from next-token prediction to an objective that predicts multiple path steps ahead and backward lets small transformers navigate mazes that next-token models of the same size cannot, and even outperform…
desk verdict Strong empirical result on MLM-U for maze navigation, but the architecture confound and loose efficiency numbers keep the headline from being fully clean. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The load-bearing object is MLM-U, a masked-language-modeling objective with a uniform masking rate. For each training example it draws a mask proportion $\mu \sim U(0,1)$, masks a random subset of the solution-path tokens, and minimizes the cross-entropy only on those masked tokens, so the model must reconstruct multiple future steps and past steps from whatever context remains. Because the mask is random, the model sees many prediction horizons and contexts rather than one fixed left-to-right factorization. The implementation pairs this objective with an encoder-decoder transformer using rotary position embeddings, which encode relative positions into attention; inference is still done autoregressively left-to-right with argmax. The paper's argument is that this variable-horizon objective, not the architecture alone, is what allows the model to plan.
What would settle it
Train an 8M-parameter next-token transformer on 20x20 DFS mazes using the exact encoder-decoder architecture, RoPE encodings, optimizer, batch size, and compute budget that the MLM-U runs use. If that next-token model also reaches near-perfect navigation accuracy, the gap is explained by architecture or tuning rather than by multi-step prediction.
Extended reading notes
Core claim
The paper's central claim is that explicitly training a transformer to predict many tokens of the solution path, both ahead and backward from random masked contexts, turns it into a reliable maze navigator, while next-token prediction leaves it short-sighted. In DFS mazes, an 8M-parameter transformer trained with MLM-U reaches 100% navigation accuracy on held-out mazes through 20x20 grid sizes, whereas the matched next-token baseline peaks at 20.6% on 20x20 and 18.8% on 30x30. In A* mazes, 8M MLM-U achieves 98.5%, 95.2%, and 85.5% accuracy on 10x10, 20x20, and 30x30 mazes, respectively, beating a 175M next-token transformer trained with extra A* search-trace supervision, which reaches 94.9%, 53.5%, and 19.3% without traces and 98.5%, 90.4%, and 70.2% with traces on those sizes. The paper also reports that MLM-U needs about 4x fewer training samples on small mazes and converges around 2x faster in GPU hours, and that scaling model size helps MLM-U most on harder mazes.
Load-bearing premise
The results assume the accuracy gap is caused by the training objective, but MLM-U and next-token runs are not fully matched: they use different architectures, and the A* baselines come from an earlier study with different compute and hyperparameters.
Editorial extensions
If this is right
- An 8M-parameter transformer trained with MLM-U achieves 100 percent navigation accuracy on held-out DFS mazes up to 20x20, while next-token training peaks at 20.6 percent on 20x20 mazes.
- On 30x30 A* mazes, 8M MLM-U reaches 85.5 percent accuracy, outperforming a 175M next-token model trained with additional A* trace supervision, which reaches 70.2 percent.
- MLM-U is roughly 4x more data-efficient on 5x5 and 10x10 mazes and about 2x faster in GPU hours to convergence on 5x5 mazes.
- Scaling MLM-U from 3M to 8M parameters raises 20x20 DFS accuracy from 85 percent to 100 percent, while next-token scaling improves more slowly.
- Next-token training overfits on held-out 10x10 mazes, while MLM-U generalizes without overfitting; the paper attributes this to the harder, variable-context objective.
Reading between the lines
- If the active ingredient is variable-horizon prediction rather than maze-specific structure, the same objective should transfer to other long-horizon sequence tasks such as code, proof, or instruction following; the paper only studies mazes.
- The uniform random masking rule is almost certainly not optimal; curriculum or confidence-based masking that focuses training on hard-but-solvable subsets could improve convergence further, and the paper notes this as future work.
- The positional-precision result suggests a testable rule for masked training: rotary position embedding precision may need to grow with sequence or grid length, so lower-precision mixed-precision training could silently cap planning performance.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper studies whether MLM-U, an objective that masks arbitrary subsets of a sequence to predict multiple tokens ahead and backward, improves transformer maze navigation relative to standard next-token prediction. The authors train parameter-matched transformers from scratch on DFS-generated and A*-generated mazes of varying sizes, and report large navigation-accuracy gains for MLM-U (e.g., 100% vs. 20.6% held-out path accuracy on 20x20 DFS mazes), as well as data-efficiency and GPU-efficiency gains. They also report that an 8M-parameter MLM-U model outperforms a 175M-parameter next-token model trained with additional A* trace supervision on 30x30 mazes.
Significance. If the claims hold, the results are significant: they suggest that the learning objective can dramatically alter a transformer's planning behavior, with practical implications for training and inference efficiency. The paper includes several useful controls: a hyperparameter sweep for the next-token baseline, two maze families with different representations, per-token accuracy measures, scaling curves, and released code. However, the central causal attribution — that the improvement is due specifically to multi-step prediction — is currently weakened by an architectural confound, and the efficiency numbers contain internal inconsistencies. With matched-architecture baselines and corrected efficiency analyses, the findings would be a solid contribution to the growing literature on training objectives for planning.
major comments (4)
- [Section 4.4 / Appendix A.2] The paper claims that MLM-U and next-token baselines are trained under 'identical settings' and that the experiments isolate 'modifying only the learning objective' (Sections 1 and 4), but this is not fully supported. Next-token DFS baselines use a decoder-only transformer, while MLM-U uses an encoder-decoder with RoPE; Appendix A.2 itself shows that architecture choice affects next-token performance. For A* mazes, the next-token numbers are imported from Lehnert et al. (2024) without matching training compute, hyperparameters, or tokenizer. Thus the observed accuracy gap could partly reflect the architectural change (e.g., bidirectional encoding and cross-attention) rather than multi-step prediction alone. Please add a matched-architecture next-token arm on at least the DFS setting (e.g., train the same encoder-decoder with a next-token loss), or explicitly temper the causal wording to describe the combined objective-architecture choice. Without this, the headline claim overstates what is controlled.
- [Section 5.1, Figure 3] The '4x more data efficient' claim is only well-defined for 5x5 mazes, where MLM-U reaches 100% accuracy with 25k training samples while next-token requires 100k. For 10x10 mazes, the text states that next-token 'reaches a peak performance of less than 50% with 75k training samples' and never reaches full accuracy, so a 4x sample-efficiency ratio cannot be computed. The claim 'suggesting MLM-U is 4x more data efficient' is therefore imprecise. Please state explicitly that the 4x ratio is for 5x5 mazes, and describe the 10x10 result separately as a ceiling difference rather than a data-efficiency ratio.
- [Section 5.1 and Abstract] The GPU-hours efficiency claim is internally inconsistent as written. The text reports that training for 2k epochs on 8 Tesla V100 GPUs takes 13.7 hours for next-token and 17.7 hours for MLM-U, and it calls the difference a '7% overhead' (the actual overhead is 29%). Combined with the reported 2.17x epoch-based speedup, the implied GPU-hour speedup is about 1.7x (13.7 / (17.7 / 2.17)), not '2x faster in terms of GPU training hours' as stated in the abstract. Please clarify how the GPU-hour speedup was computed, correct the overhead figure, and align the abstract, Section 5.1, and Figure 4.
- [Section 5.4] The reported numbers for the positional-encoding precision experiments are inconsistent. The text says that with 16-bit positional encodings MLM-U reached 50% navigation accuracy on 30x30 mazes and that 32-bit encodings 'solved 30x30 mazes perfectly,' then later states that 'by increasing the precision from 16 to 32-bits... MLM-U performance on 30x30 DFS mazes improved from 40% to 93.8%.' Please disambiguate which numbers correspond to the 100-maze training set and which to the full 100k-maze training set, and ensure the figure and text agree.
minor comments (4)
- [Section 4.1] The text says 'we generate 500k mazes across five levels of complexity' but then lists only four grid sizes (5x5, 10x10, 20x20, 30x30); Table 1 includes 15x15, so either add 15x15 to the list or change the count.
- [Section 5.1] The text says MLM-U achieves 'nearly 3x the performance' of next-token on 30x30 mazes, but Table 1 gives 93.8% vs. 18.8%, which is roughly a 5x improvement. Please correct the multiplier.
- [Abstract and Section 1] The abstract says 'converges 2x faster in terms of GPU training hours' but Section 5.1's numbers imply a smaller speedup; after correcting the computation, ensure the abstract and main text use consistent figures.
- [References] The reference list contains some incomplete entries (e.g., 'Gemma' without author names) and the in-text citation 'Ivanitskiy et al. (2023a)' appears as 'Ivanitskiy et al. (2023b)' in Section 2; please standardize the bibliography.
Circularity Check
No circularity: MLM-U's advantage over next-token prediction is measured on held-out mazes, not derived from the objective's definition; the only self-citations are to the source of the training objective and are not load-bearing.
full rationale
The paper is an empirical comparison, not a derivation. Its central claim—that MLM-U improves maze navigation over next-token prediction—is supported by measured held-out accuracies (Tables 1–2, Figures 3–6). Equation (2) defines the MLM-U loss, but the reported performance is not an algebraic consequence of that loss; it is an experimental outcome. No fitted parameter is renamed as a prediction: the only tuning is a learning-rate/weight-decay sweep on 10x10 DFS mazes (Section 4.4, Appendix B), and the headline results cover multiple maze sizes and two maze encodings, including mazes not used for model selection. The paper's self-citations are (a) adopting MLM-U from Kitouni et al. (2024a), an overlapping-author paper, and (b) citing that paper's Figure 2 for the claim that uniform masking avoids a train/inference distributional shift (Sections 3.2 and 4.3). These are methodological premises from prior work, not an imported uniqueness theorem, and the central comparison also uses independent baselines: an in-house next-token model trained from scratch and external A* baselines from Lehnert et al. (2024). The nearest thing to a concern is an experimental confound, not circularity: MLM-U uses an encoder-decoder with RoPE while the DFS next-token baseline is decoder-only (Section 4.4, Appendix A.2), and the A* baselines are not matched in compute or hyperparameters; this weakens causal attribution of the gap to the objective, but the conclusion is not equivalent to its inputs by construction.
Assumptions & free parameters
assumptions (4)
- domain assumption MLM-U's uniform masking rate over [0,1] ensures no distributional shift between training and autoregressive inference, as claimed in Kitouni et al. (2024a).
- domain assumption The encoder-decoder transformer with RoPE and tied embedding head, taken from Kitouni et al. (2024a), is an appropriate architecture for MLM-U on maze navigation.
- domain assumption The exact-match path accuracy metric, inherited from Lehnert et al. (2024), is a valid measure of navigation performance, treating any deviation from the deterministic shortest path as an error.
- domain assumption The 100k-maze deterministic A* dataset used for MLM-U is the same as the dataset used for the published Lehnert et al. baselines.
Cite this review
Pith. "Pith review of Transformers Can Navigate Mazes With Multi-Step Prediction." pith.science (2026). https://pith.science/paper/XERAXLQL
@misc{pith2026241205117,
author = {Pith},
title = {Pith review of: Transformers Can Navigate Mazes With Multi-Step Prediction},
year = {2026},
howpublished = {\url{https://pith.science/paper/XERAXLQL}},
note = {Machine review of arXiv:2412.05117}
}
read the original 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 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
Forward citations
Cited by 1 Pith paper
-
Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking
EB-Sampler dynamically unmasks multiple low-entropy tokens per function evaluation, accelerating masked diffusion model sampling by 2-3x with negligible accuracy loss.
Reference graph
Works this paper leans on
-
[5]
Gemma: Open models based on gemini research and technology.arXiv preprint arXiv: 2403.08295 ,
Gemma. Gemma: Open models based on gemini research and technology.arXiv preprint arXiv: 2403.08295 ,
-
[6]
Better & faster large language models via multi-token prediction.arXiv preprint arXiv:2404.19737 ,
Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, and Gabriel Synnaeve. Better & faster large language models via multi-token prediction.arXiv preprint arXiv:2404.19737 ,
-
[7]
doi: 10.48550/arXiv.2211.15029. Michael Igorevich Ivanitskiy, Rusheb Shah, Alex F Spies, Tilman Räuker, Dan Valentine, Can Rager, Lucia Quirke, Chris Mathwin, Guillaume Corlouer, Cecilia Diniz Behn, et al. A configurable library for generating and manipulating maze datasets. arXiv preprint arXiv:2309.10498 , 2023a. Michael Igorevich Ivanitskiy, Alex F. Sp...
-
[8]
The factorization curse: Which tokens you predict underlie the reversal curse and more, 2024a
Ouail Kitouni, Niklas Nolte, Diane Bouchacourt, Adina Williams, Mike Rabbat, and Mark Ibrahim. The factorization curse: Which tokens you predict underlie the reversal curse and more, 2024a. URLhttps://arxiv.org/abs/2406.05183. Ouail Kitouni, Niklas Nolte, James Hensman, and Bhaskar Mitra. Disk: A diffusion model for structured knowledge, 2024b. URL https:...
-
[9]
Lucas Lehnert, Sainbayar Sukhbaatar, Paul Mcvay, Michael Rabbat, and Yuandong Tian. Beyond a*: Better planning with transformers via search dynamics bootstrapping.arXiv preprint arXiv: 2402.14083 ,
-
[10]
doi: 10.48550/arXiv.2205.14217. A. Liu and A. Borisyuk. Investigating navigation strategies in the morris water maze through deep reinforcement learning. Neural Networks,
-
[12]
Simple and effective masked diffusion language models.arXiv preprint arXiv: 2406.07524,
Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models.arXiv preprint arXiv: 2406.07524,
-
[13]
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu
URLhttps://arxiv.org/abs/2406.04329. Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding,
Show all 19 references
-
[14]
Aviv Tamar, S
URLhttps://arxiv.org/abs/2104.09864. Aviv Tamar, S. Levine, P. Abbeel, Yi Wu, and G. Thomas. Value iteration networks.Neural Information Processing Systems,
-
[15]
Notably, the tokenizers for A* and DFS mazes treat cell representations differently
They depict how the prompt and response are made from maze instantiations of the A* and DFS type. Notably, the tokenizers for A* and DFS mazes treat cell representations differently. In DFS mazes each grid cell is one distinct token. This is done to avoid making the sequences ...
2024
-
[16]
Haitong Wang, Aaron Hao Tan, and Goldie Nejat
URLhttps://arxiv.org/abs/1706.03762. Haitong Wang, Aaron Hao Tan, and Goldie Nejat. NavFormer: A Transformer Architecture for Robot Target-Driven Navigation in Unknown and Dynamic Environments, February 2024a. URL http://arxiv.org/abs/2402.06838. arXiv:2402.06838 [cs]. Siwei W...
-
[17]
We choose10−4 for next token and no weight decay forMLM-U
InMLM-U, we generally don’t see overfitting and therefore also don’t need any weight decay. We choose10−4 for next token and no weight decay forMLM-U. We found training to be most stable with the AdamW optimizer with beta valuesβ1 = 0.9 and β2 = 0.999 and batch sizes of 128 an...
2024
-
[19]
In light of this view we evaluate the following ablation: InMLM-U each token in the solution path is masked with some (uniformly drawn) probability, independently of other tokens
mT ′ pred ← False end if end for ˆy1:T ← (ˆy1, ...,ˆyT ) E Miscellaneous experiments E.1 Ordered masks One of our motivations for utilizing a training scheme likeMLM-U is that such a scheme enables more explicit reasoning over tokens that are further in the future than the imm...
2024
-
[2016]
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N
doi: 10.24963/ijcai.2017/700. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need,
2017 doi
-
[2019]
Abhimanyu Dubey et al
URLhttps://arxiv.org/abs/1810.04805. Abhimanyu Dubey et al. The llama 3 herd of models.arXiv preprint arXiv: 2407.21783 ,
-
[2021]
The pitfalls of next-token prediction.arXiv preprint arXiv: 2403.06963 ,
Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls of next-token prediction.arXiv preprint arXiv: 2403.06963 ,
-
[2022]
Jacob Austin, Daniel D
doi: 10.1109/IROS47612.2022.9982133. Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces.Neural Information Processing Systems ,
2022
- [2023]
-
[2024]
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova
URLhttps://arxiv.org/abs/2309.12288. Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding,
Reviewed August 11, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.