Speculative decoding: accept, reject and resample · AI Engineering

Speculative decoding makes a large model generate text faster with no change to its output. A small draft model proposes a few tokens. The large target model checks all of them in one forward pass. A rejection rule keeps the output distribution exactly equal to the distribution of the target model.

Task

Write four functions. In this problem, p is a distribution of the target model and q is a distribution of the draft model, over a vocabulary of V tokens.

acceptance_rate(target, draft) returns \alpha = \sum_x \min(p(x), q(x)). This is the probability that the rule accepts one draft token.
residual_distribution(target, draft) returns \mathrm{norm}(\max(0, p - q)): the elementwise maximum of 0 and p - q, divided by its sum. When that sum is 0, return a copy of target as a float array.
speculative_step(target_probs, draft_probs, draft_tokens, rng) runs the verify step for \gamma draft tokens and returns the output tokens.
expected_tokens(alpha, gamma) returns the expected number of tokens of one step, (1 - \alpha^{\gamma + 1}) / (1 - \alpha). For \alpha = 1 it returns \gamma + 1.

speculative_step does these steps. The order of the random draws is part of the task.

For each draft token x_i, with i from 0 to \gamma - 1, let p_i be row i of target_probs and q_i be row i of draft_probs.
Draw u = rng.random(). When u < \min(1, p_i(x_i) / q_i(x_i)), accept x_i, add it to the output, and go to the next draft token.
Otherwise, draw one token with rng.choice(V, p=r), where r is the residual distribution of p_i and q_i. Add it to the output and return the output.
When all \gamma draft tokens are accepted, draw one bonus token with rng.choice(V, p=p_gamma) from the last row of target_probs. Add it and return the output.

Input and output

Name	Type	Shape	Meaning
`target`	float array	`(V,)`	a target distribution p
`draft`	float array	`(V,)`	a draft distribution q
`target_probs`	float array	`(gamma + 1, V)`	the target distributions at each draft position and one more
`draft_probs`	float array	`(gamma, V)`	the draft distributions at each draft position
`draft_tokens`	list of int	length `gamma`	the tokens that the draft model sampled
`rng`	`np.random.Generator`		the source of all random draws
`alpha`	float		an acceptance rate from 0 to 1
`gamma`	int		the number of draft tokens

acceptance_rate returns a Python float.
residual_distribution returns a float array of shape (V,) that sums to 1.
speculative_step returns a list of Python ints of length from 1 to gamma + 1.
expected_tokens returns a Python float.

Constraints

2 <= V <= 64 and 0 <= gamma <= 8. Each row of probabilities sums to 1.
Each draft token has q_i(x_i) > 0, because the draft model sampled it. The target probability p_i(x_i) can be 0.
Use only rng.random() and rng.choice(V, p=...), in the order above. Make no other draws.
Do not change the input arrays or the list of draft tokens.
Use numpy only.

Examples

python

target = np.array([0.5, 0.3, 0.2, 0.0])
draft = np.array([0.25, 0.25, 0.25, 0.25])
acceptance_rate(target, draft)
residual_distribution(target, draft)

text

0.7
[0.83333333 0.16666667 0.         0.        ]

python

target_probs = np.array([[0.5, 0.3, 0.2, 0.0], [0.1, 0.6, 0.2, 0.1], [0.25, 0.25, 0.25, 0.25]])
draft_probs = np.array([[0.25, 0.25, 0.25, 0.25], [0.4, 0.2, 0.2, 0.2]])
speculative_step(target_probs, draft_probs, [0, 1], np.random.default_rng(0))
speculative_step(target_probs, draft_probs, [3, 0], np.random.default_rng(0))
speculative_step(target_probs, draft_probs, [0, 0], np.random.default_rng(3))

text

[0, 1, 0]
[0]
[0, 0, 3]

python

expected_tokens(0.8, 4), expected_tokens(1.0, 4), expected_tokens(0.0, 4)

text

(3.3616, 5.0, 1.0)