Can Multi-Token Prediction Be Better? Unlock Soft-MTP with One Line of Code

Multi-Token Prediction (MTP) has become a standard recipe in frontier pretraining, but a recent study points out that its implementation is far from optimal. The researchers found that the gradient of standard MTP hides a KL self-distillation signal masked by a noisy residual—a signal that could improve speculative decoding, while the noisy residual harms next-token prediction performance. Through a simple modification called Soft-MTP, requiring just one line of code, the two signals can be decoupled, yielding universal improvements for MTP across different tasks.

The Hidden Concern of MTP: Signal and Noise in the Gradient

The core idea of MTP is to have the model predict not only the next token but also multiple future tokens. This approach has been widely adopted in frontier pretraining because it provides a richer training signal. However, this study reveals a neglected problem: the gradient of standard MTP is not purely a "useful signal."

Specifically, the gradient contains two components: one is the KL self-distillation signal, which helps the model learn more robust representations and thereby improves speculative decoding; the other is a noisy residual, which interferes with the accuracy of next-token prediction. The two are entangled, so that while MTP improves certain capabilities, it may harm basic language modeling.

The authors describe this finding as "the standard MTP gradient hides a KL self-distillation signal." The key question is how to extract the beneficial KL signal alone without introducing noise.

KL Self-Distillation is a variant of knowledge distillation: the teacher and student models share the same parameters or architecture, and the student learns by minimizing the KL divergence between its output distribution and the teacher's. In the MTP setting, the model's predictive distribution over multiple future tokens can be regarded as a kind of "soft label"; using it to constrain the next-token prediction distribution at the current position forms the self-distillation signal. This signal benefits speculative decoding because it makes the model's output probability distribution smoother and closer to the true multi-step future distribution, so that when sampled by the draft model, tokens are more likely to be accepted by the target model. Conversely, the noisy residual refers to the gradient component introduced by the MTP head that is inconsistent with the next-token prediction objective; it carries no additional information about the next token and instead biases the update direction of the shared backbone.

Speculative Decoding is a lossless technique for accelerating LLM inference: a lightweight draft model (or a head of the model itself) quickly generates several candidate tokens, and the target model then verifies these candidates in parallel, accepting tokens consistent with the target distribution and rejecting and correcting the inconsistent ones. The speedup depends directly on the "acceptance rate"—the proportion of candidate tokens accepted by the target model. MTP improves speculative decoding precisely because it trains the model to predict multiple future tokens simultaneously, bringing the draft distribution closer to the target distribution and thus raising the acceptance rate. Understanding this is essential to see why preserving the KL self-distillation signal in the MTP gradient while removing the noisy residual affects both generation quality and inference efficiency.

The Magic of One Line: A KL Gradient Recovery Trick

The research team proposes an elegant solution: by modifying the residual connection of the hidden states, the KL gradient is "recovered" from the noise. The operation requires just one line of code:

h_ntp += h_mtp.detach() - h_mtp

What this line does is, when computing the hidden state for next-token prediction, add the difference between the gradient-detached version of the MTP hidden state and the original version. This way, during backpropagation the gradient flows only through the KL self-distillation part, while the noisy residual is effectively blocked. The authors call this method Soft-MTP and regard it as almost a "free lunch"—it requires no extra computation, does not change the model architecture, yet delivers cross-task performance gains.

detach() and stop-gradient are fundamental operations in deep learning frameworks such as PyTorch: tensor.detach() returns a new tensor that shares data with the original but is detached from the computation graph, so gradients do not flow through it during backpropagation. The expression h_mtp.detach() - h_mtp is numerically always zero, so the hidden state is unchanged in the forward pass; but in the backward pass, since the h_mtp on the right side of the minus sign still retains the computation graph, the gradient flows through the MTP branch with a negated sign. This "zero forward, non-zero backward" trick is often called gradient surgery or zero-valued residual injection; it allows researchers to precisely control which gradient paths are preserved and which are cut off without changing the model's output.

Soft-MTP: Empirical Results of a Free Lunch

According to the paper, Soft-MTP outperforms standard MTP on a variety of tasks. Although the original material does not disclose specific benchmark data, the authors emphasize "improving MTP across a range of tasks," indicating that the method generalizes well.

In principle, the advantage of Soft-MTP is that it makes the MTP training objective "purer": the model can improve the acceptance rate of speculative decoding through the KL signal without sacrificing next-token prediction accuracy due to the noisy residual. This is of practical significance for large models that need to balance both generation quality and inference efficiency.

"Free lunch" here has a clear technical meaning: Soft-MTP adds no forward or backward computation, introduces no new parameters, and does not change the model architecture at inference time—all changes occur only in the gradient paths during training. In contrast, many methods for improving MTP require increasing the number of prediction heads, expanding the context window, or introducing an additional distillation teacher model, all of which incur memory and compute overhead. Soft-MTP costs almost nothing while delivering gains in both next-token prediction accuracy and speculative decoding acceptance rate—this is the basis for the authors calling it a free lunch. Note, however, that its effectiveness still depends on the training quality of the original MTP head; if the MTP head itself is underfit, the KL signal that can be separated out will also be limited.

Open-Source Resources and Reproduction Path

The research team has open-sourced the code, paper, and a HuggingFace model collection to facilitate community reproduction and further exploration:

• Code: https://github.com/JaydenTeoh/Soft-MTP
• Paper: https://arxiv.org/abs/2610.38331
• HuggingFace: https://huggingface.co/collections/JaydenTeoh/soft-multi-token-prediction

For researchers and engineers currently using MTP for pretraining, this one-line code change has an extremely low cost and is well worth trying.

Implications for Pretraining Practice

The value of this work lies not only in proposing a trick but also in reminding us that even a widely adopted "standard recipe" may have undetected flaws. The gradient decomposition perspective on MTP provides a new analytical tool for future research—by separating different training signals, the model's learning process can be controlled more precisely.

In the future, ideas similar to Soft-MTP may be extended to other multi-task training scenarios, helping researchers obtain cleaner and more controllable gradient signals while maintaining computational efficiency.

One-line Summary

One line of code decouples the KL signal from the noise in the MTP gradient, achieving a free cross-task improvement.

Paragraph Summary

This article introduces a recent study on Multi-Token Prediction (MTP). The researchers found that the gradient of standard MTP contains both a beneficial KL self-distillation signal and a harmful noisy residual: the former improves the acceptance rate of speculative decoding, while the latter harms next-token prediction accuracy. To address this, the authors propose Soft-MTP, which uses just one line of code, h_ntp += h_mtp.detach() - h_mtp, to modify the residual connection of the hidden states—the forward value is unchanged, and only the KL gradient is retained in the backward pass—thereby separating the two signals without adding computation or changing the model architecture. The paper reports that the method generally outperforms standard MTP on a variety of tasks and has open-sourced the code, paper, and HuggingFace model collection, making the change extremely low-cost for teams currently using MTP for pretraining.