Hard Research

Speculative Decoding Accept/Reject

A draft model proposes a token. A target model may assign that token a lower probability, so we cannot always keep the proposal. Compute the probability of keeping it and the replacement distribution to use if it is rejected.

This exercise covers a deterministic calculation for one proposed token. Do not sample a token, make a random accept/reject decision, or run a model.

Inputs

Implement speculative_decode_step(draft_probs, target_probs, drafted_token). The two arrays have shape (V,), with V >= 1, and contain probability distributions over the same vocabulary. They have the same floating-point dtype and device. Values are finite, nonnegative and sum to one.

Write q = draft_probs and p = target_probs. The integer drafted_token is a valid vocabulary index t, and q[t] > 0, as it must be for a token sampled from the draft distribution. Other entries of q may be zero. The inputs are probabilities, so do not apply softmax or renormalize them.

Calculation

Compute accept_prob = min(1, p[t] / q[t]).
For every vocabulary entry, compute r = maximum(p - q, 0).
If r.sum() > 0, divide r by its sum. Otherwise, return the zero vector as the residual. In this case the distributions match and every possible draft token is accepted, so no replacement is needed.

Return a dictionary with exactly two entries:

"accept_prob": an ordinary Python float. Do not round it.
"residual": the resulting array of shape (V,), preserving the input dtype and device. Compute it even when this particular token is always accepted. It depends on the full distributions, not the selected token.

Leave both input arrays unchanged.

Example

For q = [0.8, 0.2], p = [0.4, 0.6] and t = 0, the acceptance probability is 0.4 / 0.8 = 0.5. The positive differences are [0, 0.4], so the residual is [0, 1]. A rejected proposal of token 0 would therefore be replaced by token 1.

In the full sampling procedure, proposals drawn from q are accepted with this probability; rejected proposals are replaced by a draw from the residual. Together, these two steps give the target distribution p. The acceptance rule alone does not provide that guarantee.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Check a draft model's proposed token

Speculative decoding uses a cheaper draft model to propose tokens and a target model to verify them. Let q be the draft distribution and p the desired target distribution for the same prefix. For a token x sampled from q, accept it with probability

$$\min\left(1,\frac{p(x)}{q(x)}\right).$$

If rejected, sample a replacement using max(p - q, 0) over the whole vocabulary, divided by its sum. This replacement distribution supplies the probability that is missing from accepted draft samples.

Work through a two-token example

For p = [0.7, 0.3] and q = [0.4, 0.6]:

Token 0 is proposed with probability 0.4 and always accepted.
Token 1 is proposed with probability 0.6 and accepted half the time, giving accepted probability 0.3.
The remaining probability 0.3 is rejected. Since max(p - q, 0) = [0.3, 0], all replacements are token 0.

The final probabilities are therefore [0.7, 0.3], matching p. More generally, the accepted probability at a token is min(p, q), and replacement adds the missing positive difference.

Handle zero-probability cases

Use the acceptance ratio only for the proposed token, but compute the replacement distribution over every token. A token with q(x) = 0 cannot be sampled from q. If p = q, every proposal is accepted and the replacement sum is zero. Avoid normalizing that unused zero vector, or use the fallback required by the task.

This exercise isolates the verification arithmetic. Full speculative generation also verifies blocks of tokens and manages model state. See the speculative decoding paper.

Detecting runtime…