Soft-MTP: a one-line fix for multi-token prediction training

A paper submitted to arXiv on 28 October 2026 (Computer Science > Machine Learning), titled "Multi-Token Prediction Done Right", takes aim at how multi-token prediction (MTP) is trained. The abstract opens by noting that MTP is now a common ingredient in language-model pre-training, because it improves downstream performance and enables faster inference through speculative decoding.

The authors say they have found a flaw in standard MTP training. A residual component of the MTP gradient interferes with both next-token and multi-token learning, and does so increasingly as models get larger. Their analysis shows that the MTP cross-entropy gradient can be decomposed exactly into two parts. The first is a KL self-distillation term that pulls each MTP head toward the model's own next-token prediction (NTP) distribution; this term also improves speculative decoding. The second is a residual with a low signal-to-noise ratio that buries the useful KL signal.

To get rid of the residual, the paper introduces what it calls the KL gradient-recovery trick. When two distributions share a linear output head, the gradient of their KL divergence can be recovered exactly by combining their cross-entropy gradients through straight-through gradient routing. Unlike minimizing the KL directly, the trick needs no vocabulary-sized logit matrices and works with standard fused cross-entropy kernels.

That yields Soft Multi-Token Prediction (Soft-MTP). It is a one-line code change to MTP that removes the noisy residual and backpropagates only the KL self-distillation signal, reusing the NTP and MTP cross-entropy computations that already exist. The authors say Soft-MTP supports both DeepSeek- and Meta-style MTP, adds no parameters, and has no measurable overhead.

On results, the abstract compares Soft-MTP with standard MTP from 167M to 1.2B parameters. At every scale tested, Soft-MTP gives lower NTP perplexity, faster speculative decoding and higher zero-shot accuracy. Lower MTP perplexity shows up from 689M onward. At 1.2B, the authors report stronger math and code reasoning after fine-tuning. Against direct KL divergence minimization, Soft-MTP matches, if not exceeds, the results on most metrics. The authors sum it up as, in effect, a free lunch for MTP optimization, and say their code is available.

Key facts

  • The authors say the MTP cross-entropy gradient splits exactly into a useful KL self-distillation term and a low signal-to-noise residual that interferes with learning, increasingly so with scale.
  • Soft-MTP is a one-line change that backpropagates only the KL term, reusing the existing NTP and MTP cross-entropy computations, with no added parameters and no measurable overhead.
  • The KL gradient-recovery trick needs no vocabulary-sized logit matrices and works with standard fused cross-entropy kernels.
  • From 167M to 1.2B parameters, Soft-MTP beats standard MTP on NTP perplexity, speculative decoding speed and zero-shot accuracy at every scale; MTP perplexity improves from 689M onward.
  • At 1.2B the authors report stronger math and code reasoning after fine-tuning, and parity or better against direct KL minimization on most metrics.

Why it matters

MTP is described as a common part of language-model pre-training, so a flaw in how its gradient is handled would touch a lot of training runs. The paper's claim is specific: the standard objective mixes a useful KL self-distillation signal with a noisy residual, and the problem gets worse with scale. If that holds, the fix is unusually cheap, a single line, with no extra parameters and no measurable overhead according to the authors.

Who it affects

Teams that pre-train language models with MTP are the direct audience, since the method supports both DeepSeek- and Meta-style MTP. People who rely on speculative decoding for faster inference are also in scope, because the KL term the method keeps is the one the authors say improves it. The abstract does not say DeepSeek or Meta adopted or endorsed Soft-MTP; it only says the method supports both styles.

How to use it

The change is described as one line in an existing MTP setup. It reuses the NTP and MTP cross-entropy computations already being run, routes gradients with a straight-through trick, and is compatible with standard fused cross-entropy kernels, so no vocabulary-sized logit matrices have to be built. The authors say their code is available, but the abstract gives no URL for it in readable form.

How solid is it

This is an arXiv preprint submitted on 28 October 2026, and every claim here is the authors' own. No peer review or publication venue is stated. The decomposition of the gradient and the recovery trick are presented as exact. The empirical gains are reported across 167M to 1.2B parameters, but the abstract gives no numeric effect sizes, no perplexity values, speedup factors or accuracy percentages, and does not say which benchmarks or datasets were used for zero-shot accuracy or for math and code reasoning.

Risks and caveats

The tested range tops out at 1.2B parameters, so the abstract's claims say nothing about larger models, even though the flaw is said to grow with scale. Lower MTP perplexity is claimed only from 689M onward, not at every size. Soft-MTP matches or exceeds direct KL minimization on most metrics, not all of them. The "free lunch" description is the authors' framing, and the stronger reasoning result is reported only at 1.2B after fine-tuning.

“Soft-MTP is, in effect, a free lunch for MTP optimization.”

— Abstract of "Multi-Token Prediction Done Right", arXiv