Researchers Remove the Noise from LLM MTP Training with a Single Line of Code

Almost all recent large language models solve a slightly harder task during pretraining than plain next-token prediction. They predict not only the next token but also several tokens ahead. The technique is called multi-token prediction, or MTP. It serves two purposes. The model handles long-horizon tasks such as math and code better, and the extra "heads" can act as fast drafters in speculative decoding. A multi-token draft is cheap to produce, and the main model verifies it in a single pass. Researchers behind Soft-MTP found that standard training of these heads carries a noisy contaminant, and they propose a way to remove it. It takes one additional line of code.

What is wrong with standard MTP training

MTP heads are trained with standard cross-entropy against real data. The authors show that the gradient of this loss splits into two parts, and the split is exact, with no approximations.

• The first part is KL self-distillation. KL divergence measures how much one probability distribution differs from another. Here it pulls each MTP head toward the model's own prediction for the same token. This part is what raises the acceptance rate of drafts in speculative decoding: a token from a head is accepted only if it matches the main model's prediction. • The second part is the residual. It is the gradient of ordinary next-token cross-entropy against the data, arriving by a different path. It has a low signal-to-noise ratio and drowns out the useful KL signal.

Why is the residual noisy? The target token in the data is a single sample from the true distribution, a one-hot vector with a 1 in one position. The variance of such a sample does not vanish as the model grows or as training runs longer. The authors measured this directly. The residual has a markedly lower signal-to-noise ratio than the KL component, and the gap widens with model size. The residual also points against the update that the main next-token prediction task receives. In effect, it hurts both the training of the MTP heads and the training of the model itself.

One line instead of vocabulary-sized matrices

The pure KL gradient can be obtained head-on: build the teacher and student distributions and minimize the divergence between them. But with a vocabulary of tens of thousands of tokens, this means holding huge logit matrices in memory. This approach also does not work with standard fused cross-entropy kernels, which exist precisely to avoid materializing those matrices. That is why frontier models in practice train MTP heads with ordinary cross-entropy.

The authors propose the KL gradient-recovery trick. If two distributions share a linear output head, the gradient of their KL can be recovered exactly by combining two cross-entropy gradients that have already been computed. Technically, this is straight-through gradient routing: the expression hk.detach() - hk is added to the hidden state of the main head. The value does not change, so training on the main task proceeds as before. But an additional gradient flows into the MTP representation, and when summed with the ordinary MTP loss gradient, it subtracts the noisy residual. As a result, only the KL signal reaches the MTP heads and the shared trunk of the model.

This is how Soft-MTP works. It adds no new parameters and creates no noticeable overhead. It works with both DeepSeek-style MTP and Meta's variant, and it is compatible with the parallelism used in pretraining large models. The code is available on GitHub: JaydenTeoh/Soft-MTP.

What the experiments from 167M to 1.2B showed

The authors trained models from 167M to 1.2B parameters and compared Soft-MTP with standard MTP.

• Next-token perplexity is lower at every size, and the gap grows with scale. • MTP-head perplexity is lower starting from the 689M-parameter model. • Soft-MTP keeps its zero-shot accuracy advantage at every scale. • Self-speculative decoding gets faster in all tested domains. • After fine-tuning, the 1.2B-parameter model does better on math and code tasks.

The authors also compared the method with direct KL minimization, which requires explicitly computing logits. On most metrics Soft-MTP is at least as good, and in places it wins. The likely reason is that Soft-MTP still trains the shared output head with ordinary cross-entropy against the data. Under pure KL, the head would drift toward the model's own predictions.

Cost and limitations

The method involves an honest trade-off. By removing the noisy residual, Soft-MTP introduces bias: the MTP heads are pulled toward the model's own distribution rather than the true data distribution. The authors argue that the trade gets more favorable with scale: the teacher model's error shrinks, while the sampling noise remains. Besides, MTP heads are needed primarily to speed up inference, where matching the main model's prediction is exactly what is wanted.

For teams pretraining their own LLMs with MTP, the practical takeaway is simple. The change is tiny, costs nothing in memory or speed, and does not break the familiar fused kernels. The experiments are limited to models of up to 1.2B parameters, so the result still has to be verified at DeepSeek-V3 scale. But the trend of gains growing with model size looks encouraging.