Multi-token prediction carries a noisy gradient term: stripping it lifts a 1.2B model past next-token prediction on zero-shot accuracy
Synopsis
Multi-token prediction (MTP), which trains a language model to predict several future tokens, is now a common pre-training ingredient, but its standard cross-entropy gradient contains a low signal-to-noise residual that grows noisier with scale. Soft-MTP removes that residual and backpropagates only the KL self-distillation signal hidden inside it, reusing cross-entropy gradients already computed and materializing no vocabulary-sized logit matrices, as a one-line code change. On LLaMA-style models from 167M to 1.2B parameters, Soft-MTP beats standard MTP at every scale with lower next-token perplexity, faster self-speculative decoding and higher zero-shot accuracy; from 689M it also lowers MTP perplexity, and at 1.2B it gives stronger math and code reasoning after fine-tuning.
Interpretation
Multi-token prediction training has been carrying a gradient component that does not help learning, and that component can now be removed while the useful part is kept. The standard recipe trains MTP heads with cross-entropy against the data, as frontier models do; aligning MTP heads with the model's own next-token distribution previously meant explicitly minimizing a KL divergence, which requires materializing full-vocabulary teacher and student distributions. At the shared trunk, the residual accounts for 80% of the MTP gradient reaching it despite its low signal-to-noise ratio; a unit step along the residual consistently raises next-token loss, while a step along the KL component lowers it.
With the residual removed, multi-token prediction improves at once on language modeling, speculative decoding and downstream tasks, and the advantage grows with model size. Standard MTP does not improve multiple-choice accuracy over next-token prediction at small scales, consistent with earlier observations. Across the four scales from 167M to 1.2B, Soft-MTP beats standard MTP on next-token perplexity, on speculative-decoding acceptance length and speedup in five domains, and on average zero-shot accuracy; at 1.2B its average zero-shot accuracy is 57.31 against 56.80 for standard MTP and 56.75 for next-token prediction.
The implementation recovers the KL gradient from computations already performed at the hidden-state level, adding no parameters and no measurable overhead. Minimizing the KL directly requires materializing vocabulary-sized logit matrices, which is incompatible with the fused cross-entropy kernels used in pre-training. At 1.2B, Soft-MTP uses 10.28 GFLOP per token of matmul work and 107.0 GB fused-CE peak memory, the same as standard MTP, with throughput within 1%; KL-MTP uses 10.43 GFLOP and 128.1 GB.
Perspective
Experiments use LLaMA-style models from 167M to 1.2B parameters with the DeepSeek-style sequential MTP architecture and three MTP heads, trained on FineWeb-Edu for about 31.5B tokens in the main runs; Meta-style parallel MTP results are reported at 332M scale on 7.3B tokens. The authors also discuss compatibility with data parallelism, ZeRO, expert parallelism and pipeline parallelism, and identify sequence sharding as the only axis needing extra communication.
The authors frame Soft-MTP as trading bias for variance: without the residual, MTP heads are biased toward the model's own next-token distribution rather than the data distribution, a trade that becomes more favorable with scale but leaves the bias in place. The gradient signal-to-noise and gradient-share analyses are measured offline on frozen checkpoints, and the authors note this only characterizes behavior at the trunk-to-MTP-head interface, since the realized training update also depends on the optimizer. The authors also state that the long-term role of multi-token prediction in pre-training remains uncertain.