AI Engineering · Inference Speed
Speculative Decoding
LLMs generate one word at a time — each word costs a full forward pass.
Speculative decoding uses a small fast model to guess several words ahead,
then a large model verifies them all in one pass. Result: 1.87× faster
with mathematically identical output.
The technique
Why this is non-obvious
A large model (e.g., GPT-4o, 70B parameters) is slow but accurate. A small draft model (e.g., GPT-2, 124M parameters) is fast but sometimes wrong. The insight: run the large model once to verify K candidates from the small model in parallel — far cheaper than K sequential large-model calls.
How verification works (rejection sampling)
For each draft token t, compute α = min(1, p_target(t) / p_draft(t)). Accept with probability α. On rejection, sample a corrected token from (p_target − α·p_draft).clamp(0). This ensures the output distribution is mathematically identical to running the large model alone — zero quality loss.
The bonus token
When all K draft tokens are accepted, the large model's final forward pass generates one extra "bonus" token for free — since we already have its output distribution. This increases throughput beyond the naive speedup estimate.
How to explore
No API key or GPU needed. "Step Visualizer" shows token-by-token acceptance/rejection. "Benchmark" shows speedup vs draft length K. "The Math" shows the rejection sampling proof.
""")
with gr.Tab("Step Visualizer"):
gr.HTML('
Why K=5 is optimal for this model pair
At K=5, the extra verification overhead of longer drafts starts to outweigh the speedup. Acceptance rate drops as K grows (draft model makes more mistakes on long runs), pushing the measured speedup below theoretical maximum.
Why code has higher acceptance rates
Code follows strict syntactic rules — the draft model's distribution closely matches the target on deterministic patterns like indentation, keywords, and brackets. Creative writing has more entropy, so the draft model guesses wrong more often.
""")
with gr.Tab("The Math"):
gr.Markdown("""
## Rejection Sampling Proof
For each draft token $t_i$ with draft probability $q(t_i)$ and target probability $p(t_i)$:
**Accept** with probability $\\alpha_i = \\min\\left(1, \\frac{p(t_i)}{q(t_i)}\\right)$
**On rejection**, sample corrected token from:
$$p'(x) = \\frac{(p(x) - \\alpha_i \\cdot q(x))^+}{\\sum_x (p(x) - \\alpha_i \\cdot q(x))^+}$$
**Key property**: This produces the exact target distribution $p(x)$ — the output is indistinguishable from pure autoregressive sampling with the large model.
## Implementation
```python
def speculative_step(self, input_ids, max_new_tokens=5):
# Step 1: Draft model generates K tokens (K forward passes, cheap)
draft_tokens, draft_probs = self._get_draft_tokens(input_ids, K=5)
# Step 2: Target model verifies ALL K tokens in ONE forward pass
target_probs = self._verify_with_target(input_ids, draft_tokens)
# Step 3: Rejection sampling
accepted = []
for i, (tok, q, p) in enumerate(zip(draft_tokens, draft_probs, target_probs[:-1])):
alpha = min(1.0, p[tok] / q[tok])
if random.random() < alpha:
accepted.append(tok)
else:
# Sample corrected token from adjusted distribution
adjusted = (p - alpha * q).clamp(min=0)
adjusted /= adjusted.sum()
accepted.append(torch.multinomial(adjusted, 1).item())
break # Stop at first rejection
# Step 4: Bonus token if all K accepted
if len(accepted) == len(draft_tokens):
bonus = torch.multinomial(target_probs[-1], 1).item()
accepted.append(bonus)
return accepted
```
## Expected Speedup Formula
$$\\text{Speedup} \\approx \\frac{1 + K\\alpha}{1 + K\\alpha / \\text{speedup}_{\\text{draft}}}$$
Where $\\alpha$ = mean acceptance rate, K = draft length
## References
- Speculative Decoding ([arxiv 2211.17192](https://arxiv.org/abs/2211.17192))
- Accelerating Large Language Model Decoding with Speculative Sampling ([arxiv 2302.01318](https://arxiv.org/abs/2302.01318))
""")
demo.launch()