Multi-Token Prediction Done Right

Hugging Face Daily Papers Papers
multi-token-prediction language-models self-distillation speculative-decoding training-optimization

Summary

Soft-MTP 提出一种 KL 梯度恢复技巧，去除多token预测训练中的噪声残差梯度，在无额外开销下改进困惑度、投机解码与推理能力。

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.
Original Article
View Cached Full Text

Cached at: 10/30/26, 11:09 AM

Paper page - Multi-Token Prediction Done Right

Source: https://huggingface.co/papers/2610.38331

Abstract

Soft-MTP removes a noisy gradient residual from multi-token prediction via a straight-through KL recovery trick, improving perplexity, speculative decoding, and reasoning without added cost.

Multi-Token Prediction(MTP) is now a common ingredient in language-model pre-training, as it improves downstream performance and enables faster inference throughspeculative 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 MTPcross-entropy gradientcan be decomposed exactly into aKL self-distillationterm that pulls each MTP head toward the model’s ownnext-token prediction(NTP) distribution and improvesspeculative 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 theircross-entropy gradients throughstraight-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 SoftMulti-Token Prediction(Soft-MTP), a one-line code change to MTP that removes the noisy residual and backpropagates only theKL self-distillationsignal by reusing the existing NTP and MTP cross-entropy computations.Soft-MTPsupports both DeepSeek- and Meta-style MTP, adds no parameters, and incurs no measurable overhead. From 167M to 1.2B parameters,Soft-MTPimproves over standard MTP with lower NTP perplexity, fasterspeculative decodingand 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-MTPalso matches, if not exceeds, direct KL divergence minimization on most metrics.Soft-MTPis, in effect, a free lunch for MTP optimization. Our code can be found at https://github.com/JaydenTeoh/Soft-MTP.

View arXiv pageView PDFGitHub1Add to collection

Get this paper in your agent:

hf papers read 2610\.38331

Don’t have the latest CLI?curl \-LsSf https://hf\.co/cli/install\.sh \| bash

Models citing this paper12

#### JaydenTeoh/ntp-167m-model Text Generation• 0.2B• Updatedabout 4 hours ago • 1 #### JaydenTeoh/hard-mtp-167m-model Text Generation• 0.2B• Updatedabout 4 hours ago • 1 #### JaydenTeoh/soft-mtp-167m-model Text Generation• 0.2B• Updatedabout 4 hours ago #### JaydenTeoh/hard-mtp-332m-model Text Generation• 0.4B• Updatedabout 4 hours ago Browse 12 models citing this paper## Datasets citing this paper0

No dataset linking this paper

Cite arxiv.org/abs/2610.38331 in a dataset README.md to link it from this page.

Spaces citing this paper0

No Space linking this paper

Cite arxiv.org/abs/2610.38331 in a Space README.md to link it from this page.

Collections including this paper1

Similar Articles

AdaMTP: An Adaptive Training Paradigm for Multi-Token Prediction

arXiv cs.CL

This paper introduces AdaMTP, an adaptive training paradigm for multi-token prediction that dynamically aligns prediction horizons with sequence predictability using entropy-based segmentation, consistently outperforming standard MTP on math, code, and general benchmarks across three LLM backbones.

Multi-Token Residual Prediction

arXiv cs.LG

Introduces Multi-token Residual Prediction (MRP), a lightweight module for diffusion language models that enables dependency-aware multi-token denoising within a single backbone forward pass, achieving up to 1.42× lossless speedup.

@jakevin7: Recently I've been reading about GLM 5.2 and found some interesting things to share. GLM-5.2 uses MTP (Multi-Token Prediction) to accelerate inference: a lightweight "draft model" quickly predicts multiple tokens, then the main model verifies them all at once; if accepted, it skips the decoding steps.

X AI KOLs Following

GLM-5.2 adopts MTP (Multi-Token Prediction) technology to accelerate inference and fixes a training-inference discrepancy in GLM-5.1's MTP that caused KV cache mixing issues.