Transformers Can Navigate Mazes With Multi-Step Prediction

Niklas Nolte, Ouail Kitouni, Adina Williams, Michael Rabbat, Mark Ibrahim

Keywords: transformers, planning, learning objectives

TL;DR: Compared to standard next token prediction training, we find transformers trained to predict multiple steps ahead (and back) learn to navigate mazes efficiently.

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.

Primary Area: learning theory

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 4671