AIarXiv
Ranked by formulaVerdict: NotableMulti-Token Prediction Done Right
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.
Caveats
- Preprint; not yet peer reviewed.
VerdictWorth a reader's time today.
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.
The editor's rubric
- Leverage4 / 5
- A general-purpose tool used across several fields (Adam, ResNet, LoRA, next-generation sequencing).
- Magnitude3 / 5
- Large gain: roughly 2x, or a clear new state of the art on a hard, unsaturated problem.
- Evidence3 / 5
- Solid: multiple benchmarks or cohorts, ablations, fair baselines, released code or data.
- Novelty3 / 5
- A genuinely new approach to an open problem.
- Trajectory3 / 5
- A clear path to scale.
- Stakes2 / 5
- Benefits a professional community (practitioners, clinicians, engineers).
Editor’s rationale
Heuristic triage from title and abstract text only, not a reading of the paper. Cues found: method (we propose); breadth (zero/few-shot); gains (outperforms); novelty (discovery); verification (false alarm rate, code released); scale (scalable, efficient); stakes (general AI).
How the score was computed
Desks place stronger verdicts first, then use this score to order work within the same verdict.
- Merit
- 6.3 / 10
- Adjusted merit
- 5.0 / 10
- Attention
- 0%
- Freshness
- 60%