Multi-Token Prediction Done Right

Jayden Teoh

cs.LGOct 28, 2026

Multi-Token Prediction Done Right

Authors: Jayden Teoh

Organizations: MIT

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.

Figures & tables

Panel of Figure 1 — Figure 1: Soft-MTP’s illustration, pseudocode, and loss performance compared to MTP. Blue line is the one-line code change . We refer readers to Algorithm 1 for the full code implementation of Soft-MTP, including the MTP loss weighting.

Figure 2: Illustration of position-aligned NTP distributions.

Method	Output head $W_{\theta}$	MTP heads and shared trunk
Standard MTP	CE gradient	CE gradient
KL-MTP	KL gradient	KL gradient
Soft-MTP	CE gradient	KL gradient

Table 1: Gradient routing under different MTP objectives. We omit the standard NTP gradient which is unchanged across methods.

Panel of Figure 3 — Figure 3: The residual has substantially lower SNR than the KL gradient but dominates the MTP gradient signal. (a-b) Signal-to-noise ratio (SNR) of each gradient component with respect to the shared trunk and MTP head parameters, across scales. (c) Share of the MTP gradient in the shared trunk carried by each component, averaged over training steps; the residual carries $\sim$ 80% at every scale.

Panel of Figure 4 — Figure 4: The residual worsens NTP and MTP loss especially as scale increases. (a) Change in NTP loss $\Delta\mathcal{L}_{\text{NTP}}$ from a unit step of the trunk hidden state representation along each gradient component across model scales. (b-c) NTP and MTP loss of standard MTP and Soft-MTP throughout training across scales.

Size	Model	Verifier match		FineWeb-Edu		Wikipedia		Books		Science		Math
Size	Model	KL $\downarrow$	TV $\downarrow$	Speedup	Acc. Len.	Speedup	Acc. Len.	Speedup	Acc. Len.	Speedup	Acc. Len.	Speedup	Acc. Len.
167M	MTP	0.485	0.305	1.90	1.52	1.88	1.49	1.88	1.49	1.88	1.49	1.76	1.36
167M	Soft-MTP	0.470	0.300	1.92 +1.1%	1.54 +1.3%	1.89 +0.5%	1.50 +0.7%	1.91 +1.6%	1.52 +2.0%	1.90 +1.1%	1.51 +1.3%	1.81 +2.8%	1.41 +3.7%
332M	MTP	0.448	0.288	1.96	1.57	1.94	1.54	1.94	1.55	1.94	1.54	1.87	1.47
332M	Soft-MTP	0.430	0.282	1.98 +1.0%	1.60 +1.9%	1.95 +0.5%	1.55 +0.6%	1.97 +1.5%	1.58 +1.9%	1.97 +1.5%	1.58 +2.6%	1.93 +3.2%	1.53 +4.1%
689M	MTP	0.410	0.271	1.88	1.62	1.84	1.58	1.85	1.60	1.86	1.60	1.81	1.57

Table 2: Self-speculative decoding performance across domains. KL divergence / TV distance: teacher-forced divergence of the MTP distribution from the model’s own next-token verifier on FineWeb-Edu validation set (mean over the three MTP heads). Speedup: improvement in inference speed using self-speculative decoding (relative to autoregressive decoding). Acceptance Length (Acc. Len.): average accepted tokens per draft, excluding the next-token prediction which is always accepted. Green: relative improvement of Soft-MTP over MTP. We bold the better of MTP and Soft-MTP per column; Soft-MTP wins every metric at every scale.

Size	Model	FW-Edu ppl $\downarrow$	Wiki. ppl $\downarrow$	LAMB. ppl $\downarrow$	LAMB. acc $\uparrow$	PIQA acc $\uparrow$	HellaS. acc_n $\uparrow$	Wino. acc $\uparrow$	ARC-e acc $\uparrow$	ARC-c acc_n $\uparrow$	OBQA acc $\uparrow$	SciQ acc_n $\uparrow$	Avg.
167M	NTP	15.230	27.311	42.157	34.66	66.87	40.63	52.01	58.46	28.33	23.80	74.60	47.42
	MTP	15.983	29.074	43.013	33.88	66.43	39.70	51.70	59.01	26.37	23.40	74.20	46.84
	Soft-MTP	15.952	28.318	43.790	34.12	66.10	39.55	53.20	58.84	27.65	22.20	73.60	46.91
332M	NTP	13.073	22.482	26.005	38.54	69.91	46.65	53.51	63.51	31.40	24.80	78.00	50.79
	MTP	13.602	23.856	26.997	38.77	69.59	45.92	54.06	63.01	31.23	22.20	78.30	50.38
	Soft-MTP	13.551	23.249	28.566	37.45	68.72	45.58	53.59	63.89	32.17	22.20	80.10	50.46

Table 3: Downstream language modeling evaluation. We bold the stronger result between standard MTP and Soft-MTP for every model scale (NTP results are only for reference).

Model	HumanEval			MBPP			GSM8K	MATH-500
Model	pass@16	pass@32	pass@64	pass@16	pass@32	pass@64	acc	acc
NTP	28.89 (T=0.5)	33.77 (T=0.6)	39.02 (T=0.6)	47.83 (T=0.5)	53.65 (T=0.5)	58.80 (T=0.5)	47.69	24.80
MTP	31.70 (T=0.5)	36.37 (T=0.5)	40.85 (T=0.7)	48.28 (T=0.5)	54.56 (T=0.5)	59.80 (T=0.6)	48.60	23.20
Soft-MTP	31.99 (T=0.5)	37.26 (T=0.5)	43.90 (T=0.6)	47.88 (T=0.5)	54.59 (T=0.6)	61.00 (T=0.7)	49.66	25.00

Table 4: Math and code fine-tuning results . Each 1.2B model is fine-tuned with the NTP loss only, on the Python subset of Stack-Edu for code and on OpenMathInstruct-2 for math. For code, we sample with top-

p=1.0

and no top-

k

, sweep the temperature over

\{0.1,0.2,\ldots,1.0\}

, and report the best pass@

k

, with the chosen temperature indicated in gray . For math, we evaluate 0-shot with greedy decoding. We bold the best result between MTP and Soft-MTP (NTP is shown for reference).

Panel of Figure 5 — Figure 5: Effect of MTP weight and prediction horizon for 332M models trained on 7.3B tokens. Every axis is lower-better. (a-b) Varying MTP weight $\lambda_{\text{MTP}}\in\{0.25,0.5,0.75,1.0\}$ at fixed horizon $K=3$ . Soft-MTP dominates the Pareto frontier for both plots. (c-d) Varying horizon $K\in\{3,4,5,6,7,8\}$ at fixed $\lambda_{\text{MTP}}=1.0$ . Soft-MTP leads on both axes at every $K$ shown.

	Matmul FLOPs	Throughput	MFU	Peak memory (GB)
Objective	(GFLOP/token)	(tok/s)		naive CE	fused CE
NTP	8.16	62,851	51.9 %	98.3	95.5
MTP	10.28	49,783	51.7 %	120.7	107.0
KL-MTP	10.43	46,990	49.6 %	129.8	128.1
Soft-MTP	10.28	49,748	51.7 %	121.0	107.0

Table 5: Training cost by objective at the 1.2B model with

K=3

MTP heads, measured on a single NVIDIA H200 GPU. Throughput and MFU are measured at the largest micro-batch size each objective fits. Peak memory is reported at a matched micro-batch of 19. Soft-MTP’s KL recovery is free : identical FLOPs, MFU and fused-CE peak memory to standard MTP, with throughput within

0.07\%

. KL-MTP is worst on every axis —lowest throughput and MFU, and the highest peak memory—and is the only objective that does not benefit much from using fused cross-entropy kernels.

	Metric	NTP	MTP	KL-MTP	Soft-MTP
LM	FW-Edu ppl $\downarrow$	10.57	10.89	10.74	10.74
	WikiText ppl $\downarrow$	17.50	18.16	17.53	17.49
	LAMBADA ppl $\downarrow$	14.16	15.00	14.30	14.07
	Avg. accuracy $\uparrow$	56.75	56.80	57.22	57.31
Spec. dec.	KL $\downarrow$	–	0.384	0.363	0.364
	TV $\downarrow$	–	0.260	0.252	0.253

Table 6: Comparison against KL-MTP at 1.2B parameters. We bold the best of the three MTP objectives (NTP is shown for reference). Language modeling (LM): perplexity metrics and average zero-shot accuracy over lm-eval tasks. Speculative decoding (Spec. dec.): self-speculative decoding on FineWeb-Edu; acceptance length excludes the next-token prediction, which is always accepted. Fine-tuning (Math / Code): performance after NTP-only fine-tuning on math and code.

Figure 6: Effective rank of the logit map $W_{\gamma}$ . Soft-MTP stays closest to the next-token readout.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 7 — Figure 7: Validation NTP and MTP losses of standard MTP and Soft-MTP with Meta-style MTP heads at 332M.

Data
	FineWeb-Edu pre-training	Section 6.2 $K/\lambda_{\text{MTP}}$ sweeps	Math fine-tuning	Code fine-tuning
Corpus	FineWeb-Edu	FineWeb-Edu	OpenMathInstruct-2	Stack-Edu (Python)
Sequence length	2048	2048	2048	2048
Steps	60 000	14 000	8 000	40 000
Tokens	31.5B	7.34B	4.19B	21.0B
Training parameters

Appendix

Table 7: Training, model and MTP hyperparameters. The global batch is identical everywhere (

256\times 2048

tokens). Both fine-tunes are trained with next-token prediction only.

Multi-Token Prediction

Multi-token prediction trains models to forecast several future tokens at once, supporting more efficient generation and richer predictive representations.

See all 43 papers

Oct 5, 2026

LightMTP: Lightweight Latent Multi-Token Prediction

Next-token prediction (NTP) is the standard pretraining objective for large language models, yet it provides an explicit training signal only for the immediate next token, which can lead models to exploit local patterns instead of capturing longer-range structure and ideas. Multi-token prediction (MTP) addresses this by training models to predict several future tokens. However, existing MTP methods often introduce a large number of new parameters with limited improvements in downstream performance. Latent MTP approaches address this efficiency issue by encoding future tokens into a vector representation. However, these approaches usually rely on external helper models for future token encoding. We propose LightMTP, a lightweight, i.e., parameter-efficient, latent MTP approach that bootstraps the future token representations from the model's own hidden states. Our two LightMTP variants extend supervision to more future tokens without requiring the additional computational overhead of conventional MTP nor the external supervision latent MTP normally relies on. LightMTP adds at most 1% extra parameters, retains better performance on general language modeling benchmarks, and achieves similar gains in planning, coding, and reasoning.

Tamara Czinczoll, Julie Kallini, Gerard de Melo +1

Hasso Plattner Institute / University of Potsdam · Stanford University · Tel-Aviv University

Language Model Pretraining Multi-Token Prediction Sequence Prediction

Ask PDF

Oct 1, 2026

Match the Distribution, Not the Compute: Post-Training Multi-Token Prediction Heads

Multi-token prediction (MTP) improves the throughput of autoregressive generation by enabling the language model to draft multiple next tokens per forward pass, while a verification step over draft tokens ensures that token distribution of the backbone is preserved. Every open MTP-family release (MiMo-7B, DeepSeek-V3, Qwen3) trains its heads jointly with the backbone over the full pretraining run of tens of trillions of tokens, thus setting the drafter quality at pretraining time. We ask whether a lightweight post-training pass on target-generated chain-of-thought is enough to reach the same expected throughput speedup on a frozen reasoning model, and study how a serving-time system built on such a checkpoint can be optimized. We present three findings. 1) On a frozen Qwen3-8B with

K{=}3

chained MTP heads, we show that a post-training recipe with plain cross-entropy on

\approx\!2.5

B tokens reaches or exceeds the expected speedup of jointly trained MiMo-7B on math, coding and knowledge benchmarks. Our post-training recipe utilizes

10^3

10^4\times

less MTP-training tokens as compared with joint pre-training of MiMO-7B MTP baseline. 2) We propose a chain-aware relaxation of draft token verification rule that allows a bounded drift from backbone language model token distribution. We show that this relaxation lifts expected speedups by

+12

+16\%

per benchmark while preserving task accuracy. 3) We propose an adaptive controller that dynamically chooses the number of MTP heads to be engaged at inference time and demonstrate recovery of upto

11

14\%

loss in speedup using fixed maximum MTP draft length.

Prachi Badarayani, Aidan Jay, Chenghui Zhou +8

Microsoft

Multi-Token Prediction Speculative Decoding LLM Inference Acceleration Language Model Post-Training

Ask PDF

Jun 25, 2026

EntMTP: Accelerating LLM Inference with Entropy Guided Multi Token Prediction

Multi-token prediction has been shown to increase data density during training, improve downstream text-generation quality, and serves as the defacto approach for self-speculative decoding. Existing foundation and open source models that use MTP heads commit to a static tree-based attention topology throughout the entire generation sequence, meaning the speculation depth, and thus the compute required during verification, stays constant regardless of the context. This is fundamentally misaligned with the entropy patterns of natural language where low-entropy regions often support reliable multi-step drafting, while high-entropy regions require more conservative speculation. To address this, we propose Entropy-guided Multi-Token Prediction (EntMTP), a training-free scheduler that toggles between tree-based attention topologies from a set of task-specific pareto-optimal trees conditioned on a running estimate of local generation entropy. By matching speculation depth to context predictability, EntMTP maximizes expected accepted-token throughput across the full distribution of generated text without sacrificing generation quality. When evaluated across Humaneval, ShareGPT, GSM8k, and Litbench benchmarks, EntMTP consistently achieves a 1.15x speedup against Hydra and peak speedup of 1.36x against Medusa baselines respectively.

Carrie Chen

Cornell University

Multi-Token Prediction LLM Inference Scheduling Speculative Decoding LLM Inference Acceleration

Ask PDF

Abstract

Multi-Token Prediction

Explore similar work

LightMTP: Lightweight Latent Multi-Token Prediction

Match the Distribution, Not the Compute: Post-Training Multi-Token Prediction Heads

EntMTP: Accelerating LLM Inference with Entropy Guided Multi Token Prediction