Speculative Decoding and Speculative Sampling: Mathematical Foundations, Lossless Rejection Sampling, Draft Model Architectures, and Serving Economics

Large language models generate text autoregressively, predicting one token at a time by passing the full sequence through dozens of transformer layers. During inference at low batch sizes, this generation process is fundamentally memory-bandwidth bound rather than compute bound. Speculative decoding addresses this bottleneck by using a fast drafting mechanism to propose candidate token sequences and verifying them in parallel with the primary target model in a single forward pass.

When configured with modified rejection sampling, speculative decoding is mathematically lossless: the sampled output matches the exact probability distribution of the target model while yielding 2x to 3.5x wall-clock latency reductions.

The Autoregressive Inference Bottleneck

Standard autoregressive transformer generation requires loading the complete set of model weights from high-bandwidth device memory (HBM) into SRAM or register files for every single generated token.

For a model with $P$ parameters executing at batch size 1 in 16-bit precision (FP16 or BF16), generating a single token requires transferring $2P$ bytes of memory. The arithmetic intensity of this single-token forward pass is:

Arithmetic Intensity = (2 * P FLOPs) / (2 * P Bytes) = 1 FLOP / Byte

Modern datacenter accelerators exhibit theoretical compute-to-bandwidth ratios between 100 and 250 FLOPs per byte. For example, an NVIDIA H100 SXM GPU provides 1,979 TFLOPs of dense FP16 tensor core compute alongside 3.35 TB/s of memory bandwidth, yielding a balance point of approximately 590 FLOPs per byte transferred. At 1 FLOP per byte, single-token generation leaves tensor cores idle for more than 98% of the execution cycle while waiting for weight retrieval from HBM.

Speculative decoding leverages the observation that scoring $K$ candidate tokens simultaneously requires virtually the same memory bandwidth and execution latency as scoring a single token. Because all $K$ tokens can be evaluated in parallel across matrix-matrix multiplication (GEMM) operations rather than sequential matrix-vector (GEMV) operations, the target model processes multiple positions in one forward pass without increasing memory transfers.

Mathematical Formulation and Lossless Rejection Sampling

Speculative decoding was independently formalized by Leviathan et al. (2022) at Google Research and Chen et al. (2023) at DeepMind. The core mechanism pairs a target model distribution $P(x \mid x_{<t})$ with a fast draft distribution $Q(x \mid x_{<t})$.

The Sampling Protocol

Given a prompt prefix $x_{<t}$ and a lookahead speculation window of $\gamma$ tokens:

1. Draft Phase: The draft system autoregressively generates $\gamma$ speculative tokens:

x_1, x_2, ..., x_\gamma ~ Q(x)

1. Parallel Verification Phase: The target model executes a single parallel forward pass over the full draft sequence, computing true probability distributions $P(x \mid x_{<t+i-1})$ for all positions $i \in \{1, 2, \dots, \gamma+1\}$.
2. Rejection Sampling Step: For each speculative token $i$ from $1$ to $\gamma$:

• Sample a uniform random variable $r_i \sim U(0, 1)$.
• Evaluate the acceptance condition:

r_i <= min(1, P(x_i | x_{<i}) / Q(x_i | x_{<i}))

• If accepted, the token $x_i$ is retained in the accepted prefix.
• If rejected at index $k \le \gamma$, the rejection loop terminates. A corrected replacement token is sampled from the adjusted residual distribution:

P'(x) = max(0, P(x) - Q(x)) / \sum_y max(0, P(y) - Q(y))

• All subsequent speculative tokens $x_{k+1}, \dots, x_\gamma$ are discarded.

1. Bonus Token Generation: If all $\gamma$ draft tokens are accepted, an additional $(\gamma+1)$-th token is sampled directly from the target model distribution $P(x \mid x_{\le \gamma})$ at zero incremental latency cost, since its logits were already computed in the parallel verification forward pass.

Proof of Distributional Equivalence

The rejection sampling scheme ensures that the generated sequence distribution under speculative decoding is identical to standard sampling from the target model.

For any candidate token $x$, the probability of emission under speculative sampling $P_{\text{spec}}(x)$ is the sum of two mutually exclusive events: acceptance through the draft branch and recovery through the rejection branch:

P_spec(x) = Q(x) * min(1, P(x) / Q(x)) + (1 - \alpha) * P'(x)

Here, $\alpha$ denotes the total probability of draft acceptance:

\alpha = \sum_y Q(y) * min(1, P(y) / Q(y)) = \sum_y min(P(y), Q(y))

The recovery distribution $P'(x)$ evaluates to:

P'(x) = (P(x) - min(P(x), Q(x))) / (1 - \alpha) = max(0, P(x) - Q(x)) / (1 - \alpha)

Substituting $P'(x)$ and $\alpha$ back into the emission equation:

P_spec(x) = min(P(x), Q(x)) + (1 - \alpha) * [ (P(x) - min(P(x), Q(x))) / (1 - \alpha) ] P_spec(x) = min(P(x), Q(x)) + P(x) - min(P(x), Q(x)) P_spec(x) = P(x)

The resulting token distribution exactly equals the target model distribution $P(x)$ across all vocabulary elements without approximation error or distributional drift.

Acceptance Rate and Latency Speedup Mechanics

The acceptance probability at each step is directly governed by the statistical proximity between draft and target distributions:

\alpha = 1 - D_TV(P, Q) = 1 - (1/2) \sum_x |P(x) - Q(x)|

where $D_{\text{TV}}(P, Q)$ represents the total variation distance between the two distributions.

Assuming independent and identically distributed acceptance across positions with average acceptance rate $\alpha$, the expected number of accepted tokens $\mathbb{E}[N]$ per verification round with lookahead $\gamma$ is:

\mathbb{E}[N] = (1 - \alpha^{\gamma+1}) / (1 - \alpha)

If $\alpha = 0.8$ and $\gamma = 5$, the expected token yield per target forward pass is 3.67 tokens. If $\alpha = 0.9$, the expected yield rises to 4.69 tokens.

Wall-Clock Speedup Factor

Let $t_{\text{target}}$ be the execution time of one forward pass of the target model, and $t_{\text{draft}}$ be the execution time of one forward pass of the draft model. Defining the relative cost ratio $c = t_{\text{draft}} / t_{\text{target}}$, the theoretical wall-clock speedup $S$ is given by:

S = \mathbb{E}[N] / (\gamma * c + 1)

Achieving net wall-clock acceleration requires two conditions:

1. High acceptance rate: The draft model must closely align with the target distribution, keeping $\alpha > 0.65$.
2. Low draft overhead: The per-token drafting cost $c$ must remain small, typically $c \le 0.05$ to $0.10$.

If $c = 0.05$, $\gamma = 5$, and $\alpha = 0.8$, the expected speedup is:

S = 3.67 / (5 * 0.05 + 1) = 3.67 / 1.25 = 2.94x

Draft Architectures and Taxonomy

Deployments employ four primary architectural paradigms for draft generation:

1. Independent Small Draft Models

A smaller autoregressive model (such as a 1B to 7B parameter model drafting for a 70B parameter target) generates the draft sequence. Both models must share an identical tokenizer and vocabulary mapping to avoid token fragmentation misalignment.

• Advantages: Completely independent architectures; no modifications to the target model.
• Trade-offs: Consumes separate GPU memory for draft weights and KV-cache; cross-model distribution shifts can reduce $\alpha$ on out-of-distribution prompts.

2. Multi-Head Speculation (Medusa)

Introduced by Cai et al. (2024), Medusa eliminates the separate draft model by appending multiple parameter-efficient feedforward heads directly on top of the final transformer layer of the target model.

• Head Architecture: Head $k$ predicts the token distribution at offset $t+k$ given the current hidden state $h_t$.
• Tree-Based Attention: Medusa generates a tree of candidate tokens across top-$k$ predictions from each head. A custom tree attention mask evaluates all candidate paths concurrently in a single forward verification pass.
• Advantages: Zero additional model loading; simple single-engine deployment; no vocabulary alignment issues.

3. Feature-Level Autoregressive Drafting (EAGLE and EAGLE-2)

Proposed by Li et al. (2024), EAGLE generates speculative tokens by autoregressively predicting the target model's top-layer hidden feature representations rather than token distributions directly.

• Mechanism: A lightweight single-layer transformer ingests the target model's second-to-top hidden state combined with shifted token embeddings, capturing contextual representations with lower semantic variance.
• Performance: EAGLE achieves higher acceptance rates ($\alpha > 0.85$ across standard benchmarks) than token-level drafting, delivering 2.5x to 3.5x speedups on instruction-tuned models.

4. Non-Parametric N-Gram and Prompt Lookup Decoding

For structured, repetitive, or retrieval-heavy generation (such as code editing, JSON extraction, and RAG continuations), candidates can be drawn directly from the input prompt or previously generated text via exact n-gram matching without any neural drafting model.

• Mechanism: The engine searches for matching sequence prefixes in the prompt context and speculates the following $K$ tokens.
• Overhead: Near-zero FLOP overhead; highly effective when outputs heavily quote input context.

System Implementation and KV-Cache Dynamics

Integrating speculative decoding into production inference engines (such as vLLM, TensorRT-LLM, and SGLang) requires dedicated memory management mechanics.

Rollback Management in Linear Speculation

When the target model rejects token $k$ during linear speculation, the KV-cache entries corresponding to positions $k+1 \dots \gamma$ must be invalidated. In paged memory architectures, the engine truncates the logical sequence length and returns unused physical KV blocks to the free pool without reallocating active sequence memory.

Tree Attention Masking

When using tree speculation (e.g., Medusa or EAGLE), candidates branch into a directed tree of possible continuations. Instead of running separate verification passes for each branch, engines construct a 2D causal tree attention mask:

Candidate Tree:
       [Root: A]
       /        \
   [B1]          [B2]
   /  \          /  \
[C1]  [C2]    [C3]  [C4]

The tree attention mask permits token $C_1$ to attend only to its causal ancestors (Root $A$ and parent $B_1$), preventing cross-branch attention leakage while computing logits for all 7 candidate paths in one unified GEMM kernel.

Serving Economics and Batching Trade-Offs

Speculative decoding provides dramatic acceleration under specific serving regimes, but introduces distinct computational trade-offs:

The Batch Size Crossover

Speculative decoding is optimized for memory-bandwidth bound execution (batch sizes 1 through 8). As concurrent batch sizes increase beyond 32 or 64 sequences, aggregate matrix operations transition from memory-bound GEMVs to compute-bound GEMMs.

In compute-bound regimes:

• Hardware tensor cores are already fully utilized.
• Verifying $\gamma$ tokens requires $\gamma$ times more FLOPs, increasing target forward pass latency.
• Draft model execution competes with the target model for tensor core compute and GPU memory bandwidth.
• Net throughput (tokens per second per GPU) plateaus or degrades relative to non-speculative continuous batching.

Production Deployment Guidelines

• Interactive Chat and Coding: Highly effective. At batch size 1 to 4, time-per-output-token (TPOT) drops by 50% to 70%, improving user-perceived responsiveness.
• High-Throughput Offline Batching: Ineffective. Standard continuous batching without speculation achieves higher overall token throughput per dollar by dedicating all compute capacity to serving concurrent requests.
• Agentic Loops and Reasoning Chains: Highly effective. Extended test-time reasoning traces and multi-step tool calls generate thousands of sequential tokens per request where single-stream latency is the primary operational bottleneck.

Sources

• Leviathan, Y., Kalman, M., & Matias, Y. (2022). Fast Inference from Transformers via Speculative Decoding. Google Research. arXiv:2211.17192.
• Chen, C., Borgeaud, S., Irving, G., Lespiau, J. B., Sifre, L., & Jumper, J. (2023). Accelerating Large Language Model Decoding with Speculative Sampling. DeepMind. arXiv:2302.01318.
• Cai, T., Li, Y., Geng, Z., Peng, H., Lee, J. D., Chen, D., & Dao, T. (2024). Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads. arXiv:2401.10774.
• Li, Y., Wei, F., Zhang, C., & Zhang, H. (2024). EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty. ICML 2024. arXiv:2401.15077.