Speculative Decoding: Lossless Multi-Token Generation

First check the one-token correction mechanism. Even though some samples come from the draft model, the accepted-or-repaired output should match the target distribution. Then simulate the accepted-prefix length that creates the latency upside.

import numpy as np

target_p = np.array([0.45, 0.30, 0.15, 0.10])
draft_q = np.array([0.30, 0.45, 0.20, 0.05])
rng = np.random.default_rng(0)

def residual_distribution(p, q):
    repair_mass = np.maximum(0.0, p - q)
    return repair_mass / repair_mass.sum()

def speculative_one_token(p, q):
    draft = rng.choice(len(q), p=q)
    accept_prob = min(1.0, p[draft] / q[draft])
    if rng.random() < accept_prob:
        return draft, True
    repair = rng.choice(len(p), p=residual_distribution(p, q))
    return repair, False

samples = np.zeros_like(target_p)
accepted = 0
for _ in range(200_000):
    token, was_accepted = speculative_one_token(target_p, draft_q)
    samples[token] += 1
    accepted += int(was_accepted)

print("target distribution:   ", np.round(target_p, 3))
print("speculative samples:   ", np.round(samples / samples.sum(), 3))
print("draft-token acceptance:", round(accepted / samples.sum(), 3))

def expected_accepted(alpha, k):
    # Expected length of consecutive acceptances (prefix), capped at k.
    if alpha == 1.0:
        return float(k)
    return alpha * (1.0 - alpha**k) / (1.0 - alpha)

def simulate(alpha, k, trials=50000, seed=0):
    rng = np.random.RandomState(seed)
    total = 0
    for _ in range(trials):
        a = 0
        for _ in range(k):
            if rng.rand() < alpha:
                a += 1
            else:
                break
        total += a
    return total / trials

for alpha in [0.3, 0.6, 0.85]:
    for k in [2, 4, 8]:
        theo = expected_accepted(alpha, k)
        mc = simulate(alpha, k)
        print(f"alpha={alpha:.2f} k={k:>2}  E[accepted] theo={theo:.3f}  mc={mc:.3f}")