🤖 machine learning

Variational Speculative Decoding: Rethinking Draft Training from Token Likelihood to Sequence Acceptance

This paper introduces Variational Speculative Decoding (VSD), a novel framework that reformulates draft training as variational inference to maximize target-model sequence acceptance, thereby achieving significant inference speedups over existing methods by resolving the training-decoding discrepancy through an Expectation-Maximization procedure with adaptive weighting and regularization.

Original authors: Xiandong Zou, Jianshu Li, Jing Huang, Pan Zhou

Published 2026-08-13

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Xiandong Zou, Jianshu Li, Jing Huang, Pan Zhou

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a super-smart robot how to write a story. This robot, known as a Large Language Model (LLM), is incredibly talented but also incredibly slow. It works like a person writing a sentence one word at a time, pausing after every single word to think about what comes next. This "stop-and-think" process is safe and accurate, but it makes the robot feel like it's wading through molasses.

To speed things up, scientists invented a clever trick called "speculative decoding." Think of it like a team of two: a fast, energetic junior assistant (the "draft model") and a wise, slow supervisor (the "target model"). The junior assistant tries to guess the next few words of the story very quickly. The supervisor then checks these guesses. If the guesses are right, the robot accepts them all at once, skipping the slow "think" time. If a guess is wrong, the supervisor fixes it, and they try again. The goal is to get the junior assistant to guess long, correct chains of words so the supervisor can approve them in big batches, making the whole process fly.

However, there's a catch. The junior assistant is usually trained by showing it the "perfect" answer and asking it to memorize that single path. But in the real game, the supervisor doesn't just look at one perfect path; it looks at a whole tree of possibilities, picking the best ones from a crowd. It turns out that training the assistant to only follow one straight line makes it terrible at guessing the right path when the supervisor is actually looking at a whole forest of options. This mismatch between how the assistant is trained and how it actually plays the game has been holding back the speed of these AI robots.

The Paper's Big Idea: A New Way to Train the Assistant

In this paper, the researchers propose a new training method called Variational Speculative Decoding (VSD). They realized that the old way of training the draft model was like teaching a race car driver to only drive on a single, straight track, when the actual race involves navigating a complex, winding course with many possible turns.

The Problem with the Old Way
The authors explain that current methods train the draft model to be "greedy." This means the model is taught to always pick the single most likely next word, creating one straight line of text. But during the actual game, the system doesn't just check one line; it generates a "tree" of many different possible word combinations and then picks the best ones. The paper shows that this creates a mismatch: the model is trained to be perfect at one specific path, but the game requires it to be good at many different paths. In fact, their experiments showed that about 30% of the time, the "perfect" path the model was trained on gets thrown away during the game because it wasn't actually the best choice among the many options.

The New Solution: Learning from the Whole Forest
VSD changes the rules of training. Instead of just memorizing one straight line, the new method treats the draft model as a "variational inference" problem. In plain English, this means the model is taught to understand the whole forest of possibilities, not just one tree. It learns to guess a set of paths that are likely to be accepted by the supervisor, rather than just guessing the single "best" word.

To do this, the researchers use a clever two-step process (called an Expectation-Maximization procedure):

The Guessing Phase (E-step): The model generates many different possible story paths. A special "oracle" (a smart filter) checks these paths and keeps the ones that look like they would pass the supervisor's test, while tossing out the bad ones.
The Learning Phase (M-step): The model learns from these "kept" paths. But here's the twist: the researchers added two special tools to make learning stable and smart.
- Adaptive Rejection Weighting (ARW): This acts like a coach who knows when the student is struggling. If the model is making lots of bad guesses, the coach ignores the noise and focuses on the few good ones. If the model is doing well, the coach pays attention to the mistakes to help it improve faster.
- Confidence-Aware Regularization (CAR): This prevents the model from being "overconfident." If the model is 99% sure about a wrong guess, this tool gives it a big penalty. It teaches the model to be humble and explore other options rather than stubbornly sticking to a wrong idea.

What They Found
The researchers tested this new method on several different AI models, including text-only models (like LLaMA) and models that can see images (like LLaVA). They found that VSD consistently made the AI faster without making it dumber.

Speed: On text models, VSD made the AI up to 9.6% faster than the previous best method (EAGLE-3).
Acceptance: The "assistant" was able to get more words approved at once. On average, it accepted 6 to 7 tokens (chunks of text) in a single go, compared to the 5 to 6 accepted by the old methods.
Multimodal Magic: It worked just as well on image-and-text models, improving speed by 7.9% over the best existing method for those models (ViSpec).

Why It Matters
The paper proves that by changing how we train the "assistant" to think about the whole game instead of just one move, we can significantly speed up AI. The authors show mathematically that their method is guaranteed to improve the "expected acceptance length" (how many words get approved at once). They didn't just guess this; they measured it across many different tasks, from writing code to solving math problems and answering questions about pictures.

In short, VSD fixes the training mismatch by teaching the draft model to be a better explorer of possibilities, ensuring that when it makes a guess, it's a guess that the supervisor is actually likely to accept. This makes the AI feel less like it's wading through molasses and more like it's sprinting.

Technical Summary: Variational Speculative Decoding (VSD)

1. Problem Statement: Training-Decoding Distributional Discrepancy

Current speculative decoding methods accelerate Large Language Model (LLM) and Multimodal LLM (MLLM) inference by using a lightweight draft model to propose tokens, which a target model verifies in parallel. However, a fundamental training-decoding discrepancy persists:

Deterministic Training vs. Stochastic Decoding: Existing draft models are typically trained via cross-entropy to predict a single "greedy" token trajectory. In contrast, practical decoding involves stochastic multi-path sampling (e.g., tree-based exploration), ranking candidates, and verifying multiple paths.
Token-Level Likelihood vs. Path-Level Utility: Training optimizes the likelihood of individual tokens along one path. Decoding, however, makes decisions at the path level, where a draft path's value depends on its cumulative confidence and ability to survive verification over a sequence.
The Consequence: Standard training collapses the draft distribution toward a single deterministic path. Empirical analysis in the paper shows that roughly 30% of training-time greedy paths are pruned during draft-tree construction, and the final accepted path coincides with the greedy path in only ~36% of cases. This misalignment limits the achievable inference speedup.

2. Methodology: Variational Speculative Decoding (VSD)

The authors propose Variational Speculative Decoding (VSD), a framework that reformulates draft training as a variational inference problem over latent draft paths.

2.1. Variational Formulation

VSD treats the draft path $z$ as a latent proposal. The goal is to maximize the marginal probability that a path is valid (i.e., accepted by the target model).

Path-Level Validity: For a draft path $z = (z_1, \dots, z_\ell)$ , the validity probability $\kappa(x, z)$ is defined as the joint probability that all tokens are accepted: $\kappa(x, z) = \prod_{i=1}^\ell \alpha_i(x, z_{<i})$ , where $\alpha_i$ is the token-level acceptance probability.
Evidence Lower Bound (ELBO): The authors derive an ELBO for the log-marginal probability of acceptance, $\log Z_\theta(x)$ $lo g Z_{θ} (x)$ :
$\mathcal{L}_{VSD}(\psi; x) = \mathbb{E}_{z \sim q_\psi}[\log \kappa(x, z)] - D_{KL}(q_\psi(z|x) \parallel p_\theta(z|x))$
- The first term encourages the draft model to generate paths with high acceptance probability (path-level utility).
- The second term (KL divergence) acts as a regularizer to keep the draft distribution close to the target model's distribution, preventing degenerate solutions.

2.2. Optimization via Expectation-Maximization (EM)

Since the posterior over valid paths is intractable, VSD optimizes the ELBO using a Monte Carlo-based EM procedure:

E-Step (Oracle-Filtered Sampling): The draft model generates a batch of candidate paths. An "oracle" (based on the target model's greedy path) filters these candidates, retaining high-utility paths and substituting rejected ones with corrective samples from the target policy. This constructs an empirical approximation of the valid-path posterior.
M-Step (Weighted Likelihood Maximization): The draft model parameters are updated to maximize the likelihood of the accepted samples from the E-step. To ensure stability and efficiency, two mechanisms are introduced:
1. Adaptive Rejection Weighting (ARW): A control variate coefficient dynamically reweights rejected samples based on the draft model's current reliability. This reduces gradient variance, preventing unstable updates when the draft model is weak.
2. Confidence-Aware Regularization (CAR): This penalizes the draft model for assigning high confidence to invalid paths (hard-negative mining). It reweights rejected proposals based on the draft model's confidence, discouraging overconfident errors that increase verification costs.

3. Key Contributions

Novel Framework: VSD is the first framework to explicitly address the training-decoding distributional discrepancy by formulating draft training as variational inference over latent paths, rather than single-token prediction.
Theoretical Guarantees:
- Theorem 1: Proves that maximizing the VSD objective increases a provable lower bound on the expected accepted length.
- Theorem 2: Demonstrates that VSD provides a strictly tighter lower bound on expected accepted length compared to traditional KL-only training (which only minimizes $D_{KL}(q_\psi \parallel p_\theta)$ ).
- Theorem 3: Characterizes the optimal draft distribution at each EM step as the step-level valid-path posterior.
Algorithmic Innovations: The integration of ARW and CAR stabilizes the training of path-level objectives, managing the high variance inherent in Monte Carlo sampling of draft trees.

4. Experimental Results

The authors evaluated VSD across diverse LLMs (LLaMA-3.1/3.3, Vicuna, DeepSeek-R1) and MLLMs (LLaVA-1.5) on benchmarks including MT-Bench, HumanEval, GSM8K, and various visual tasks (VQAv2, ChartQA, etc.).

Performance Gains:
- LLMs: VSD achieves up to a 9.6% speedup over the state-of-the-art EAGLE-3 and improves acceptance length by 7.6% on average.
- MLLMs: VSD achieves up to a 7.9% speedup over the best ViSpec baseline and improves acceptance length by 7.9%.
Compatibility: VSD is modular and can be applied to enhance existing methods (e.g., EAGLE-3, MSD, ViSpec, GRIFFIN, HASS), consistently yielding improvements regardless of the underlying draft architecture.
Ablation Studies:
- Increasing the number of latent proposals ( $S$ ) consistently improves performance, confirming VSD's ability to leverage multi-path exploration.
- Both ARW and CAR contribute significantly to performance; removing them leads to reduced speedup and acceptance length due to higher variance and overconfident errors.
- VSD shows greater relative gains on tasks requiring long-horizon consistency (e.g., code generation, math reasoning) and when the baseline greedy acceptance rate is lower.

5. Significance and Claims

The paper claims that VSD represents a principled shift in speculative decoding training. By moving from token-level likelihood to sequence acceptance, VSD aligns the draft model's objective with the actual decoding process (multi-path verification).

Theoretical Alignment: The work establishes a direct mathematical link between the variational bound and wall-clock speedup, theoretically guaranteeing efficiency improvements over traditional token-level training.
Practical Impact: The results demonstrate that VSD is a complementary, drop-in enhancement that significantly boosts inference efficiency for both text and multimodal models without compromising output quality (lossless).
Limitations: The authors note that scaling the number of latent proposals ( $S$ ) beyond 40 is computationally constrained in their current setup, though they expect further gains with larger $S$ . They also identify extending VSD to other modalities (e.g., speech) as future work.