Implement the speculative-decoding rejection sampler and prove it emits the target's distribution

mediumFree drillimplement emitted_distribution

Before you code

If the draft and target distributions disagree, would accepting every draft token preserve the target distribution? State what your sampler must preserve.

Review the prerequisite concept

The concept page shows its own access level. The statement and visible tests below are free.

At one draft position the draft model samples a token from its distribution $q$ , and the target model's distribution is $p$ . The verifier accepts the drafted token $x$ with probability $\min(1, p(x)/q(x))$ . On rejection it samples a replacement from the residual $(p - q)_+$ , normalized. Write the function that returns the exact law of the emitted token, not a sample of it.

emitted_distribution(p, q) -> (dist, alpha)
  p, q   [V] probability vectors (non-negative, each sums to 1; q may be 0 where p is not)
  dist   [V] probability that the verifier emits each token
  alpha  scalar, the probability that the drafted token is accepted

A correct sampler is lossless: dist equals p for every pair, and alpha is one minus the total-variation distance between p and q. When q == p the residual has no mass, and the function must still return a finite answer.

The visible tests check losslessness on random pairs. The full suite adds an identical draft, tokens the draft never proposes, tokens the target never wants, and the acceptance probability.

YOUR IMPLEMENTATION · PYTHON 3.14 + NUMPY

Free drill: the visible tests run for everyone, the edge cases and the solution for any signed-in reader.

The visible tests

import numpy as np


def _pairs(seed, n):
    rng = np.random.default_rng(seed)
    for _ in range(n):
        V = int(rng.integers(2, 40))
        yield rng.dirichlet(np.ones(V)), rng.dirichlet(np.ones(V) * 0.3)


def test_emits_the_target_distribution(ns):
    worst = 0.0
    for p, q in _pairs(0, 300):
        dist, _ = ns["emitted_distribution"](p, q)
        worst = max(worst, float(np.abs(np.asarray(dist) - p).max()))
    assert worst < 1e-12, f"emitted distribution differs from the target by up to {worst:.3g}"


def test_output_is_a_distribution(ns):
    for p, q in _pairs(1, 50):
        dist, _ = ns["emitted_distribution"](p, q)
        dist = np.asarray(dist)
        assert (dist >= -1e-15).all(), "negative probability"
        assert abs(dist.sum() - 1) < 1e-12, f"sums to {dist.sum():.6f}, not 1"

The reasoning behind this component: the companion page.