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.

---

Question 1

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?

As we state in Section 3.2 of our paper, the interleaving pattern is flexible—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.

To 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.

These 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.

Moreover, 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.

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.

This 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—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.

---

Question 2

the offset $d$ is sampled uniformly and FIXED during training, i.e. for different register token, they have the same $d$ parameter?

First, 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.

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.

As 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}$.

Regarding 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.

Relevant 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.

---

Question 3

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.

We 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).

Through 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.

Moreover, 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.

Recent 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.

Could the author provide some numerical experiments on pretraining? 1b or 7b models are desirable.

While 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—particularly its adaptability and scalability compared to existing multi-token methods.

We thank the reviewer for their time and feedback, and we hope to have addressed their concerns.

---

[1] Kaplan et al.: "Scaling laws for neural language models." arXiv preprint 2020.

[2] Hoffman et al.: "Training compute-optimal large language models." arXiv preprint 2022.

[3] Gloeckle et al.: "Better & Faster Large Language Models via Multi-token Prediction." ICML 2024.

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.

---

Question 1

How does the method compare to other types of MTP like sequential MTP from DeepSeek V3?

We 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. Using 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).

The 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.

As 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).

Besides 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].

We appreciate the reviewer's feedback and will include these experiments in the final version of the paper.

---

Question 2

What is the influence of the ratio of $L_{ntp}$ and $L_{reg}$?

We analyzed the impact of the auxiliary loss coefficient $a$ (controlling the relative weight between $L_{ntp}$ and $L_{reg}$), while fixing $d_{max}=4$. We finetuned the Gemma 2B model on GSM8K training set and the 1M-GSM split.

The 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).

We 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.

Notably, 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. We will include these ablation experiments and insights in the final version of the paper.

We appreciate the reviewer's time and hope we have addressed their concerns.

---

[1] Liu et al., "DeepSeek-V3 Technical Report", arXiv preprint 2024.

[2] Gloeckle et al., "Better & Faster Large Language Models via Multi-token Prediction." ICML 2024.

We thank the reviewer for their time and feedback.

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.

---

Summary

This paper introduces a method termed "Computational Halting," which aims to enhance the mathematical reasoning capabilities of large language models (LLMs).

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.

We 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.

In 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.

Based 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.

---

Originality : A Reimplementation of Prior Work

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".

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.

We kindly argue that the reviewer's statement is invalid. Our proposed method, MuToR, is orthogonal to the mentioned previous work [1], since:

1. 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.

2. 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".

3. We propose a method for multi-token prediction in autoregressive transformer models. Our work introduces a novel multi-token prediction training objective—associated with the learnable tokens—and investigate its effectiveness compared to the baselines. However, the authors in [1] investigate whether—and in what settings—can a Language Model benefit from increased computations. They do not use any auxiliary training loss for their learnable tokens.

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.

---

Quality : Critical Technical and Methodological Flaws

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.

As 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.

The paper seems to conflate its method with other, mechanistically distinct "special token" approaches, revealing a shallow understanding of the literature.

As 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.

---

Quality: Weak and Misleading Evaluation

The paper's primary claim of improvement is based on a comparison to the base Llama 3 8B Instruct model.

The experiments are confined to a single model scale (8B parameters).

In 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. We list our experiments here:

1. Supervised Finetuning of Pretrained LLM, using Gemma 2B and Llama 3 8B (the base model),in both mathematical reasoning and abstractive summarization benchmarks.
2. Parameter Efficient Finetuning, using Gemma 2B model.
3. Autoregressive Image Generation, where we pretrain a LlamaGen transformer model on ImageNet 256x256.
4. Experiments on the synthetic star-graph task, where we finetune a pretrained GPT2-Large model.

The gains reported in this paper are dwarfed by the improvements from a standard, data-centric fine-tuning approach.

Our 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. Thus, 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.

The paper fails to perform the necessary ablation studies to validate its claims.

In 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.

---

Questions

1.

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.

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.

2.

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.

As we already stated (see our arguments above), our method is orthogonal to StdPT_PauseFT and PausePT_PauseFT (both variants proposed in [1]).

3.

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?

As 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.

Regarding 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).

4.

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.

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.

We 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.

---

[1] Goyal et al., "Think before you speak: Training Language Models With Pause Tokens.", ICLR 2024.

[2] Darcet et al., "Vision Transformers Need Registers.", ICLR 2024.

[3] Toshniwal et al., "OpenMathInstruct-2: Accelerating AI for Math with Massive Open-Source Instruction Data.", ICLR 2025

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.

---

Question 1

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.

We 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.

To 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×—that means that we even used more compute for the baseline methods.

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.

We 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.

Regarding 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.

---

Question 2

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.

As the reviewer correctly points out, FlashAttention does not currently support custom attention masks, which prevents us from using it directly with MuToR. To 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.

Using 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:

• 0.35s—with MuToR + FlexAttention (sequence_length=512, custom mask).
• 0.34s—with standard FlashAttention-2 (sequence_length=512, assuming a causal mask).
• 0.16s—with standard FlashAttention-2 (sequence_lenth=256, Next-Token baseline).

This demonstrates that MuToR’s training speed penalty—due to the custom mask—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.

We thank the reviewer for their time and we hope that their concerns might be resolved.

---

[1] Kaplan et al., "Scaling laws for neural language models." arXiv preprint 2020.

[2] Hoffman et al., "Training compute-optimal large language models." arXiv preprint 2022.

[3] Dong et al., "Flex Attention: A programming model for generating optimized attention kernels." arXiv preprint 2024.

1. The authors are missing a practical trade-off that is at least as important as the claim MuToR will help maximize performance from a given dataset.

especially in scenarios where data is limited—a common real-world constraint in many academic and industrial applications.

The dataset sizes MuToR operates on in your experiments are tiny (100M tokens) compared to what data is available for grade-school math etc. Any reasonable user given the knowledge the MuToR costs X% more (either wall-clock time or FLOPs) should ask themselves, should I pay for X% more cost by using MuToR or by simply training on more data. This is a completely natural and reasonable comparison. Your experiments do not approach the scale of data being limited. If your argument is "Use MuToR when training data is limited" instead of "Use MuToR", you should say so because they are vastly different arguments.

2. Regarding the additional training costs due to sequence length. The authors will be aware that especially in Math-style questions, the answer can be extremely long (Reasoning/Chain-of-thought etc.) so once this method is scaled beyond 512 tokens (which is again very small), the additional training cost of MuToR will become significant and the authors can do a more rigorous FLOPs calculation for their chose models to work out when the sequence length term in compute will become dominant (it will be relatively low for their small models because the FFNs are not that large).