Speculative decoding

The Sequential Generation Bottleneck

WHAT is the problem?

Standard autoregressive generation produces one token at a time:

p(x_t | x_{<t}) = \text{softmax}(W_{\text{out}} h_t)

For a sequence of length $T$ , this requires $T$ sequential forward passes. Each pass:

Loads model weights from memory (billions of parameters)
Has latency dominated by memory bandwidth, not FLOPs
Cannot be parallelized across the sequence dimension

The inefficiency: Modern GPUs have teraflops of compute but spend most time waiting for memory I/O. The model is "idling" between tokens.

How Speculative Decoding Works

The Two-Model Setup

Draft model $M_q$ : Small, fast model (e.g., 1B parameters, distilled from target)
Target model $M_p$ : Large, slow model (e.g., 70B parameters, the one we actually want)

Key property: $M_q$ is 5-10× faster per token but slightly less accurate.

The Algorithm (Step-by-Step)

Formula Speculative Decoding Protocol

Input: Prefix $x_{<t}$ , draft model $M_q$ , target model $M_p$ , lookahead $K$

Step 1: Draft phase (parallel speculation)

For i = 0 to K-1:
    x̃_{t+i} ~ q(· | x_{<t+i})  # Sample from draft model

Generates $K$ candidate tokens: $\tilde{x}_t, \tilde{x}_{t+1}, \ldots, \tilde{x}_{t+K-1}$ Step 2: Verification phase (single parallel pass)

Run target model on $[x_{<t}, \tilde{x}_t, \tilde{x}_{t+1}, \ldots, \tilde{x}_{t+K-1}]$
Get probabilities: $p(\cdot | x_{<t+i})$ for $i = 0, \ldots, K-1$ Step 3: Acceptance sampling For each candidate $\tilde{x}_{t+i}$ in order:

Accept with probability: min(1, p(x̃_{t+i} | x_{<t+i}) / q(x̃_{t+i} | x_{<t+i}))
If rejected:
    Resample from adjusted distribution: max(0, p - q) / Z
    Break (reject all subsequent candidates)
If all accepted:
    Sample bonus token from p(· | x_{<t+K})

WHY this works:

Lossless: The acceptance sampling ensures output distribution exactly matches $M_p$ alone
Fast: When draft is correct (often 60-80% of tokens), we accept $K$ tokens in one pass
Bounded: Worst case is rejecting everything = same as standard decoding (no slowdown)

Derivation of Acceptance Criterion

Goal: Match target distribution $p(x)$ using proposals from draft $q(x)$ .

Start with rejection sampling:

Sample $\tilde{x} \sim q(x)$
Accept with probability $\alpha = \min(1, p(\tilde{x}) / q(\tilde{x}))$

WHY this gives $p(x)$ :

P(\text{accept } x) = q(x) \cdot \frac{p(x)}{q(x)} = p(x) \quad \text{when } p(x) \leq q(x)

When $p(x) > q(x)$ , we reject with probability $1- p(x)/q(x)$ and resample from:

p_{\text{resample}}(x) = \frac{\max(0,\, p(x) - q(x))}{\int \max(0,\, p(y) - q(y))\, dy}

Combined distribution:

P_{\text{final}}(x) = q(x) \min\left(1, \frac{p(x)}{q(x)}\right) + \left(1 - \sum_y q(y) \min\left(1, \frac{p(y)}{q(y)}\right)\right) \cdot p_{\text{resample}}(x) = p(x)

The math guarantees output matches the target model exactly.

Worked Example: 4-Token Speculation

| Position | Draft Token | $q$ (draft) | $p$ (target) | Accept? | |----------|-------------|------------|------| | t | "Par" | 0.7 | 0.8 | ✓ (min(1, 0.8/0.7) = 1) | | t+1 | "is" | 0.5 | 0.6 | ✓ (min(1, 0.6/0.5) = 1) | | t+2 | "," | 0.4 | 0.2 | ✗ (min(1, 0.2/0.4) = 0.5, coin flip fails) |

Outcome:

Accept "Par" and "is" (2 tokens in one target pass!)
Reject "," and resample from $\max(0, p - q)$ distribution at position t+2
Likely samples "." (which had $p=0.5, q=0.1$ )
Final: Generated 3 tokens with 1 draft pass + 1 target pass (vs. 3 passes normally)

WHY this step: Position t+2 fails because target model strongly prefers "." over ",", so draft's confidence in "," is misleading. The adjusted distribution removes draft's bias.

Speedup Analysis

Expected Acceptance Rate

Let $\alpha$ = probability a draft token is accepted (empirically 0.6-0.8 for good draft models).

Tokens accepted per speculation:

\mathbb{E}[\text{tokens}] = \sum_{i=0}^{K-1} (i+1)\,\alpha^{i}(1-\alpha) + (K+1)\,\alpha^K = \frac{1 - \alpha^{K+1}}{1 - \alpha}

Derivation (careful with indices):

Accept exactly $i$ tokens then reject (for $0 \le i < K$ ): the first $i$ candidates are accepted (probability $\alpha^i$ ) and the $(i+1)$ -th is rejected (probability $1-\alpha$ ). We then resample one token from the adjusted distribution, giving $i+1$ output tokens. So the term is $(i+1)\,\alpha^i(1-\alpha)$ .
Accept all $K$ candidates (probability $\alpha^K$ ): we additionally sample a bonus token, giving $K+1$ output tokens. So the term is $(K+1)\,\alpha^K$ .
Summing and simplifying the geometric series gives the closed form $\dfrac{1 - \alpha^{K+1}}{1 - \alpha}$ .

WHY speedup is less than $K$ : Rejections waste the target pass. The maximum speedup (as $\alpha \to 1$ , so output rate $\to K+1$ ) is

\text{Speedup}_{\max} = \frac{c_p (K+1)}{K c_q + c_p},

which equals $K+1$ only in the idealized limit $c_q \to 0$ (a free draft model). With a nonzero draft cost the draft phase itself takes time, so the true ceiling is below $K+1$ . As $\alpha \to 0$ , speedup $\to c_p/(K c_q + c_p) \le 1$ (no gain, possibly slight overhead from the draft passes).

Optimizing the Draft Model

WHAT makes a good draft model?

High agreement ( $\alpha$ ): Minimize distribution mismatch with target
Low latency: Must be much faster than target (typically >5× speedup)
Same vocabulary: Token spaces must align exactly

HOW to create one:

WHY distillation works best: Explicitly optimizes for distribution match, not just task accuracy. A model with 95% task accuracy might have only 60% token-level agreement, while a distilled model can achieve 75%+ agreement at lower task accuracy.

Common Mistakes

Practical Implementation Details

Batching Speculation

Challenge: Different sequences in a batch may accept different numbers of tokens.

Solution: Process in "sync points":

Draft all sequences with $K$ tokens
Target verifies all
Each sequence accepts 0 to $K$ tokens independently
Pad shorter sequences, continue from their current position

WHY this works: Target verification is still a single batched forward pass. The asynchrony is handled in post-processing.

Temperature and Sampling

With temperature $\tau$ :

Draft samples from $q_\tau(x) = \text{softmax}(\text{logits}_q / \tau)$
Target uses $p_\tau(x) = \text{softmax}(\text{logits}_p / \tau)$
Acceptance criterion unchanged: $\min(1, p_\tau / q_\tau)$

Top-k, nucleus sampling: Apply after acceptance sampling on the adjusted distribution.

Extensions and Variants

Medusa decoding

Uses multiple "draft heads" on the target model itself (adds $K$ small prediction heads at intermediate layers). Trades some memory for not needing a separate draft model.

Speculative RAG

For retrieval-augmented generation, draft model proposes both tokens and retrieval queries. Target verifies both.

Multi-candidate speculation

Draft model proposes a tree of $K$ candidates with multiple branches. Target verifies all paths. Handles uncertainty better when $\alpha$ is low.

Recall Feynman Explanation (ELI12)

Imagine you're writing an essay and you have a younger sibling who writes really fast but makes some mistakes. Here's the trick:

Your sibling writes the next 4 sentences quickly (the "draft")
You read all 4 sentences at once and check each one
If a sentence is good, you keep it. If it's wrong, you stop there, fix it, and have your sibling start again from that point Why is this faster than you writing alone? Because reading and checking 4 sentences takes almost the same time as writing 1 sentence yourself. If your sibling gets even 2 out of 4 right, you're saving time!

In AI, the "you" is a giant smart model (slow), the "sibling" is a small quick model (fast but less accurate), and "sentences" are tokens. The small model guesses ahead, the big model checks everything in one shot. Most guesses are right, so we skip most of the slow generation steps.

Connections

Model Quantization — Complementary: quantize both draft and target for further speedup
Flash Attention — Orthogonal memory optimization; both can be combined
Beam Search — Alternative decoding strategy; speculative decoding compatible with beam search
Knowledge Distillation — Used to train the draft model
Transformer Architecture — Understanding self-attention explains why memory-bound
Distributed Training — Speculative decoding reduces inference latency; distributed training reduces time

#flashcards/ai-ml

What is speculative decoding? :: A lossless inference acceleration technique where a fast draft model generates $K$ candidate tokens in parallel, and a slow target model verifies all candidates in a single forward pass using rejection sampling to maintain exact target distribution.

Why is autoregressive decoding slow?

Each token requires a sequential forward pass through the model, which is memory-bound (loading billions of parameters) rather than compute-bound, leaving GPU compute resources idle between tokens.

What is the acceptance criterion in speculative decoding?

Accept candidate token

\tilde{x}

with probability

\min(1, p(\tilde{x})/q(\tilde{x}))

where

p

is target model probability and

q

is draft model probability. This ensures output matches target distribution exactly.

What is the expected number of output tokens per speculation?

\mathbb{E}[\text{tokens}] = \frac{1 - \alpha^{K+1}}{1 - \alpha}

where

\alpha

is the probability a draft token is accepted and

K

is the lookahead length. It equals the sum

\sum_{i=0}^{K-1}(i+1)\alpha^i(1-\alpha) + (K+1)\alpha^K

, ranging from 1 (worst case) to

K+1

(best case).

What makes a good draft model for speculative decoding?

(1) High agreement rate

\alpha

with target model (60-80%), (2) Much faster than target (>5× speedup), (3) Same vocabulary and valid probability distributions. Architecture can differ from target.

Why doesn't speculative decoding change the output distribution?

The rejection sampling in Step 3 mathematically guarantees the output matches the target model's distribution exactly, even though we use draft model proposals. It's "trust but verify" with statistical correction.

What is the maximum theoretical speedup, accounting for draft cost?

\text{Speedup}_{\max} = \frac{c_p(K+1)}{K c_q + c_p}

where

c_q, c_p

are draft/target per-token costs. This equals

K+1

only in the idealized limit

c_q \to 0

; a nonzero draft cost lowers the ceiling.

What is the optimal lookahead

K

in speculative decoding?

Typically

K=4

8

. Larger

K

increases wasted computation when early rejections occur. Optimal

K

balances expected accepted tokens against cost:

K^* = \arg\max_K \frac{\mathbb{E}[\text{accepted}]}{K \cdot c_q + c_p}

How does speculative decoding handle different temperatures?

Both draft and target apply temperature

\tau

to their logits before sampling:

q_\tau = \text{softmax}(\text{logits}_q/\tau)

p_\tau = \text{softmax}(\text{logits}_p/\tau)

. Acceptance criterion uses these temperature-adjusted probabilities. Top-k/nucleus applied after acceptance.

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Chalo isko simple tarike se samajhte hain. Socho aapke paas ek bahut bada aur intelligent model hai (target model, jaise 70B parameters), jo har token generate karne mein slow hai kyunki har baar poore model ke weights ko memory se load karna padta hai. Isko autoregressive generation kehte hain—ek time pe ek hi token, aur ye memory-bound hota hai, matlab GPU ka compute idle rehta hai bas memory I/O ka wait karta rehta hai. Yehi asli bottleneck hai—hamare paas teraflops ka compute hai par hum uska sirf ~2% use kar paate hain.

Ab speculative decoding ka jugaad ye hai: ek chhota, fast draft model (jaise 1B parameters) use karo jo jaldi-jaldi kuch tokens guess kar deta hai—jaise ek intern jaldi se sentence sketch kar deta hai. Fir bada target model in saare guessed tokens ko ek hi parallel forward pass mein check kar leta hai—jaise expert editor ek nazar mein poora draft approve kar deta hai. Kyunki draft model aksar 60-80% tokens sahi guess kar leta hai, hum ek hi mehnge pass mein kai tokens accept kar lete hain, aur bahut saare slow sequential steps skip ho jaate hain. Isliye speed bahut badh jaati hai.

Sabse important baat—ye technique lossless hai. Acceptance sampling ka jo maths hai (min(1, p/q) probability se accept karna, warna adjusted distribution se resample karna) wo mathematically guarantee karta hai ki final output bilkul wahi distribution follow karega jo akela target model deta. Matlab quality mein zero compromise, sirf speed ka fayda. Aur worst case bhi safe hai—agar saare guesses galat ho jaayein, tab bhi ye normal decoding jitna hi time lega, usse slow nahi. Isliye ye real-world mein LLM inference fast karne ke liye ek bahut powerful aur practical trick hai.

The Sequential Generation Bottleneck

WHAT is the problem?

How Speculative Decoding Works

The Two-Model Setup

The Algorithm (Step-by-Step)

Derivation of Acceptance Criterion

Worked Example: 4-Token Speculation

Speedup Analysis

Expected Acceptance Rate

Optimizing the Draft Model

WHAT makes a good draft model?

HOW to create one:

Common Mistakes

Practical Implementation Details

Batching Speculation

Temperature and Sampling

Extensions and Variants

Medusa decoding

Speculative RAG

Multi-candidate speculation

Connections

Concept Map

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Scaling & Efficient Architectures

Connections