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

LightMTP: Lightweight Latent Multi-Token Prediction
Oct 5, 2026Tamara Czinczoll, Julie Kallini, Gerard de Melo +1Language Model Pretraining Multi-Token Prediction
Hasso Plattner Institute / University of Potsdam · Stanford University · Tel-Aviv University
Match the Distribution, Not the Compute: Post-Training Multi-Token Prediction Heads
Oct 1, 2026Prachi Badarayani, Aidan Jay, Chenghui Zhou +8Multi-Token Prediction Speculative Decoding
Microsoft
EntMTP: Accelerating LLM Inference with Entropy Guided Multi Token Prediction
Jun 25, 2026Carrie ChenMulti-Token Prediction LLM Inference Scheduling
Cornell University

Oct 5, 2026J/K move · Enter open · S save

LightMTP: Lightweight Latent Multi-Token Prediction

Tamara Czinczoll, Julie Kallini, Gerard de Melo, Chen Shani

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

Language Model Pretraining Multi-Token Prediction Sequence Prediction

Open paper AskPDF arXiv

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.

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