Speculative Decoding Math: Algorithms & Speedup Limits › Code Implementation
Language AI HandbookAll chapters Speculative Decoding Math: Algorithms & Speedup Limits › Code ImplementationOne section · Sections in this chapter7 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement the speculative decoding mathematics to see these concepts in action. The implementation serves two purposes: it verifies our mathematical derivations empirically, and it provides a clean reference implementation that shows how the theoretical concepts translate into actual code.
import numpy as np
rng = np.random.default_rng(42)First, we'll implement the acceptance criterion for a single token:
from typing import Tuple
import numpy as np
def acceptance_probability(p: np.ndarray, q: np.ndarray) -> float:
"""
Calculate the expected acceptance probability alpha.
Args:
p: Target model distribution over vocabulary
q: Draft model distribution over vocabulary
Returns:
Expected acceptance probability
"""
return np.sum(np.minimum(p, q))
def sample_with_rejection(
p: np.ndarray, q: np.ndarray, draft_token: int
) -> Tuple[int, bool]:
"""
Apply acceptance criterion to a draft token.
Args:
p: Target model distribution
q: Draft model distribution
draft_token: Token proposed by draft model
Returns:
Tuple of (final_token, was_accepted)
"""
# Calculate acceptance probability for this specific token
accept_prob = min(1.0, p[draft_token] / q[draft_token])
# Sample uniform and compare
u = rng.uniform()
if u < accept_prob:
return draft_token, True
else:
# Sample from residual distribution
residual = np.maximum(0, p - q)
residual = residual / residual.sum() # Normalize
new_token = rng.choice(len(p), p=residual)
return new_token, FalseLet's verify that this produces the correct distribution:
# Create example distributions
vocab_size = 10
p = np.array([0.3, 0.25, 0.15, 0.1, 0.08, 0.05, 0.03, 0.02, 0.01, 0.01])
q = np.array([0.2, 0.2, 0.2, 0.15, 0.1, 0.05, 0.04, 0.03, 0.02, 0.01])
# Run many samples to verify distribution
n_samples = 100000
final_tokens = []
for _ in range(n_samples):
# Draft model proposes a token
draft_token = rng.choice(vocab_size, p=q)
# Apply acceptance criterion
final_token, _ = sample_with_rejection(p, q, draft_token)
final_tokens.append(final_token)
# Calculate empirical distribution
empirical_dist = np.bincount(final_tokens, minlength=vocab_size) / n_samplesTarget distribution p(x): [0.3 0.25 0.15 0.1 0.08 0.05 0.03 0.02 0.01 0.01] Draft distribution q(x): [0.2 0.2 0.2 0.15 0.1 0.05 0.04 0.03 0.02 0.01] Empirical distribution: [0.298 0.25 0.152 0.101 0.08 0.049 0.03 0.02 0.01 0.01 ] Max deviation from target: 0.0018
The empirical distribution closely matches the target, confirming our acceptance criterion preserves the correct distribution. With 100,000 samples, the maximum deviation from target probabilities is well within statistical noise bounds. This provides strong empirical evidence that the theoretical guarantee holds in practice.
Now let's calculate expected speedup for different configurations:
def expected_tokens(alpha: float, k: int) -> float:
"""Calculate expected number of tokens per iteration."""
return (1 - alpha ** (k + 1)) / (1 - alpha)
def speedup(alpha: float, k: int, c: float) -> float:
"""
Calculate speedup ratio.
Args:
alpha: Acceptance probability
k: Draft length
c: Cost ratio (target time / draft time)
Returns:
Speedup factor
"""
expected_n = expected_tokens(alpha, k)
iteration_cost = 1 + k / c
return expected_n / iteration_costLet's visualize how speedup varies with draft length for different acceptance rates:
k_values = np.arange(1, 16)
c = 20 # Target model is 20x slower than draft
alphas = [0.5, 0.6, 0.7, 0.8, 0.9]
# Calculate speedups for each alpha
speedup_results = {}
for alpha in alphas:
speedup_results[alpha] = [speedup(alpha, k, c) for k in k_values]
The plot shows that optimal draft length increases with acceptance probability. In this configuration, peaks near 13 draft tokens, while peaks near 3. The exact optimum changes with the draft-to-target cost ratio, so the broader lesson is to lengthen drafts only while the additional accepted-token yield exceeds their generation cost.
Let's find the optimal for different scenarios:
def find_optimal_k(
alpha: float, c: float, max_k: int = 20
) -> Tuple[int, float]:
"""Find the optimal draft length and corresponding speedup."""
best_k = 1
best_speedup = speedup(alpha, 1, c)
for k in range(2, max_k + 1):
s = speedup(alpha, k, c)
if s > best_speedup:
best_speedup = s
best_k = k
return best_k, best_speedupconfigs = []
for alpha in [0.6, 0.7, 0.8, 0.9]:
for c in [10, 20, 50]:
opt_k, opt_speedup = find_optimal_k(alpha, c)
configs.append((alpha, c, opt_k, opt_speedup))Optimal draft length for different configurations: -------------------------------------------------- α c Optimal k Speedup -------------------------------------------------- 0.6 10 3 1.67x 0.6 20 4 1.92x 0.6 50 6 2.17x 0.7 10 4 1.98x 0.7 20 6 2.35x 0.7 50 8 2.76x 0.8 10 6 2.47x 0.8 20 8 3.09x 0.8 50 11 3.82x 0.9 10 10 3.43x 0.9 20 13 4.67x 0.9 50 19 6.37x
The results confirm our analysis: higher acceptance rates and faster draft models (larger ) support longer draft sequences and achieve greater speedup. Notice that the optimal increases with both and , which matches our theoretical understanding of the diminishing returns structure.
Finally, let's visualize the acceptance rate as a function of distribution similarity:
# Generate random distribution pairs with varying divergence
alphas_sim = []
kl_divs = []
n_pairs = 200
for _ in range(n_pairs):
# Create base distribution
base = rng.dirichlet(np.ones(20))
# Add noise to create draft distribution with varying divergence
noise_scale = rng.uniform(0, 2)
noise = rng.dirichlet(np.ones(20) * (1 / (noise_scale + 0.1)))
mix = rng.uniform(0.3, 1.0)
q_sim = mix * base + (1 - mix) * noise
q_sim = q_sim / q_sim.sum()
p_sim = base
# Calculate alpha and KL divergence
alpha_sim = np.sum(np.minimum(p_sim, q_sim))
# KL divergence (with smoothing to avoid log(0))
eps = 1e-10
kl = np.sum(p_sim * np.log((p_sim + eps) / (q_sim + eps)))
alphas_sim.append(alpha_sim)
kl_divs.append(kl)
The relationship between distribution divergence and acceptance probability is clear: as the draft model's distribution diverges from the target, acceptance rates drop substantially. This shows the importance of choosing draft models that approximate the target well.