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

Raghavv Goel; Mukul Gagrani; Mingu Lee; Chris Lott

Source-linked AI summary

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

Raghavv Goel, Mukul Gagrani, Mingu Lee, Chris Lott

arXiv:2603.17942v2Mar 18, 2026cs.CL

TL;DR

Existing multi-token prediction methods often require auxiliary training, modified weights, or draft models, limiting their simplicity and suitability for constrained settings. ESP instead probes a frozen LLM with embedding-space mask tokens, builds and prunes speculative token trees, and verifies predictions in parallel. The paper reports consistent gains over training-free baselines, while theoretical and empirical analyses support the alignment underlying its predictions.

Problem
Existing multi-token prediction approaches often rely on auxiliary training, modified model weights, or external draft models, creating engineering, parameter, and memory burdens.
Method
ESP probes a frozen LLM with on-the-fly embedding-space mask tokens, organizes parallel proposals into dynamically expanded and pruned token trees, and verifies them against the base model.
Results
ESP consistently outperforms training-free baselines across models and tasks, improving acceptance, forward-pass efficiency, and token throughput.
Takeaways & Limitations
Training-free probing can support efficient speculative inference without draft models, N-gram caches, offline token trees, or retraining.
Takeaways & Limitations
Under matched block complexity, STAND outperforms ESP on coding, RAG, and summarization for LLaMA3.1-8B-Instruct because of its n-gram copy advantage.

Abstract

from arXiv · show

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.

1. Introduction

ESP presents a training-free, single-model approach to multi-token prediction that probes frozen LLMs in embedding space, organizes proposals with adaptive token trees, and verifies them losslessly. Across diverse tasks and model families, it improves acceptance, forward-pass efficiency, and throughput without auxiliary models or retraining.

Method: ESP supports mean-of-prompt and token-embedding initialization, with mean-prompt initialization performing best across model families.
Results: ESP consistently outperforms training-free baselines across SpecBench tasks, sampling temperatures, and LLaMA3 and Qwen3 model families.It achieves higher acceptance rates, fewer forward passes, and improved token throughput.
Results: ESP performs strongly on open-ended and constrained tasks and is particularly suitable for compute-limited settings such as edge devices.
Method: ESP uses mask-token probing in a base model’s embedding space to enable multi-token generation without retraining or external draft models.
Method: Dynamic speculative token-tree expansion adapts token paths to predicted probabilities, removing the need for manually designed tree structures.
Efficiency: Efficient static-tree attention masks and position updates improve throughput for fixed tree structures.The implementation is designed for GPU-friendly decoding.
Analysis: Cosine-similarity alignment between mask-token and true-token representations supports correct-token inclusion in Top-K predictions across decoder layers.

2. Background

ESP injects dynamically computed mask tokens into a frozen autoregressive model to elicit predictions for multiple future positions. The resulting speculative tokens are jointly verified against the base model, while branching and pruning organize alternative futures.

Mask-token probing: A frozen autoregressive model produces next-token logits and distributions from a prompt sequence.
Mask-token probing: ESP appends dynamically computed embedding-space mask tokens to the prompt without modifying model parameters.
Mask-token probing: Each mask token is designed to elicit a prediction of a different future token position.
Verification: Predicted tokens are verified against the base model’s own next-token distribution through simultaneous verification.
Verification: Accepted predictions are appended to the prefix and used to verify subsequent predictions, yielding a lossless speculative-decoding procedure.
Tree construction: Multiple candidate futures are organized through tree branching, with dynamic expansion and pruning used to improve proposal diversity.

3. Methods

ESP enables training-free multi-token prediction by injecting dynamically generated mask tokens into a frozen LLM, constructing and pruning speculative token trees, and verifying proposals in parallel. Its design combines embedding-space probing, representation-alignment analysis, lossless verification, and compute-aware tree construction.

Mask Token Injection: ESP injects dynamically generated mask tokens into a frozen LLM to elicit multiple future-token predictions without auxiliary models or fine-tuning.The mask tokens are synthesized in embedding space and appended to the prompt.
Mask Token Injection: Mask tokens are initialized from the mean prompt embedding and updated with generated-token embeddings during decoding.This prompt-dependent initialization is intended to align the probes with the prompt distribution.
Why Mask Tokens Enable Multi-Token Prediction: Decoder layers align mask-token representations with valid-token states, making higher final-state similarity increase the likelihood that the correct token enters Top-K.Accepted tokens reach about 0.45 average cosine similarity after layer 15, whereas rejected tokens plateau near 0.35.
Logit-based Prediction and Verification: ESP samples Top-K mask-token logits to form a speculative tree, expands only Top-1 tokens at each depth, and verifies proposals with the same base model.Verification is lossless because tokens are accepted only when they match the base model’s next-token prediction.
Logit-based Prediction and Verification: Block complexity counts the last accepted token, future-token tree nodes, and associated mask tokens processed in parallel.Methods are compared under matched block complexity because larger trees can improve acceptance while increasing latency and compute cost.
Dynamic Tree Construction: Dynamic tree construction ranks trajectories by cumulative probability and retains the Top-B −1 paths, adapting branch exploration to model uncertainty.Tree pruning improves average token acceptance by up to 4%.

4. Experiments

ESP is evaluated on diverse SpecBench tasks using LLaMA3 and Qwen3 models, with lossless sample matching and block complexities of 10, 30, and 60. It generally improves acceptance, throughput, and implementation efficiency, while task characteristics affect the best tree configuration.

Experimental Setup: ESP is evaluated on LLaMA3 and Qwen3 models across SpecBench tasks including summarization, translation, writing, coding, retrieval, and mathematics.Experiments use exact-match sample matching for lossless generation and report average acceptance length and speedup ratio.
Results: ESP achieves the strongest overall speed–quality trade-off across models and budgets while maintaining lossless generation.The evaluation compares average accepted tokens and speedup over autoregressive decoding at BC = 30 and BC = 60.
Results: Up to 12% higher τ than LADE on LLaMA3 and 13–18% higher τ than STAND on Qwen3 are reported, with up to 42% fewer model invocations at BC = 60.ESP also records the best speedup ratios across the evaluated models and budgets.
Results: ESP achieves the highest τ and speedup across most tasks and block-complexity settings for LLaMA3.1-8B-Instruct and Qwen3-32B.On LLaMA3.1-8B-Instruct, its peak reported result is τ=1.81 with 1.43× speedup on math/reasoning at BC=60.
Results: Coding, retrieval, and summarization are exceptions on LLaMA3.1-8B-Instruct, where STAND outperforms ESP through its n-gram copy advantage.ESP nevertheless performs well at low block complexity, including an average 1.22× speedup at BC=10 for LLaMA3.1-8B-Instruct.
Ablations: Dynamic branching performs best with two mask tokens at BC=30 and BC=60, while one or two masks outperform three at fixed BC=60.Open-ended tasks favor single-mask wider trees, whereas closed-ended tasks favor two-mask deeper trees.
Implementation: The efficient tree-attention and positional-update implementation improves throughput by about 22% for LLaMA3.2-3B-Instruct and 21% for LLaMA3.1-8B-Instruct.At BC=60, gains reach 29–30% for LLaMA3.1-8B-Instruct and 27–28% for LLaMA3.2-3B-Instruct.
Mask Embeddings: Mean prompt-embedding initialization performs best across LLaMA3 variants, while initialization outside the embedding-table distribution causes a minor performance drop.The results show that mask-token initialization affects training-free multi-token prediction quality.

5. Related Work

Prior multi-token and speculative decoding methods often modify model architectures, train auxiliary components, or rely on caching and heuristic matching. ESP instead uses mask-token probing and dynamic tree construction without those additions.

Interpretability: Future Lens provides a mechanistic rationale for probing by showing that future-token information can be linearly decoded from single hidden states.Related work also formalizes hypotheses about whether transformers plan ahead.
Multi-Token Prediction: Prior MTP methods add decoding heads, sampler modules, independent output heads, or fine-tune base-model parameters.These approaches differ from ESP’s probing-based design.
Speculative Decoding: Speculative decoding extensions use draft tokens, token trees, reused target-model layers, or modified attention mechanisms for parallel verification.The cited approaches often require architectural changes or auxiliary models, unlike ESP.
Training-Free Acceleration: Training-free acceleration methods such as LADE, PLD, and STAND rely on caching or heuristic matching, whereas ESP uses probing for dynamic tree construction.ESP thereby avoids the memory overhead associated with those mechanisms.

6. Conclusion

The paper presents training-free multi-token prediction through embedding-space mask probing, dynamic token-tree expansion, and efficient attention implementation. Experiments and analysis support speed and efficiency gains across tasks and models while identifying task-dependent tree preferences.

Conclusion: ESP uses embedding-space mask tokens and dynamic token-tree expansion for efficient speculative inference without draft models, N-gram caches, or offline token trees.The framework maintains lossless generation while reducing model forward calls.
Conclusion: A single mask token can provide substantial block-efficiency gains, while dynamic branching with multiple masks further improves deeper tree configurations.Open-ended tasks favor wider trees, whereas closed-ended tasks prefer deeper, more focused trees.
Conclusion: Efficient attention-mask and position-ID implementations reduce runtime overhead and improve throughput.These implementation changes complement the probing and tree-expansion framework.

Impact Statement

ESP targets inference speedups without retraining, auxiliary networks, or architectural modifications. Its lossless efficiency and low resource requirements are presented as relevant to constrained deployments, while its probing behavior also informs interpretability.

Impact: ESP enables inference-time speedups without model retraining, auxiliary networks, or architectural modifications.The framework uses mask-token probing and dynamic token-tree expansion.
Impact: Lossless generation and reduced storage, memory, and power requirements make ESP relevant to mobile, embedded, and edge deployments.The stated scope is compute-constrained environments where heavyweight speculative-decoding frameworks are impractical.
Impact: The method also offers interpretability insights into latent multi-step predictive structure in decoder representations.This extends the paper’s impact beyond inference efficiency.

A. Intuition for Multi-Token Prediction Using Mask Tokens

ESP probes a frozen LLM with mask-token representations to predict future tokens, relying on hidden-state alignment to place correct continuations in Top-K candidates. The method supports training-free parallel proposals, while its theoretical analysis connects cosine similarity with acceptance.

Intuition: ESP compares mask-token and next-true-token hidden states when both predict the same future position.The analysis is performed in hindsight because the next-true token is unavailable during inference.
Theoretical connection: Higher cosine similarity between mask and next-true hidden states places the next-next true token in the mask token’s Top-K predictions.The guarantee holds above a threshold δ∗, and larger Top-K sets can permit a smaller threshold.
Open question: The authors leave the mechanism that injects relevant information into mask-token states for future work.This is presented as an unresolved question behind the observed higher cosine similarity.
Design rationale: The method is training-free and head-free, using embedding-space mask probing rather than auxiliary heads, modified weights, or external draft models.This design distinguishes ESP from training-based MTP and draft-model speculative decoding approaches.
Tree construction: ESP organizes mask-token predictions into token trees and uses Top-1 expansion when two or more mask tokens are employed.Top-1 expansion retains only the highest-probability token at each depth for child-node generation.

E. Efficient Tree Attention Mask and Position ID Construction

ESP reduces tree-attention overhead by exploiting fixed tree structures and reusing position IDs and attention-mask structure across generation steps. This implementation improves token throughput and model-call efficiency across models and tasks.

Bottleneck: Custom tree-attention mask generation is latency-intensive because it sequentially iterates over all branches of the token tree.The branching hierarchy must be reflected in both position IDs and attention masks during input construction.
Static implementation: Static tree structures enable GPU-friendly attention-mask and positional-index updates during simultaneous generation and verification.The input token order is fixed for a given block complexity, allowing reuse across passes.
Position IDs: Position IDs are updated by uniformly shifting the previous tensor by the number of generated tokens.This avoids recomputing the full position-index structure at every generation step.
Attention masks: Attention masks gain BE columns filled with zeros, which can be prepended to the previous mask as the key-value cache grows.This produces the mask for the current pass while preserving the evolving cache structure.
Throughput results: 4% and 19.6% token-rate increases are reported for LLaMA3.2-3B-Instruct at BC = 30 and BC = 60 with one mask token.For LLaMA3.1-8B-Instruct, the corresponding increases are 28% and 22.4%.
Evaluation: The method achieves the highest reduction in model forward passes, benefiting compute- and energy-constrained devices.Additional evaluations report higher block efficiency across most SpecBench tasks, with specified exceptions for retrieval and summarization.

G.5. Comparison with Training-based Method

ESP is compared with the training-based EAGLE3 method in terms of speedup and training cost. The comparison highlights a plug-and-play advantage for training-free methods.

Training cost: EAGLE3 requires 200–300 GPU hours of training per model, whereas the other compared methods require 0 GPU hours.The table reports speedup ratios over autoregressive decoding and marks unavailable official EAGLE3 checkpoints as N/A.
Practical trade-off: Training-free methods can be used in a plug-and-play manner without additional compute cost.The comparison notes that EAGLE3’s training time may increase for larger model sizes.

G.6. Impact of Number of Mask Tokens

The best number of mask tokens depends on block complexity and task type. Open-ended tasks favor shallower, wider trees in some settings, while more constrained tasks can benefit from deeper, narrower trees.

BC = 60: At BC = 60, roleplay and question answering achieve higher block efficiency with one mask token, while coding, summarization, RAG, and math/reasoning favor two.Writing shows mixed behavior across mask-token configurations.
BC = 30: At BC = 30, one mask token consistently outperforms two mask tokens.The two-mask configuration uses dynamic tree expansion, and differences within 0.01 are treated as similar.

G.7. Impact of tree-pruner

The lightweight tree-pruner improves block efficiency in multi-mask settings by removing redundant branches while preserving accuracy. With a single mask token, it maintains block efficiency without degradation.

Up to 3% higher BE for LLaMA3.2-3B-Instruct and up to 4% for LLaMA3.1-8B-Instruct when using two mask tokens.The pruner introduces no additional computational overhead.
With a single mask token, the pruner maintains BE without degradation.
The pruner removes repeated tokens when the second mask’s Top-K predictions include the first mask’s Top-1 token.This enables more diverse branching without sacrificing accuracy.