Soft-MTP: Improving Multi-Token Prediction with One Line of Code

Multi-Token Prediction (MTP) has become a standard component of frontier large-model pretraining, but a new study points out that a KL self-distillation signal masked by noise is hidden within the gradients of standard MTP. The researchers' proposed Soft-MTP can cleanly extract this signal through a one-line code change, thereby improving MTP's performance across a variety of tasks for free.

The Hidden Problem of MTP: The KL Signal Is Drowned Out by Noise

The core idea of MTP is to have the model predict tokens at multiple future positions while predicting the next token. This mechanism has been widely used in frontier pretraining and directly helps with speculative decoding—because the model can generate multiple candidate tokens at once, decoding efficiency improves.

But the authors of this study found that the gradients of standard MTP are not "clean." It actually mixes two parts: one is the KL self-distillation signal, which is exactly what we hope to gain from MTP to improve speculative decoding; the other is a noise residual, which instead harms standard next-token prediction (NTP) performance. In other words, MTP's benefits and side effects are bundled into the same gradient and cannot be treated separately during training.

According to the authors' introduction on Reddit, this discovery is the starting point of the entire work: since the gradient contains both the useful KL signal and the harmful residual, can the KL gradient be "fished out" separately without changing the overall training framework?

The KL self-distillation signal refers to the following: in MTP, the model uses the same network to output predictions for multiple future positions simultaneously, and the distribution differences among these predictions can be regarded as a kind of "self-distillation"—the model tries to make the prediction distributions at farther positions consistent with those at nearer positions, thereby conveying richer supervisory information. This consistency constraint is mathematically equivalent to minimizing the KL divergence between prediction distributions, hence it is called the KL self-distillation signal. Unlike traditional knowledge distillation (using a large model to guide a small model), here the teacher and student are different output heads of the same model, which falls under the category of "self-distillation." This signal is believed to help improve the model's ability to model long-range dependencies and directly improve the acceptance rate of speculative decoding.

The noise residual refers to the remaining part of the MTP gradient that is orthogonal to the KL signal. The authors found that this part of the gradient does not correspond to any beneficial learning objective, and instead interferes with the standard next-token prediction (NTP) loss, leading to a decline in basic language modeling performance. The residual may originate from cumulative errors introduced in multi-step prediction, distribution mismatch caused by position shifts, and gradient conflicts among different prediction heads. Separating the KL signal from the noise residual is precisely the core motivation of Soft-MTP.

The One-Line KL Gradient Recovery Trick

The authors' answer is a technique called "KL gradient recovery," whose implementation is surprisingly concise—only a one-line code change is needed:

h_ntp += h_mtp.detach() - h_mtp

What this line of code does is keep the numerical value unchanged during the forward pass (because h_mtp.detach() - h_mtp is zero in the forward pass), but change the direction of gradient flow during backpropagation, thereby separating the KL self-distillation gradient from the noisy MTP gradient. The authors specifically mention in the post that this is what they consider the "coolest part" of the entire paper.

The engineering value of this approach lies in the fact that it does not require redesigning the model architecture or adding an extra training stage; existing pretraining pipelines that use MTP can apply it directly. For teams already running large-scale pretraining, the cost of changing one line is extremely low.

The gradient-blocking mechanism of .detach(): In PyTorch, tensor.detach() returns a new tensor disconnected from the computation graph. Its value is the same as the original tensor, but during backpropagation gradients are not passed back to the original tensor. Therefore, h_mtp.detach() - h_mtp is always zero in the forward pass and does not change the output value; but in the backward pass, the h_mtp on the right side of the minus sign receives a negative gradient, while the h_mtp.detach() on the left side is treated as a constant. Adding this term to h_ntp is equivalent to injecting the gradient produced by the MTP head into the NTP path in the opposite direction during backpropagation, thereby canceling the part of the MTP gradient that is not the KL signal. This "zero forward, nonzero backward" trick is often called gradient surgery or a variant of the straight-through estimator in the literature. Its elegance lies in introducing no additional parameters or loss terms and achieving selective routing of gradients solely through computation-graph operations.

Speculative decoding is an inference acceleration technique: a small model (or the shallow layers of a model) quickly generates multiple candidate tokens, which are then verified in parallel by a large model, thereby reducing the number of serial decoding steps. MTP is naturally suited to speculative decoding because the model itself can output multiple future tokens at once. But if the noise residual in the MTP gradient damages the quality of the base model, the acceptance rate of speculative decoding will also decline accordingly. By purifying the gradient, Soft-MTP enables the MTP head to focus more on improving the consistency of multi-step prediction, thereby indirectly safeguarding the gains from speculative decoding.

The Effectiveness and Scope of Soft-MTP

The authors named this method Soft-MTP and call it a "free lunch"—it can improve MTP's performance on a variety of tasks. Although the post does not elaborate on specific benchmark numbers, judging from the problem formulation, the benefits of Soft-MTP are mainly reflected in two directions: first, preserving MTP's gains for speculative decoding; second, avoiding the drag of the noise residual on next-token prediction.

This positioning is worth noting. Previously, MTP was regarded as a standard recipe for pretraining, and the community focused more on "whether to add MTP," whereas Soft-MTP pushes the question to "how to make MTP exhibit only its good side." For deployment scenarios that rely on speculative decoding for inference acceleration, this means continuing to enjoy the efficiency benefits brought by multi-token prediction without sacrificing basic language modeling quality.

The authors have also open-sourced the code, paper, and HuggingFace model collection, with relevant links provided in the post to facilitate reproduction and integration by researchers.

The meaning of "free lunch" in the machine learning context: It usually refers to methods that stably improve performance without increasing computational cost, changing the model architecture, or introducing additional hyperparameters. Soft-MTP fits this description because it modifies only one line of forward-pass code, and the training overhead is exactly the same as standard MTP. However, it should be noted that "free" does not mean "risk-free": changing gradient routing may affect training stability, especially when the loss surface is steep in the early stage of training. The authors' suggestion in the Reddit discussion to start with small-scale experiments to validate it is precisely out of this consideration. In addition, the benefits of Soft-MTP depend on the premise that the MTP head and the NTP head share the backbone network; if MTP uses independent parameters or a separated architecture, the applicability of this trick needs to be reevaluated.

Community Response and Follow-up Observations

Judging from the Reddit discussion, the appeal of this work mainly comes from two things: the extreme simplicity of the change and the clarity of the problem insight. Being able to separate gradient signals with one line of code—this "accomplish a great deal with little effort" style usually spreads quickly in engineering communities.

However, the post itself is more of a results announcement and does not yet provide detailed side-by-side comparisons with other MTP variants. The robustness of Soft-MTP across different model scales and different training-data mixtures, as well as whether the KL gradient recovery trick can be generalized to other multi-objective training scenarios, still awaits verification through the full paper and subsequent reproductions.

For teams currently using MTP, a pragmatic approach is to first validate the effect of this one-line change in small-scale training, and then decide whether to introduce it into the main training pipeline.

One-line Summary

Soft-MTP uses one line of code to separate the KL signal from the noise residual in MTP gradients, improving multi-token prediction for free.

Paragraph Summary

Soft-MTP is an improvement method for Multi-Token Prediction (MTP). The researchers found that the gradients of standard MTP mix a beneficial KL self-distillation signal with a noise residual that harms next-token prediction, making it impossible to separate MTP's benefits from its side effects. To address this, the authors propose a "KL gradient recovery" trick that requires only one line of code, h_ntp += h_mtp.detach() - h_mtp, to separate a pure KL gradient during backpropagation while keeping the forward value unchanged. This method requires no changes to the model architecture or additional training stages and can be directly embedded into existing MTP pretraining pipelines. The authors claim that Soft-MTP achieves a "free lunch" on a variety of tasks, preserving MTP's gains for speculative decoding while avoiding the drag of the noise residual on basic language modeling. The code, paper, and HuggingFace models have all been open-sourced, but robustness across model scales remains to be verified.