Multi-Token Prediction Needs Registers
Authors: Anastasios Gerontopoulos
Spyros Gidaris
Nikos Komodakis
Venue: NeurIPS 2025 poster ()
Keywords: autoregressive transformers
next-token prediction
multi-token prediction
register tokens
Primary area: deep_learning

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.

Multi-Token Prediction Needs Registers
Anastasios Gerontopoulos
Spyros Gidaris²
Nikos Komodakis 1,3,4
Archimedes, Athena Research Center University of Crete "IACM-Forth
valeo.ai
Abstract
Multi-token prediction has emerged as promising objective for improving lan- guage model pretraining, but its benefits have not consistently generalized to other settings such as fine-tuning.
this paper, we propose MuToR, a simple and effec- tive 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: introduces only negligi- ble number of additional parameters, requires no architectural ceranne--nnsuriii compatibility with off-the-shelf pretrained language models-ann remains aligned with the next-token pretraining objective, making it especially well-suited for su- pervised fine-tuning. Moreover, 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 available at |ttps//gtthub.com/mssssg//tuu
Introduction
Autoregressive Transformer architectures have become the cornerstone of modern Large Language Models (LLMs), enabling unprecedented capabilities across wide range of natural language process- ing tasks [Achiam et al. 2023 Liu et al. [2024]. Their success has also extended to domains such as
image generation [Esser al. 2021 Sun et al. 2024] and multimodal models (Alayrac et al. 2022
Liu et al. 2023). These models are primarily trained using a simple yet effective approach: next-token prediction with teacher forcing. By supplying ground truth tokens as context, teacher forcing enables fully parallelized computation via masked self-attention, thus accelerating and stabilizing training.
However, next-token prediction with teacher forcing has notable limitations. Models trained this way often focus on short-term patterns while struggling with harder, long-range decisions. Prior work has shown that this setup can lead to shortcut learning, where valuable supervision diminishes BBachaan and Nagarajan] 2024]. and that it underperforms on tasks requiring planning (Bubeck] These findings strongly suggest the need for training objectives that go beyond standard next-token prediction.
al. |lll||22222 2023].
To address these limitations, multi-token prediction training Qi at.
2020 Gloeckle et al. 2024
et al. 2024) has emerged as a promising alternative. Rather than predicting just one token at a time, the model learns to predict multiple future tokens at each position. Recent implementations achieve
this through additional transformer output heads; Gloeckle et al. 2024) employs parallel heads, one
for each future token position, while Liu et al. 2024 uses sequential heads. Crucially, this approach is used only during training, as its primary goal to provide more informative learning signal rather than to speed up inference. Compared to standard teacher forcing, multi-toke tokin prediction encourages the model to develop internal "planning" representations, and mitigates overfitting to local patterns. Experiments by Gloeckle et al. (2024) demonstrate that it leads to improved generative performance and increased data efficiency.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Motivated by these findings, we explore whether a more effective approach for multi- prediction can furthe enhance autoregressive transformers.
We propose a simple but powerful modification: instead of adding extra transformer layers for future tokens interleaved between the regular tokens.
token prediction, introduce register Each register token is assigned a randomly sampled offset d, and the model is trained to predict the token d steps ahead (rather than just the next token) for these tokens. The original next-token prediction objective remains unchanged for all regular tokens.
These register tokens are used exclusively during training to propagate richer supervisory signal through the model. At inference time, they are discarded to preserve generation speed. This is made possible by carefully designed attention mask: register tokens are allowed to attend only to preceding regular tokens-enable entirely blind register tokens in both directions. This ensures full compatibility with standard autoregressive inference while encouraging the model to internalize forward-looking, multi-step planning during training.
them to learn predictive representations -while regular tokens
Compared to approaches that add output heads or transformer blocks, our register-based method offers several key advantages:
No architectural changes:
Only small number of additional trainable parameters are introduced through register embeddings, while the core transformer layers remain untouched and no extra output heads are required.
Fine-tuning compatibility:
Our method is particularly well-suited for fine-tuning pretrained LLMs (e.g.,Llama Llama GGrattafiori et al. 2024], Gemma [Gemma Team et al. 2024)) ). introduces minimal parameter overhead, preserves original attention patterns for regular tokens, and uses carefully selected position ids and attention masks for register tokens boingngbringing multi-token prediction
closer the next-token pretraining setup. In contrast, previous methods [Gloeckle et al. 2024
Liu et al. [2024] rely on separate transformer heads, adding many new parameters that must be trained from scratch, making them less effective in fine-tuning scenarios.
Scalable prediction horizons: Because the number of register tokens remains fixed regardless of the offiet the training cost independent of the prediction horizon, which can be scaled arbitrarily. Register tokens thus provide greater flexibility in future token prediction. For example, autoregressive image generation, our method naturally extends to predicting tokens in a two-dimensional neighborhood a capability not easily achieved by adding output heads.
Overall, our work delivers the following contributions:
We introduce MuToR (Multi-Toker prediction with Registers), a novel multi-token prediction method that employs trainable, interleaved register tokens tasked with predicting multiple future targets. MuToR enables scalable prediction horizons with minimal additional parameters and seamlessly integrates with any pretrained autoregressive language model without architectural modifications.
Through experiments on language modeling benchmarks, we validate the effectiveness of MuToR in both supervised fine-tuning and parameter-efficient fine-tuning (PEFT) settings, consistently surpassing standard fine-tuning baselines under equivalent training compute.
We further demonstrate the versatility of MuToR by applying it to autoregressive image gener- ation a pretraining setting, where improves performance over standard teacher-forcing highlighting its broader potential across diverse domains and training settings.
2
Related Work
Limitations of Next-Token Prediction Bachmann and Nagarajan (2024) introduce a path-finding task on star-graphs to highlight key limitations of standard next-token prediction (i.e., teacher forcing). Their findings reveal that next-token prediction encourages shortcuts, making the underlying task intractable and reducing validation accuracy to random random-guessing levels. Interestingly, this "cheating" behavior can be mitigated by multi-token prediction with lookahead embeddings, as in Monea et al. [023]. While the task is an extreme case, alternative transformer architectures can solve it [Frydenlund] 2024), reinforcing the intuition that tasks requiring planning might need tailored training
objectives [Bubeck et al. 2023

---

Several works have explored decoding multiple future tokens. primarily to accelerate inference rather than improve generation quality [Stern al. 2018 Monea et al. 2023 Li et al., 2024 Cai et al. 2024]. In contrast, Qi et al. [2220] propose multi- token prediction pre-training objective that enhances performance on some sequence-to-sequence tasks. However, their method scales poorly to large decoder models because their multi-stream attention mechanism becomes computationally expensive as prediction depth increases. More Pecently.Gloeckle Gloeckle et al. (2024) proposed an architecture with multiple parallel decoding heads for multi-token prediction, leading to better generative performance (mostly in coding tasks). Liu et al. 2024) modified this approach to sequential decoding heads instead. Both studies suggest that multi-token prediction helps by providing richer supervision, better information sharing between tokens, and implicit planning in hidden states. However, these benefits mainly observed during pretraining, with limited success in fine-tuning scenarios-a gap our work addresses.
Multi-token and Lookahead Prediction
A less related line of research trains autoregressive models on permuted sequences Yang et al. 2019 Pannatier et al., 2024 Kakogeorgiou et al. 2024 Yu et al. 2024 Pang et al., [024). By predicting next tokens in shuffled order, these methods force the model to recover distant dependencies in the original sequence. While conceptually interesting, they differ significantly from our approach.
Trainable Dummy/ / Register Tokens Recent work investigates the use of trainable tokens as transformer inputs, revealing emergent properties. [Burtsev et al. 2020) prepend dummy tokens to prompts to induce working memory, but observe minimal gains. Goyal et al. [2024) demonstrate that appended trainable tokens can improve performance by increasing operations between tokens, thus encouraging deeper "thinking". Pfau et al. (022) study the conditions under which language model can leverage the extra computation provided by the dummy tokens. Related efforts incorporate
planning tokens trained to predict latent reasoning steps (Wang al. 2023 Yin et al. 2024]. In
the vision domain, Darcet et al. [2024] employ register tokens during pretraining to prevent the appearance of high-norm artifacts. Similarly to these works,
learnable tokens but also
associate them with an auxiliary multi-token prediction objective to create a denser training signal.
More related to our work, Monea et al. 2023 appends lookahead tokens to the input sequence and trains them (while freezing the base model) to enable parallel decoding for speculative sampling. Bachmann and Nagarajan (2024) adapt this idea to path-finding on star-graphs, by simultaneously predicting multiple answer tokens from a prefix. While effective for these specific graph problems, the approach's 's complete "teacherless" training paradigm and parallel inference requirements make it fundamentally unsuitable for the broader range generative tasks we consider. Unlike these methods, approach only uses register tokens during training to enhance supervision, without modifying the inference procedure or requiring additional compute at test time.
Method
3.1 Preliminaries
Next-Token Prediction Building on the foundational work ohannon [1448] 1951], next-token pre- diction remains the core objective for autoregressive language models. Given sequence (#1,..., ......... 22). the model is trained to predict the next token r++1 by maximizing the joint probability under left-to- right factorization:
P .T) = II P( xgs))
(I)
where 7 is the sequence length. For model Po and dataset D, the training objective will be to minimize the the expected negative log-likelihood loss over the dataset:
[-
Cotp ED
Σ log Po (x (gg))
Multi-Token Prediction In contrast, Gloeckle et al. (2024) propose predicting multiple future tokens per position. Their objective minimizes:
Cmpp = Ep
Σ log Po Pa xCS)
(3)
3

---

Targete 32
Transformar
Transi ornee
VS
Inpats
#2
Position ids
(1 (1++4) -1 (2
- 4)+4)
Next-token prediction
Multi-tozen prediction with registers
Figure I: Next-token prediction vs. Multi-token prediction with registers (MuToR), The trans- former block represents any decoder-only autoregressive model, with colored lines indicating per- mitted attention connections between tokens. Left: Standard next-token prediction, where each It predicts 21+1 conditioned on preceding tokens. Right: MuToR interleaves register tokens ra to predict tokens d steps ahead (Iz+a), conditioned only previous regular tokens. Register tokens assigned position ids (e.g., + d for ra targeting T+++a) that mimic next-token prediction. Regular tokens follow the standard next-token prediction formulation, unaffected by the registers.
where we denote the maximum prediction horizon as dmaax to align notation with our method. Recent implementations utilize additional transformer heads either parallel Gloockle al. 2024 or sequential Liu et al. 2024
3.2 Our approach
In this work, we introduce MuToR, an alternative multi-token prediction method (illustrated in Figure T). Our approach inserts interleaved learnable tokens-termed registers, following DDrcet et al. 2024] - training sequences, where each register is tasked with predicting future token at uniformly sampled offset d. By optimizing this auxiliary prediction objective alongside the primary next-token prediction task, the model benefits from richer supervisory signals, which enhances the quality of learned representations.
Let (I1, ..... 2T) be training quence. We augment x by interleaving register tokens Ta. yielding:
Register Tokens
F4
F2
Fd
d' d (dd, ................................ 17 IT);
where each ra predicts the future token at offset By default, all ra share a single learnable embedding, adding minimal train- able parameters, while the target offset d is specified via the register's 's position id (detailed later). The augmented sequence x' is processed by the transformer Pp. which shares all compo- nents enieddinbs,,,ggs, layers, and prediction head betwenn the regular tokens I, and register tokens Td-
AF2
Figure MuToR's attention mask. Each cell indicates whether the row can attend to the column.
Attention Masking We modify the causal attention mask satisfy two conditions: (i) regular tokens cannot attend to any register tokens, and (ii) register tokens cannot attend other register tokens in the sequence (see aeeFFigure>> 2}. This pattern preserves the standard next-token prediction objective (Equation 2) for regular tokens, as their representations remain unaffected by registers. As result, the registers can be discarded during inference.
Multi-token Prediction with Registers Each register token ra inserted after Ia predicts the future token Ftta with offset d sampled uniformly per sequence from (1...., dmax), where dmax
'The interlcaving pattern is flexible-c.g. registers may sparse or appear consccutively. For simplicity, fixed d per sequence, though registers with different offsets could be mixed.

---

determines the maximum prediction horizon. The auxiliary register loss over dataset D is:
reg = ED Σ log Po xge,
(5)
Our attention masking ensures predictions depend only preceding regular tokens, excluding other registers. This approach enables flexible prediction horizons while maintaining parameter efficiency, all register tokens regardlsss of share single learnable embedding.
Position Embeddings for Registers Since we do not use specialized heads for future token prediction, we encode prediction offsets through positional bias. We utilize the token indices in the original (register-fire -feee) sequence to compute positional embeddings. While regular tokens 2₁ keep their natural positions each register ra inserted after 3y (predicting #2++d) receives position t d (see position ids in FFiuure 1). This matches the position id of the regular token that would normally predict Fe+d under standard next-token prediction.
Our implementation focuses on RoPE [Su et al. 2024), the dominant positional encoding modern language [Touvron et al. 2023b and autoregressive image models (Sun et al. 2024). However, our position manipulation works with any embedding scheme (sinusoidal, relative, etc.), requiring no architectural changes while effectively encoding prediction offsets through positional bias.
Overall Training Loss We jointly optimize the standard next-token prediction loss Cutp and the auxiliary register loss Lreg. The overall loss combines these two objectives through weighted sum:
Cmtp = (1 - -a) + aLreg
(6)
where (0, 1) controls the relative contribution of each loss term.
Inference During inference, we discard register tokens entirely, leaving the model's computational graph and latency unchanged. This is made possible by our attention masking strategy, which prevents regular tokens from ever attending to registers during training. Unlike prior approaches (Goyal et al.
(2024), Pfau et al. (2024) ) that inserted tokens increase inference computation, method maintains the standard autoregressive process without modification.
3.3 Adaptation in Language Modeling
A key application of our method is supervised fine-tuning of pre- trained language models. For generative tasks (e.g., mathematical reasoning), where datasets contain (prefix, answer) pairs-with an- swer sequences potentially including chain-of-thought tokens-we we interleave register tokens only within the answer sequence. This aligns with standard practice where prefix predictions are excluded from loss computation. Beyond this adaptation, the method fol- lows the implementation described in Isubsection 3.2 without mod- ification.
21
#
X4
Zg
F2
x
3.4 Adaptation in Autoregressive Image Generation
Figure 3: The 2D neighborhood of possible prediction targets (depicted in red) for register token. The register ra inserted after 27, and dmax 2D is set to 3.
Autoregressive transformers achieve strong performance on image generation by modeling discrete visual tokens [Esssr et al. 2021 Sun et al. 2024]. A VQ- VAE tokenizer [an Den Oord ct al. 2017 first encodes an image into 2D token grid zn e x which then flattened into a ID sequence x via raster-scan ordering.
The model then learns autoregressive prediction conditioned e (either class labels or captions).
We adapt MuToR to images by modifying the offset sampling to respect the 2D image structure. For each sequence xr we sample a 2D offset pair (da, d, with both d, and d,, drawn uniformly from (1...., dmax 2D). excluding (da, (d,)) (1, 1), it denotes the image token after which the register inserted (see Figure 3p. We then compute the rasterized offset as d = (dh - 1) + d,
in subsection 3.2 each register token ra predicts the token d steps ahead in the sequence, with all other
5

---

components (attention masking, positional embeddings, and loss) implemented identically- excepp for the additional conditioning on c. Each register thus predicts one of de mD possible future tokens.
This 2D extension enriches the training signal by capturing spatial dependencies inherent in visual data, while requiring minimal architectural changes. Unlike prior multi-token prediction approaches that would require multiple additional prediction heads (one for each possible 2D offiset), MuToR achieves this capability with negligible parameter overhead.
4 Results
4.1 Language Modeling
Experimental Setup We focus mathematical reasoning with chain-of-thought and abstractive
summarization, two challenging generative tasks that provide a rigorous testbed for our method.
To evaluate performance, we fine-tune two pretrained decoder-only language models: Gemma 2B Gemma Team et al.) 2024] and Llama 3 8B [Grattafiori et al. [022]. Our experiments target three widely used mathematical reasoning benchmarks: GSM8K Cobbe et al. 2021], MATH500
Lightman et al. 2023), and AQUA-RAT [Ling et al. 2017]. For fine-tuning, we use curated subsets from OpenMathInstruuee
[oshniwal al., 2025]. a high-quality dataset derived from GSM8K and MATH training samples. Specifically, we filter the IM and 2M splits to isolate grade school (GSM-style) or MATH-style problems, referring to them as /M-GSM, 2M-GSM, and IM-MATH. We also fine-tune on the original GSM8K and AQUA-RAT training splits. Additional details about the experimental setup are provided in Appendix A.I As for summarization, we target the following benchmarks: SAMSum [Gliwa et al., [219)] and DialogSum [Chen et al. [021]. For mathematical
tasks, we measure exact-match accuracy; for summarization, we use ROUGE scores Lin 2004).
Baselines We mainly consider two baselines: (i) Next-Token, the standard fine-tuning recipe using the next-token prediction objective, and (ii) Multi- [Gloeckle et al. 2024], adapted for fine-tuning by adding dmax 1 auxiliary prediction heads and applying loss- loightiig strategy for these heads, similar to our method. To ensure a fair comparison, for both Multi-Toker and MuToR, tune the number of predicted future tokens, dmaxx alongside the auxiliary loss coefficient. Implementation details are provided in Appendix A.
Comparative Results Tables [ndd]] 2 present our results, showing only the best configurations for Multi-Token and our MuToR method. In mathematical reasoning (Table D. our approach consistently outperforms both baselines, including Multi-Token which introduces a substantial number of additional trainable parameters (see Table T. Moreover, MuToR's gains are preserved across varying training set sizes, demonstrating its effectiveness even in settings with high-quality fine- tuning data. In contrast, Multi- 's benefits seem to diminish with larger models or more data. For summarization (Table 2). our MuToR method improves all ROUGE scores, achieving superior results over Multi-Token, demonstrating broad applicability in sequence-to-sequees generative tasks.
Comparison with DeepSeek's Sequential Multi-Token Prediction We further compare MuToR against the multi-token prediction variant (Liu et al. 2024], that uses sequential heads to predict up to damax tokens ahead while maintaining the causal chain of prediction. The results, presented in Table 16 (Appendix C.I, indicate that DS-Multi tellti--oooeee Toker, offers only marginal improvements over the parallel Mfulti-Tokee variant Glloeckle et al. (004). , despite its greater architectural complexity. Importantly, MuToR outperforms both methods, demonstrating better adaptability on the supervised fine-tuning setting. As both multi-token prediction variants yield similar performance, we stick with Multi-Toker for most of the experiments and ablations that are included in this paper.
Matching the Training-Time Compute Our method increases sequence length during training (and thus training compute) by inserting register tokens. To ensure gains are not due to higher compute alone, we train the baselines for more epochs (roughly doubling compute), thus matching or even
For brevity, we refer to the sequential multi token prediction as DS-Multa --okee,, since it was first used when training Deepseck-V? V3.

---

Downstream accuracy (%) in mathematical reasoning benchmarks. The subheaders refer to the training split used each experiment. Results for Gemma 2B are averaged over seeded runs. MuToR offers consistent improvement over both standard Next-Tokee and Multi-Token.
Table 1:
GSMSK
MATH500
AQUA-RAT
Model
Method
GSM8K
IM-GSM
2M-GSM
IM-MATH
AQUA-RAT
Next Token
38.87
66.09
69.02
26.73
40.16
Gemma 2B
Token
40.66
66.69
69.02
26.87
38.45
MuToR (ours)
42.10
68.33
70.56
28.13
41.73
Next Text-Tooke Token
66.41
85.74
87.11
41.4
Ilama3 8B
Tokenn
66.56
85.67
86.35
42.6
MuToR (ours)
67.85
87.11
87.64
43.4
Table 2: Experimental results for fine- fine-tuning Gemma 2B on abstractive summarization benchmarks. We select the checkpoint with the higher ROUGE-L in the validation set, and report ROUGE scores on the test set.
SAMSum
DialogSum
Method
ROUGE-1
ROUGE-2
ROUGE-L
ROUGE-1
ROUGE-2
ROUGE-L
Next-Token 51.47
27.29
43.23
47.23
20.91
38.77
Multi-Token
51.90
27.44
43.50
47.98
21.23
39.25
MuToR (ours) 52.32
28.08
44.09
48.22
21.71
39.48
exceeding the compute of MuToR. As shown in Tables 12 and 13 (Appendix Bb, extending training does not yield further improvements for the basclines and, in some instances, even degrade performance due to mild overfitting. These results suggest that integrating MuToR into the fine-tuning pipeline offers a more effective approach to leveraging increased training compute, when the available data is fixed common real-world constraint in many academic and industrial applications.
Matching Compute with Additional Data Through our experiments, we focus on the data-constrained set- ting both for practical reasons (limited resources) and methodological clarity (ensuring fair comparison with baselines). However, also investigate the possibility of matching MuToR's increased training-time compute by leveraging additional data for the Next-Token base- line. The result (presented in Table shows that MuToR can surpass gains from additional data. While a thor- ough exploration of the data-abundant regime is left for future work, this preliminary finding highlights MuToR 's potential beyond data-constrained scenarios. More details are provided in the Appendix B.
Table 3: Downstream accuracy (%)
GSMSK. IM-GSM-extra split corresponds to matching MuToR's training compute by utilizing additional fine-tuning data.
Method & Data split
GSM8K
Next-Token (IM-GSM)
66.09
Next-Token (IM-GSM-extra)
66.69
MuToR (IM-GSM)
68.33
Integration in Cerameter-Efficiien Fine-tuning We test our method with LoRA Hu et al. 2022]: for both Next-Tokee and our MuToR method, we apply rank-32 adapters all linear layers (approximately 39M trainable parameters for Gemma 2B; register tokens in MuToR add a negligible number of parameters). As shown Table our LORA-MuTOR approach improves accuracy over standard LoRA fine-tuning across both training splits. Interestingly, MORA-MuTE matches or even exceeds the full fine-tuning Next-Token performance, demonstrating its utility in PEFT setups. In contrast, the Multi-Token approach is less compatible with PEFT settings, as it requires training several additional transformer layers from scratch.
Impact of Maximum Lookahead Offset The key hyperparameter dmax in our MuToR method controls how many tokens ahead the registers predict. Larger values offer richer supervision but increase task difficulty. Experiments (Table 4 show that dmax = 4 is optimal for this particular setting (in general the optimal value may depend on the training data and the downstream task). Notably, as dmax increases, Multi- Token's performance degrades, barely beating Next-Token
7

---

Table 4: Downstream accuracy (%) with respect the maximum offset dmax. using Gemma 2B. #Add. Param. denotes the additional trainable parameters for Multi-Tokeen and MuToR (in up- proximation).
Table 5; Downstream accuracy (%) in a PEFT scenario, using Gemma 2B and LoRA.
GSM8K
Method
GSMSK
IM-GSM
Full fine-tuning Next-Token
38.87
66,09
#Add. Puram.
GSM8K
dmax
Method
LoRa-Next: -oken
36.34
66. 11
GSM8K
IM-GSM
LoRA-MuToR (ours)
38.59
68.11
Next- Token
38.87
66.09
110M
Multi Taken
40.66
66.69
Table 6: Ablation regarding the register embed- dings, using Gemma 2B and dmax = 4.
2
(ours)
67.15
220M
Multi-Toke Token
40.59
66.36
3
NuToR (ours)
41.60
68.01
Register embedding
GSMSK
330M
Multi Token
39.78
65.53
2K
NuToR (ours)
42.10
68.33
GSM8K
IM-GSM
550M
Kulti-Token
40.23
65.55
Same
42.10
68.33
2K
NuToR (ours)
39.90
68.16
Different
41.85
68.18
when using the IM-GSM training split. In comparison, MuToR's gains maintained across dmaa values and training split sizes.
Shared vs. Different Register Embeddings Per Offset We test whether having different register embeddings per offset improves performance. As seen in Table shared embeddings (Same), which the default, perform slightly better than distinct ones (Different), suggesting our positional encoding scheme provides sufficient offset information.
4.2
Autoregressive Image Generation
Experimental Setup
We train LlamaGen-B (111M parameters; Sun et al.
2024) on ImageNet
[Deng et al. 2009] at 256x256 resolution, using ADM's preprocessing pipeline Dhariwal and Nichol
2021]
The dataset is pre-tokenized using VQ-VAE tokenizer from Sun et al. [0244]
We compare three approaches: (1) Next-Token a standard next-token prediction baseline; (2)
MuToR-1D 1D with ID offsets (as in language modeling); and (3) MuToR with 2D offsets (described
subsectionnn 3.4). Implementation details are included in Appendix A.2
For evaluation, we generate 50,000 images using classifier-free guidance (scale=2.0) (Ho and Sali? mans 2022] and compute FID [Heusel et al. [2017]. IS (Salimans et al.) [016]]. Precision, and Recall Kynkäänniemi et al. 2019 using TensorFlow scripts from Dhariwal and Nichol [2021]
Results Table shows that both MuToR variants consistently outperform Next -okke in FID and IS across different training iterations. Notably, when comparing under similar training-time compute, the MuToR-20 2D variant at 100K steps surpasses the Next -oken model at 200K steps, demonstrating both improved performance and faster convergence.
Extending the offset to 2D The 2D extension in MuToR 2D significantly improves performance by leveraging spatial dependencies in the image data (Table 7). despite requiring prediction of much more possible future tokens. This demonstrates that the 2D formulation effectively captures structural patterns, providing richer training signal.
Scaling down the number of registers during pretraining To reduce computational costs while maintaining performance, we investigate using fewer register tokens in MuToR-2D. Specifically, we test inserting only 80 randomly placed registers per image, increasing the sequence length by just 30%. As shown in Table 8) this sparse version achieves very similar performance to the full setup (with 256 registers) while requiring less computation. These results demonstrate that substantial performance gains can be obtained with relatively few register tokens and only a modest increase in training compute. They also highlight untapped potential MuToR's design, particularly regarding optimal register placement and sparsity strategies, which merit further investigation.
8

---

Table 7: Conditional generation performance on Imagenet 256 256 (cfg scale 2.0). Both dmax and d duax 2D are set to 4, for MuToR-1D and MuToR-22 respectively.
Iter.
Method
FID
IS
Pre. f
Rec.
Next-Token
7.71
146.5
0.830
0.439
100K
MuToR-1D
7.01
155.6
0.828
0.438
MuToR-2D
6,57
163.2
0.831
0.445
Text-Tokeet
6,83
158.4
0.833
0.443
200K
MuToR-1D
6,43
163.0
0,836
0.444
MuToR-2D
5.65
183.5
0,842
0.448
6.18
171.5
0.841
0.443
360K
MuToR-1D
5.79
178.3
0.841
0.441
MuToR-2D
5.09
195.3
0.839
0.457
Maximum Offset Analysis We examine how expanding the prediction neighborhood in MuToR-2D affects performance by varying dmax_2D. 2D. which determines how many future tokens each register must predict. Using 80 randomly placed registers per sequence, test different dmax__D values. achieves optimal performance, while increasing it leads to degrading
Table shows that 2D results. In this optimal setup, each register predicts up 15 future tokens prediction horizon that would require 14 additional transformer heads in prior multi-token approaches, making them computationally impractical. This demonstrates MuToR-2D's unique ability to effectively leverage long-range predictions while maintaining training efficiency.
Table 9: Ablation using MuToR 2D, with varying dmax 2D and 80 registers inserted in random positions. Targets denote the number of future tokens (d2 max 2D
Table Ablation using MuToR-2D. 2D, dmax 2D 4, and varying number of registers (# Reg.) inserted in ran-
1) that the
registers predict during training (see Figure 3
dom positions.
Iter.
diaa 2D
# Targets
FID
IS
# Iter.
# Reg.
FID
IS
t
7.20
150.8
100K
6.87
162.1
80
6.87
162.1
100K
5
7.08
154.8
256
6.57
163.2
6.19
169.0
80
5.89
175.6
200K
200K
15
5.89
175.6
256
5.65
183.5
5
24
6.21
166.2
80
5.03
195.6
360K
3
5.56
178.3
256
5.09
195.3
360K
4
15
5.03
195.6
5
24
5.53
179.2
4.3 Synthetic Data
Next-Token
MuToR
GPT2-Large
We further evaluate MuToR on the star- star-graph path- finding problem [Bachmann and Nagarajan] 2024).
Accuracy
which highlights limitations of the next-token predic- tion objective. Consider a directed graph G (n,,). where n denotes the number of paths emanating from the start node, Ustart and / denotes their length. Given an end node, uend. the model must identify the unique path from Wstart to Wend-
G(2.5)
G(5.5)
G(10.5)
G(20.5)
G(2110)
Figure 4: Solve rate (%) finetuned GPT2-L model on different star graph configurations.
Despite the task's 's simplicity, transformer models trained by teacher forcing fail to solve it, due to shortcut learning and the loss of meaningful training signal. We thus investigate whether the look- ahead induced by MuToR's objective can provide necessary supervision to learn underlying task. To this end, we experiment with fine-tuning pretrained GPT2 model [Radford et al. 2019 following the setup from Bachmann and Nagarajan 2024]. The relevant implementation details are provided in the Appendix A.3

---

Results
The results are presented in Figure 4 MuToR solves the task across various graph configu- rations, effectively overcoming the "cheating phenomenon" " that causes standard teacher foreing to fail. These findings indicate that MuToR can be particularly effective in scenarios where some form of shortcut learning applies, recovering valuable training signal.
Conclusion
We introduced MuToR, simple yet effective approach to multi-token prediction that leverages interleaved, trainable register tokens to predict future targets. MuToR enables scalable prediction horizons with minimal parameter overhead and is fully compatible with existing pretrained model architectures, allowing seamless integration into standard fine-tuning tinettuning pipelines. Empirical results demonstrate that MuToR consistently improves performance in both language modeling tasks-such as mathematical reasoning and summarization and autoregressive image generation, highlighting its versatility across modalities. This positions MuToR promising foundation for using roken-based lookahead mechanisms to propagate richer supervisory signals during training.
Limitations It is worth noting that MuToR currently uses uniformly interleaved or randomly po- sitioned register Sokens-Strategies that may not align optimally with the structure or semantics of specific tasks. While this simple placement scheme has proven effective across modalities, it leaves room for substantial improvement. By learning or adapting the placement of register to- kens potentially guided by model uncertainty task-specific priors MuToR could deliver more targeted supervision with fewer auxiliary tokens, further enhancing efficiency and performance.
Acknowledgments and Disclosure of Funding
This work has been partially supported by project MIS 5154714 of the National Recovery and Resilience Plan Greece 2.0 funded by the European Union under the NextGenerationEU Program. Hardware resources were granted with the support of GRNET.
References
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et Gpt-4 technical report. arXiv preprint 2:xXx:::333.87744 2023.
Jean-Baptiste Alayrac. Jeff Donahue. Pauline Luc, Antoine Miech, Iain Barr, Yana Hasson, Karel Lenc, Arthur Mensch, Katherine Millican, Malcolm Reynolds, et al. Flamingo: a visual language model for few-shot learning. In NeurlPS, 2022.
Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls of next-token prediction. In Proceedings of the 41st International Conference on Machine Learning, pages 2296-2318, 2024.
Sébastien Bubeck, Varun Chadrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, et al. Sparks of artificial general intelligence: Early experiments with gpt-4, 2023.
Mikhail Burtsev, Yuri Kuratov, Anton Peganov, and Grigory V Sapunov. Memory transformer. arXiv preprint arXiv:2006. 1527, 2020.
Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D Lee, Deming Chen, and Tri Dao. Medusa: Simple 11m inference acceleration framework with multiple decoding heads. In Proceedings of the 41st International Conference Machine Learning, pages 5209-5235, 2024.
Yulong Chen, Yang Liu, Liang Chen, and Yue Zhang. Dialogsum: A real-life scenario dialogue summarization dataset. In Findings the Association for Computational Linguistics: ACL-IJCNLP 2021, pages 5062-5074, 2021.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Rejichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint 2rXiv:2110.14166 2021.
10

---

Timothée Darcet. Maxime Oquab, Julien Mairal, and Piotr Bojanowski. Vision transformers need registers. In The Twelfth International Conference on Learning Representations, 2024.
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 leee, 2009.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-8794, 2021.
Li Dong, Nan Yang, Wenhui Wang, Furu Wei, Xiaodong Liu, Yu Wang, Jianfeng Gao, Ming Zhou, and Hsiao- Wuen Hon. Unified language model pre-training for natural language understanding and generation. Advances in neural information processing systems, 32, 2019.
Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high- high-resolution image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 12873 11883,, 2021.
Arvid Frydenlund. The mystery of the pathological path-star task for language models. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 12493- 12516, 2024.
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.08225, 2024.
Bogdan Gliwa, Iwona Mochol, Maciej Biesek, and Aleksander Wawer. Samsum corpus: A human- annotated dialogue dataset for abstractive summarization. Proceedings of the 2nd Workshop on New Frontiers in Summarization, pages 70-79, 2019.
Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, and Gabriel Synnaeve. Better & faster large language models via multi-token prediction. Proceedings of the 41st International Conference on Machine Learning, pages 15706-15734, 2024.
Sachin Goyal, Ziwei Ji, Ankit Singh Rawat, Aditya Krishna Menon, Sanjiv Kumar, and Vaishnavh Nagarajan. Think before you speak: Training language models with pause tokens. In The Twelfth International Conference on Learning Representations, 2024.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, er al. The llama 3 herd of models. arXiv preprint arXiv:2407.2178. 2024.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Mcasuring mathematical problem solving with the math dataset. In Thirty-fifth Conference Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598 2022.
Jordan Hoffmann. Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 15556, 2022.
Edward J Hu, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2022.

---

Ioannis Kakogeorgiou, Spyros Gidaris, Konstantinos Karantzalos, and Nikos Komodakis. Spot: Self- training with patch-order permutation for object-centric learning with autoregressive transformers. In Proceedings of the IEEE/CVE Conference Computer Vision and Puttern Recognition, pages 22776-22786, 2024.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.0836/ 2020.
Dawid Kopiczko, Tijmen Blankevoort, and Yuki M Asano. Bitune: Bidirectional instruction-tuning. arXiv preprint arXhv:2405.14862 2024.
Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. Advances in neural information processing systems, 32, 2019.
Yuhui Li. Fangyun Wei, Chao Zhang, and Hongyang Zhang. Eagle: speculative sampling requires rethinking feature uncertainty. In Proceedings of the 41st International Conference on Machine Learning, pages
2024.
Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. In The Twelfth International Conference on Learning Representations, 2023.
Chin-Yev Lin. Rouge: A package for automatic evaluation of summaries. In Text summarization branches out, pages 74, 2004.
Wang Ling, Dani Yogatama, Chris Dyer, and Phil Blunsom. Program induction by rationale gen- eration: Learning to solve and explain algebraic word problems. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1; Long Papers), pages 158-167, 2017.
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 1:X1v:2412.19437, 2024.
Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee. Visual instruction tuning. In NeurIPS, 2023.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. 1rXiv:1711.05101, 2017.
arXiv preprint
Giovanni Monea, Armand Joulin, and Edouard Grave. Pass: Parallel speculative sampling. arXiv preprint arXiv:2311.1358., 2023.
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.00827,, 2024.
Arnaud Pannatier, Evann Courdier, and François Fleuret. o-gpts: A new approach to autoregressive models. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pages 143-159. Springer, 2024.
Jacob Pfau, William Merrill, and Samuel R Bowman. Let's think dot by dot: Hidden computation in transformer language models. First Conference on Language Modeling, 2024.
Weizhen Qi, Yu Yan, Yeyun Gong, Dayiheng Liu, Nan Duan, Jiusheng Chen, Ruofei Zhang, and Ming Zhou. Prophetnet: Predicting future n-gram for wetpreeee---eeueeneeee--eeee the Association for Computational Linguistics: EMNLP 2020, pages 2401-2410 2020.
Findings of
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog. 1(8):9, 2019,
12

---

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140):1-61 2020.
Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
Claude E Shannon. A mathematical theory of communication. The Bell system technical journal, 27 (3):379-423, 1948.
Claude E Shannon. Prediction and entropy of printed english. Bell system technical journal, 30(1): 50-64, 1951.
Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit. Blockwise parallel decoding for deep autore- gressive models. Advances in Neural Information Processing Systems, 31, 2018.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
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 preprini arXiv:2406.06525, 2024.
Shubham Toshniwal, Wei Du, Ivan Moshkov, Branislav Kisacanin, Alexan Ayrapetyan, and Igor Gitman. Openmathinstruct-2:: Accelerating ai for math with massive open-source instruction data. In The Thirteenth International Conference on Learning Representations, 2025.
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.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023b.
Aaron Van Den Oord, Oriol Vinyals, al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
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, 2023.
Zhilin Yang, Zihang Dai, Yiming Yang. Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. Advances in neural information processing systems, 32, 2019.
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.
Qihang Yu, Ju He, Xueqing Deng, Xiaohui Shen, and Liang-Chich Chen. Randomized autoregressive visual generation. arXiv preprint arXiv:441.00077, 2024.
13

---

A
Implementation Details
In this section, we provide the necessary training details, as well as the hyperparameters used in our experiments.
A.1
Language Modeling
A.1.1 Mathematical Reasoning
Datasets GSM8K comprises approximately 8.7K grade-school level math problems, with around 1.3K of them forming the test set. On the other hand, MATH500 test set consists of 500 problems. uniformly sampled from the original MATH test set [Hendrycks al., [022]. covering more advanced and diverse mathematical domains, Finally, AQUA-RAT includes multiple-choice math problems with chain-of-though solutions, split among the training set (~97K samples), the validation (~250 samples) and the test set (~250 samples).
As mentioned in subsection 4.1 we finetune our base models on standard downstream train- ing datasets, such as GSM8K and AQUA-RAT, as well as on curated subsets derived from Daphnnaatlnttruutt22 The utilized training splits are listed below, including information about the employed filtering and the amount of samples:
GSM8K training set (~ 7.4K samples),
IM-GSM split (~152K samples), obtained by filtering OpenMaahhnsttrct225 1M split for grade school math-like problems,
2M-GSM split (~277K samples), similarly derived from OpenMatthnsttuut-22s 2M split using the same grade-school filtering criteria,
IM-MATH split (~200K samples), constructed by filtering OpenMathlnsttuutt22s IM split for MATH-style problems and randomly sampling 200K of them,
AQUA-RAT training (~~9K samples).
Our filtering depends on the source dataset of each sample, which is included OpenMaathlsstrutt22'' metadata.
Training details Throughout all experiments and for all methods, we use bidirectional attention among the prefix tokens, proposed in previous works [Dong et al., 2019 Raffel et al. 2020
Kopiczko et al. 2024]. since it benefits prefix-asswe tasks. We finetune all models for 5 epochs, using AdamW W optimizer [Loshchilov and Hutter [017] without weight decay and batch size of 10. We also employ a learning rate scheduler with linear decay and warmup, setting the peak learning rate be 5e-5 for Gemma 2B and 2e-5 for Llama 3 8B. Both language models loaded and finetuned in b precision. To match our available resources, we filter out training sequences that are longer than 512 token Moreover, all experiments with the 2B model are conducted using three random seeds to ensure statistical robustness. This setting is not replicated for the 8B model though, due to the substantial computational resources that are required.
All experiments with Gemma 2B model are run using a single A100 GPU and gradient accumulation. For the experiments with the Llama 3 8B model, utilize 5 x A100 and Fully Sharded Data Parallelism (FSDP).
Evaluation During evaluation, assess performance by using greedy decoding and applying exact match techniques to verify the correctness of the generated answer. For experiments on GSM8K and MATH500, the best- perffrrnnnn checkpoint from each run is then used for comparisons. For experiments on AQUA-RAY we choose the checkpoint with the highest accuracy on the provided validation set. Since register tokens are not used during inference, the inference process identical across both our method and the bascline approaches.
'In IM-MATH, we keep sequences up to 768 tokens.
14

---

A.1.2 Abstractive summarization
Datasets In experiments, we target SAMSum and Sim,ogSum, two widely used dialogue summarization benchmarks. SAMSum consists of approximately 16K messenger--iks conversations with summaries. They are split between the training set (~14K samples), the validation set (~818 samples) and the test set (~819 samples). On the other hand, DialogSum contains approximately pairs, focusing more daily-life formal conversations, such business negotiations. These pairs are split between the training set (~12.5K samples), the validation set (~500 samples) and the test set (~1.5K samples).
14K conversation-sumnsn
Training details Across all experiments, finetune the model for 3 training epochs, following the same setup (optimization details, weights' precision), with the mathematical reasoning experiments. We filter out training sequences that longer than 768 tokens. Since these experiments are conducted with the 2B model, we utilize a single A100 GPU and gradient accumulation.
Evaluation We calculate ROUGE scores against the ground truth reference summaries, select the best checkpoint with respect ROUGE-L on the validation set, and report ROUGE-1, ROUGE-2 and ROUGE-L scores on the test set.
A.1.3 Best performing configurations
In Tables I0 and [1 provide the best performing configurations (dmax and for MuToR, across all the training splits that were utilized in our experiments. For Multi-Token after carefully tuming the same hyperparameters, we found that the optimal configuration across all experiments is dmax = 2, = 0.1.
Table 10: MuToR's best performing hyper- parameters for Gemma 2B model.
Table 11: MuToR's best performing hyper- parameters for Llama 8B model.
Training data dmax
GSM8K
0.3
Training data
dmax
IM-GSM
0.3
GSM8K
4
0.3
2M-GSM
0.1
IM-GSM
4
0.3
IM-MATH
0.3
2M-GSM
6
0.1
AQUA- RAT
4
0.3
1M-MATH
4
0.3
SAMSum
0.5
DialogSum
0.5
A.2 Autoregressive Image Generation
Training details In the autoregressive image generation experiments, we pretrain LlamaGen-B model (approximately 11 parameters). Specifically, the model itself is an autoregressive decoder- only transformer, similar to Llama [Touvron et al., 2023a]. It utilizes RoPE as positional bias, extended to two dimensions. The model also involves learnable class embeddings, which are used as conditionals for image generation. During training, the class embeddings are dropped with probability of 0, to enable the use of classifier-free guidance at inference time.
To tokenize the image patches, we use a VQ-VAE tokenizer provided by Sun et al. 2024 with codebook size equal 16384 and embedding dimension equal 8. Unlike LlamaGen, we pre- tokenize the dataset with ADM's preprocessing scheme [Dhariwal and Nichol 2021 resulting in two crops per image.
Both Token baseline and MuToR are trained for 360K update steps, using AdamW optimizer with 3g 0.9, 88 0.95 and weight decay 0.05. We employ constant learning rate, equal to 0.0004, and a batch size of 1024. All experiments are run using 8 x H100 GPUs and the Distributed Data Parallel (DDP) framework.
For MuToR's implementation, we empirically tune the loss coefficient a to 0.5, so that Lreg has equal contribution with Catp (Equation 6). This indicates that the auxiliary loss provides valuable
15

---

supervision that the model leverages during pretraining, to improve its learned representations. In the experiments that fewer registers are inserted, we sample random positions for each training sample.
Evaluation To benchmark generative performance, sample 50,000 images at 256 256 resolu- tion, setting temperature = 1.0 and fixed random seed for fair comparison. We use classifier free reported
guidance [Ho and Salimans [022] with scale 2.0, which original LlamaGinn. Thus, the logit ly is formed as such: lg = lu + 2(L. L.), where l, denotes the unconditional logit (with the class embedding dropped) and le denotes the conditional logit. Then,
optimal value in the
the generated samples are used to calculate the performance metrics, using the Tensorflow scripts
from Dhariwal and Nichol [2021]
A.3 Synthetic Data
We setup our experiments using the official implementation from Bachmann and Nagarajan [2224] (regarding the model's 8 architecture and optimization). All model configurations (Next-Token and MuToR) are trained for sufficient number of epochs, using a single A100 GPU.
MuToR's implementation for the star graphs problem follows Subsection 3.3] with a task-specific modification; we sample the offset d from (m, thus excluding the next-token from the register's 's prediction. In this way, we prevent the registers from learning the "Clever Hans Cheat", and enable them to focus on planning look-ahead predictions. In these experiments, we set dmax = 4 and 0.5 for the graphs with path length 5 and dmax 6, 0.3 for graphs with path length 10, which are more challenging.
dmax)]
B Matching the Training Compute
As mentioned earlier, MuToR inserts additional tokens in the training sequence, thereby increasing its length and the overall training compute. To estimate the total training-time compute, we follow the widely used approximation:
6~6xNXL
(7)
where C' denotes the compute (in FLOPs), N denotes the model's parameters, and D denotes the total number of tokens processed during training. This formulation is supported by both empirical studies and theoretical analyses in prior scaling literature [Kaplan et al. 2020 Hoffmann et al. 2022
While the self-attention mechanism has a theoretical O(N2) complexity, in practice the compute cost is dominated by the feedforward layers (unless the sequence length I. becomes disproportionally larger than the embedding dimension), making the total training compute effectively linear in sequence length. This approximation holds for the Large Language Models and the sequence lengths used in our experiments, since the tipping points for these particular models lie far beyond their supported context length.
We further validate this approximation empirically by measuring training wall-clock time. In our IM-GSM experiments, MuToR requires 1.4x the wall-clock time of the Next-Token baseline for the same number of epochs. This aligns with expectations, especially since MuToR only interleaves register tokens into the answer portion, not the prefix, so the effective sequence length increase is less than 2x. As a result, doubling the number epochs for the baselines leads to equal or higher total compute in the data-constrained setting that we explore.
B.1 Language Modeling
Tables 12 and 14 report the results of fine-tuning baseline methods for twice the number of training epochs than MuToR, in order investigate whether improved downstream performance can be derived from increasing the training compute. Due to our limited computational resources, we keep the data fixed in these experiments, thus simulating data-constrained setting.
Matching the Training Compute via Additional Data In Table 3 (subsection 4.1). we report results where the baselines are trained with additional data to match MuToR's training compute.
As tipping points we refer to the sequence lengths which attention cost equals the feedforward cost.
16

---

Table 12:
Gemma 2B: downstream accuracy (%) in mathematical reasoning benchmarks across different number of training epochs. The subheaders refer to the training split used each experiment. Results are averaged across three seeded runs.
GSM8K
MATH500
Method
# Epochs
GSM8K
IM-GSM
IM-MATH
5
38.87
66.09
26.73
Next-Token
10
37.91
65.07
27.07
5
40.66
66.69
26.87
Multi-
10
39.98
64.92
26.73
MuToR (ours)
5
42.10
68.33
28.13
Table 13: Gemma 2B: ROUGE metrics comparison across different number of training epochs. We select the checkpoint with the higher ROUGE-L in the validation set, and report ROUGE scores on the test set.
Dataset
Method
#Epochs
ROUGE-1
ROUGE-2
ROUGE-L
3
51.47
27.29
43.23
5
51.67
27.64
43.32
SAMSum
3
51.90
27.44
43.50
Multi-Token
5
51.04
26.35
42.48
MuToR (ours)
3
52.32
28.08
44.09
3
47.23
20.91
38.77
Next Token
5
46.92
20.39
38.25
DialogSum
47.98
21.23
39.25
Multi-Tokes
5
47.46
20.76
38.97
MuToR (ours)
3
48.22
21.71
39.48
Table 14: Llama 3 8B: downstream accuracy (%) in mathematical reasoning benchmarks across different number of training epochs. The subheaders refer to the training split used in each experiment.
GSM8K
Method
# Epochs
GSM8K
IM-GSM
66.41
85.74
Next-Token
10
64.74
85.74
66.56
85.67
Multi-Ton Token
65.20
84.98
MuToR (ours)
67.85
87.05
Specifically, we construct IM-GSM-extra split by augmenting IM-GSM with samples from 2M- GSM, resulting in approximately 1.4x more training samples-consister with MuToR's 1.4x higher training time. The reported accuracies are obtained with Gemma 2B and are averaged over three seeded runs. All other training details follow subsecttio
follow/subsection A.l
B.2 Autoregressive Image Generation
As discussed (subsection 4.2, MuToR-2D demonstrates substantial performance gains even when using only 80 register tokens. In this configuration, training compute increases by approximately 30%. To ensure that our performance gains cannot be matched by increasing the pretraining iterations, also train the baseline Next-Tokee for 500K deeps-correspondiing
similar increase in
17

---

computational cost. As shown in Table 15 MuToR-22 2D delivers significantly better generation quality under comparable compute constrants.
Table 15: Conditional generation performance with respect to training iterations. We compare against the MuToR-2D configuration using 80 randomly placed registers, which results in approximately a 30% increase in training compute.
Method
# Iter.
FID
IS
360K
6.18
171.5
Next-Token
500K
5.82
175.8
MuToR-20 2D (ours)
360K
5.03
196.6
Additional Experiments
C.1 Comparison with Sequential Multi-Token Prediction
to the fine-tuning task by employing dinax 1 additional transformer layers. Our implementation closely follows Liu et al. (2024). and all training nyperparameters are identical with those described in (suuseection A. We fine-tune Gemma 2B with MMMMMlltoookeee on mathematical reasoning datasets, including GSM8K, IM-GSM and IM-MATH. The hyperparameters dmax and are empirically tuned to 2 and 0.1 respectively.
We adapt the sequential multi-token prediction method
Table rejepotss results averaged over three seeded runs. Notably, DS-Multi Token performs on par with the parallel Multi -oken baseline [Gloeckle et al. [202].. while requiring more additional trainable parameters. MuToR consistently outperforms both methods, using only negligible extra parameters, thus suggesting its superiority when integrated in the supervised fine-tuning stage.
Table 16: Downstream accuracy (%) in mathematical reasoning benchmarks, using Gemma-2B The subheaders refer the training split used in each experiment. All results are averaged over 3 seeded runs.
GSM8K
MATH500
Method
# Add. Params.
GSM8K
IM-GSM
IM-MATH
Next-Token
38.87
66.09
26.73
Mukei
110M
40.66
66.69
26.87
DS-Multi-Toker
118M
40.61
66.97
27.00
MuToR (ours)
2K
42.10
68.33
28.13
Table 17: Effect of varying the loss coefficient a, while fixing dmax 4. The results obtained from single run experiments.
GSMSK
GSM8K
IM-GSM
0.5
40.78
67.39
0.4
40.94
68.00
0.3
42.00
68.38
0.2
41.24
67.55
0.1
40.33
67.32
0
38.43
66.33
(lext-Token)
18

---

C.2 Impact of the Loss Coefficient
We analyze the effect of the auxiliary loss coefficient a, which controls the relative weight be- tween Cntp and Creg. In these experiments. we fine-tune Gemma 2B on mathematical reasoning datasets, while fixing dmax 4. As shown in Table setting 0.3 yields optimal downstream performance, with both higher and lower values leading in mild degradation.
This trend is expected, prediction supervisory signals. However, the optimal value for likely depends on factors like the model, task complexity, data distribution, and dmaxx For example, in autoregressive image generation, setting 0.5 achieves better generation quality, likely because the multi-token prediction objective captures rich spatial dependencies.
optimal balance should exist between next-token and multi-token
D Broader Impact
Improving large language models can enhance their accuracy, efficiency, and safety, enabling more reliable applications across education, healthcare, and communication. However, important to remain mindful of potential risks such as misuse bias amplification and to continue developing responsible deployment practices.
19

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The abstract and introduction accurately reflect the proposed method (MuToR), the main contributions, and the results from our experimental evaluation, which are described in detail in section 4
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss limitations of our work in section
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
20

---

Answer: [NA]
Justification:
This paper does not include theoretical results that require proofs.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
The paper includes a detailed Appendix section (Appendix A), which contains all the necessary information to reproduce our experimental results. In this section, we describe the training splits, optimal hyperparameter values (for both the baselines and our method), and other key details related training and evaluation. We believe this information ensures the reproducibility of our work. Additionally, our method is explained clearly in section of the main paper.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibiliity In the case of closed-source models, it may be that access to the model is limited in
21

---

some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification:
Upon acceptance, we will release our code publicly.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed run reproduce the results. See the NeurlPS code and data submission guidelines (https
ctcaps.cc/public// http://wieeeeeeeee Swweeddueeeodeuubmmicce ssionPolicy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The paper describes experimental setup in detail in Appendix A including training splits, hyperparameters, and training procedures. We also provide the evaluation methodology in the same section.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
report results three seeded runs for statistical ro- bustness. However, due to computational constraints, we could not run multiple seeded trials for the image generation experiments and the Llama 3 8B language modeling experiments.
Justification: Tables
Guidelines:
The answer NA means that the paper does not include experiments.
22

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative eтror rates).
show in tables or
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide hardware and computational resource details Appendix A
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: We have read the code of ethics and made sure that the paper conforms it in every aspect.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We discuss broader impact in Appendix D
23

---

Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairesss considerations (e.g., dep.oymmly of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
The paper proposes novel method for training autoregressive transformers. Neither the method itself nor the standard benchmark data used appear to pose high risk for misuse, thus not necessitating specific release safeguards beyond standard open-source practices.
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We properly cite all models, datasets, and code implementations used to run our experiments and evaluate models.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.

---

The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The primary new asset is the implementation code for the proposed method, which we plan to release upon acceptance.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve research with human subjects, therefore IRB approval is not applicable.
25

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Answer: [Yes]
Justification: We propose methodologies that can be used for improving the performance of Large Language Models.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non non-standard components.
Please refer to our LLM policy http::://eerrips.scc/ooeeeeeeeee for what should or should not be described.
26


Reviews:
[{"review_id": "8gzoID4beA", "reviewer": "Reviewer_zeH6", "summary": "This work proposes a novel method to perform multi token prediction in language model and image generation settings. The proposed method is to be applied during training only. Inference does not utilize any multi-token prediction functionality despite being used during training. The authors argue that the standard next-token prediction used in auto-regressive generation is not sufficient to instill long-range planning abilities into a model. The authors note that existing work which proposes multi-token prediction during AR generation generally requires some small but non-negligible architecture modifications and additional parameters (e.g. additional language heads to each additional token to be predicted at a given position.\n\nThe authors propose MuToR, which interleaves learnable \"register\" tokens with usual tokens from data during training. The register token is tasked to predict a future token. The training setup using attention masking to ensure that a register token as position t, predicting the token at t + t', can only attend to non-register tokens at positions <t (i.e. a register token is blind to all other register tokens). This method ensures that during inference the register setup can be discarded and standard AR generation can be used. Positional embeddings are used to indicate the position of the to-be predicted token. This limits the number of additional parameters (the authors use a single learnt register embedding) to the hidden dimension of the model (~10^3) making the method easily usable with pretrained models (where subsequently added new tokens is relatively commonplace now) and with parameter-efficient fine-tuning methods.\n\nThe authors conduct experiments with MuToR on language modeling tasks (Math) and autoregressive image generation extending their method naturally to 2D grids of tokens. They show that their method outperforms a previous multi-token prediction method (which adds additional language model heads) and standard next-token prediction only.", "questions": "1. Please address the query outlined under \"Weaknesses\" regarding the \"Matching the Training-Time Compute\".\n2. Are the authors still about to use FlashAttention with MuToR? If not, what is the penalty in terms of training speed for the experiments conducted with respect to standard next-token prediction.", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 4, "strengths_and_weaknesses": "## Strengths\n1. Novelty\nI think the proposed method is highly intuitive (the concept of having a special token(s) to predict further into the future rather than adding specialized model parameters), very simple and seems to require only small adjustments to basic model training (attention masking, one added token). There have been some prior works on multi-token prediction, however this work appears to be orthogonal to this work in that it comes at the problem of multi-token prediction is a very novel way. Additionally the fact that this method does not impact inference in any way (for which rigid but highly optimized libraries now exist) is a bonus.\n\n2. Clarity\nExcept for specific places (see below), I find the paper very clear (which makes sense given how simple and intuitive the proposed method is). The figures are easily understandable for anyone who understands the basics of attention in sequence models.\n\n3. Experiments\nWith certain exceptions (see below), the experiments are good in that they cover multiple modalities, compare to the most related methods such that understanding the benefits of the proposed method is easy for the reader and the key parameters are ablated (Table 3 and 5). The method appears to be effective compared to the related methods.\n\n## Weaknesses\n1. Matching the Training-Time Compute\nThis experiment needs to additional rigor to be convincing.\nFirst, it appears the authors alternate between data tokens and register tokens roughly doubling the sequence length per training sample compared to standard next-token prediction. If one wishes to match compute, given the nature of attention (compute scales quadratically with sequence length), one should more than double the number of epochs for a fair comparison.\nSecond, given the relatively small amounts of data the authors are using (e.g. at most 100M data tokens for 1M-GSM), repeated epochs of the same data should definitely not be used for these kinds of experiments, instead _additional_ data should be used. I personally given the experimental setup for this part, do not find it convincing.\n\n2. Training time due to custom masking\nThe novel method the authors propose require a custom attention mask. As far as I can see, this means standard optimizations in language model training (such as FlashAttention{1,2,3}) would not be able to be used with MuToR. This is a significant weakness of the method and is not addressed by the authors in terms of training time. At scale this could become a large issue.\n\n3. Minor points on clarity\nFirst, the authors should make it clear that the method currently uses a register token in every other position, thus the sequence length approximately doubles.\nSecond, some of the extra details on how the datasets are curated (which are in the SM) need to be put into the main paper like the maximum token length, these will caveat that the results are valid at short sequence lengths only.", "was_revised": true, "final_justification": "I thank the authors for their responses. I believe this discussion period has helped clarify some of the limitations of their novel method which should be given to the reader. I maintain my score given the originality of the work and I expect the authors to include some of the limitations and caveats of their work revealed during the discussion period.", "rebuttal": "We thank the reviewer for their positive and thoughtful feedback. We are pleased that they find MuToR intuitive, novel, and effective. Below we address the concerns that they raised.\n\n---\n## **Question 1**\n> Matching the Training-Time Compute This experiment needs to additional rigor to be convincing. First, it appears the authors alternate between data tokens and register tokens roughly doubling the sequence length per training sample compared to standard next-token prediction. If one wishes to match compute, given the nature of attention (compute scales quadratically with sequence length), one should more than double the number of epochs for a fair comparison. \n\nWe estimate training compute using the widely accepted formula: $C \\approx 6 \\times N \\times D$, where $C$ denotes the total training compute (in FLOPs), $N$ denotes the model's parameters, and $D$ denotes the total number of tokens processed during training. This formulation is supported by both empirical studies and theoretical analysis in prior scaling literature [1,2]. While the self-attention mechanism has a theoretical $O(L^2)$ complexity, in practice, the compute cost is dominated by the feedforward layers (unless the sequence length $L$ becomes disproportionally larger than the embedding dimension), making the total training compute effectively linear in sequence length. **Following this, doubling the number of tokens (as our method does) corresponds to approximately doubling compute, which is why we doubled the baseline epochs to match total tokens.**\n\nTo further validate this, we measured actual training time. In our 1M-GSM experiments, **MuToR requires 1.4× the wallclock time of the Next-Token baseline (for the same number of epochs)**. This is consistent with expectations, especially since MuToR only interleaves register tokens into the answer portion, not the prefix, so the effective sequence length increase is less than 2×&mdash;**that means that we even used more compute for the baseline methods.** \n\n\n\n>Second, given the relatively small amounts of data the authors are using (e.g. at most 100M data tokens for 1M-GSM), repeated epochs of the same data should definitely not be used for these kinds of experiments, instead additional data should be used. I personally given the experimental setup for this part, do not find it convincing.\n\nWe kindly disagree with the reviewer's claim that repeated epochs of the same data should not be used in such experiments. Using a fixed dataset and allowing repeated epochs is a widely accepted experimental setup, especially in scenarios where data is limited—a common real-world constraint in many academic and industrial applications. By running these experiments, we validate that **MuToR provides a more effective way to leverage additional training compute, when the amount of data remains fixed**. This insight may be useful for practitioners, as **MuToR can aid in getting the maximum performance out of specific datasets**, a highly relevant and realistic scenario.\n\n\nRegarding the reviewer's suggestion for using additional data, we respectfully argue that doing so would confound the experimental comparison, as improvements could no longer be attributed solely to our proposed method or training compute. By controlling for both compute and data, we isolate the contribution of the multi-token prediction objective—ensuring a clean, interpretable comparison with the baselines.\n\n---\n## **Question 2**\n> Are the authors still about to use FlashAttention with MuToR? If not, what is the penalty in terms of training speed for the experiments conducted with respect to standard next-token prediction.\n\nAs the reviewer correctly points out, FlashAttention does not currently support custom attention masks, which prevents us from using it directly with MuToR. \nTo address this, we leverage FlexAttention [3], a recently introduced API by the PyTorch team that supports flexible masking while still compiling to fused, FlashAttention-style kernels.\n\nUsing FlexAttention, MuToR achieves near parity with FlashAttention-2. For example, in our autoregressive image generation experiments (360K iterations, global batch size 1024), the average training step latency was:\n\n- **0.35s**&mdash;with MuToR + FlexAttention (sequence_length=512, custom mask).\n- **0.34s**&mdash;with standard FlashAttention-2 (sequence_length=512, assuming a causal mask).\n- **0.16s**&mdash;with standard FlashAttention-2 (*sequence_lenth=256, Next-Token baseline*).\n\n\n\nThis demonstrates that MuToR’s training speed penalty&mdash;due to the custom mask&mdash;is minimal when using FlexAttention. In comparison with the standard Next-Token baseline, the training latency roughly doubles, approximating the expected linear scaling due to increased sequence length.\n\nWe thank the reviewer for their time and we hope that their concerns might be resolved.\n\n\n\n\n---\n[1] Kaplan et al., \"Scaling laws for neural language models.\" arXiv preprint 2020.\n\n[2] Hoffman et al., \"Training compute-optimal large language models.\" arXiv preprint 2022.\n\n[3] Dong et al., \"Flex Attention: A programming model for generating optimized attention kernels.\" arXiv preprint 2024."}, {"review_id": "Yg3aazntip", "reviewer": "Reviewer_foD8", "summary": "This paper introduces a method termed \"Computational Halting,\" which aims to enhance the mathematical reasoning capabilities of large language models (LLMs). The core technique involves augmenting input prompts with a sequence of special, learnable \"halting tokens\" during the fine-tuning process of a pre-trained model. The authors apply this method to the Llama 3 8B Instruct model, fine-tuning it on a dataset of mathematical problems. The central hypothesis is that these halting tokens provide the model with additional, sequential computational steps—effectively allowing it to \"think\" or halt to process information—before generating a final answer. The authors present results on the GSM8K and MATH benchmarks, claiming that their method leads to performance improvements over the baseline Llama 3 8B Instruct model and thereby validates their approach.", "questions": "1. **On Novelty and Relation to Prior Work:** The proposed \"computational halting\" mechanism appears identical to the \"pause-finetuning\" on a standard pre-trained model (`StdPT_PauseFT`) described in Goyal et al. (2024). Could you please clarify the novel technical contributions of your work beyond what was presented by Goyal et al.? A simple rebranding of a prior technique is insufficient.  \n\n   \n\n2. **On the Choice of Training Regime:** The central finding of Goyal et al. was that pause tokens require inclusion in pre-training (`PausePT_PauseFT`) to be effective, as fine-tuning alone (`StdPT_PauseFT`) yields minimal or negative results. Your work uses the latter, less effective approach. Could you provide a strong justification for this design choice? An ablation study that directly compares the performance of   \n\n   \n\n   `StdPT_PauseFT` (your method) with `PausePT_PauseFT` on Llama 3 would be necessary to validate this decision.\n\n3. **On Baselines and State-of-the-Art Comparison:** Your evaluation shows an improvement over the base Llama 3 8B Instruct model. However, fine-tuning this same model on the public `OpenMathInstruct-2` dataset yields a MATH score of 67.8% , far exceeding the performance you report. How do you position the significance of your method's contribution in light of these SOTA results, especially given that your method increases inference latency while the data-centric approach does not? A convincing submission must include this SOTA comparison.  \n\n   \n\n4. **On Reproducibility:** To allow for verification of your results, could you please provide a detailed description of your fine-tuning dataset, including its source, size, and composition? Furthermore, please provide all hyperparameters used for training, including the learning rate, batch size, number of training epochs, and the specific number of \"halting tokens\" used during fine-tuning and inference.", "limitations": "No, the authors have not adequately addressed the limitations of their work. A revised version should include a dedicated section with constructive discussion on the following points:\n\n- **Performance vs. Latency Trade-off:** The paper completely ignores the inference latency overhead introduced by the halting tokens. The community is actively developing methods to *accelerate* inference, such as speculative decoding. Any method that deliberately slows down inference must rigorously justify this trade-off with a corresponding significant improvement in performance, which is not demonstrated here. Measurements of this latency overhead are essential.  \n\n  \n\n- **Generalization vs. Overfitting:** The authors should discuss the risk that the model is not learning a generalizable \"thinking\" skill but is instead overfitting to a simple correlation in the fine-tuning data (i.e., learning that seeing N special tokens is associated with a certain type of math problem).\n\n- **Mechanism Transparency:** The work provides no analysis of *how* the halting tokens are being used by the model. A deeper scientific contribution would involve investigating the internal mechanisms. For example, do attention patterns show a meaningful flow of information through these tokens? Does their utility change with the complexity of the problem? Without such analysis, the method remains a black box.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 1, "originality": 1, "strengths_and_weaknesses": "## Strengths\n\nThe primary strength of this work lies in its focus on a critical and challenging research area. Improving the multi-step, logical reasoning abilities of LLMs, particularly in formal domains like mathematics, remains a key frontier in artificial intelligence research. The paper correctly identifies the limitations of standard autoregressive models, which allocate a fixed amount of computation per generated token, and seeks to address this shortcoming.\n\nThe underlying idea of providing a model with more \"thinking time\" is conceptually sound and aligns with an emerging research direction that explores dynamic computation and decoupling token generation from computational steps. This line of inquiry, which questions the rigid structure of the standard Transformer, is a promising avenue for developing more powerful and capable models.  \n\n\n\n\n\n## Weaknesses\n\n### Originality: A Reimplementation of Prior Work\n\n\n\nThe most significant weakness of this paper is its profound lack of originality. The proposed \"Computational Halting\" method is not novel; rather, it appears to be a direct and unacknowledged reimplementation of a specific experimental variant from Goyal et al. (2024), \"Think before you speak: Training language models with pause tokens\".  \n\n\n\nGoyal et al. introduce the concept of \"pause tokens\" (`<pause>`), which are learnable tokens appended to an input prefix to delay the model's output, thereby allowing for additional computation. The methodology described in this paper—fine-tuning a pre-existing model by appending special tokens to the input—is identical to the \"Standard Pretraining and Pause-Finetuning (StdPT_PauseFT)\" variant explicitly detailed and analyzed by Goyal et al.. The rebranding of this technique as \"Computational Halting\" does not constitute a novel contribution. The failure to cite, discuss, and differentiate from this highly relevant prior work represents a severe gap in the literature review and positions the paper's contribution as derivative.  \n\n\n\n\n\n### Quality: Critical Technical and Methodological Flaws\n\n\n\nThe paper's technical execution is fundamentally flawed because it ignores the central findings of the very work it replicates.\n\n1. **Ignoring the Necessity of Pre-training:** The core technical conclusion of Goyal et al. is that pause tokens are most effective when integrated during *both* pre-training and fine-tuning (`PausePT_PauseFT`). The authors of that study demonstrated that simply introducing these tokens during fine-tuning (`StdPT_PauseFT`)—the exact approach taken in this paper—yields \"fewer and milder gains, and sometimes even a clear drop in performance\". The rationale is that standard pre-training ingrains a strong architectural bias towards immediate next-token prediction, conditioning the model to be \"quick.\" Without being exposed to delay-based computation during pre-training, the model does not learn to effectively utilize the additional computational pathways offered by the pause tokens during fine-tuning. This paper's methodology is therefore predicated on a technique that has already been shown to be suboptimal, and it provides no justification for this choice, suggesting a failure to understand the key technical lessons from the most relevant prior work.  \n\n   \n\n2. **Conceptual Confusion:** The paper seems to conflate its method with other, mechanistically distinct \"special token\" approaches, revealing a shallow understanding of the literature. For instance, the concept is unrelated to \"register tokens\" in Vision Transformers, which serve as a memory sink to absorb high-norm outlier activations and clean up attention maps, rather than extending sequential computation. Similarly, it is the conceptual opposite of multi-token prediction heads used for speculative decoding, which aim to   *parallelize* future predictions to *reduce* latency, whereas this method *serializes* extra computation and *increases* latency.  \n\n   \n\n\n\n### Quality: Weak and Misleading Evaluation\n\n\n\nThe experimental evaluation is insufficient and presents a misleading picture of the method's efficacy.\n\n1. **Inadequate Baselines:** The paper's primary claim of improvement is based on a comparison to the base Llama 3 8B Instruct model. While a gain over this baseline may be observed, it is an inappropriate and misleading benchmark for the task of mathematical reasoning. The current state-of-the-art for this specific model involves fine-tuning on high-quality, domain-specific data. For example, fine-tuning the same Llama 3 8B base model on the `OpenMathInstruct-2` dataset achieves a score of **67.8% on the MATH benchmark** and **91.7% on GSM8K**. The base Llama 3 8B Instruct model scores only 30.0% on MATH. The gains reported in this paper are dwarfed by the improvements from a standard, data-centric fine-tuning approach. This context completely undermines the significance of the paper's results; the proposed method is a complex, latency-increasing technique that is vastly outperformed by standard fine-tuning on a public dataset.  \n\n   \n\n2. **Lack of Essential Ablation Studies:** The paper fails to perform the necessary ablation studies to validate its claims. A rigorous analysis would, at a minimum, investigate the effect of varying the number of \"halting tokens,\" a key hyperparameter of the method. More critically, the authors do not provide the essential comparison between their   \n\n   \n\n   `StdPT_PauseFT` approach and the `PausePT_PauseFT` regime, which is necessary to determine if the technique has any merit at all.\n\n3. **Limited Scope:** The experiments are confined to a single model scale (8B parameters). As shown in other work on architectural modifications, performance benefits can vary significantly with model scale. Without testing on larger models (e.g., Llama 3 70B), any claims of generalizability are entirely unsupported.  \n\n   \n\n\n\n### Clarity and Reproducibility\n\n\n\nThe paper lacks the necessary details for the research community to reproduce or verify its findings. Key omissions include:\n\n- A precise description of the fine-tuning dataset used.\n\n- A complete list of training hyperparameters, such as learning rate, batch size, number of epochs, and optimizer settings.  \n\n  \n\n- The specific number of \"halting tokens\" used during fine-tuning and inference. Given the documented difficulty in reproducing official benchmark scores even with some of these details , this lack of transparency is a serious flaw.", "was_revised": true, "final_justification": "After carefully reviewing the authors' rebuttal and cross-referencing it with the comments from other reviewers, I have gained a better understanding of the paper's motivation and technical contributions. Although some concerns remain, the authors have satisfactorily addressed the majority of my initial issues. Consequently, I have decided to increase my rating.", "rebuttal": "We thank the reviewer for their time and feedback. \n\n**However, we believe that the reviewer has misunderstood our core contribution, as well as the scope of our work.** Below we shall answer to each of their specific criticisms.\n\n---\n## **Summary**\n> This paper introduces a method termed \"Computational Halting,\" which aims to enhance the mathematical reasoning capabilities of large language models (LLMs). \n\n>   The central hypothesis is that these halting tokens provide the model with additional, sequential computational steps—effectively allowing it to \"think\" or halt to process information—before generating a final answer.\n\nWe respectfully argue that **this summary of our work is misleading and inaccurate.** The reviewer claims that we propose \"Computational Halting\", a method to insert learnable tokens during the finetuning stage, in order to increase the inference computation and boost downstream performance. \n\nIn contrast, our submission introduces MuToR, **a novel multi-token prediction method that uses interleaved, learnable tokens to predict future targets during training and enrich supervision for autoregressive transformers. Our work does not utilize the learnable tokens during inference, but it completely discards them.**\n\nBased on this, our central hypothesis is in fact orthogonal with what the reviewer describes in this summary. **In fact, the term \"Computational Halting\" is not used in our submission, so we sincerely wonder why the reviewer insists on this**.\n\n---\n\n## **Originality : A Reimplementation of Prior Work**\n\n>The most significant weakness of this paper is its profound lack of originality. The proposed \"Computational Halting\" method is not novel; rather, it appears to be a direct and unacknowledged reimplementation of a specific experimental variant from Goyal et al. (2024), \"Think before you speak: Training language models with pause tokens\".\n\n> The failure to cite, discuss, and differentiate from this highly relevant prior work represents a severe gap in the literature review and positions the paper's contribution as derivative.\n\nWe kindly argue that the reviewer's statement is invalid. **Our proposed method, MuToR, is orthogonal to the mentioned previous work [1]**, since:\n\n1. **We use interleaved learnable register tokens during training to predict future targets in the sequence**, thus enhancing the supervisory signal and propagating richer gradients through the model's layers.  **In contrast, the authors of [1] suggest appending the learnable tokens to the prompt to leverage the extra computation during inference.**\n\n\n2. **Our carefully designed attention mask and positional embeddings enable us to completely discard the registers at inference time.** In contrast, the proposed method in [1] primarily depends on the extra operations between tokens on inference, **to simulate \"thinking time\"**.\n\n3. We propose a method for multi-token prediction in autoregressive transformer models. **Our work introduces a novel multi-token prediction training objective&mdash;associated with the learnable tokens&mdash;and investigate its effectiveness compared to the baselines.**\nHowever, the authors in [1] investigate whether&mdash;and in what settings&mdash;can a Language Model benefit from increased computations. **They do not use any auxiliary training loss for their learnable tokens**. \n\n\n**We also point to our Related Work (Section 2), where the reviewer can find a proper citation and acknowledgement of this previous work**, as well as a comprehensive literature review in the field. \n\n---\n\n## **Quality : Critical Technical and Methodological Flaws**\n\n\n> Without being exposed to delay-based computation during pre-training, the model does not learn to effectively utilize the additional computational pathways offered by the pause tokens during fine-tuning.\n\nAs we have already stated, **our proposed method does not involve any delay-based computation, so this criticism is not applicable**. Regardless of that, our experimental evaluation (see Table 1 and 2) indicates that **applying MuToR directly in the finetuning stage of pretrained Language Models results in significant performance gains**.\n\n> The paper seems to conflate its method with other, mechanistically distinct \"special token\" approaches, revealing a shallow understanding of the literature. \n\nAs we clearly state in Section 3, we adopt only the terminology (\"register tokens\") from [2]. **We respectfully point to our Related Work section, where we have thoroughly discussed the relevant literature, acknowledged their contributions and carefully differentiated our approach from previous works**.\n\n---\n\n## **Quality: Weak and Misleading Evaluation**\n\n> The paper's primary claim of improvement is based on a comparison to the base Llama 3 8B Instruct model.\n\n>The experiments are confined to a single model scale (8B parameters).\n\nIn section 4 of our paper, we provide a thorough experimental evaluation of our method. **In contrast with the reviewer's claim, our evaluation is not limited to the Llama 3 8B Instruct model (in fact, this specific model variant is not used in our work!), but we present results in both finetuning and pretraining settings, using pretrained LLMs of different size (for finetuning), and  validating MuToR's effectiveness in image generation as well.** \nWe list our experiments here:\n1. Supervised Finetuning of Pretrained LLM, using Gemma 2B and Llama 3 8B (the base model),in both mathematical reasoning and abstractive summarization benchmarks.\n2. Parameter Efficient Finetuning, using Gemma 2B model.\n3. Autoregressive Image Generation, where we pretrain a LlamaGen transformer model on ImageNet 256x256.\n4. Experiments on the synthetic star-graph task, where we finetune a pretrained GPT2-Large model.\n\n\n\n> The gains reported in this paper are dwarfed by the improvements from a standard, data-centric fine-tuning approach.\n\nOur experiments involve finite amounts of data, which are far less than the large-scale dataset needed to obtain the state-of-the-art performance mentioned by the reviewer [3]. These results are **not comparable to our experiments**, since they were obtained by finetuning the model on **~14M training samples**, while our finetuning splits are typically **smaller than 200K samples**.\nThus, we respectfully argue that our evaluation is not misleading; all the implementation details are included in the appendices, and our setup guarantees a fair comparison between methods.\n\n> The paper fails to perform the necessary ablation studies to validate its claims. \n\nIn section 4 of our paper, we include a comprehensive ablation study (**as acknowledged by the other reviewers**), providing key insights and understanding on what drives MuToR's gains. \n\n---\n\n## **Questions**\n\n#### **1.**\n> Could you please clarify the novel technical contributions of your work beyond what was presented by Goyal et al.? A simple rebranding of a prior technique is insufficient.\n\n**We have already analyzed our key differences with the work presented in [1], and thoroughly explained why our proposed method is orthogonal to theirs.** Please see \"Originality: A Reimplementation of Prior Work\" in this response.\n\n#### **2.**\n> Could you provide a strong justification for this design choice? An ablation study that directly compares the performance of StdPT_PauseFT (your method) with PausePT_PauseFT on Llama 3 would be necessary to validate this decision.\n\nAs we already stated (see our arguments above), **our method is orthogonal to StdPT_PauseFT and PausePT_PauseFT (both variants proposed in [1])**.\n\n#### **3.**\n> How do you position the significance of your method's contribution in light of these SOTA results, especially given that your method increases inference latency while the data-centric approach does not?\n\nAs mentioned before, **our method does not introduce inference latency at all. This is also clearly stated in Section 3.2 of the paper (see Inference), and acknowledged by other reviewers as well.** \n\nRegarding the significance of our method, **the results that the reviewer mentions are in fact not comparable to ours, since the cited work [3] uses a much larger amount of finetuning data**. We again argue that the scope of our work is neither to reach state-of-the-art performance in mathematical reasoning, nor to outperform models that have been finetuned on far larger amounts of data. **On the contrary, the significance of our work is proven through the performance gains obtained with respect to our baselines, and through the fact that MuToR offers specific advantages over the existing multi-token prediction methods (see Section 1 of the paper).**\n\n#### **4.**\n\n> On Reproducibility: To allow for verification of your results, could you please provide a detailed description of your fine-tuning dataset, including its source, size, and composition? Furthermore, please provide all hyperparameters used for training, including the learning rate, batch size, number of training epochs, and the specific number of \"halting tokens\" used during fine-tuning and inference.\n\n**All of these key details are already included, in the supplementary material (technical appendices) of our submission.** We also plan to publicly release the code for our method, to ensure reproducibility.\n\n\n\n\nWe thank the reviewer for their time. We provide detailed answers to their criticisms in order to resolve any misunderstandings regarding our work. **We sincerely hope that the reviewer will be able to re-assess our submission in the light of our counter-arguments**.\n\n\n---\n[1] Goyal et al., \"Think before you speak: Training Language Models With Pause Tokens.\", ICLR 2024.\n\n[2] Darcet et al., \"Vision Transformers Need Registers.\", ICLR 2024.\n\n[3] Toshniwal et al., \"OpenMathInstruct-2: Accelerating AI for Math with Massive Open-Source Instruction Data.\", ICLR 2025"}, {"review_id": "yyXscDs0zS", "reviewer": "Reviewer_nf2t", "summary": "The paper introduces MuToR (Multi-Token Prediction with Registers), a method that enhances autoregressive models by inserting learnable register tokens during training. Each register token is added in between of the two normal tokens, and predicts a future token d steps ahead. The proposed approach show advantage in both language model fine-tuning and autoregressive image generation.", "questions": "1. The proposed method is interesting yet a lot of operations are un-explanable to me. For example, why do you want to insert registers after each token and use it to predict the future token after d steps? If I insert register for maybe every two normal tokens, what would be the difference in performance? Also it is unclear to me why a shared embedding is sufficient, in particular Table 5 shows that different register embedding gives worse performance, which is countering my intuition. Could the authors comment on these?\n\n2. Training compute is only modestly increased by register token insertion, but the real cost when scaling to large batches, long contexts, or high $d_{max}$ isn't fully quantified. Also, I want to confirm that the offset $d$ is sampled uniformly and FIXED during training, i.e. for different register token, they have the same $d$ parameter?\n\n3. The proposed method is only evaluated on LLM finetuning tasks. To my knowledge, next token prediction is current the standard paradigm for pretraining. In finetuning, the final prediction target is the true label/answer, which is not exactly where we could see the full benefit of replacing next token prediction by multi toekn prediction. Could the author provide some numerical experiments on pretraining? 1b or 7b models are desirable.", "limitations": "Yes.\n\nAlso I don't see any potential negative societal impact of this work.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths: The proposed approach is both novel and simple, and it is versatile for both language and vision models. MuToR consistently outperforms both standard next-token and prior multi-token methods in both full and parameter-efficient finetuning (e.g., LoRA). Thoughtful exploration of hyperparameters like d_max, number of registers, embedding sharing, and 2D extensions provides a clearer understanding of what drives MuToR’s gains.\n\nWeakenss (please respond to the questions section directly): Lack of justification of some operations, for example it is unclear why registers are inserted uniformly or randomly, and a shared embedding is sufficient; Training compute is only modestly increased by register token insertion, but the real cost when scaling to large batches, long contexts, or high d_{max} isn't fully quantified; There is no pretraining experiment for the proposed method.", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback. The reviewer acknowledges the novelty, the effectiveness and the versatility of our proposed method. We address specific concerns and questions below.\n\n\n---\n## **Question 1**\n> For example, why do you want to insert registers after each token and use it to predict the future token after d steps? If I insert register for maybe every two normal tokens, what would be the difference in performance?\n\nAs we state in Section 3.2 of our paper, the **interleaving pattern is flexible**&mdash;we can either place registers after each token, or adopt a sparser strategy. In the language modeling experiments (LLM finetuning), we chose to fully interleave the answer sequences. This formulation provides a **\"denser\" multi-token prediction signal at every token position.** \n\nTo justify this design choice, we ran experiments to compare the performance gain obtained by the two settings; *the default setting when we fully interleave the answer sequence*, and *a \"sparser\" alternative where we insert one register for each two normal tokens*, as the reviewer suggests. We finetuned Gemma 2B on both GSM8K training set and 1M-GSM split. In the following table, we report the average accuracy over 3 seeded runs.\n\n\n| Method            | GSM8K         |  1M-GSM                      |\n|---                               | ---         |  ---            |\n| *Next-Token (baseline)*            |*38.87*    | *66.09*         |\n| MuToR-sparser variant                  | 40.33       |  67.48          |\n| **MuToR-fully interleaving (default)** |**42.10**    | **68.33**       |\n \n\nThese results suggest that the sparser interleaving pattern achieves lower performance than the fully interleaving setting, although it still outperforms the vanilla Next-Token Baseline. Therefore, we adopted it as the default setting throughout our finetuning experiments.\n\nMoreover, in the paper we also explored a \"sparser\" alternative during pretraining for the autoregressive image generation task where we placed the registers in random positions within each sequence. Our results (Table 7) indicate that while using a fewer amount of registers during training (~30% of the sequence length), the model performs on par with the default scenario. This showcases an effective way of using MuToR in pretraining settings, with only a modest increase in total training compute.\n\n> Also it is unclear to me why a shared embedding is sufficient, in particular Table 5 shows that different register embedding gives worse performance, which is countering my intuition.\n\nThis is indeed a valid question. Our initial intuition was that different register embeddings would further boost downstream performance, but the results presented in Table 5 (in our paper) indicate otherwise, so we used a shared embedding through all experiments. We believe that the reason behind this lies in our positional encoding scheme&mdash;the offset information is encoded effectively through the position modification for the register tokens, and this makes the additional embeddings redundant. This design is clearly explained in Section 3.2 of our submission.\n\n---\n## **Question 2**\n\n> the offset $d$ is sampled uniformly and FIXED during training, i.e. for different register token, they have the same $d$ parameter?\n\nFirst, let us clarify the proposed design regarding the prediction offset. As we state in Section 3.2, the offset $d$ is sampled uniformly and fixed per sequence, i.e. all registers within the same sequence have the same $d$, but different sequences within the same batch have different offsets. This formulation is not strict though; we could also insert registers with different offsets within the same sequence. \n> Training compute is only modestly increased by register token insertion, but the real cost when scaling to large batches, long contexts, or high $d_{max} isn't fully quantified.\n\nAs we state in Section 1 (Introduction), our register-based multi-token prediction method enables us to **scale $d_{max}$ arbitrarily, without increasing the total training cost**, since the number of inserted registers does not depend on $d_{max}$. \n\nRegarding the reviewer's concern about the scaling cost, we recall the widely used approximation for the total training compute: $C \\approx 6 \\times N \\times D$, where $C$ denotes the total training compute (in FLOPs), $N$ denotes the model's parameters, and $D$ denotes the total number of tokens processed during training. \n\nRelevant literature on scaling [1],[2] suggests that the total compute needed grows linearly with the total number of tokens, and we observe similar trends through all of our experiments. Therefore, from the perspective of compute, we do not think that our method will have a problem scaling to larger batches or sequences, due to this linear scaling. In fact, our experiments in autoregressive image generation (pretraining setting) were conducted with a large batch size (128 per GPU), and we generally observed the aforementioned scaling trend.\n\n---\n## **Question 3**\n\n> The proposed method is only evaluated on LLM finetuning tasks. To my knowledge, next token prediction is current the standard paradigm for pretraining. In finetuning, the final prediction target is the true label/answer, which is not exactly where we could see the full benefit of replacing next token prediction by multi toekn prediction.\n\nWe respectfully clarify that the next-token prediction objective is the standard paradigm not only for pretraining, but also for supervised finetuning in generative tasks, which involve prefix-answer pairs (with the answer sequences potentially including chain-of-thought tokens).\n\nThrough our experiments, we experimented with mathematical reasoning benchmarks, which include the complete and analytical chain-of-thought reasoning trace within the answer part. We also evaluate our method in the abstractive summarization task, which also requires generating the answer autoregressively.\n\nMoreover, we evaluated our method on the autoregressive image generation task (section 4.2), in a challenging pretraining setting, leading to improved generation quality and faster convergence.\n\nRecent work [3] also recommends evaluating multi-token prediction on generative benchmarks (those that require autoregressive generation of the answer sequence) rather than likelihood-based benchmarks (such as multiple-choice). We therefore think that the benchmarks used in our experiments provide a rigorous and insightful testbed for evaluating our method and its benefits.\n\n> Could the author provide some numerical experiments on pretraining? 1b or 7b models are desirable.\n\nWhile we agree that large-scale pretraining (1B or 7B LLMs) would further validate our method, our computational resources did not allow for training such models from scratch.  We prioritized experiments that could highlight MuToR’s key advantages&mdash;particularly its adaptability and scalability compared to existing multi-token methods.\n\n\nWe thank the reviewer for their time and feedback, and we hope to have addressed their concerns.\n\n---\n[1] Kaplan et al.: \"Scaling laws for neural language models.\" arXiv preprint 2020.\n\n[2] Hoffman et al.: \"Training compute-optimal large language models.\" arXiv preprint 2022.\n\n[3] Gloeckle et al.: \"Better & Faster Large Language Models via Multi-token Prediction.\" ICML 2024."}, {"review_id": "NH7CpsyAE4", "reviewer": "Reviewer_sXdq", "summary": "The paper introduces multi-token prediction with registers (MuToR). The method offers advantages including (1) no architectural changes (2) fine-tuning compatibility and (3) scalable prediction horizons. The authors validate MuToR across various settings, including Supervised Fine-Tuning (SFT), Parameter-Efficient Fine-Tuning (PEFT), and image auto-regressive generation. Additionally, ablation studies are conducted to analyze the influence of different components, such as the use of same versus different registers and the maximum prediction distance $d_{max}$.", "questions": "1. How does the method compare to other types of MTP like sequential MTP from DeepSeek V3?\n2. What is the influence of the ratio of $L_{ntp}$ and $L_{reg}$?", "limitations": "yes", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "# Strengths\n1. MuToR is a simple yet effective method, demonstrating better performance over both parallel MTP and NTP across diverse settings, including SFT, PEFT, and auto-regressive image generation.\n2. The ablation experiments are comprehensive, effectively analyzing the influence of individual components and yielding notable observations, such as the finding that using identical registers is more effective than using different ones.\n3. The paper is clearly and neatly written.\n\n# Weakness\n1. The baselines seems to be limited. Why not compared to the sequential MTP from DeepSeek V3? I think the method is also mentioned in related works.", "was_revised": true, "final_justification": "The rebuttal addressed my concerns. Therefore, I maintain my rating as Borderline accept.", "rebuttal": "We thank the reviewer for their positive and constructive feedback.   We are pleased that they found our method to be \"simple yet effective\" and that they appreciated our ablation analysis and writing style. Below we address the concerns they raised.\n\n---\n## **Question 1**\n> How does the method compare to other types of MTP like sequential MTP from DeepSeek V3?\n\nWe did not include a comparison with DeepSeek-V3’s sequential MTP[1] in the original submission due to the complexity of adapting its sequential-head architecture for finetuning pretrained LLMs. Following the reviewer’s suggestion, we have now implemented and evaluated this baseline.\nUsing the Gemma 2B model, we conducted experiments on the GSM8K training set, the 1M-GSM split and the 1M-MATH split, tuning the key hyperparameters $d_{max}=2$ (maximum prediction depth) and $a=0.1$ (auxiliary loss coefficient).\n\nThe results can be seen in the table below; we report average test accuracy (%) over 3 seeded runs, and include the other baselines as well. The headers refer to the finetuning split used. \n\n| Method                         |# Add. Params.| GSM8K       |  1M-GSM     | 1M-MATH |\n|---                             |---           | ---         |  ---        | ---     |\n| *Next-Token*                   |-             |  38.87      |  66.09      | 26.73   |\n| *Multi-Token (Gloeckle et al. [2])* |110M         |  40.66     |  66.69       | 26.87   |\n| Multi-Token (DeepSeek-V3 [1])       |118M          |  40.61      |  66.97      | 27.0    |\n| **MuToR (ours)**               |2K            | **42.10**   |  **68.33**  |**28.13**|\n\n\nAs seen in the results, the sequential MTP achieves similar accuracies with the parallel Multi-Token method, but at the cost of greater complexity. Notably, our MuToR method outperforms both of them, suggesting that it is better suited for the downstream finetuning of pretrained Language Models.  In addition, it does so by relying on negligible trainable parameters when comparing to the multi-token baselines(see column 2).\n\nBesides that, the increased complexity of the sequential MTP method (especially as $d_{max}$ scales) limits its applicability, while our method can seamlessly integrate within models without requiring architectural changes. This is particularly significant in domains such as autoregressive image generation, where scaling $d_{max}$ (thus predicting further into the future) would require lots of additional transformer heads for the existing methods [1],[2].\n\nWe appreciate the reviewer's feedback and will include these experiments in the final version of the paper.\n\n---\n## **Question 2**\n> What is the influence of the ratio of $L_{ntp}$ and $L_{reg}$?\n\nWe analyzed the impact of the auxiliary loss coefficient $a$ (controlling the relative weight between $L_{ntp}$ and $L_{reg}$), while fixing $d_{max}=4$.\nWe finetuned the Gemma 2B model on GSM8K training set and the 1M-GSM split.\n\nThe results are presented in the table below. We report test set accuracy (%) over *1 seeded run* (due to the substantial computational resources needed to replicate each experiment 3 times).\n\n| a         | GSM8K       |  1M-GSM     |\n|---        | ---         |  ---        |\n| 0.5       | 40.78       | 67.39       |\n| 0.4       | 40.94       | 68.00       |\n| **0.3**   |  **42.0**   | **68.38**   |\n| 0.2       | 41.24       | 67.55       |   \n| 0.1       |  40.33      | 67.32       |\n| *0 (Next-Token)*| *38.87* | *66.09*   |\n\nWe observe that optimal downstream performance is achieved when $a=0.3$, with both higher and lower values leading to mild degradation. These results are expected; an optimal balance should exist between next-token and multi-token prediction supervisory signals.  \n\nNotably, the optimal $a$ depends on factors like the pretrained model, task complexity, data distribution, and $d_{max}$. For example, in autoregressive image generation, we found that $a=0.5$ achieves better generation quality, likely because the multi-token prediction objective captures rich spatial dependencies. \nWe will include these ablation experiments and insights in the final version of the paper.\n\nWe appreciate the reviewer's time and hope we have addressed their concerns.\n\n\n---\n[1] Liu et al., \"DeepSeek-V3 Technical Report\", arXiv preprint 2024.\n\n[2] Gloeckle et al., \"Better & Faster Large Language Models via Multi-token Prediction.\" ICML 2024."}]

Author rebuttal:

Meta-review: In this work, the authors propose MuToR with the main idea being to augment the prompt during training with some learnable 'register' tokens, whose task is to predict a future token. This way, they are able to do multi-token prediction instead of the traditional next-token ones, effectively giving their models better planning capabilities. The authors show experiments in several setting from mathematical problems to autoregressive image generation, validating their approach.

Before the rebuttal, the paper received mixed scores with 1 Accept, 2 Borderline Accepts and 1 Borderline Reject. Two reviewers praise the paper's novelty and three of the reviewers find the experimental section quite good. Several reviewers also find the paper easy to read and well-written. All three positive reviewers find only mild weaknesses in the paper, and those weaknesses are properly addressed in the rebuttal.

On the other hand, reviewer foD8 has a somehow strange argument saying that the paper introduces "computational halting", a term never used in the paper. They also mention that the paper is a reimplementation of Goyal et al. (2024) paper, something that the authors strongly push against. Having read both papers, I tend to agree with the authors in this point. The reviewer also has some arguments on the evaluation mismatch (e.g., in Llama 3 8B Instruct) or in utilization of additional computational pathways which do not stand. Despite the relative initial harsh review, the reviewer finds the authors arguments in the rebuttal convincing and decides to increase their score to Borderline Accept.

At this stage, all four reviewers unanimously recommend the paper to be accepted, with the reviewer zeH6 finding the paper quite strong. After reading the reviews and the rebuttal, I agree with the reviewers about the quality of the paper and find all the weaknesses in the reviews completely addressed in the rebuttal. I urge the authors to integrate their strong rebuttal in the camera-version of the paper. Congratulations to the authors!
Decision: Accept (poster)