Multi-Token Prediction Done Right

Multi-token prediction (MTP) trains language models to predict several future tokens, potentially improving model quality and enabling faster speculative decoding. Standard MTP produces a gradient with two parts: a useful self-distillation signal that moves future-token heads toward the model’s next-token distribution, and a noisy residual that obscures that signal and interferes with both next-token and multi-token learning. The interference grows with model scale. The proposed KL gradient-recovery method reconstructs the gradient of a KL divergence from cross-entropy gradients when both distributions share a linear output head, using straight-through gradient routing. It avoids creating vocabulary-sized logit matrices and works with fused cross-entropy kernels.

Soft-MTP applies this method as a one-line change, reusing existing next-token and multi-token computations while backpropagating only the useful self-distillation signal. It adds no parameters or measurable overhead and supports both DeepSeek- and Meta-style MTP. In experiments from 167 million to 1.2 billion parameters, it improves next-token perplexity, speculative-decoding speed, and zero-shot accuracy at every tested scale. Multi-token perplexity improves from 689 million parameters onward, and fine-tuning at 1.2 billion parameters produces stronger math and code reasoning. Soft-MTP also matches or exceeds direct KL-divergence minimization on most measures, making gradient cleanup a practical way to improve MTP without additional model or runtime cost.