Transformers Can Navigate Mazes With Multi-Step Prediction

We’re thrilled to see the reviewer appreciate the robustness of our controlled experiments, originality in exploring the role of learning objectives in maze navigation, and clarity of our presentation.

Limited Novelty in Training Objective

Please refer to the general response on the goal of this paper.

We agree similar objectives have been explored in the reinforcement learning, which we acknowledge in our related works. However, reinforcement learning methods have not played a role in the pretraining of foundation models, and our goal is to showcase limitations of Next Token Prediction (NTP). NTP remains the defacto objective for large language model training. Here we show in controlled scientific experiments that NTP is not well-suited to mutli-step navigation tasks. We believe these findings would benefit the research community by providing clear scientific evidence for the advantages of alternative training objectives. Based on your findings, we’ve also added additional wording to make this contribution clear in Section 6 and highlight the strengths of reinforcement learning approaches in the related work in Section 2.

Limited Baseline Comparisons

Thank you for the suggestion to include additional baselines, including approaches from the reinforcement learning literature. While we don’t think these are directly comparable to our approach given the differences in architecture and training setup, specifically in the representation of the mazes. However, we do of course acknowledge their superiority (see for instance https://arxiv.org/pdf/2406.08404) and we can see RL as a gold standard to which we hope to rise with general (discrete) foundation models.

Baselines from prior related work

Otherwise, in addition to the baselines, we compare MLM-U performance in a rigorous manner against both prior work (Lehnert), including a very challenging baseline that uses extra supervision from A* search traces, and our own experiments with NTP. We consider the more challenging DFS setting which produces longer solution paths, where we perform a sweep over hyperparamaters for the next token baseline and consider multiple architectures (see Appendix A.2).

Scalability Concerns

We agree scalability of model size and understanding the role of hyperparameters are important. In Figure 6, we explore model size finding that MLM-U does benefit from scaling model size. When we scale from size 3 to 8M parameters, we find performance improves by 15% to perfectly solve 20x20 mazes. Regarding hyperparameters, we’d like to point out a potential misunderstanding: MLM-U does not have any masking ratio hyperparameters; instead, MLM-U uniformly samples masking ratios per sample with no additional hyperparameters involved. We do perform ablations of the other hyperparameters, in Section 5.3. We also provide an additional experiment in our general response exploring an alternative masking strategy of predicting only in the forward direction, which we show does not work as well as the uniform masking procedure in MLM-U used throughout the paper. We do also highlight several limitations of MLM-U in Section 6, such as the limitation regarding the input context window of transformers.

Insufficient Analysis of Failure Cases

Fortunately, MLM-U for mazes up to size X perfectly solves all held-out mazes. In the most complex setting we studied, size 30x30 and 20x20, we did observe failures to solve the mazes in X% of cases. We provide a qualitative analysis of mazes along with their solutions in Appendix C (as shown in Figure 10 and Figure 12). We also describe a failure due to insufficient precision, where we found X. To address some of these challenges, we’d be excited to explore curriculum learning approaches as well as specialized tokenization approaches to better represent world states.

Other Questions

Can you provide results or discuss potential outcomes for applying MLM-U to tasks beyond maze navigation?

We’re excited about the potential for MLM-U to tasks beyond maze navigation. We see maze navigation as a first test bed for demonstrating the capacity of MLM-U to predict multiple steps ahead. We’ve conducted early experiments on graph traversal problems from Stargraph where MLM-U also shows great promise for navigation capabilities. In principle, MLM-U is a general purpose objective that would work with any data that can be turned into a sequence of inputs for a transformer, which is quite a broad set. We would be excited to explore MLM-U in real world planning scenarios. We’ve highlighted this more explicitly in our conclusion section 6

1. In this paper, maze accuracy always refers to an exact token match of the entire path. We also provide additional token-accuracies in Appendix to gauge paths that are approximately accurate, but found both measures to provide comparable results. Based on your feedback, we’ve also made this explicitly by updating the caption of Table 2.

2. Please refer to appendix D.2 and let us know if any details remain unclear.

3. Please refer to section 4.3 and appendix D.1. We are not sure we understand the question "What are the sets of uniform masks used for MLM-U". To reiterate here, MLM-U works as follows: 1. draw a number m uniformly in [0,1]. 2. Each token in the solution path is masked with probability m.

4. We decode autoregressively, just like NTP. Therefore we don't need to know the trajectory length. This is not necessarily (or rather likely not) the most effective way of utilizing the model at inference time, but we wanted to keep it the same as in NTP to make this study more controlled. We wanted to make sure that the performance gains come from training, not from inference. See L285ff.

5. We don't have exact numbers here. About half of the failure modes had parsing errors. The other failures often had consistent paths, but more often than not they were longer than the shortest one. In conclusion, less than 25%.

6. Good question. The failure modes usually correspond to the longest paths, which might be a result of the autoregressive decoding with finite error probability. Note for instance Figure 10 vs 11, which were randomly picked failures in both NTP and MLM-U. You can see the difference in average path length. We also noticed that the models are very good at keeping paths consistent. It is rare that a model does an illegal step.

7. It takes a long time to train these models. They are quite deep and the 30x30 maze representations have context lengths of 3k+ tokens, scaling quadratically with maze size. We found ourselves compute limited when attempting to train on 40x40 mazes, for instance.

8. You are right, this would be interesting, but that takes even longer to train (see 7.) and is therefore deferred to future work. As you see, we instead explored the other direction and checked what happens if you lower the amount of data, see Fig 3.

We appreciate your effort in carefully reviewing this paper. We are excited that you consider our experimental results impressive and our writing well-structured. Please also refer to the general response.

Regarding Weaknesses

Some of the key setup and explanation details are missing from the paper, for example: What is the input representation of the mazes (does it match the Lehnert paper? Or are there modifications)? What is the definition of "navigation accuracy" (is this the same as A* solution accuracy in the Leonhart paper)

The setup is the same as the deterministic case in the Lehnert paper. Since it is deterministic, we can employ the "exact match" metric in all experiments on A* mazes. We do describe this exact match criteria in Appendix C (L718), but we agree that that should be described more specifically in the main text as well. Based on your feedback, we now include an expanded description in the main text in the caption of Table 2 in the updated paper. Please let us know if you have any further questions regarding the setup. We’d be glad to clarify or update the draft accordingly.

I think that primarily focusing on maze navigation is interesting, but perhaps too limited.

We agree maze navigation is of course one setting to study multi-step planning capabilities. We admit that this task does not cover imperfect information cases, but otherwise we consider it a pretty good perfect-information planning proxy. We found that proving such a superiority in maze navigation is in itself worth publishing, since the standard next token prediction objective falls short quite significantly. As reviewer ZjXN noted, we also explore the benefits of scaling MLM-U: “The scaling results are impressive, suggesting that this is a promising direction of research for larger models, and more complex tasks,” which we’re eager to further explore in follow-up work.

MLM-U is applied to full trajectories, which while it does encourage predicting multiple steps in the past, seems to be significantly harder to apply during the inference phase.

We believe this may come down to a misunderstanding of the difference between inference and training for MLM-U. MLM-U is at least trained to predict tokens further ahead (and backward) of its context, at least with some finite probability. At inference time however, we generate tokens in an autoregressive manner identical to next token prediction as described in Section 4.3. We agree there is a question of whether forward and backward token prediction in training is actually what drives the performance is of course unknown to us. We also ran a naive experiment where we mask all future tokens from some point onwards, but found this did not perform as well as multi-token predictions both ahead and back (as described in the general response). We hope this additional experiment as well as the distinction between inference and training for MLM-U have fully answered your question. Please let us know if you have further questions.

The paper has no baselines beyond next-token prediction

Fair point. Note however that our main goal was to point out the flaw in NTP. We see MLM-U as what one would converge to if one took Gloeckle to its extreme.

The results are drawn from a single training run on each of the methods. It would be good to see the variance of the data, particularly in Figures 3 and 4.

We agree it’s important to account for the variability across runs, thank you for pointing this out. Based on your feedback we’ve update Figure 4 to show the standard error about the mean for both next token prediction and MLM-U. Fortunately, the results across multiple seeds confirm MLM-U has stable convergence, whereas next token prediction overfits with continued training across the multiple seeds we ran. We hope this addresses your concern regarding variability and provides evidence for the robustness of our findings.

Key details We added Figure 14 (from lehnert) and 15 (adapted) to address your concern. As a summary here: The prompts work exactly like in Lehner, without the additional A* trace supervision in the prompts. Specifically, a sequence looks like: <BOS> <maze repr> <start node> <goal node> <insert path here> <EOS>. Sampling temperature is 0 (argmax) (see Enc-Dec description in appendix), given that we are looking for exact matches.

Limited experiments Thank you for bringing this potential misunderstanding to our attention. Our goal is not to overclaim by suggesting we’ve solved the open-challenge of multi-step planning. We choose maze navigation as an exemplary setting, studied by multiple prior works, requiring multi-step planning where standard next token training still struggles. We are excited about the potential of applying alternative training strategies for other domains. However in this work we found it worthy of note for other researchers in the field, and therefore worthy of publication, to find such a striking difference in performance for next token prediction vs a masking objective for maze navigation.

Training vs Inference Re policy shift: Just to clarify the procedure: Training is like BERT (so indeed masking), but uniformly sampled masking rate instead of fixed 15%. Inference is exactly like standard AR. In L288ff, we explain that this is not off-policy, referring to Kitouni et al (2024). One intuition that might help: Training like we do is (in expectation) the same objective as training a model in the usual fashion (left to right causal), but instead of always having the mask be causal (tokens are functions of previous tokens), sampling uniformly a permutation of that left to right order to condition tokens instead. This is referred to as permutation language modelling, with XLNet as an example.

Baselines We don’t believe methods such as A* are appropriate baselines for this work for two reasons: 1) search algorithms such as A* and DFS are used to generate the original maze solutions used for training. 2) our goal is not to compete with dedicated search algorithms; we make no claims to this effect. Our goal is to isolate the effect of learning objectives for transformers in maze navigation.

Variance We’re glad you appreciate our effort to add standard error to Figure 4. While of course including standard errors for as many figures as possible is our ambition, in this case we prioritized based on your feedback doing so for Figure 4. In the 5x5 setting, we found the error bounds for MLM-U < 1% for overall maze navigation path accuracy and observed similarly small margins for other settings. Given this relatively small deviation across runs in contrast to the large margins by which MLM-U outperforms next token training: 100% versus 45.2% (10x10), 100% versus 24.4% (for 15x5) and 100% versus 20.6% (for 20x20), and 93.8% versus 18.8% (for 30x30) as shown in Table 1, the performance differences, as noted by Reviewer vUWK, constitute strong empirical evidence for the benefits of training with MLM-U for maze navigation.

1. Thanks for the clarification.
2. This doesn't actually answer my question - Given that these are sequence models, I would expect at least somewhere in the paper to show the actual sequences passed to the model? How do the prompts work? What are the sampling temperatures? etc. All of these details are quite important, and not discussed at all in the design of the mazes. Similar to Figure 1 in Lehnert et al, these details should be discussed.
3. Thank you for clarifying these details.
4. See my answer in the other comment - is there not off-policy shift when decoding auto-regressively here?
5. Thanks for clarifying, it would be good to report this in the paper.
6. Thanks for clarifying, it would be good to report this in the paper.
7. Given that compute intensivity is a concern, what is the advantage of using such an approach over a compute efficient approach (such as A*, which can easily and optimally solve any 50x50 maze)? It seems like this should be discussed deeper in the limitations.
8. Thanks for clarifying.

We’re thrilled the reviewer agrees planning abilities of transformers is a “timely” and important topic. We’re also glad the reviewer found our experimental setting to be fair, providing good evidence for the benefits of alternative learning objectives for better maze navigation.

Limited novelty While this paper does not propose MLM-U specifically, we believe our contribution is valuable to the research community in providing clear scientific evidence to the benefit of factorization agnostic training for overcoming one of the main limitations of today’s transformers (as you noted as other reviewers): multi-step planning. As reviewer ZjxN noted, “no existing work has clearly demonstrated the effects of MLM-U, or MLM-based training approaches on maze prediction tasks.” Our empirical experiments carefully cover two distinct maze solving tasks, across a range of maze grid size complexities and model sizes. We believe these findings would benefit the research community by providing clear scientific evidence for the advantages of alternative training objectives to overcome transformers’ limitations with regard to multi-step planning.

Please also refer to the general statement about the goal of the paper.

Questions

Did the experiments on 30x30 mazes leverage the traces of 30x30 mazes for training [..]?

In our MLM-U baselines, we do not use supervision from any search traces in training. MLM-U is trained only on maze solutions. All models are then evaluated on held-out mazes of the same size seen during training. What we found remarkable is that MLM-U outperformed the baseline from prior work that in addition to the same training data used for MLM-U also uses supervision from A* search traces. MLM-U outperformed this baseline that uses additional supervision by more than 15% for the most complex 30x30 maze setting as shown in Table 2.

It would be great to study the generalization of MLM-U in comparison with next token in terms of the task difficulty.

In our experiments, following the standard setup from prior work such as Lehnert et al. we train and evaluate all methods on mazes of the same grid size. We agree it would be interesting to explore generalization of models trained on one maze size to solve mazes of other grid sizes. Based on your suggestion we ran additional experiments to this effect. We now include these results in a new Appendix Section E.2 showing that while MLM-U is more sensitive to modifications of the tokenizer, when the maze is embedded in the same tokenization space MLM-U does generalize significantly better than next token to solve mazes of grid sizes other than those used for training.

In Section 5.4, why would the precision of the positional encoding especially matter in this MLM-U training paradigm?

We hypothesize precision of positional encodings also likely also matters in next token prediction, but to a lesser extent than it does for MLM-U. Intuitively, this is related to the difference training objective: next token training aims at local predictions that do not require predicting long-range steps in maze states. On the other hand, MLM-U leverages higher precision positional encodings as predictions can span the entire range of maze. More explicitly, next token prediction has an additional positional bias via the causal mask, which as shown in prior work (NoPE paper https://arxiv.org/pdf/2404.12224) can remove the need for positional encodings entirely with next token prediction. MLM-U on the other hand—due to its multiple token ahead and backward predictions—has no causal masking. Instead the positional bias in MLM-U comes from the explicit encoding.

In this work, the Transformer models are trained from scratch using NTP or MLM-U objectives. However, I believe the contributions of the work would be much greater if the authors conduct experiments with LLMs in the finetuning stage

Thank you for the suggestion! We agree it would be interesting to assess how well MLM-U can serve as a finetuning objective for maze navigation. Our central scientific aim is to assess in a controlled, apples-to-apples comparison the effect of learning objective in maze navigation, an exemplary tasks requiring multi-step planning. In order to do so, we chose to remove as many confounding factors such as the effect pretraining as well as the possible distribution shift in pre-post training arising from finetuning with a different objective than was used in pretraining. Although finetuning falls outside the scope of the scientific claims we target in this work, se do agree exploring alternatives to next token prediction in post-training is also an interesting and practically valuable contribution given the rise of open-source models.

Thank you for your time reviewing our paper. We are happy that you appreciate our detailed analysis of the comparison between next token prediction and MLM-U. We respond to your feedback and specific questions below:

Some implementation details are not clear to me. Does MLM-U use an encoder-decoder architecture while next-token prediction is used for decoder-only architecture?

We’d be happy to clarify the choice of architectures. We found that encoder-decoder architecture works better for MLM-U (and base this choice on prior work). We found decoder-only architecture works better for the next token baseline. In order to give both objectives a fair chance in the competition, we use the best model for each objective. To ensure our results are not the result of the choice of architecture, we ablate the choice of model choice in Appendix A.2. where we explore both encoder-decoder and decoder architecture and choose the stronger baseline for comparisons in the main paper. Please let us know if you have any other questions regarding our experimental setup.

emphasis on the role of RoPE positional encodings.

Thank you for pointing this out. We agree our presentation could have been better to appropriately emphasize the core message of this section. While we do use RoPE, we do not seek to emphasize this implementation of positional encoding. We see how the title of Section 5.4 is confusing, thank you for making us aware! Based on your feedback, we’ve revised the title of the section in the manuscript to appropriately emphasize the goal of the section. The point we are trying to make in that section is that the positional encoding might have more of an impact in MLM-U like objectives than in NTP, for two reasons: 1. in NTP, one does not strictly need an additional positional bias on top of the causal mask (see the NoPE paper) and 2. since small changes (like the precision used in calculating it) make a big difference (see Fig 7).

The novelty of the MLM-U objective might be small.

We agree MLM-U isn’t the only objective to explore multi-token prediction masking objectives as detailed in our general response. The value of contribution is in the application of such an objective to the open challenge of multi-step planning in transformers. As noted by reviewer ZjxN, “no existing work has clearly demonstrated the effects of MLM-U, or MLM-based training approaches on maze prediction tasks.” We show via controlled experiments the benefits such an objective can bring in overcoming the multi-step planning limitations of transformers for maze navigation as well as the promise of scaling such an approach.

In an effort to alleviate some of the concerns about the diversity of experiments, we introduce a couple of additional experiment results that can be found in appendix E and in the main text of the revised manuscript. We summarize them here:

1. We include a new version of Fig 4 showing the standard error across the mean for three random seeds as suggested by reviewer ZjxN. The new results suggest MLM-U has stable convergence across multiple seeds, whereas next token prediction overfits with continued training across the multiple seeds we ran.

2. Another reviewer criticized that we don't provide additional baselines like the solution provided in Bachmann et al (2024), and another reviewer hinted that they think that the benefits of MLM-U in maze solving do not come from increased "foresight", but rather from "justification" and "hindsight". To investigate both concerns, we conduct the following experiment: Instead of training with MLM-U where all tokens are masked independently of each other with some probability, we uniformly pick a position in the solution path and mask all tokens to the right of this position. Then we predict all of those tokens as a function of the context to the left of the chosen position. This relates closely to Bachmann et al, and since the mask is always "oriented forward" and we don't predict tokens "backwards". However, we find that this method is far inferior to MLM-U in the 10x10 A* maze setting that we tested. The maximum per-token accuracy observed is 73%, with less than 4% full path accuracy. This result indicates that future work is needed to understand the failure modes of looking far ahead, but maybe it sheds some light on the concerns of the reviewers.

3. One reviewer was interested in seeing whether models can generalize beyond the exact maze distribution they were trained on, specifically when changing the maze size. We have no hope of that generalization to larger mazes than the training distribution, because of the problems of transformers in length generalization. We can however achieve non-trivial accuracy when going to smaller mazes. We evaluate 10x10 DFS mazes with models trained on 20x20 mazes. Interestingly, we find that next token training can somewhat generalize to smaller mazes without any changes except adjusting tokenization, whereas MLM-U only performs well when embedding the 10x10 maze into a quadrant of a random 20x20 maze. Without embedding, MLM-U has 0% accuracy and 100% with the embedding trick. Next token training achieves 21% without embedding and 29% with it, see also appendix E.2 in the revision.

We revised the manuscript to clarify our contribution and aspects of our experimental setup.

With these revisions thanks to the thorough feedback from reviewers, we’re confident the updated manuscript would be a valuable contribution to the research community showing the promise moving away from next token prediction in training transformers for improved mutli-step planning in tasks such as maze navigation. We remain available to incorporate any further feedback or to answer any remaining questions any reviewers may have.