Efficient Training-Free Multi-Token Prediction via Embedding-Space Probing

Raghavv Goel; Mukul Gagrani; Mingu Lee; Chris Lott

cs.CLMar 18, 2026

Efficient Training-Free Multi-Token Prediction via Embedding-Space Probing

Authors: Raghavv Goel, Mukul Gagrani, Mingu Lee, Chris Lott

Organizations: Qualcomm AI Research

Abstract

Large Language Models (LLMs) possess latent multi-token prediction (MTP) abilities despite being trained only for next-token generation. We introduce ESP (Embedding-Space Probing), a simple and training-free MTP method that probes an LLM using on-the-fly mask tokens drawn from its embedding space, enabling parallel future-token prediction without modifying weights or relying on draft models. ESP constructs a speculative token tree by sampling Top-K candidates from mask-token logits and applies a lightweight pruning rule to retain high-probability continuations. During generation, predictions are verified in parallel, yielding lossless decoding while significantly reducing model calls and increasing token throughput. ESP consistently outperforms existing training-free baselines, improving acceptance length by 7-11% over LADE on LLaMA3 and 7-8% on Qwen3, and increasing throughput by up to 15-19% over the strongest baseline. Finally, we provide theoretical insight and empirical evidence showing that decoder layers naturally align mask-token representations with next-token states, enabling accurate multi-step prediction without retraining or auxiliary models.

Figures & tables

Figure 1: (Left) Standard next-token prediction setup for autoregressive models, (middle) multi-token prediction during prefill-stage by probing mask tokens which are appended to prompt tokens, (right) multi-token prediction with parallel verification and generation. Mask tokens are associated with last generated token ( $x_{s}$ ) and future tokens ( $\hat{x}_{s+1},\hat{x}_{s+2}$ ) through custom tree attention mask.

Figure 2: We use Dolly-Databricks (creative-writing) ( Conover et al., 2023 ) samples (100) to measure average cosine similarity across layers for mask and true-future token hidden-states. For Llama3.2-3B-Instruct , higher cosine similarity in later layers (15 onwards) correlates with token acceptance (green), while lower similarity correlates with rejection (red).

Figure 3: Mask tokens are present for each input token: last generated (blue) and future (orange) tokens, when processed by model all tokens are flattened and mask tokens are placed at the end with appropriate position indices.

Model	BC=30								BC=60
	PLD		STAND		LADE		ESP		PLD		STAND		LADE		ESP
	$\tau$	S/R	$\tau$	S/R	$\tau$	S/R	$\tau$	S/R	$\tau$	S/R	$\tau$	S/R	$\tau$	S/R	$\tau$	S/R
LLaMA3.2-3B-Instruct	1.43	1.19 $\times$	1.55	1.01 $\times$	1.43	1.01 $\times$	1.56	1.22 $\times$	1.43	1.19 $\times$	1.62	1.07 $\times$	1.57	1.10 $\times$	1.63	1.22 $\times$
LLaMA3.1-8B-Instruct	1.44	1.23 $\times$	1.58	1.10 $\times$	1.45	1.06 $\times$	1.63	1.35 $\times$	1.44	1.23 $\times$	1.64	1.14 $\times$	1.60	1.14 $\times$	1.71	1.38 $\times$
Qwen3-8B	1.31	1.12 $\times$	1.42	1.02 $\times$	1.55	1.12 $\times$	1.63	1.37 $\times$	1.31	1.12 $\times$	1.48	1.06 $\times$	1.73	1.21 $\times$	1.74	1.43 $\times$
Qwen3-32B	1.29	1.09 $\times$	1.42	1.10 $\times$	1.52	1.22 $\times$	1.60	1.43 $\times$	1.29	1.09 $\times$	1.48	1.13 $\times$	1.69	1.31 $\times$	1.70	1.48 $\times$

Table 1: Comparison of multi-token prediction average acceptance length (

\tau

) and speedup ratio (S/R) over autoregressive decoding across models and methods averaged on Spec-bench tasks for block complexity

BC=30

and

BC=60

. All results on a single NVIDIA H100 GPU with 256-token outputs. PLD uses a fixed depth-10 tree (BC

\approx

10) regardless of budget; its values are shown for reference under both BC columns.

Figure 4: Evaluation on Spec-Bench using LLaMA3.1-8B-Instruct across block complexities (BC = 10, 30, 60). Our method (green) consistently achieves the highest average accepted tokens across most tasks and BC settings. WRIT=writing, ROLE=roleplay, CODE=coding, TRANS=translation, SUMM=summarization, M/R=math/reasoning, HUM.=humanities

Figure 5: Evaluation on SpecBench using Qwen3-32B across block complexities (BC = 10, 30, 60). Our method (green) consistently achieves the highest average accepted token across most tasks and BC settings.

Model	dynamic	[7,2]	[5,4]	[3,6]	dynamic	[15,4]	[12,7]	[10,9]	[8,11]	[6,13]
	BC=30 ( $m_{1},m_{2}$ )				BC=60 ( $m_{1},m_{2}$ )
LLaMA3.2-3B-Instruct	1.506	1.504	1.476	1.417	1.630	1.631	1.620	1.602	1.574	1.540
LLaMA3.1-8B-Instruct	1.571	1.570	1.540	1.479	1.712	1.708	1.701	1.687	1.665	1.624

Table 2: Average acceptance length (

\tau

) for BC

=30

and BC

=60

for LLaMA3.2-3B / LLaMA3.1-8B-Instruct with two mask tokens (

m_{1},m_{2}

) across different branch configurations. Dynamic uses adaptive branching per input; static configs use fixed trees. The best

\tau

is bold ; second best is underlined .

Model	( $m_{1}$ )	( $m_{1},m_{2}$ )	( $m_{1},m_{2},m_{3}$ )
	$[29]$	$[15,4]$	$[7,5,3]$
LLaMA3.2-3B-Instruct	1.65	1.63	1.51
LLaMA3.1-8B-Instruct	1.73	1.71	1.57

Table 3: Impact of number of mask tokens at BC=60. Branch widths are set in decreasing order across mask tokens. Bold = best

\tau

Method	$\mathbf{m_{1}}$ (30)	$\mathbf{m_{1},m_{2}}$ (30)	$\mathbf{m_{1}}$ (60)	$\mathbf{m_{1},m_{2}}$ (60)
LLaMA3.2-3B-Instruct
$\tau$	1.56	1.50	1.65	1.63
Naive	0.99 $\times$	0.99 $\times$	0.95 $\times$	0.96 $\times$
Efficient	1.22 $\times$	1.08 $\times$	1.20 $\times$	1.22 $\times$
LLaMA3.1-8B-Instruct
$\tau$	1.63	1.57	1.73	1.71

Table 4: Speedup over AR for naive and efficient implementations across mask configurations single and two mask tokens. For two mask tokens (

m_{1},m_{2}

[7,2]

for BC=30 and

[15,4]

for BC=60.

	LLaMA3.2-3B-Instruct			LLaMA3.1-8B-Instruct
Method	$\mathbf{m_{1}}$ (10)	$\mathbf{m_{1}}$ (30)	$\mathbf{m_{1},m_{2}}$ (60)	$\mathbf{m_{1}}$ (10)	$\mathbf{m_{1}}$ (30)	$\mathbf{m_{1},m_{2}}$ (60)
Last K (hard init)	1.36	1.53	1.62	1.38	1.56	1.67
Sample	1.39	1.57	1.65	1.41	1.60	1.69
(embedding distribution)	1.39	1.57	1.65	1.41	1.60	1.69
Mean (soft init)	1.41	1.59	1.67	1.42	1.62	1.71

Table 5:

\tau

for different mask token initializations for LLaMA3.2-3B-Instruct and LLaMA3.1-8B-Instruct. Last K sets

m_{i}=\mathbf{e}_{t-k-i}

, Sample draws

m_{i}\sim\mathcal{N}(\mu,\sigma^{2}I)

where

\mu,\sigma

are the mean and standard deviation of the full embedding table, and Mean (used in our main method).

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Example of dynamic token tree expansion for block complexity (BC) = 30 using two mask tokens. As the tree expands, each child node inherits the probability of its parent, resulting in a multiplicative score. We denote the last accepted token (root) as $0$ , the Top- $9$ tokens from the output logits of the first mask token as $\hat{t}_{1,1:9}$ , and the Top- $8$ tokens from the second mask token as $\hat{t}_{2,1:8}$ .

Figure 7: Attention mask and position index for the current (top) and next (bottom) LLM forward pass in our method with static branch configuration [3,2] . The tokens in blue are part of key-value cache, green token is the last generated token which will be inputted to the model to generate next token, orange tokens are future tokens, and purple tokens are mask tokens. Position IDs of each token is mentioned above it. The structure of the attention mask remains unchanged except for the addition of BE columns filled with zeros (assuming non-attended positions are represented by

-\infty

). Similarly, the position indices are preserved in order but uniformly incremented by BE .

Model	BC=30				BC=60
	PLD	STAND	LADE	ESP	PLD	STAND	LADE	ESP
LLaMA3.2-3B-Instruct	30.07	35.48	30.07	35.90	30.07	38.27	36.31	38.65
LLaMA3.1-8B-Instruct	30.56	36.71	31.03	38.65	30.56	39.02	37.50	41.52
Qwen3-8B	23.66	29.58	35.48	38.65	23.66	32.43	42.20	42.53
Qwen3-32B	22.48	29.58	34.21	37.50	22.48	32.43	40.83	41.18

Appendix

Table 6: Comparison of multi-token prediction reduction in model forward calls performance across models and methods averaged on Spec-bench tasks for block complexity

BC=30

and

BC=60

Panel of Figure 8 — Figure 8: Block efficiency across SpecBench tasks for LLaMA3.2-3B-Instruct and QWen3-8B at block complexity (BC) = 60. Our method consistently outperforms baselines across most tasks, demonstrating strong performance in both open-ended and closed-ended settings.

Method	BC	WRIT	ROLE	CODE	TRANS	SUMM	QA	RAG	M/R	AVG
PLD	30	1.53	1.10	1.24	1.22	1.62	1.15	1.77	1.46	1.39
STAND	30	1.31	1.17	1.31	1.20	1.55	1.18	1.53	1.43	1.33
LADE	30	1.37	1.27	1.39	1.33	1.38	1.39	1.39	1.52	1.38
OURS	30	1.51	1.46	1.53	1.52	1.57	1.44	1.49	1.61	1.52
STAND	60	1.33	1.22	1.33	1.22	1.61	1.22	1.58	1.46	1.37
LADE	60	1.49	1.38	1.56	1.42	1.50	1.47	1.55	1.64	1.50

Appendix

Table 7: Block-efficiency with temperature=

1.0

for Llama3.2-3B-Instruct across SpecBench tasks for BC = 30 and BC = 60. Task acronyms: WRIT (writing), ROLE (roleplay), CODE (coding), TRANS (translation), SUMM (summarization), QA (question answering), RAG (retrieval-augmented generation), M/R (math/reasoning).

Method	BC	WRIT	ROLE	CODE	TRANS	SUMM	QA	RAG	M/R	AVG
PLD	30	1.40	1.12	1.25	1.20	1.51	1.14	1.48	1.40	1.31
STAND	30	1.24	1.14	1.34	1.25	1.56	1.17	1.50	1.39	1.32
LADE	30	1.22	1.16	1.31	1.28	1.26	1.21	1.25	1.32	1.25
OURS	30	1.52	1.48	1.59	1.56	1.61	1.53	1.55	1.65	1.56
STAND	60	1.26	1.17	1.35	1.28	1.60	1.21	1.56	1.47	1.36
LADE	60	1.25	1.22	1.38	1.37	1.41	1.28	1.34	1.46	1.34

Appendix

Table 8: Block-efficiency with temperature=

1.0

for LLama3.1-8B-Instruct across tasks for BC = 30 and BC = 60. Task acronyms: WRIT (writing), ROLE (roleplay), CODE (coding), TRANS (translation), SUMM (summarization), QA (question answering), RAG (retrieval-augmented generation), M/R (math/reasoning).

Model	$m_{1}$	$m_{1},m_{2}$	$m_{1},m_{2},m_{3}$
	(BC=30)	(BC=60)	(BC=120)
LLaMA3.2-3B-Instruct	1.561	1.631	1.669
LLaMA3.1-8B-Instruct	1.632	1.708	1.756

Appendix

Table 9: Impact of number of mask tokens on average acceptance length (

\tau

). Each additional mask doubles the budget: BC=30 (

m_{1}

), BC=60 (

m_{1},m_{2}

), BC=120 (

m_{1},m_{2},m_{3}

). All configs use static tree expansion.

Model	BC	PLD	STAND	LADE	Ours	GPU hrs	EAGLE3	GPU hrs
						(training)		(training)
LLaMA-3.1-8B-Instruct	30	1.24	1.10	1.06	1.35	0	2.36	$\sim$ 200–300
LLaMA-3.1-8B-Instruct	60	1.24	1.15	1.14	1.38	0	2.61	$\sim$ 200–300
Qwen3-32B	30	1.07	1.07	1.17	1.43	0	N/A (no official checkpoint)
Qwen3-32B	60	1.07	1.11	1.24	1.48	0	N/A (no official checkpoint)

Appendix

Table 10: Speedup ratio (S/R) over autoregressive decoding. EAGLE3 requires 200–300 GPU hours of training; all other methods are training-free (0 GPU hours). N/A = no official EAGLE3 checkpoint available.

Config	WRIT	ROLE	CODE	TRANS	SUMM	QA	RAG	M/R
Llama3.2-3B-Instruct
$m_{1}$	1.56	1.54	1.57	1.56	1.64	1.55	1.70	1.69
$m_{1},m_{2}$	1.52	1.51	1.53	1.53	1.59	1.50	1.69	1.67
Llama3.1-8B-Instruct
$m_{1}$	1.58	1.56	1.62	1.62	1.65	1.59	1.73	1.70
$m_{1},m_{2}$	1.54	1.45	1.58	1.59	1.62	1.52	1.71	1.68

Appendix

Table 11: Block-efficiency across different tasks for BC=

\mathbf{30}

with different mask tokens probed (single mask:

m_{1}

, two mask tokens:

m_{1},m_{2}

). Task acronyms are WRIT: writing, ROLE: roleplay, CODE: coding, TRANS: translation, SUMM: summarization, M/R: math/reasoning

Config	WRIT	ROLE	CODE	TRANS	SUMM	QA	RAG	M/R
Llama3.2-3B-Instruct
$m_{1}$	1.64	1.65	1.64	1.65	1.72	1.64	1.78	1.78
$m_{1},m_{2}$	1.62	1.61	1.67	1.64	1.74	1.62	1.83	1.81
Llama3.1-8B-Instruct
$m_{1}$	1.67	1.64	1.73	1.70	1.74	1.69	1.80	1.78
$m_{1},m_{2}$	1.69	1.60	1.75	1.71	1.75	1.66	1.87	1.83

Appendix

Table 12: Block-efficiency across different tasks for BC=

\mathbf{60}

comparing different mask tokens probed. Task acronyms are WRIT: writing, ROLE: roleplay, CODE: coding, TRANS: translation, SUMM: summarization, M/R: math/reasoning

Method	$\mathbf{m_{1}}$ (10)	$\mathbf{m_{1}}$ (30)	$\mathbf{m_{1},m_{2}}$ (30)	$\mathbf{m_{1},m_{2}}$ (60)
LLaMA3.2-3B-Instruct
w/o tree pruning	1.41	1.59	1.50	1.63
w/ tree pruning	1.41	1.59	1.55	1.67
LLaMA3.1-8B-Instruct
w/o tree pruning	1.42	1.62	1.51	1.66
w/ tree pruning	1.42	1.62	1.57	1.71

Appendix

Table 13: Impact of tree pruning for different LLaMA models

Method	WRIT	ROLE	CODE	TRANS	SUMM	QA	RAG	M/R	AVG
$\mu+{\epsilon}\sigma$	1.60	1.58	1.65	1.61	1.72	1.59	1.64	1.79	1.65
$\mu+5\sigma$	1.56	1.55	1.60	1.60	1.67	1.55	1.61	1.75	1.61
$\mu+10\sigma$	1.54	1.55	1.61	1.60	1.69	1.55	1.61	1.75	1.61

Appendix

Table 14: Llama3.2-3B-Instruct average block efficiency (BE) across Spec-bench categories with BC=60. Task acronyms are WRIT: writing, ROLE: roleplay, CODE: coding, TRANS: translation, SUMM: summarization, M/R: math/reasoning

Model	BC	$\lambda=0.01$	$\lambda=0.1$ (default)	$\lambda=0.5$
LLaMA3.2-3B-Instruct	30	1.53	1.59	1.43
LLaMA3.2-3B-Instruct	60	1.63	1.67	1.52
LLaMA3.1-8B-Instruct	30	1.60	1.62	1.47
LLaMA3.1-8B-Instruct	60	1.69	1.71	1.57

Appendix

Table 15: Impact of EMA update rate

\lambda

for mask token embeddings on average acceptance length (

\tau

) at BC=30 and BC=60. Default

\lambda=0.1

consistently performs best; higher

\lambda

degrades performance more than lower

\lambda

, suggesting the method is robust to slower updates but sensitive to over-adaptation.

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 42 papers

Feb 27, 2026

Distribution-Aware Companding Quantization of Large Language Models

Large language models such as GPT and Llama are trained with a next-token prediction loss. In this work, we suggest that training language models to predict multiple future tokens at once results in higher sample efficiency. More specifically, at each position in the training corpus, we ask the model to predict the following n tokens using n independent output heads, operating on top of a shared model trunk. Considering multi-token prediction as an auxiliary training task, we measure improved downstream capabilities with no overhead in training time for both code and natural language models. The method is increasingly useful for larger model sizes and keeps its appeal when training for multiple epochs. Gains are especially pronounced on generative benchmarks like coding, where our models consistently outperform strong baselines by several percentage points. Our 13B parameter models solves 12 % more problems on HumanEval and 17 % more on MBPP than comparable next-token models. Experiments on small algorithmic tasks demonstrate that multi-token prediction is favorable for the development of induction heads and algorithmic reasoning capabilities. As an additional benefit, models trained with 4-token prediction are up to 3X times faster at inference, even with large batch sizes.

Athul Radhakrishnan, Siddhant Mohan, Mahima Sachdeva

LLM Training Multi-Token Prediction Efficient Language Model 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

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

Abstract

Multi-Token Prediction

Explore similar work

Distribution-Aware Companding Quantization of Large Language Models

EntMTP: Accelerating LLM Inference with Entropy Guided Multi Token Prediction

LightMTP: Lightweight Latent Multi-Token Prediction