Papers
arxiv:2610.38331

Multi-Token Prediction Done Right

Published on Oct 28
· Submitted by
Jayden Teoh
on Oct 30
Authors:

Abstract

Soft-MTP removes a noisy gradient residual from multi-token prediction via a straight-through KL recovery trick, improving perplexity, speculative decoding, and reasoning without added cost.

Multi-Token Prediction (MTP) is now a common ingredient in language-model pre-training, as it improves downstream performance and enables faster inference through speculative decoding. We identify a flaw in standard MTP training, in which a residual component of the MTP gradient interferes with both next-token and multi-token learning, increasingly so with scale. Specifically, the MTP cross-entropy gradient can be decomposed exactly into a KL self-distillation term that pulls each MTP head toward the model's own next-token prediction (NTP) distribution and improves speculative decoding, and a residual that has a low signal-to-noise ratio and buries the useful KL signal. We introduce 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, this trick requires no vocabulary-sized logit matrices and is compatible with standard fused cross-entropy kernels. This enables Soft Multi-Token Prediction (Soft-MTP), a one-line code change to MTP that removes the noisy residual and backpropagates only the KL self-distillation signal by reusing the existing NTP and MTP cross-entropy computations. Soft-MTP supports both DeepSeek- and Meta-style MTP, adds no parameters, and incurs no measurable overhead. From 167M to 1.2B parameters, Soft-MTP improves over standard MTP with lower NTP perplexity, faster speculative decoding and higher zero-shot accuracy at every scale, lower MTP perplexity from 689M onward, and stronger math and code reasoning after fine-tuning at 1.2B. Soft-MTP also matches, if not exceeds, direct KL divergence minimization on most metrics. Soft-MTP is, in effect, a free lunch for MTP optimization. Our code can be found at https://github.com/JaydenTeoh/Soft-MTP.

Community

Paper author Paper submitter

We show the standard multi-token prediction (MTP) gradient hides a KL self-distillation gradient signal behind a noisy residual, and introduce the KL gradient–recovery trick to extract it. This is a free lunch that improves MTP across a range of tasks.

Really elegant result: splitting the MTP cross-entropy gradient into a KL self-distillation term plus a residual, then recovering only the KL part through straight-through routing on the shared head, as a one-line change that still works with fused cross-entropy kernels. The SNR analysis stood out to me, with the residual carrying about 80% of the MTP gradient reaching the trunk at roughly 4x lower SNR, and the gap widening from 4.1x at 167M to 4.6x at 1.2B. Given that trend, do you expect the gains to keep growing beyond 1.2B, or does the KL signal eventually saturate once the MTP heads track the NTP distribution closely? Curious too whether the 1.8-2.2% longer acceptance length holds up with more than K = 3 draft heads.

Sign up or log in to comment

Get this paper in your agent:

hf papers read 2610.38331
Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash

Models citing this paper 12

Browse 12 models citing this paper

Datasets citing this paper 0

No dataset linking this paper

Cite arxiv.org/abs/2610.38331 in a dataset README.md to link it from this page.

Spaces citing this paper 0

No Space linking this paper

Cite arxiv.org/abs/2610.38331 in a Space README.md to link it from this page.

Collections including this paper 4