Hugo's Cyber Heritage

Shine or Die

Speculative Decoding

Posted on Edited on In
Symbols count in article: 198 Reading time ≈ 1 mins.

Learning about Speculative Decoding

Speculative Decoding Study Notes

1. The Essence of LLM Inference: Autoregression

The generation process of an LLM is a standard autoregressive chain:

Therefore it must:

  • Generate tokens one step at a time
  • Run a full Transformer forward pass for every single token generated
  • Low GPU utilization (serialization bottleneck)

The root cause of slow inference:

Every single generated token requires one complete forward pass of the large model.

2. KV Cache: Solving a Problem That Cannot Be Solved During Training

What KV Cache Is For:

Cache the Key / Value of past tokens:

  • Avoid recomputing K/V
  • Reduce attention complexity from O(N²) → O(N)
  • Speeds up attention, but does not reduce the number of forward passes

What KV Cache Cannot Solve:

  • A new Q must be computed every time a token is generated
  • The FFN cannot be cached at all
  • Every layer still has to compute attention(Q, K_cache, V_cache)
  • Number of forward passes still = number of tokens

Why Can’t Training Use KV Cache?

Training computes over the whole sequence in parallel; every token’s Q/K/V appears for the first time, so there is no “historical” information to cache.

3. What Does Inference Acceleration Really Need to Optimize?

All acceleration techniques share exactly one goal:

Reduce the number of full forward passes of the large model.

Because the FFN is where the main FLOPs go during inference (60%–75%); attention has already been greatly accelerated by KV Cache.

4. The Core Idea of Speculative Decoding

Speculative decoding exploits two properties of Transformers:

  1. The large model cannot generate multiple tokens in parallel (autoregressive constraint)
  2. But a Transformer can verify all positions of a sequence in parallel at once (parallel forward over the sequence)

So we can:

  1. Use a small model to quickly generate a run of tokens (the draft)
  2. Use one forward pass of the large model to verify the entire draft
  3. Use rejection sampling to decide which tokens can be accepted
  4. Accepted tokens go straight into the real output, reducing the number of large-model forward passes

5. Why Can a Small Model “Guess” the Large Model?

Sources:

  • Distillation: the small model q learns the distribution of the large model p
  • Language distributions are highly concentrated:
    the large and small models overlap heavily on the top 1–5 candidate tokens

The more accurately the small model guesses, the higher the acceptance rate → the greater the speedup.

6. Why Can the Large Model Verify Multiple Tokens at Once?

Because the Transformer architecture natively supports:

  • A parallel forward pass over the entire sequence
  • The causal mask guarantees each position only looks to its left
  • The logits of every position can be output simultaneously

Therefore:

One forward pass of the large model gives all the probabilities for y₁,y₂,…,yₖ.

But:

  • It is not “predicting” multiple tokens in parallel
  • It is only “verifying” the draft sequence in parallel

7. The Rejection Sampling Mechanism: Guaranteeing an Unbiased Output Distribution

The small model generates draft tokens (y_i).

The large model computes the true probability:

The small model’s probability:

Acceptance probability α:

Sample a random number u:

  • (u < \alpha_i) → accept
  • (u \ge \alpha_i) → reject and fall back to the large model’s autoregressive generation

This mechanism guarantees:

No matter how bad the small model is, the final sampling distribution is always = the large model’s true distribution p (unbiased).

8. The Complete Speculative Decoding Pipeline (Schematic)

Suppose the draft length k = 4

Step 1: The Small Model Generates the Draft

1
draft = [y1, y2, y3, y4]

Step 2: One Forward Pass of the Large Model

Input sequence = X + draft

The large model computes at once:

  • p(y1|X)
  • p(y2|X,y1)
  • p(y3|X,y1,y2)
  • p(y4|X,y1,y2,y3)

Step 3: Check One by One with Rejection Sampling

If all are accepted → “swallow” k tokens in one go
If y3 is rejected → accept y1,y2, and y3 is regenerated by the large model

Step 4: Repeat

9. Speculative Decoding Pseudocode

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
X = initial_context
while not finished:
    # Step 1: small model draft
    draft = S.generate_k_tokens(X)

    # Step 2: large model verify (one forward)
    logits_L = L.forward(X + draft)

    accepted = []
    for i, y in enumerate(draft):
        p = softmax(logits_L[i])[y]
        q = S.prob(y)

        alpha = p / (c * q)
        u = random()

        if u < alpha:
            accepted.append(y)
        else:
            # fallback: large model generates from here
            y_L = L.sample_token(X + accepted)
            accepted.append(y_L)
            break

    X = X + accepted

10. Why Doesn’t This Break Autoregressive Causality?

Because the causal mask guarantees:

  • The hidden state of y₂ can only use X + y₁
  • y₃ can only use X + y₁,y₂
  • …

Future tokens never leak information.

11. Pros and Cons of Speculative Decoding

Pros

  • Huge inference speedup (2×~4×, or even higher)
  • Output quality is completely unchanged (unbiased)
  • Fits modern inference architectures with KV Cache + batching
  • Already adopted by mainstream frameworks such as GPT-4, Llama, Qwen, and vLLM

Cons

  • Requires training or distilling a small model
  • The more accurate the small model → the greater the speedup
  • A small model that is too poor lowers the speedup ratio
  • Not suitable for tasks with very few output tokens (e.g. classification)

12. The Essence of Speculative Decoding in One Sentence

The small model predicts a draft, the large model verifies it in parallel with one forward pass, rejection sampling guarantees the distribution is unbiased, and thus the number of large-model forward passes is reduced, achieving faster inference.

Translated from the Chinese original.

中文