Multi-Token Prediction Done Right

Multi-Token Prediction Done Right. By Jayden Teoh, from MIT. Published on the twenty-eighth of October.

Multi-token prediction, or MTP, is now a common ingredient in language-model pre-training. It improves downstream performance, and it enables faster inference through speculative decoding. The paper identifies a flaw in how it is usually trained. A residual component of the MTP gradient interferes with both next-token learning and multi-token learning, and the interference grows with scale.

The diagnosis comes from an exact decomposition. The MTP cross-entropy gradient splits into two parts. One is a KL self-distillation term, which pulls each MTP head toward the model's own next-token prediction distribution, the NTP distribution. That term is useful: it is the part that improves speculative decoding. The other part is a residual with a low signal-to-noise ratio, and it buries the useful KL signal underneath it.

If only the KL term is wanted, the obvious move would be to minimize the KL divergence directly. The paper instead 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, using straight-through gradient routing.

That matters for cost. Unlike direct KL minimization, the trick needs no vocabulary-sized logit matrices, and it stays compatible with standard fused cross-entropy kernels.

The trick makes possible Soft Multi-Token Prediction, or Soft-MTP. It is described as a one-line code change to MTP. It removes the noisy residual and backpropagates only the KL self-distillation signal, and it does so by reusing the next-token and multi-token cross-entropy computations that training already performs.

The practical profile follows from that reuse. Soft-MTP supports both the DeepSeek and the Meta styles of MTP. It adds no parameters. And it incurs no measurable overhead.

The experiments span models from one hundred and sixty-seven million to one point two billion parameters. At every scale in that range, Soft-MTP improves over standard MTP on three counts: lower next-token perplexity, faster speculative decoding, and higher zero-shot accuracy.

Two further gains are scale-dependent. MTP perplexity itself is lower from six hundred and eighty-nine million parameters onward. And at one point two billion parameters, after fine-tuning, math and code reasoning are stronger.

Against the alternative it set aside, direct KL divergence minimization, Soft-MTP matches it, if not exceeds it, on most metrics, while avoiding that approach's large logit matrices. The author's own summary is that Soft-MTP is, in effect, a free lunch for MTP optimization.

The argument runs in a straight line. Write the multi-token gradient out exactly, and it splits into two signals sharing one channel: a useful pull toward the model's own next-token predictions, and a noisy residual whose interference increases with scale. Keep the first, drop the second, and find a way to compute the first from quantities the training loop already has. Each step is small, and the last one is what makes the whole thing cost nothing.