Neural Mastery
← Practice
Speculative Decoding Rejection Samplermedium

Speculative Decoding Rejection Sampler

medium Inside an LLM 20–25 min

NumPy allowed · plain Python earns 10 bonus points

LLM step 1/4

Logits become a next-token distribution

Softmax turns a model score vector into probabilities.
Open lesson →
01prompt
02tokens
03logits
04probabilities
attention→next token
mechanism65%
layer20%
pattern10%
model5%

Softmax(logits) → probability distribution

Accelerate autoregressive inference without changing the target distribution using Speculative Decoding verification.

Task

Implement verify_speculative_tokens(draft_tokens, draft_probs, target_probs, random_samples). For each step $i$, accept draft token $x$ if $u_i le min(1, P_{target}(x) / P_{draft}(x))$. If accepted, append to accepted_tokens. If rejected: resample from the adjusted distribution $P'(y) = ext{normalize}(max(0, P_{target}(y) - P_{draft}(y)))$ choosing the argmax token as bonus_token, and terminate verification. If all $K$ tokens are accepted and target_probs contains a step $K$, sample bonus_token as $argmax P_{target}(K)$. Return dict with accepted_tokens, bonus_token, and acceptance_count.

Function Signature

verify_speculative_tokens(draft_tokens: list[int], draft_probs: list[list[float]], target_probs: list[list[float]], random_samples: list[float]) -> dict

Examples

Example 1: All tokens accepted with bonus token
Input: draft_tokens = [2, 1] draft_probs = [[0.1, 0.2, 0.7], [0.2, 0.6, 0.2]] target_probs = [[0.05, 0.15, 0.8], [0.1, 0.7, 0.2], [0.5, 0.3, 0.2]] random_samples = [0.5, 0.4]
Output: {'accepted_tokens': [2, 1], 'bonus_token': 0, 'acceptance_count': 2}
Example 2: Second token rejected with resample
Input: draft_tokens = [0, 1] draft_probs = [[0.8, 0.2], [0.9, 0.1]] target_probs = [[0.9, 0.1], [0.1, 0.9]] random_samples = [0.5, 0.8]
Output: {'accepted_tokens': [0, 1], 'bonus_token': None, 'acceptance_count': 2}
Example 3: Input length mismatch error
Input: draft_tokens = [0] draft_probs = [[0.5, 0.5]] target_probs = [[0.5, 0.5]] random_samples = [0.1, 0.2]
Output: Raises ValueError

Constraints

  • Mismatched input lengths must raise ValueError.
  • Rejection must immediately stop verification of remaining draft tokens.
  • Argmax tie-breaking must prefer higher probability, then lower token index.
  • A plain-Python solution earns 10 bonus points.
Python3Saved ✓
def verify_speculative_tokens(draft_tokens: list[int], draft_probs: list[list[float]], target_probs: list[list[float]], random_samples: list[float]) -> dict:
"""Verify draft tokens via speculative decoding rejection sampling.
Args:
draft_tokens: Sequence of K tokens proposed by draft model.
draft_probs: Draft model probabilities for each proposed step.
target_probs: Target model probabilities for each step (length >= K).
random_samples: Uniform random samples in [0, 1) for acceptance test.
Returns:
dict with 'accepted_tokens', 'bonus_token', and 'acceptance_count'.
"""
# Your implementation here
pass

Case 1: All tokens accepted with bonus token
Input: draft_tokens = [2, 1] draft_probs = [[0.1, 0.2, 0.7], [0.2, 0.6, 0.2]] target_probs = [[0.05, 0.15, 0.8], [0.1, 0.7, 0.2], [0.5, 0.3, 0.2]] random_samples = [0.5, 0.4]
Expected: {'accepted_tokens': [2, 1], 'bonus_token': 0, 'acceptance_count': 2}
Case 2: Second token rejected with resample
Input: draft_tokens = [0, 1] draft_probs = [[0.8, 0.2], [0.9, 0.1]] target_probs = [[0.9, 0.1], [0.1, 0.9]] random_samples = [0.5, 0.8]
Expected: {'accepted_tokens': [0, 1], 'bonus_token': None, 'acceptance_count': 2}
Case 3: Input length mismatch error
Input: draft_tokens = [0] draft_probs = [[0.5, 0.5]] target_probs = [[0.5, 0.5]] random_samples = [0.1, 0.2]
Expected: Raises ValueError