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.
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.
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.
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.
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.
Click a step above to visualize it.
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.
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.
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
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)
- Accelerating Large Language Model Decoding with Speculative Sampling (arxiv 2302.01318)