Speculative Decoding

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. The Autoregressive Bottleneck
5. Speculative Decoding Algorithm
6. Acceptance-Rejection Proof
7. Expected Speedup Analysis
8. Draft Model Selection
9. Tree-Based Speculation (SpecInfer)
10. Medusa: Multi-Head Decoding
11. Common Pitfalls
12. Summary
13. Exercises

Learning Objectives

1. Derive the acceptance-rejection scheme for speculative decoding.
2. Prove that speculative decoding produces samples from the exact target distribution.
3. Compute the expected number of accepted tokens for a given acceptance rate.
4. Explain tree-based speculation and its theoretical advantages.
5. Describe Medusa's approach to self-speculative decoding.

Notation

• $p(\cdot)$ — target (large) model distribution
• $q(\cdot)$ — draft (small) model distribution
• $K$ — number of speculative tokens (draft length)
• $\alpha$ — average acceptance rate: $\mathbb{E}[\min(1, p(x)/q(x))]$
• $c_p, c_q$ — per-token cost of target and draft models

Core Intuition

Autoregressive generation is slow because each token requires a full forward pass through the large model — sequentially, one at a time. But the large model can verify multiple tokens in parallel (same cost as generating one). Speculative decoding exploits this: a cheap draft model guesses $K$ tokens, and the large model verifies all $K$ in one pass, accepting those that match its distribution.

The Autoregressive Bottleneck

Standard generation: Each step requires loading all model weights from HBM for a single token. For a 70B model in FP16:

• Weights: 140 GB loaded per token.
• At 2 TB/s bandwidth: ~70 ms per token minimum (memory-bound).
• Generating 1000 tokens: ~70 seconds.

Key observation: The large model processes batch-size-1 inputs inefficiently. Whether it processes 1 token or 8 tokens, the weights must be loaded once — so verifying 8 tokens costs roughly the same as generating 1.

Speculative Decoding Algorithm

Input: Target model $p$, draft model $q$, draft length $K$, prefix tokens $x_{1:t}$.

Algorithm:

1. Draft phase: Generate $K$ tokens from draft model: $\hat{x}_{t+1}, \ldots, \hat{x}_{t+K} \sim q(\cdot | x_{1:t}, \hat{x}_{1:i-1})$.

2. Verify phase: Run target model on all $K+1$ positions simultaneously: Compute $p(x | x_{1:t}, \hat{x}_{1:i-1})$ for $i = 1, \ldots, K+1$.

3. Accept/Reject: For each position $i = 1, \ldots, K$:

   • Sample $u \sim \text{Uniform}(0, 1)$.
   • If $u < \frac{p(\hat{x}_{t+i})}{q(\hat{x}_{t+i})}$: accept $\hat{x}_{t+i}$.
   • Else: reject. Sample correction token from adjusted distribution. Stop.

4. Bonus token: If all $K$ tokens accepted, sample one additional token from $p$ at position $K+1$.

Acceptance-Rejection Proof

Theorem. The output of speculative decoding is distributed exactly as $p$.

Proof for position $i$: The probability of outputting token $x$ is:

Case 1: $p(x) \geq q(x)$. Draft proposes $x$ with probability $q(x)$, always accepted (ratio $\geq 1$). Rejection correction also samples $x$ with probability $\frac{p(x) - q(x)}{Z}$ where $Z = \sum_y\max(0, p(y)-q(y))$.

Total: $q(x) + (1-\alpha) \cdot \frac{p(x)-q(x)}{Z}$. Since $(1-\alpha) = Z$ (the total rejection mass), this equals $q(x) + p(x) - q(x) = p(x)$. $\checkmark$

Case 2: $p(x) < q(x)$. Accepted with probability $p(x)/q(x)$. Contribution: $q(x) \cdot p(x)/q(x) = p(x)$. Not sampled from correction (numerator is 0). Total: $p(x)$. $\checkmark$ $\blacksquare$

Expected Speedup Analysis

Expected accepted tokens per round: If each token has independent acceptance probability $\alpha$:

For $K \to \infty$: $\mathbb{E}[\text{accepted}] = \alpha/(1-\alpha)$.

Tokens per round (including bonus): $\mathbb{E}[\text{tokens}] = (1-\alpha^{K+1})/(1-\alpha)$.

Speedup over standard decoding:

where cost per round = 1 target forward pass + $K$ draft forward passes.

Example: $\alpha = 0.8, K = 5, c_q/c_p = 0.05$:

• Expected tokens: $(1 - 0.8^6)/(1-0.8) = 3.69$.
• Cost: $1 + 5 \times 0.05 = 1.25$.
• Speedup: $3.69/1.25 = 2.95\times$.

Draft Model Selection

Requirements:

• Fast (much smaller than target): $c_q \ll c_p$.
• Accurate (high $\alpha$): distribution close to target.
• Same vocabulary as target.

Options:

• Smaller version of same family: LLaMA-7B drafting for LLaMA-70B ($\alpha \approx 0.7$–$0.8$).
• Quantized target: INT4 version of the same model ($\alpha \approx 0.85$–$0.9$).
• Distilled model: Trained to mimic target distribution.
• n-gram model: For highly predictable tokens (boilerplate, formatting).

Tradeoff: Better draft → higher $\alpha$ → more accepted tokens, but also higher $c_q$.

Tree-Based Speculation (SpecInfer)

Instead of a single draft sequence, generate a tree of possibilities:

Structure: At each position, the draft model generates top-$k$ continuations, forming a tree of width $k$ and depth $K$.

Verification: The target model verifies all paths simultaneously (batch all leaf-to-root paths).

Advantage: Even if one path is rejected, another branch may be accepted. Higher expected acceptance for the same draft cost.

Expected tokens: Significantly more than linear speculation for the same total draft tokens generated.

Medusa: Multi-Head Decoding

Self-speculative decoding without a separate draft model:

Architecture: Add $K$ additional "prediction heads" to the target model. Head $k$ predicts the token at position $t+k$ given the hidden state at position $t$:

Training: Only the heads are trained (target model frozen). Each head is a small MLP.

Verification: Generate candidates from all heads, form a tree, verify with one target forward pass.

Advantage: No separate draft model; heads add minimal latency; works with any base model.

Common Pitfalls

Pitfall 1. Using too large a draft model. If $c_q/c_p > 0.2$, the overhead of running the draft model erodes the speedup. Optimal when draft is 10–50x smaller.

Pitfall 2. Setting $K$ too high when $\alpha$ is low. For $\alpha = 0.5$: average accepted = 1 token regardless of $K > 3$. Large $K$ wastes draft compute.

Pitfall 3. Forgetting that speculative decoding is lossless. The output distribution is exactly $p$ — there's no quality-speed tradeoff (unlike quantization or pruning).

Summary

• Speculative decoding: draft $K$ tokens cheaply, verify in one target pass.
• Lossless: outputs exact target distribution (provably).
• Expected speedup: $2$–$3\times$ for well-matched draft models.
• Tree-based speculation increases acceptance probability.
• Medusa: self-speculative with additional prediction heads.
• Optimal $K$ depends on $\alpha$: higher acceptance → larger $K$ worthwhile.

Exercises

Exercise 1. Compute the expected number of accepted tokens for $\alpha = 0.7$ and $K = 4, 8, 16$.

Exercise 2. Prove that the correction distribution $\max(0, p(x) - q(x)) / Z$ is a valid probability distribution.

Exercise 3. For $\alpha = 0.9$ and a draft model 20x faster than target: find the optimal $K$ that maximizes tokens/second.

Exercise 4. Explain why speculative decoding provides more speedup for memory-bound inference than compute-bound inference.

Exercise 5. Design a tree-based speculation strategy with branching factor 3 and depth 4. How many candidate sequences are verified?