REVIEW 3 major objections 6 minor 2 cited by
Multi-Token Prediction Needs Registers
T0 review · 3 major / 6 minor · reviewed 2026-08-15 · deepseek-v4-flash
Pith's one-line read MuToR claims that interleaving trainable register tokens to predict tokens d steps ahead delivers multi-token prediction benefits during training with negligible parameter overhead and no inference change.
desk verdict A clean training-only register mechanism for multi-token prediction with consistent but modest gains; worth engaging, though the compute control and missing baselines need work. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The load-bearing object is the register token r_d: one learnable embedding inserted after x_t, given position id t+d-1, and trained to predict x_{t+d}. The offset d is sampled uniformly per sequence from {1,...,d_max}, so registers share a single embedding and the offset is encoded through positional bias such as rotary position embeddings. A custom attention mask ensures regular tokens cannot attend to registers and registers cannot attend to other registers, keeping the next-token loss and inference graph unchanged. Because there is one register per position regardless of d_max, the prediction horizon can be scaled almost for free; in the 2D image variant a sampled offset pair (d_h,d_w) maps to a rasterized offset and lets one register cover an entire spatial neighborhood.
What would settle it
Train MuToR and a next-token baseline with identical total regular-token processing (equal optimizer steps scaled to match sequence length) on a fixed fine-tuning set; if next-token then matches or beats MuToR, the claimed advantage collapses. A sharper control is to insert the same register tokens but remove their auxiliary loss: equal performance would show the extra positions, not the lookahead objective, are responsible.
Extended reading notes
Core claim
The central claim is that an auxiliary lookahead objective can be implemented by interleaving a single shared learnable register token after regular tokens, giving each register the task of predicting the token d steps ahead. The register is assigned position id t+d-1 so that its prediction looks like ordinary next-token prediction, and a custom attention mask lets registers attend only to preceding regular tokens while regular tokens never see registers. That mask preserves the next-token objective exactly and makes registers discardable at inference. Empirical results show MuToR outperforming next-token and multi-token-head baselines on mathematical reasoning and summarization after fine-tuning, improving LoRA fine-tuning to near full-fine-tuning level, and, in image pretraining, reaching better FID at 100K steps than the next-token model at 200K steps.
Load-bearing premise
The load-bearing premise is that running baselines for extra epochs faithfully matches the extra training compute that register tokens consume; if doubled-epoch baselines overfit or misrepresent the cost of longer sequences, the reported gains could come from training budget rather than from the register objective.
Editorial extensions
If this is right
- In fine-tuning, MuToR can be applied to an off-the-shelf pretrained LLM without architecture changes, so lookahead supervision becomes a training-only addition to the standard recipe.
- LoRA fine-tuning with registers reaches or exceeds full-fine-tuning next-token performance, so parameter-efficient pipelines could get stronger models at lower cost.
- In autoregressive image generation, one register embedding covers a whole 2D prediction neighborhood that would otherwise require many separate prediction heads.
- Because the number of registers is fixed as d_max grows, extending the prediction horizon does not add training cost.
- The star-graph results indicate the objective can recover useful supervision in settings where teacher forcing collapses into shortcut learning.
Reading between the lines
- The attention-mask pattern is effectively an invisible auxiliary slot: the same mechanism could carry other training-only objectives, such as predicting latent features, correctness labels, or auxiliary modalities, without changing inference.
- The sparse-register experiment suggests placement is a resource; learning where to insert registers per task could produce similar gains with less compute than uniform interleaving.
- If the gain is caused by supervision density rather than architecture, the same register-token idea should transfer to encoder-decoder and non-autoregressive training, though the paper does not test this.
- The 2D offset formulation suggests registers can encode spatial structure generically, which may open a route to structured lookahead in non-image domains such as graphs or molecules.
Signed reviews
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes MuToR, a training-only multi-token prediction objective that interleaves learnable register tokens into the input sequence. Each register token is assigned a random future offset d and is trained with a cross-entropy loss to predict the token d steps ahead, while regular tokens retain the ordinary next-token loss and are masked so they never attend to register tokens. At inference the register tokens are discarded. The authors fine-tune Gemma 2B and Llama 3 8B on mathematical reasoning and dialogue summarization, apply the method with LoRA, pretrain LlamaGen-B on ImageNet with 1D and 2D offsets, and evaluate on the star-graph path-finding task of Bachmann and Nagarajan. They report consistent improvements over next-token and multi-token baselines with negligible additional parameters, and claim that the improvements are not due to higher training compute because baselines trained for twice as many epochs do not improve.
Significance. If the empirical claims survive a properly FLOP-matched comparison, MuToR is a genuinely useful training-time technique: it requires no architectural change, adds only a small embedding, is compatible with off-the-shelf pretrained models and PEFT, and supports arbitrarily large prediction horizons without extra output heads. The paper is strong on breadth: it covers full fine-tuning, LoRA, image pretraining, and a synthetic planning task, and it ablates the maximum offset, register sparsity, and shared versus distinct embeddings. The central weakness is that the compute-matched controls in Section 4.1 and Appendix B.1 do not actually match compute, and many of the headline gains are small and reported without uncertainty. If the authors add a genuine FLOP-matched next-token baseline and provide error bars, the contribution would be convincing.
major comments (3)
- [Section 4.1, Eq. (4), Tables 11-12] The compute-matched controls do not actually match compute. Eq. (4) inserts a register token after every regular token except the last, so a MuToR training sequence has length about 2T. Causal attention FLOPs grow quadratically with sequence length; even counting non-attention layers, a MuToR step is at least about twice as expensive as a Next-Token step on the original length T, and in attention-dominated regimes it is about four times as expensive. Training the baselines for only twice as many epochs, as in Tables 11 and 12, therefore leaves MuToR with substantially more total FLOPs. The extra-epoch control is also confounded by repeated data, which produces the overfitting visible in those tables. The same issue affects the image experiments: the text around Table 6 calls 100K MuToR-2D steps versus 200K Next-Token steps 'similar training-time compute' without accounting for the doubled sequence length. Consequently, the gains in Tables 1, 2, 4, and 6 could be attributable to additional training FLOPs rather than to the register objective. Please provide a true FLOP-matched or compute-matched baseline (e.g., Next-Token trained for enough steps to match MuToR's measured FLOPs, or with correspondingly increased sequence length and batch size), and report measured FLOPs or wall-clock time for all compared runs.
- [Tables 1-6 and Appendix A.1.1] Many of the headline improvements are small and are reported without uncertainty. Gemma 2B numbers are stated to be averaged over three seeds but no standard deviations or significance tests are given; Llama 3 8B results are from a single seed, as acknowledged in Appendix A.1.1; and the image experiments in Table 6 appear to be single runs. In this regime, differences such as GSM8K 38.87 versus 42.10, or 1M-GSM 66.09 versus 68.33, may be within seed variance, and the Llama 3 8B gains are particularly vulnerable because there is only one run. Please report per-seed values, standard deviations, and significance tests, or add multiple seeds for the Llama 3 8B and image-generation runs.
- [Section 4.1 and Tables 1-3] The multi-token baseline is described only as 'adapted for finetuning by adding dmax-1 auxiliary prediction heads and applying a loss-weighting strategy,' with no details on head initialization, whether the heads are trained from scratch, or the loss-weight schedule; the reported added parameters (110M-550M) are marked 'in approximation.' A baseline with poorly tuned or under-trained heads could make MuToR look better than it is. Please specify the baseline architecture and tuning procedure, and consider reporting an additional multi-token baseline with a comparable parameter budget or comparable per-step FLOPs.
minor comments (6)
- [Section 3.1] The phrase 'minimize the the expected negative log-likelihood loss' contains a duplicated article and should be corrected.
- [Section 3.2, Eq. (4)] The displayed augmented sequence ends with x_T and contains T-1 register tokens; clarify whether a register is ever inserted after the final token, and state how the register loss handles offsets with t+d > T, i.e., whether such targets are masked out.
- [Section 3.2, Position Embeddings] For d=1, a register after x_t receives position t, which is identical to x_t's position; please explain why this does not cause a RoPE collision or ambiguity, since both tokens are present in the same sequence.
- [Section 4.2 and Table 6] Specify whether all MuToR image runs use 256 registers per image and how the checkpoints at 100K, 200K, and 360K iterations are selected; the main text says 'at 100K steps' but it is unclear if these are final checkpoints or best-validation checkpoints.
- [Appendix A.1.1 and Tables 9-10] Please report the search ranges for dmax and alpha, not just the best configurations, and correct the typo 'uniformally' in the dataset description.
- [Figure 4, Section 4.3] Add error bars or report the number of runs for the star-graph solve rates, as this task is known to be sensitive to initialization.
Circularity Check
No significant circularity: MuToR's empirical claims are tested against external benchmarks rather than derived from its own fitted quantities.
full rationale
The paper makes no claimed prediction that reduces by construction to a fitted input. Its central contribution is a training-time objective defined in Equations 4 through 6, and its value is assessed by downstream accuracy on GSM8K, MATH500, AQUA-RAT, SAMSum, DialogSum, and ImageNet FID/IS, all external benchmarks. The only tuned quantities, the loss weight a and the maximum offset d_max, are standard hyperparameters selected on validation data and reported with ablations in Tables 3, 8, 9, and 10; they are not renamed as predictions. The paper uses the term 'registers' following Darcet et al. (2024), but the method adds a distinct multi-token prediction objective, so the borrowed term is not load-bearing. Cited prior work by Gloeckle et al. and Liu et al. provides external context and is not the sole support for the claimed mechanism. The compute-matching controls in Section 4.1 and Appendix B.1 are presented as empirical checks, and any objection that they understate MuToR's per-step FLOPs is a possible confound in the experimental comparison, not a circularity: the paper does not define MuToR's success in terms of those controls. The limitation acknowledged in the Conclusion, namely that register placement is not optimized, is an honest scope restriction and does not introduce a circular step. Overall, no step in the derivation chain equates an output with an input by construction.
Assumptions & free parameters
free parameters (3)
- Auxiliary loss weight (alpha) =
0.1 to 0.5 per dataset (Tables 9-10); 0.5 for images
- Maximum offset (dmax) =
3, 4, or 6 for language; dmax_2D=4 for images
- Number of register tokens =
256 (full) or 80 (sparse) in image experiments
assumptions (4)
- domain assumption A register's position id t+d-1 encodes the offset d through rotary position embeddings, so the shared register embedding can learn to predict the token d steps ahead.
- domain assumption Regular tokens attend only to regular tokens, so discarding registers at inference leaves the model's behavior identical to the base model.
- domain assumption Training baselines for extra epochs is a valid compute-matched control for the longer register-augmented sequences.
- domain assumption The auxiliary loss at register positions propagates useful gradients to the shared transformer layers and improves regular-token representations.
invented entities (1)
-
Register tokens (r_d)
Cite this review
Pith. "Pith review of Multi-Token Prediction Needs Registers." pith.science (2026). https://pith.science/paper/STJSUF57
@misc{pith2026250510518,
author = {Pith},
title = {Pith review of: Multi-Token Prediction Needs Registers},
year = {2026},
howpublished = {\url{https://pith.science/paper/STJSUF57}},
note = {Machine review of arXiv:2505.10518}
}
read the original abstract
Multi-token prediction has emerged as a promising objective for improving language model pretraining, but its benefits have not consistently generalized to other settings such as fine-tuning. In this paper, we propose MuToR, a simple and effective approach to multi-token prediction that interleaves learnable register tokens into the input sequence, each tasked with predicting future targets. Compared to existing methods, MuToR offers several key advantages: it introduces only a negligible number of additional parameters, requires no architectural changes--ensuring compatibility with off-the-shelf pretrained language models--and remains aligned with the next-token pretraining objective, making it especially well-suited for supervised fine-tuning. Moreover, it naturally supports scalable prediction horizons. We demonstrate the effectiveness and versatility of MuToR across a range of use cases, including supervised fine-tuning, parameter-efficient fine-tuning (PEFT), and pretraining, on challenging generative tasks in both language and vision domains. Our code will be available at: https://github.com/nasosger/MuToR.
Figures
Figures from the paper (1 more)
Forward citations
Cited by 2 Pith papers
-
AdaMTP: An Adaptive Training Paradigm for Multi-Token Prediction
AdaMTP uses entropy-based segmentation to adaptively mask multi-token prediction losses, improving quality and speed over fixed-horizon multi-token prediction.
-
Distilled Pretraining: A modern lens of Data, In-Context Learning and Test-Time Scaling
Distilled pretraining improves test-time scaling via generation diversity but impairs induction-head-based in-context learning, with the trade-off explained by a bigram model analysis.
Reference graph
Works this paper leans on
-
[1]
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774,
-
[5]
Imagenet: A large-scale hierarchical image database
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248–255. Ieee,
2009
-
[6]
The mystery of the pathological path-star task for language models
Arvid Frydenlund. The mystery of the pathological path-star task for language models. InProceedings of the 2024 Conference on Empirical Methods in Natural Language Processing , pages 12493– 12516,
work page 2024
-
[7]
Gemma: Open models based on gemini research and technology
10 Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based on gemini research and technology. arXiv preprint arXiv:2403.08295,
-
[8]
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783,
-
[10]
Bitune: Bidirectional instruction-tuning
Dawid J Kopiczko, Tijmen Blankevoort, and Yuki M Asano. Bitune: Bidirectional instruction-tuning. arXiv preprint arXiv:2405.14862,
-
[11]
11 Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437,
-
[12]
Decoupled weight decay regularization
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101,
Show all 24 references
-
[13]
Pass: Parallel speculative sampling
Giovanni Monea, Armand Joulin, and Edouard Grave. Pass: Parallel speculative sampling. arXiv preprint arXiv:2311.13581,
-
[14]
Randar: Decoder-only autoregressive visual generation in random orders
Ziqi Pang, Tianyuan Zhang, Fujun Luan, Yunze Man, Hao Tan, Kai Zhang, William T Freeman, and Yu-Xiong Wang. Randar: Decoder-only autoregressive visual generation in random orders. arXiv preprint arXiv:2412.01827,
-
[15]
Prophetnet: Predicting future n-gram for sequence-to-sequencepre-training
Weizhen Qi, Yu Yan, Yeyun Gong, Dayiheng Liu, Nan Duan, Jiusheng Chen, Ruofei Zhang, and Ming Zhou. Prophetnet: Predicting future n-gram for sequence-to-sequencepre-training. In Findings of the Association for Computational Linguistics: EMNLP 2020, pages 2401–2410,
2020
-
[16]
Autoregressive model beats diffusion: Llama for scalable image generation
Peize Sun, Yi Jiang, Shoufa Chen, Shilong Zhang, Bingyue Peng, Ping Luo, and Zehuan Yuan. Autoregressive model beats diffusion: Llama for scalable image generation. arXiv preprint arXiv:2406.06525,
-
[18]
Guiding language model reasoning with planning tokens
Xinyi Wang, Lucas Caccia, Oleksiy Ostapenko, Xingdi Yuan, William Yang Wang, and Alessan- dro Sordoni. Guiding language model reasoning with planning tokens. arXiv preprint arXiv:2310.05707,
-
[20]
Randomized autoregressive visual generation
Qihang Yu, Ju He, Xueqing Deng, Xiaohui Shen, and Liang-Chieh Chen. Randomized autoregressive visual generation. arXiv preprint arXiv:2411.00776,
-
[21]
In order to facilitate reproducibility, we plan to release the code for MuToR’s implementation as well
13 A Implementation Details In this section, we provide all the training details, as well as the hyperparameters used in our experi- ments. In order to facilitate reproducibility, we plan to release the code for MuToR’s implementation as well. A.1 Language Modeling A.1.1 Mathe...
2021
-
[22]
Both language models are loaded and finetuned in bfloat16 precision
We also employ a learning rate scheduler with linear decay and warmup, setting the peak learning rate to be 5e-5 for Gemma 2B and 2e-5 for Llama 3 8B. Both language models are loaded and finetuned in bfloat16 precision. To match our available resources, we filter out training ...
2024
-
[23]
Both Next-Token baseline and MuToR are trained for 360K update steps, using AdamW optimizer withβ1 = 0.9,β 2 = 0.95 and weight decay = 0.05
Unlike LlamaGen, we pre- tokenize the dataset with the ADM’s preprocessing scheme [Dhariwal and Nichol, 2021], resulting in two crops per image. Both Next-Token baseline and MuToR are trained for 360K update steps, using AdamW optimizer withβ1 = 0.9,β 2 = 0.95 and weight decay...
2021
-
[1024]
This indicates that the auxiliary loss provides valuable supervision that the model leverages during pretraining, to improve its learned representations
For MuToR’s implementation, we empirically tune the loss coefficienta to be 0.5, so thatLreg has equal contribution withLntp ( Equation 6). This indicates that the auxiliary loss provides valuable supervision that the model leverages during pretraining, to improve its learned ...
2022
-
[2017]
Classifier-free diffusion guidance
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598,
-
[2019]
Semformer: Transformer language models with semantic planning
Yongjing Yin, Junran Ding, Kai Song, and Yue Zhang. Semformer: Transformer language models with semantic planning. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 18669–18680,
2024
-
[2021]
Training verifiers to solve math word problems
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168,
-
[2023]
Memory transformer
Mikhail S Burtsev, Yuri Kuratov, Anton Peganov, and Grigory V Sapunov. Memory transformer. arXiv preprint arXiv:2006.11527,
2006 arXiv
-
[2024]
Dialogsum: A real-life scenario dialogue summarization dataset
Yulong Chen, Yang Liu, Liang Chen, and Yue Zhang. Dialogsum: A real-life scenario dialogue summarization dataset. In Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021, pages 5062–5074,
2021
-
[2025]
Llama: Open and efficient foundation language models
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023a. 12 Hugo Touvron,...
Reviewed August 15, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.