{"id":"4e0621c8-4cbd-4035-8633-905a86e11ff5","arxiv_id":"2412.05117","paper_version":2,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":0,"one_line_summary":"A multi-step masked prediction objective, MLM-U, lets small transformers solve mazes that next-token models, and even much larger A*-supervised next-token models, cannot.","lead":"This paper shows that training transformers to predict several steps of a maze solution at once, rather than only the next step, makes them far better at solving new mazes. The finding suggests that changing the training objective could be a simple way to improve long-horizon planning in language models.","discovery_kind":"new_application","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Architecture confound: MLM-U is always trained with an encoder-decoder + RoPE, while the next-token DFS baseline is decoder-only and the A* baselines come from Lehnert et al. unmatched; the 'only the objective changed' claim is not yet isolated. A matched-architecture next-token arm would settle it.","rationale":"The reader's weakest assumption and my primary concern align: the objectives are not compared under fully matched architecture and training conditions. This is the most load-bearing issue because the paper's central claim is causal ('modifying only the learning objective'), and the evidence for that causal claim is strongest when all else is held equal. The DFS comparison is partially protected by Appendix A.2, where the authors show the decoder-only next-token baseline is stronger than the encoder-decoder next-token baseline; if they had used the encoder-decoder for next-token, MLM-U would likely look even better. However, that ablation does not test the opposite arm: MLM-U with a decoder-only architecture. The A* comparison is more exposed because it relies entirely on published numbers from Lehnert et al. with no control over architecture or compute. The observed effect is large and probably real, but the paper's own framing overstates the isolation of the objective. I therefore do not recommend changing the reader's CONDITIONAL verdict; the concern supports the existing request for matched-architecture evidence rather than overturning the result. Secondary issues, such as the inconsistent GPU-hour accounting (13.7 vs 17.7 hours is a 29% overhead, not '7%') and the poorly defined '4x data efficiency' ratio, are real but do not threaten the navigation accuracy claim as directly as the architecture confound does.","tokens_in":14943,"tokens_out":9161,"duration_ms":96840,"concrete_test":"Train a parameter-matched next-token transformer with the exact encoder-decoder architecture, RoPE precision, tokenizer, and training schedule used for MLM-U on the same DFS (Table 1) and deterministic A* (Table 2) datasets, sweeping learning rate and weight decay as in Appendix B. If this next-token encoder-decoder model still saturates near 20-50% on 10x10-30x30 DFS mazes while MLM-U reaches ~100% on 20x20 and ~94% on 30x30, the architecture confound is not the driver and the central claim stands. If the gap narrows substantially, the headline should be reworded from 'learning objective' to 'objective plus architecture.'","verdict_should_be":"UNCHANGED","load_bearing_attack":"The central claim is that the learning objective, not architecture, drives the large navigation gap. But the comparisons never hold the architecture fixed. MLM-U is trained with the Kitouni et al. encoder-decoder and RoPE (Section 4.3, Algorithm 1), while the DFS next-token baseline uses a decoder-only transformer, chosen because Appendix A.2 shows it outperforms the encoder-decoder for next-token training. The A* comparison in Section 5.2 imports Lehnert et al. (2024) baselines without matching architecture, tokenizer, training compute, or hyperparameters. This leaves open the possibility that part of MLM-U's advantage comes from the encoder-decoder's bidirectional encoding and cross-attention rather than from explicit multi-step prediction. Appendix A.2 mitigates the concern for DFS by showing decoder-only is the stronger next-token baseline, but it does not eliminate it: no next-token run uses the MLM-U architecture, and no MLM-U run uses a decoder-only architecture. The paper's phrasing 'identical settings' (Section 4) and 'modifying only the learning objective' (Section 4) is therefore stronger than the evidence currently supports. The concern is load-bearing because it targets the causal attribution at the heart of the paper, not the secondary efficiency claims.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","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.","tokens_in":15192,"tokens_out":6004,"duration_ms":53524,"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":[{"comment":"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":"Section 4.4 / Appendix A.2"},{"comment":"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":"Section 5.1, Figure 3"},{"comment":"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":"Section 5.1 and Abstract"},{"comment":"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.","section":"Section 5.4"}],"minor_comments":[{"comment":"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":"Section 4.1"},{"comment":"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.","section":"Section 5.1"},{"comment":"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.","section":"Abstract and Section 1"},{"comment":"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.","section":"References"}],"recommendation":"major_revision","confidential_remarks":"The manuscript is an empirical study from a strong industrial group, and the code is released. The central qualitative finding is likely to be robust, but the paper's strongest claims ('modifying only the learning objective', 4x data efficiency, 2x GPU-time speedup) are not yet supported by the evidence as presented. The architectural confound is the most important issue; the other numbers require straightforward corrections. I would support a revised version that adds a matched-architecture next-token baseline and fixes the efficiency-statistics inconsistencies. Note that the comparison to Lehnert et al. is a same-group baseline, but this is acceptable if clearly labeled as an external baseline."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"First, the headline result is real: an 8M transformer trained with MLM-U solves held-out DFS mazes up to 20x20 perfectly and reaches 93.8% on 30x30, where next-token training sits in the teens and twenties. On Lehnert A* mazes it beats a 175M next-token model with A* trace supervision. That is a large effect, and the paper documents it carefully across maze types, sizes, model scales, token-wise and full-path accuracy, and failure modes.\n\nWhat's new: this is the first systematic application of the MLM-U objective (from the authors' own earlier paper) to maze navigation, and the comparison against A*-supervised next-token transformers is new. The hyperparameter sweep for the next-token baseline, the per-token results, the ordered-mask ablation in Appendix E.1, and the positional-encoding precision finding (16-bit vs 32-bit RoPE makes a large difference on larger mazes) are all useful. Code is released.\n\nSoft spots in order of severity. The architecture confound is the main one. MLM-U is always trained with an encoder-decoder plus RoPE; the DFS next-token baseline is decoder-only (chosen because Appendix A.2 shows it is stronger than the encoder-decoder for next-token), and the A* baselines are imported from Lehnert et al. without matching compute or hyperparameters. The paper repeatedly says 'identical settings' and 'modifying only the learning objective,' but that is not what the experiments do. Appendix A.2 mitigates part of the concern, but there is no arm that trains next-token on the MLM-U architecture or MLM-U on a decoder-only. Without that, part of the gap could be credited to bidirectional context and cross-attention rather than multi-step prediction.\n\nSecond, the efficiency claims are looser than the prose. '4x more data efficient' is only well-defined on 5x5 where both objectives reach perfect accuracy; on 10x10 next-token never gets there, so the ratio is undefined. The GPU-hours text is inconsistent: the paper says 13.7 vs 17.7 hours for the same epoch count but then claims ~2x faster convergence; presumably that factors in the lower epoch count to convergence, but the derivation isn't shown. Main tables lack error bars, though some figures report standard error.\n\nThird, the comparison with Lehnert et al. is a comparison against published numbers, not matched runs. That's fine for a sanity check but shouldn't anchor the 'outperform A*-supervised' claim.\n\nNone of these kill the core effect. The paper is honest about its limitations, including MLM-U's failure on smaller mazes out of the box (Appendix E.2) and its failure cases on A* mazes where it predicts valid alternate shortest paths.\n\nWho should read this: anyone working on training objectives for planning or on transformer reasoning. It's a solid empirical contribution that deserves serious review, but the next-token baseline needs to be run in the MLM-U architecture before the 'only the objective changed' claim stands. I'd send it to a good venue with that demand and with sharper efficiency numbers.","headline":"Strong empirical result on MLM-U for maze navigation, but the architecture confound and loose efficiency numbers keep the headline from being fully clean.","tokens_in":15760,"tokens_out":3379,"would_cite":false,"duration_ms":33136,"reading_group":"maybe","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"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…","keywords":["maze navigation","next token prediction","MLM-U","multi-step prediction","transformers","long-term planning","masked language modeling","learning objectives"],"falsifier":"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.","tokens_in":14729,"feed_emoji":"🧭","tokens_out":7914,"duration_ms":73162,"temperature":0.7,"pith_summary":"Transformers trained to predict the next token have no way to practice looking several moves ahead, and prior work shows they take shortcuts on navigation. This paper tests whether replacing the objective with MLM-U, which masks random chunks of the path and predicts many steps forward and backward, makes transformers better maze solvers. Across DFS and A* mazes from 5x5 to 30x30, an 8M-parameter MLM-U transformer outperforms next-token baselines, reaching perfect navigation on all held-out DFS mazes up to 20x20 and beating a 175M next-token model that is given A* search-trace supervision. The paper also reports 4x better sample efficiency and roughly 2x faster convergence in GPU hours in settings where both objectives can solve the task. If the finding holds, multi-step masked prediction is a lever for planning in transformers that does not require extra supervision or larger models.","feed_headline":"Multi-step training lets transformers solve mazes next-token can't","feed_subtitle":"An 8M-parameter model solves all held-out 20x20 mazes, where next-token training peaks near 20 percent accuracy.","key_machinery":"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.","core_discovery":"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.","pith_inferences":["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."],"forward_implications":["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."],"supporting_citations":[{"why":"Identifies the Clever-Hans cheat and other shortcuts that next-token prediction takes on graph navigation, motivating the search for a multi-step objective.","marker":"Bachmann & Nagarajan (2024)"},{"why":"Introduces the MLM-U objective and the encoder-decoder training setup that this paper applies to maze navigation.","marker":"Kitouni et al. (2024a)"},{"why":"Provides the A* maze dataset, the deterministic shortest-path ground truth, and the next-token baselines with and without A* trace supervision.","marker":"Lehnert et al. (2024)"},{"why":"Shows that predicting multiple future tokens with multiple heads can improve large language models, motivating explicit multi-step prediction.","marker":"Gloeckle et al. (2024)"},{"why":"Supplies the randomized DFS maze generation method and the text representation used for the DFS maze experiments.","marker":"Ivanitskiy et al. (2023a)"},{"why":"Shows that transformers trained on maze navigation can learn internal maze states, providing the baseline next-token maze training setup.","marker":"Ivanitskiy et al. (2023b)"},{"why":"Defines the rotary position embeddings whose floating-point precision the paper finds critical for MLM-U on larger mazes.","marker":"Su et al. (2023)"}],"fun_headline_variants":["Multi-step prediction lets small transformers beat huge ones on mazes","Next-token training fails mazes; multi-step succeeds","Multi-step objective makes transformers 4x more sample-efficient","8M transformer solves all 20x20 mazes with multi-step training"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"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.","fun_headline_variants_meta":{"raw":{"variants":["Multi-step prediction lets small transformers beat huge ones on mazes","Next-token training fails mazes; multi-step succeeds","Multi-step objective makes transformers 4x more sample-efficient","8M transformer solves all 20x20 mazes with multi-step training"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000983,"raw_usage":{"total_tokens":4237,"prompt_tokens":1077,"completion_tokens":3160,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":693,"completion_tokens_details":{"reasoning_tokens":3100}},"tokens_in":693,"tokens_out":3160,"duration_ms":24820,"temperature":1.0,"reasoning_tokens":3100,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-11T20:50:13.678393+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"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.","supporting_citations":[],"review_version":1}