Multi-Token Prediction Done Right

AI summary

Overview

• Research area: Machine learning / large language model pre-training objectives (multi-token prediction, self-distillation, speculative decoding).
• Technical level: Advanced — the core argument rests on an exact algebraic decomposition of cross-entropy gradients, KL divergence, and straight-through gradient routing.
• Scope: A single-author paper that diagnoses a flaw in standard Multi-Token Prediction training, derives a zero-cost fix ("Soft-MTP"), and validates it with LLaMA-style models from 167M to 1.2B parameters.

What This Paper Is About

Multi-Token Prediction (MTP) — training a model to predict several future tokens with extra heads — is now standard in frontier pre-training (the paper cites DeepSeek-V3/4, Nemotron 3, GLM-5, and Kimi K3), because it improves downstream performance and enables faster inference via speculative decoding. The author shows that the standard MTP cross-entropy gradient contains two parts: a useful KL self-distillation signal that pulls each MTP head toward the model's own next-token distribution, and a noisy "residual" that buries that signal and interferes with learning. The goal is to keep the useful part and discard the rest without paying the memory cost of explicitly computing a KL divergence over the full vocabulary.

Key Contributions

1. An exact decomposition of the MTP gradient. The hard MTP gradient with respect to the k-th MTP hidden state is shown to split exactly into g_KL (a KL self-distillation term toward the position-aligned NTP distribution q_t^(k)) and g_res (a residual equal to the aligned NTP cross-entropy gradient against the data).
2. 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 — no vocabulary-sized logit matrices, no custom kernel, and compatible with standard fused cross-entropy kernels.
3. Soft Multi-Token Prediction (Soft-MTP). A one-line code change that injects -g_res at the MTP hidden states via a numerically zero augmentation of the NTP hidden state, so only g_KL flows into the MTP heads and shared trunk. It supports both DeepSeek-style and Meta-style MTP, adds no parameters, and incurs no measurable overhead. Unlike direct KL minimization, it still trains the shared output head with cross-entropy.
4. A bias–variance explanation plus a full empirical study. The residual is decomposed into an NTP-error term (bias, which shrinks with scale) and one-hot sampling noise (irreducible variance), predicting that the trade becomes more favorable as models grow — which the experiments confirm.

Main Findings

• The residual is far noisier than the KL term. Measured over M = 256 independent batches (4 sequences of length 2048 each), the residual carries roughly 4× lower signal-to-noise ratio than the KL component, and the gap widens with scale — from 4.1× at 167M to 4.6× at 1.2B.
• The noisy part dominates the update. Despite its low SNR, the residual carries about 80% of the MTP gradient reaching the shared trunk at every scale, which is why the full hard-MTP gradient has lower SNR than the KL component alone.
• The residual actively fights NTP. A unit step along the KL component consistently decreases NTP loss, while a unit step along the residual consistently increases it, throughout training and across scales.
• Lower loss at scale. Soft-MTP reaches lower NTP loss than standard MTP, with the gap widening as scale increases; from 689M parameters onward it also reaches lower MTP loss.
• Better and faster speculative decoding. Evaluated on FineWeb-Edu, Wikipedia, Books, Science, and Math (256 prompts of length 512, 512-token continuations, drafting with K = 3 heads and verifying against the model's own next-token head), Soft-MTP wins every metric at every scale. Averaged over the five domains it improves acceptance length by 1.8–2.2% and speedup by 1.4–1.5% at each scale, with lower KL divergence and TV distance to the verifier — e.g. at 1.2B, KL drops from 0.384 to 0.364 and TV from 0.260 to 0.253.
• Higher zero-shot accuracy. Average accuracy over the reported benchmarks improves by roughly 0.1 points at 167M and 332M, and roughly 0.5 points at 689M and 1.2B, where Soft-MTP also overtakes the NTP baseline. At 1.2B it posts the best average accuracy (57.31 vs 56.80 for MTP and 56.75 for NTP) and the best WikiText (17.489 vs 18.156) and LAMBADA (14.073 vs 14.996) perplexity.
• Gains are not uniform on every individual benchmark. At 332M, for example, Soft-MTP's LAMBADA perplexity (28.566) is worse than standard MTP's (26.997), and at 167M the NTP-only baseline still has the lowest FineWeb-Edu perplexity (15.230).
• Stronger reasoning after fine-tuning. At 1.2B, after discarding the MTP heads and fine-tuning with NTP loss only, Soft-MTP leads standard MTP in seven of eight reported cells and is the only objective that never falls below the NTP baseline: +3.1 points on HumanEval at pass@64 and +1.2 on MBPP at pass@64, +1.1 on GSM8K and +1.8 on MATH-500.
• Explicit KL buys little. Retaining hard cross-entropy supervision at the LM head while propagating the recovered KL gradient only through the MTP heads and trunk matches or exceeds direct KL minimization on most metrics, at no extra cost.

Methodology in Plain English

The author starts from algebra rather than architecture. Because the MTP heads and the next-token head share the same linear output layer, the gradient of the MTP cross-entropy loss can be split cleanly into "pull toward the model's own next-token prediction" and "pull toward the observed token." The second part is the noisy one: it contains both the teacher's own error and the irreducible randomness of sampling a single one-hot target.

Because the two parts are additive, the noisy part can be cancelled by adding an equal and opposite gradient. The trick is to build a second copy of the next-token hidden state that is numerically identical to the real one (so the forward pass and the NTP loss are unchanged) but whose gradient path is routed, straight-through-estimator style, into the MTP hidden states. Backpropagating the combined loss then delivers exactly the KL gradient to the MTP heads and trunk, using only the cross-entropy computations already being performed. The output head itself still receives ordinary cross-entropy gradients, so it stays fitted to the data.

Empirically, the author pre-trains LLaMA-style models from 167M to 1.2B parameters on FineWeb-Edu for 60,000 steps (batch size 0.5M tokens, sequence length 2048, about 31.5B tokens, roughly 1.3× Chinchilla-optimal for the largest model), using the Muon optimizer under the update-RMS rescaling of Liu et al. (2025a). The MTP setup is DeepSeek-style sequential with K = 3 heads, each at 3.2% of the trunk's non-embedding parameters, and λ_MTP = 1.0. Comparisons cover gradient SNR and gradient share, loss curves, speculative decoding, zero-shot benchmarks via the LM Evaluation Harness, and math/code fine-tuning (OpenMathInstruct-2 with 4.2B tokens; the Python subset of Stack-Edu with 21.0B tokens; GSM8K, MATH-500, HumanEval, MBPP via the BigCode Evaluation Harness). A Meta-style MTP variant is checked in an appendix.

Why This Matters

Impact on research. The paper reframes a widely deployed training recipe: standard MTP is not merely "more cross-entropy on shifted tokens," it already contains an implicit self-distillation signal that is being drowned out. Showing that a KL gradient can be recovered exactly from existing cross-entropy gradients — with no extra logits and no custom kernel — offers a general recipe for any setup where a teacher and student share a linear head, sidestepping the memory barrier that has kept explicit self-distillation out of frontier pre-training.

Real-world applications:

• Faster LLM inference serving, where MTP heads act as drafters for speculative decoding and small acceptance-length gains translate directly into throughput.
• Code assistants and code-completion models, where the paper reports the largest fine-tuning gains at larger sample budgets (HumanEval pass@64).
• Mathematical reasoning models, which the paper evaluates through GSM8K and MATH-500 after fine-tuning.
• Reinforcement-learning post-training pipelines, where the paper notes that generating rollouts is often the main bottleneck and better drafters help.

Industry relevance. MTP is already in production-scale training stacks (DeepSeek-V3/4, Nemotron 3, GLM-5, Kimi K3). A one-line change with no added parameters and no measurable overhead, compatible with fused cross-entropy kernels and existing parallelism, is unusually cheap to adopt. It also connects to post-hoc alignment efforts such as Nemotron 3 Ultra's "MTP Boosting" phase, which pursues the same goal after pre-training.

Future Directions

• Test at frontier scale. Evidence is from 167M to 1.2B parameters; whether the 4× SNR gap and the widening trend continue at much larger sizes is an open empirical question.
• Combine with existing distillation variants. The paper compares against top-k logit distillation, argmax-based approaches, and cached teacher-gradient methods only in discussion; combining Soft-MTP with those, or with a temperature-scaled forward KL as in MTP Boosting, is untested.
• Probe the bias–variance tradeoff systematically. Soft-MTP deliberately biases MTP heads toward the NTP distribution and away from the data distribution; the paper argues this is desirable for acceptance rate but does not characterize where the bias could hurt, or how it interacts with the λ_MTP weighting beyond the sweep reported.
• Extend beyond MTP. The KL gradient–recovery trick is stated for any two distributions sharing a linear output head, but only MTP settings (DeepSeek-style and Meta-style) are evaluated.

Target Audience

Machine learning researchers and engineers working on LLM pre-training objectives, speculative decoding, or knowledge distillation — particularly those who already run MTP heads and want a low-cost improvement. The paper assumes comfort with cross-entropy gradients, KL divergence, softmax, and straight-through estimators, so readers without that background will get the empirical story but not the derivation. Practitioners who only want the implementation can go straight to the one-line change and the appendix code.

Authors’ abstract

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.