Transformers Can Navigate Mazes With Multi-Step Prediction (2412.05117v2)

Abstract: Despite their remarkable success in LLMing, 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

• The paper introduces the MLM-U training objective that predicts multiple steps to enhance long-horizon planning in maze navigation.
• It achieves 93.8% accuracy on 30x30 DFS mazes, dramatically surpassing the 18.8% accuracy of standard NT models.
• MLM-U also improves data and computational efficiency by reducing training samples by 75% and GPU hours by 50%.

Essay on "Transformers Can Navigate Mazes With Multi-Step Prediction"

The paper "Transformers Can Navigate Mazes With Multi-Step Prediction" presents an exploration of transformer training objectives as a means to enhance the model's long-term planning capabilities, using maze navigation tasks as a test case. The authors propose substituting the conventional next-token (NT) prediction task with a Multi-Step Masked LLMing (MLM-U) training objective to improve efficiency and effectiveness in solving mazes.

Overview of Proposed Methodology

The central hypothesis of the paper is that the standard NT prediction, despite its triumph in LLMing, fails to endow transformers with the necessary foresight to plan extended sequences effectively. This is particularly pressing in tasks like maze navigation, where multi-step planning is essential. The paper introduces the MLM-U objective as an alternative that requires the prediction of multiple steps forward and backward, thus invoking more robust planning strategies.

Transformers configured for MLM-U training engage with input sequences by masking arbitrary subsets and learning to predict these masked tokens using the surrounding context. The uniform randomization of mask positions is intended to compel the model to develop a more global understanding of sequences, thus better supporting complex planning tasks.

Experimental Setup

To validate their approach, the researchers trained transformer models on mazes generated through two distinct methodologies: Depth-First Search (DFS) and A* algorithms. Transformer models with matching parameters were trained either under the standard or MLM-U objectives across several maze complexities. The training considered factors such as navigation accuracy, data efficiency, and computational efficiency in terms of GPU hours to convergence.

Results and Analysis

The MLM-U objective demonstrated superior performance over next-token prediction across the board, with results indicating significant advantages in maze navigation accuracy, particularly as task complexity increased. Notably, an 8M parameter transformer trained on MLM-U solved 30x30 DFS mazes with 93.8% accuracy compared to the 18.8% of an NT transformer. Even more substantial was its performance against larger NT models that utilized supervision from A* search traces.

Beyond navigation accuracy, the MLM-U objective showed a marked increase in data efficiency; it required a quarter of the training samples compared to NT to achieve comparable results on simpler mazes. Computational efficiency was another area of improvement, as MLM-U achieved convergence more rapidly, illustrated by a 2× reduction in training GPU hours for similar tasks.

Importantly, the results underscore the pertinence of model scale, particularly for more challenging maze complexities. Scaling the size of MLM-U trained models correlated with enhanced navigation performance, again surpassing equivalently scaled NT models.

Discussion and Implications

The paper delivers a salient critique of NT prediction objectives while providing empirical evidence for a promising, generalizable strategy through MLM-U. As the complexity of navigation tasks grows, the limitations of fixed-token prediction objectives become apparent. Therefore, emphasizing objectives that encourage broader sequence reasoning emerges as a compelling advancement for the field, with potential implications extending to broader applications requiring long-horizon planning.

Given these findings, future research should explore understanding the interactions between MLM-U objectives and positional encodings, as these seem crucial, especially in larger maze environments. The implications of these objectives could extend beyond navigation to include tasks like strategic gameplay, automated planning in robotics, and other domains that benefit from improved anticipatory computing.

The optimistic results from MLB-U signify a crucial step toward enhancing transformers' planning capabilities, offering a fertile ground for further exploration into adaptive training goals that extend beyond immediate next-token prediction. As such, this paper adds a valuable contribution to the ongoing research into expanding the utility and adaptability of transformer models in complex environments.