Transformers Can Navigate Mazes With Multi-Step Prediction

TL;DR

Transformers trained to predict only the next token struggle with long-horizon planning, as demonstrated in maze navigation. The authors propose MLM-$\mathcal{U}$, a diffusion-like objective that masks arbitrary subsets of the input and requires predicting multiple steps ahead and backwards, formalized as $L_{MLM-\mathcal{U}} = - \mathbb{E}_{\mu \in \mathcal{U}} \log P_{\theta}(m_{\mu} X | m_{\mu}^C X)$. Across DFS and A* mazes and multiple model scales, MLM-$\mathcal{U}$ substantially improves navigation accuracy, data efficiency (roughly 4x), and training efficiency (roughly 2x) compared to standard next-token training, and can match or exceed larger autoregressive models even with A* supervision. These results highlight the critical role of learning objectives in enabling long-horizon planning in transformers and suggest MLM-$\mathcal{U}$ as a practical avenue for more capable planning in real-world tasks.

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

Figures (14)

• Figure 1: MLM-$\mathcal{U}$ predicts multiple steps ahead and backward. Standard autoregressive training only (explicitly) predicts the next step. We compare 8M parameter transformer models trained with autoregressive next token prediction versus MLM-$\mathcal{U}$ training objectives. Maze complexity is defined in terms of the maze grid size.
• Figure 2: Left: Path lengths, measured by number of traversed cells, of A* and DFS mazes for maze sizes 10x10, 20x20 and 30x30 on the validation dataset. Error bars show the standard deviation. Middle: Example 10x10 A* maze Right: Example 10x10 DFS maze. Both are real randomly selected examples illustrating the difference between encoding walls in cells (A*) versus edges with longer paths (DFS).
• Figure 3: Training Data Sample Efficiency. We compare 8M parameter model next token versus MLM-$\mathcal{U}$ held-out accuracy as we vary the number of mazes seen during training. On the left, for 5x5 mazes which both learning objectives can solve, MLM-$\mathcal{U}$ is $4\times$ more data efficient. On the right, for 10x10 mazes we see MLM-$\mathcal{U}$ converges to perfectly solve 10x10 mazes with 25k training samples, where next token performance peaks below 50% accuracy.
• Figure 4: Training efficiency of next token vs. MLM-$\mathcal{U}$ on 5x5 mazes. While both models are able to perfectly solve held-out 5x5 mazes, MLM-$\mathcal{U}$ does so 2.03x more quickly relative to next token. The shaded region shows the standard error across the mean over three random seeds. We also observe overfitting for next token training past 200k training steps whereas MLM-U accuracy remains at near perfect accuracy.On the right, we show the number of GPU hours needed for each training objective to converge.
• Figure 5: Comparing convergence rates of next token and MLM-U on 10x10 mazes. Left is training accuracy; right is navigation accuracy on held-out mazes.