Speculative Decoding & Draft-Verify Inference

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. The Inference Bottleneck
5. Speculative Decoding Algorithm
6. Acceptance-Rejection Proof
7. Draft Model Design
8. Medusa: Multi-Head Speculation
9. EAGLE: Feature-Level Drafting
10. Self-Speculative Decoding
11. Common Pitfalls
12. Summary
13. Exercises

Learning Objectives

1. Explain why autoregressive decoding is memory-bandwidth bound.
2. Derive the acceptance probability for speculative decoding.
3. Prove that speculative decoding produces the exact same distribution as standard decoding.
4. Compare Medusa, EAGLE, and draft-model approaches.
5. Compute expected speedup as a function of acceptance rate.

Notation

• $p(x_t | x_{<t})$ — target model distribution
• $q(x_t | x_{<t})$ — draft model distribution
• $\gamma$ — number of speculative tokens per step
• $\alpha$ — acceptance rate

Core Intuition

Autoregressive LLMs generate one token per forward pass — but each pass is memory-bandwidth bound (loading 70B parameters to generate ONE token). Speculative decoding uses a small "draft" model to propose multiple tokens quickly, then the large model VERIFIES them all in ONE parallel pass. If most draft tokens are accepted, we get multiple tokens per large-model pass — a 2-3x speedup with ZERO quality loss.

The Inference Bottleneck

Why LLM inference is slow:

• Each token requires loading ALL parameters from memory.
• For 70B in BF16: 140GB must be read from GPU memory per token.
• On H100 (3.35 TB/s bandwidth): min 42ms per token → max 24 tok/s.
• Arithmetic intensity is tiny: most time is memory reads, not computation.

Key insight: A single forward pass can process MULTIPLE tokens in parallel (prefill). It's the SEQUENTIAL nature of autoregressive decoding that causes the bottleneck.

Speculative decoding exploits this: Verify $\gamma$ tokens in one pass (same cost as generating 1 token, since we're bandwidth-bound either way).

Speculative Decoding Algorithm

Leviathan et al. (2023) & Chen et al. (2023):

1. Draft: Use small model $q$ to generate $\gamma$ candidate tokens: $\tilde{x}_1, \ldots, \tilde{x}_\gamma$.
2. Verify: Run target model $p$ on all $\gamma$ tokens in ONE parallel forward pass.
3. Accept/Reject: For each position $t$ in order:
   • If $p(x_t) \geq q(x_t)$: always accept.
   • If $p(x_t) < q(x_t)$: accept with probability $p(x_t)/q(x_t)$.
   • On first rejection: sample correction token from adjusted distribution.
4. Repeat from the accepted prefix.

Expected tokens per step:

where $\alpha$ is the average acceptance rate.

Acceptance-Rejection Proof

Theorem: Speculative decoding produces tokens from EXACTLY the target distribution $p$.

Proof sketch: At each position, define:

• Accept $\tilde{x}_t$ with probability $\min(1, p(\tilde{x}_t)/q(\tilde{x}_t))$.
• On rejection, resample from: $p'(x) = \text{norm}(\max(0, p(x) - q(x)))$.

The combined distribution:

This is a standard rejection sampling argument. The output is PROVABLY identical to sampling from $p$ — no approximation.

Draft Model Design

Requirements: Fast + high acceptance rate (close to target).

Options:

• Smaller model of same family: LLaMA-7B drafts for LLaMA-70B. Acceptance rate: 70-80%.
• Quantized target: INT4 target as draft for BF16 target. Acceptance: 85-90%.
• Distilled student: Small model distilled from target. Best acceptance.
• N-gram model: Extremely fast but low acceptance (40-50%).

Optimal $\gamma$ (speculation length):

Typically $\gamma = 3$-$8$ depending on acceptance rate.

Medusa: Multi-Head Speculation

Cai et al. (2024): Add multiple prediction heads to the target model:

• Head 1: predicts token at position $t+1$ (standard).
• Head 2: predicts token at position $t+2$.
• Head 3: predicts token at position $t+3$.
• ...

NO separate draft model. The target model itself predicts multiple future tokens.

Tree attention: Verify multiple candidate sequences simultaneously using tree-structured attention.

Advantage: No draft model needed; single model in memory; simpler deployment.

Speedup: 2-3x on typical tasks.

EAGLE: Feature-Level Drafting

Li et al. (2024): Draft at the FEATURE level (hidden states), not token level:

1. Use a lightweight network to predict future hidden states.
2. These predicted hidden states go through the LM head to get token candidates.
3. Verify with the real model.

Advantage over Medusa: Higher acceptance rate because feature-level prediction captures more context than independent per-position heads.

Speedup: 2.5-3.5x (higher than Medusa for code and structured text).

Self-Speculative Decoding

Draft using the model itself (no separate model):

Layer-skipping: Use early layers as the draft model (skip layers 20-32, use layers 1-19 as a "cheap" forward pass).

Early exit: If confidence is high after layer 12, use that prediction as draft.

Advantage: No separate model; no additional memory; works with any model.

Limitation: Lower acceptance rate than a properly trained draft model.

Common Pitfalls

Pitfall 1. Using a draft model from a different tokenizer. Draft and target MUST share the same tokenizer — otherwise token-level acceptance makes no sense.

Pitfall 2. Speculation too long ($\gamma$ too large). If acceptance rate is 60%, then expected accepted tokens for $\gamma=10$ is only 2.5. The time spent drafting 10 tokens is wasted. Adapt $\gamma$ based on observed acceptance.

Pitfall 3. Assuming speculative decoding helps with batched inference. With large batch sizes, inference becomes compute-bound (not memory-bound), and speculative decoding provides no benefit. It's primarily useful for low-batch (interactive) scenarios.

Summary

• Problem: Autoregressive inference is memory-bandwidth bound (one token at a time).
• Speculative decoding: Draft multiple tokens cheaply; verify in one pass.
• Provably lossless: Output distribution is EXACTLY the target distribution.
• Speedup: 2-3x for interactive inference (single-user, low batch).
• Medusa: Multi-head self-speculation; no draft model needed.
• EAGLE: Feature-level drafting; highest acceptance rates.
• Batch caveat: Only helps in memory-bound (low batch) regime.

Exercises

Exercise 1. For acceptance rate $\alpha = 0.75$ and $\gamma = 5$: compute the expected tokens per verification step.

Exercise 2. Prove that the rejection sampling scheme produces exactly the target distribution (full proof).

Exercise 3. Compare speculative decoding speedup at batch sizes 1, 4, 16, 64. At what batch size does it stop helping?

Exercise 4. Design a Medusa head for GPT-2 (768-dim hidden state). Specify architecture, training data, and expected acceptance rates per head.

Exercise 5. Compute the optimal $\gamma$ for a setup where: draft model takes 5ms per token, verification takes 50ms regardless of $\gamma$, and acceptance rate is 0.8.