All papersLLMs & Language ModelsMIT

Multi-Token Prediction Done Right

/ Paper arXiv 2610.38331

Key numbers

1.2BFrom 167M to 1.2B parameters

From the paperCopied from the abstract, not written by AI.

Abstract

From the paperThe authors' own words.

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.

From the full paper

From the paperCopied from the paper itself by our code.

Reading the paper on arXiv…

Where it fits

The most cited papers it builds on, and the papers that build on it.

/ Builds on 3 most cited

  1. 201742.1k citedDecoupled Weight Decay Regularization
  2. 201931.3k citedLanguage Models are Unsupervised Multitask Learners
  3. 202322.6k citedLLaMA: Open and Efficient Foundation Language Models ↗

/ Built on by None yet

No papers cite it yet. We check again every week.

From Semantic Scholar, updated weekly.

Try it

Code, models and demos that use this paper.

Models:
JaydenTeoh/ntp-167m-model ↗12 models on Hugging Face

Ask this paper

/ Ask a questionAnswers come from the full paper and name the section

    / Your notes

    Select text in the abstract to highlight it or add a note. They stay in this browser.

      Read the PDF