Speculative Decoding: Mathematical Foundations, Distribution Preservation Proofs, Tree-Structured Verification, and Memory-Bandwidth Amortization

Autoregressive large language model (LLM) generation suffers from an acute hardware efficiency mismatch during inference. While the prefill phase (processing the input prompt) processes tokens in parallel and achieves high arithmetic intensity on modern matrix accelerators, the decode phase (generating text token-by-token) is fundamentally memory-bandwidth bound. At small batch sizes, each generated token requires transferring the model's entire multi-billion-parameter weight matrix from High-Bandwidth Memory (HBM) into on-chip Static Random-Access Memory (SRAM) and register files to execute a single matrix-vector product.

Speculative decoding, formalized independently by Leviathan, Kalman, and Matias (2022) and Chen et al. (2023), breaks this sequential memory bottleneck. By employing a fast, lightweight draft mechanism (such as a smaller companion model or specialized prediction heads) to propose sequences of candidate tokens and verifying them concurrently in a single forward pass of the larger target model, speculative decoding achieves 2x to 4x latency reductions. Crucially, through an exact speculative rejection sampling criterion, speculative decoding guarantees that the output token distribution remains identical to sampling directly from the target model.

The Memory-Wall Bottleneck in Autoregressive Generation

To understand the mechanics and necessity of speculative decoding, one must evaluate autoregressive inference through the Roofline model on modern GPU accelerators.

Arithmetic Intensity and the Roofline Model

The theoretical execution performance $P$ (in FLOPs per second) on an accelerator is governed by:

$$P = \min\left(P_{\text{peak}}, I \times B_{\text{peak}}\right)$$

where $P_{\text{peak}}$ represents peak floating-point compute capacity (FLOP/s), $B_{\text{peak}}$ denotes peak memory bandwidth (bytes/s), and $Arithmetic\ Intensity\ (I)$ represents the ratio of computational operations to bytes transferred from global memory:

$$I = \frac{\text{Floating-Point Operations (FLOPs)}}{\text{Memory Transferred (Bytes)}}$$

The machine balance threshold $I_{\text{threshold}} = P_{\text{peak}} / B_{\text{peak}}$ defines the transition between memory-bandwidth bound and compute-bound regimes.

Consider an NVIDIA H100 SXM5 GPU:

• Peak 16-bit Tensor Core compute ($P_{\text{peak}}$): 989.4 TFLOP/s (dense FP16/BF16)
• Peak HBM3 bandwidth ($B_{\text{peak}}$): 3.35 TB/s
• Machine balance threshold ($I_{\text{threshold}}$): $\approx 295.3 \text{ FLOPs/byte}$

The Arithmetic Intensity of Single-Token Decoding

For a standard Transformer decoder model with parameter count $N_{\text{params}}$ operating in 16-bit precision (2 bytes per parameter), computing a forward pass on a batch of $B$ tokens requires:

• Compute: $2 \times N_{\text{params}} \times B \text{ FLOPs}$ (from linear projection layers, ignoring minor attention terms for short contexts)
• Memory Movement: $2 \times N_{\text{params}} \text{ bytes}$ to read weights from HBM, plus key-value cache access $2 \times L \times 2 \times d_{\text{model}} \times S \text{ bytes}$ (where $L$ is layer depth and $S$ is context length).

When generating a single token for a single sequence ($B = 1$):

$$I_{\text{decode}} = \frac{2 \times N_{\text{params}} \times 1}{2 \times N_{\text{params}}} = 1 \text{ FLOP/byte}$$

Because $I_{\text{decode}} = 1 \ll 295.3 \text{ FLOPs/byte}$, single-token autoregressive decoding utilizes less than 0.5% of the accelerator's available arithmetic execution units. The GPU cores spend over 99% of execution cycles waiting for weights to travel across the memory bus.

Prefill vs. Verification Arithmetic Intensity

When verifying a sequence of $K$ candidate tokens concurrently, the target model processes all $K$ tokens in parallel within the same forward pass:

$$I_{\text{verify}} = \frac{2 \times N_{\text{params}} \times K}{2 \times N_{\text{params}}} = K \text{ FLOPs/byte}$$

For $K = 5$, arithmetic intensity increases by $5\times$ while loading the exact same parameter weights from HBM once. Because the latency to read weights dominates the execution time, evaluating $K$ tokens simultaneously in parallel takes roughly the same wall-clock time as evaluating a single token. Speculative decoding exploits this slack capacity to convert wasted memory-read cycles into useful token verification.

Mathematical Foundations of Speculative Sampling

The primary theoretical achievement of speculative decoding is proving that acceleration does not require output approximation. The generated text follows the exact target distribution $P(x)$ without altering token probabilities or sampling entropy.

Problem Formulation

Let $\mathcal{V}$ denote the vocabulary space.

• Target model distribution: $P(x \mid x_{<t}) = P(x)$
• Draft model distribution: $Q(x \mid x_{<t}) = Q(x)$

At step $t$, the draft model generates a sequence of $\gamma$ speculative tokens sequentially:

$$\tilde{x}_1 \sim Q(x \mid x_{<t}), \quad \tilde{x}_2 \sim Q(x \mid x_{<t}, \tilde{x}_1), \quad \dots, \quad \tilde{x}_\gamma \sim Q(x \mid x_{<t}, \tilde{x}_1, \dots, \tilde{x}_{\gamma-1})$$

The target model then performs a single parallel forward pass over the concatenated sequence $[x_{<t}, \tilde{x}_1, \dots, \tilde{x}_\gamma]$, obtaining the exact conditional probability distributions $P(x \mid x_{<t}, \tilde{x}_1, \dots, \tilde{x}_{i-1})$ for all $i \in \{1, \dots, \gamma\}$ simultaneously.

The Modified Rejection Sampling Algorithm

For each candidate token $\tilde{x}_i$ at position $i \in \{1, \dots, \gamma\}$, we evaluate the speculative acceptance probability $\alpha(\tilde{x}_i)$:

$$\alpha(\tilde{x}_i) = \min\left(1, \frac{P(\tilde{x}_i)}{Q(\tilde{x}_i)}\right)$$

We sample a uniform random variable $r_i \sim \mathcal{U}(0, 1)$:

1. Acceptance: If $r_i \le \alpha(\tilde{x}_i)$, the candidate token $\tilde{x}_i$ is accepted: $x_i = \tilde{x}_i$. The verification proceeds to token $i+1$.
2. Rejection: If $r_i > \alpha(\tilde{x}_i)$, the candidate token $\tilde{x}_i$ is rejected. The loop terminates immediately. A replacement token $x_i$ is sampled from the adjusted residual distribution $P'(x)$:

$$P'(x) = \frac{\max\left(0, P(x) - Q(x)\right)}{\sum_{y \in \mathcal{V}} \max\left(0, P(y) - Q(y)\right)}$$

All subsequent draft tokens $\tilde{x}_{i+1}, \dots, \tilde{x}_\gamma$ are discarded.

1. Bonus Token: If all $\gamma$ candidate tokens are accepted, the target model's forward pass has already computed $P(x \mid x_{<t}, \tilde{x}_1, \dots, \tilde{x}_\gamma)$ for the $(\gamma+1)$-th position. A bonus token $x_{\gamma+1}$ is sampled directly from $P(x \mid x_{<t}, \tilde{x}_1, \dots, \tilde{x}_\gamma)$ at zero additional marginal inference cost.

Thus, each speculative round produces between $1$ and $\gamma+1$ valid tokens.

Mathematical Proof of Exact Distribution Preservation

We now prove that sampling a token $X$ through this speculative rejection mechanism yields the identical distribution as sampling directly from the target model: $\mathbb{P}(X = x) = P(x)$ for all $x \in \mathcal{V}$.

Proof

Let $X$ be the random variable representing the token chosen at position $i$. A token $x$ can be produced via two mutually exclusive events:

1. $x$ was proposed by the draft model ($X_{\text{draft}} = x$) and subsequently accepted.
2. The draft model proposed some token $y$, which was rejected, and $x$ was sampled from the residual distribution $P'$.

The total probability is:

$$\mathbb{P}(X = x) = \mathbb{P}(\text{Draft proposes } x \text{ and accepts}) + \mathbb{P}(\text{Draft is rejected}) \cdot P'(x)$$

First, evaluate the probability of proposing $x$ and accepting it:

$$\mathbb{P}(\text{Draft proposes } x \text{ and accepts}) = Q(x) \cdot \alpha(x) = Q(x) \cdot \min\left(1, \frac{P(x)}{Q(x)}\right) = \min\left(Q(x), P(x)\right)$$

Second, evaluate the total acceptance rate $\beta$, defined as the marginal probability that any proposed draft token is accepted:

$$\beta = \sum_{y \in \mathcal{V}} Q(y) \alpha(y) = \sum_{y \in \mathcal{V}} \min\left(Q(y), P(y)\right)$$

Consequently, the total probability of rejection is:

$$1 - \beta = 1 - \sum_{y \in \mathcal{V}} \min\left(Q(y), P(y)\right)$$

Using the algebraic identity $\max(0, a - b) = a - \min(a, b)$, we have:

$$\sum_{y \in \mathcal{V}} \max\left(0, P(y) - Q(y)\right) = \sum_{y \in \mathcal{V}} P(y) - \sum_{y \in \mathcal{V}} \min\left(Q(y), P(y)\right) = 1 - \beta$$

Substituting this into the definition of the residual distribution $P'(x)$:

$$P'(x) = \frac{\max\left(0, P(x) - Q(x)\right)}{1 - \beta}$$

Now, substitute these expressions back into the total probability formula:

$$\mathbb{P}(X = x) = \min\left(Q(x), P(x)\right) + (1 - \beta) \cdot \frac{\max\left(0, P(x) - Q(x)\right)}{1 - \beta}$$

$$\mathbb{P}(X = x) = \min\left(Q(x), P(x)\right) + \max\left(0, P(x) - Q(x)\right)$$

Using the identity $\min(a, b) + \max(0, a - b) = a$:

$$\mathbb{P}(X = x) = P(x)$$

This completes the proof. The output distribution of speculative decoding is mathematically identical to the target model distribution.

Greedy Decoding Equivalence

For greedy decoding (temperature $T = 0$), the target and draft distributions collapse to Dirac delta distributions:

$$P(x) = \begin{cases} 1 & \text{if } x = \arg\max_{v} P(v) \\ 0 & \text{otherwise} \end{cases}, \quad Q(x) = \begin{cases} 1 & \text{if } x = \arg\max_{v} Q(v) \\ 0 & \text{otherwise} \end{cases}$$

The acceptance criterion simplifies to deterministic equality:

$$\alpha(\tilde{x}) = 1 \iff \arg\max_{v} P(v) = \arg\max_{v} Q(v)$$

If the draft model's greedy prediction matches the target model's greedy prediction, it is accepted; otherwise, execution rolls back and emits the target model's greedy token.

Theoretical Speedup and Performance Economics

The acceleration achieved by speculative decoding depends on three interacting variables: the draft acceptance rate, the draft model latency, and the speculation window length.

Expected Accepted Tokens per Speculative Round

Assume an independent and identically distributed per-token acceptance rate $\beta \in [0, 1]$ across a speculation horizon of $\gamma$ draft tokens.

The probability of accepting exactly $k$ draft tokens (where $0 \le k \le \gamma$) follows a geometric truncated distribution:

• For $0 \le k < \gamma$: Probability of accepting $k$ tokens and rejecting on token $k+1$ is $\beta^k (1 - \beta)$. Each such event yields $k + 1$ total emitted tokens (including the corrective token from $P'$).
• For $k = \gamma$: Probability of accepting all $\gamma$ tokens is $\beta^\gamma$. This event yields $\gamma + 1$ total emitted tokens (including the bonus token).

The expected number of emitted tokens per round $\mathbb{E}[N]$ is:

$$\mathbb{E}[N] = \sum_{k=0}^{\gamma-1} (k + 1) \beta^k (1 - \beta) + (\gamma + 1) \beta^\gamma = \frac{1 - \beta^{\gamma+1}}{1 - \beta}$$

As $\beta \to 1$, $\mathbb{E}[N] \to \gamma + 1$. When $\beta = 0.8$ and $\gamma = 5$:

$$\mathbb{E}[N] = \frac{1 - 0.8^6}{1 - 0.8} = \frac{1 - 0.262144}{0.2} = 3.689 \text{ tokens/round}$$

Wall-Clock Speedup Formulation

Let:

• $t_{\text{target}}$ = execution time for one forward pass of the target model
• $t_{\text{draft}}$ = execution time for one forward pass of the draft model
• $c = t_{\text{draft}} / t_{\text{target}}$ = relative cost ratio ($0 < c \ll 1$)

One round of speculative decoding with speculation length $\gamma$ takes:

$$T_{\text{spec}} = \gamma \cdot t_{\text{draft}} + t_{\text{target}} = (\gamma \cdot c + 1) t_{\text{target}}$$

Standard autoregressive decoding requires $t_{\text{target}}$ for every single token, taking $\mathbb{E}[N] \cdot t_{\text{target}}$ to produce $\mathbb{E}[N]$ tokens. The theoretical speedup factor $S$ is:

$$S = \frac{\text{Standard Latency}}{\text{Speculative Latency}} = \frac{\mathbb{E}[N] \cdot t_{\text{target}}}{(\gamma \cdot c + 1) t_{\text{target}}} = \frac{1 - \beta^{\gamma+1}}{(1 - \beta)(\gamma \cdot c + 1)}$$

Table: Theoretical Speedup as a Function of Acceptance Rate (beta) and Cost Ratio (c) for gamma = 5
-------------------------------------------------------------------------------------------------
Acceptance (beta) | Cost Ratio (c = 0.05) | Cost Ratio (c = 0.10) | Cost Ratio (c = 0.15)
-------------------------------------------------------------------------------------------------
0.50              | 1.57x                 | 1.31x                 | 1.12x
0.70              | 2.37x                 | 1.97x                 | 1.69x
0.80              | 2.95x                 | 2.46x                 | 2.11x
0.90              | 3.75x                 | 3.12x                 | 2.68x
0.95              | 4.27x                 | 3.56x                 | 3.05x
-------------------------------------------------------------------------------------------------

Optimal Speculation Horizon ($\gamma^*$)

Differentiating the speedup equation with respect to $\gamma$ reveals that an optimal speculation horizon $\gamma^*$ exists for any pair $(\beta, c)$:

• If $\gamma$ is too small, the system underutilizes the target model's parallel verification capacity.
• If $\gamma$ is too large, the marginal probability $\beta^\gamma$ of accepting distant tokens approaches zero, but the draft model cost $\gamma \cdot c$ continues to accumulate linearly, degrading overall throughput.

In production systems such as vLLM and SGLang, dynamic speculative schedulers track the running acceptance rate $\hat{\beta}$ per sequence and adaptively resize $\gamma \in [3, 7]$ on a per-step basis.

Tree-Structured Speculation Topologies

Standard speculative decoding uses a linear sequence of candidate tokens. However, linear speculation suffers from sequential dependency: if the second candidate token is rejected, all subsequent tokens ($\tilde{x}_3, \dots, \tilde{x}_\gamma$) are immediately invalidated, regardless of their intrinsic quality.

Tree-based speculation constructs a branched tree of hypotheses, allowing the target model to evaluate multiple candidate execution paths in parallel.

SpecInfer and Tree-Based Verification

SpecInfer (Miao et al., 2023) introduced tree-structured speculation. The draft model generates a prefix tree (trie) of candidate tokens with branching factor $b$ and depth $d$.

To verify a tree topology with $M$ total nodes in a single forward pass of the target model, SpecInfer constructs a specialized causal tree 2D attention mask $A_{\text{tree}} \in \{0, 1\}^{M \times M}$:

$$A_{\text{tree}}(i, j) = \begin{cases} 1 & \text{if node } j \text{ is an ancestor of node } i \text{ (or } j = i \text{)} \\ 0 & \text{otherwise} \end{cases}$$

This causal tree mask prevents attention leakage across distinct sibling branches while allowing each token to attend to its exact historical prefix path. The target model evaluates all $M$ nodes in a single forward kernel call, selects the longest valid accepted path, and rolls back all unselected branches.

Medusa: Multi-Head Speculation Without Companion Models

Medusa (Cai et al., 2024) eliminates the requirement for a separate draft model by appending $K$ lightweight decoding heads (single-layer feed-forward networks) directly to the final hidden state of the target model:

$$H = \text{TransformerLayer}_L(x_{\le t})$$

$$\tilde{x}_{t+1} = \text{Head}_0(H), \quad \tilde{x}_{t+2} = \text{Head}_1(H), \quad \dots, \quad \tilde{x}_{t+K} = \text{Head}_K(H)$$

Medusa heads generate top-$k$ predictions at each positional offset simultaneously. The system takes the Cartesian product of the top candidates, filters them into a fixed tree structure (e.g., 64 candidate paths), and validates them via tree-attention masking in the next target step. Because Medusa heads run concurrently on the target model's existing activations, draft generation latency $t_{\text{draft}} \approx 0$.

EAGLE and EAGLE-2: Feature-Level Extrapolation

EAGLE (Li et al., 2024) and EAGLE-2 (Li et al., 2024) observe that language modeling is substantially more predictable in feature representation space than in discrete token space.

EAGLE feeds the target model's second-to-last layer hidden states into a single transformer decoder layer, predicting top-level feature vectors autoregressively before projecting to token logits. EAGLE-2 incorporates dynamic draft trees that condition their branching factor on contextual draft confidence (entropy of draft heads), achieving acceptance rates exceeding 85% on coding and structured data benchmarks.

Drafting Paradigms: A Comparative Analysis

Modern inference engines utilize four primary speculative drafting architectures:

Speculative Drafting Architectures
----------------------------------------------------------------------------------------
1. Small Companion Model (Draft Model)
   - Mechanism: Separate smaller model from the same architecture family (e.g. Llama-3-8B drafting for Llama-3-70B).
   - Advantages: High acceptance fidelity, shares token vocabulary and tokenizer.
   - Disadvantages: Requires additional GPU memory allocation for draft weights and draft KV cache.

2. Multi-Head Predictions (Medusa / Hydra)
   - Mechanism: Additional linear / MLP heads attached to target model's final hidden state.
   - Advantages: Zero additional weight memory overhead; no inter-model synchronization.
   - Disadvantages: Requires dedicated supervised fine-tuning of heads; lower acceptance on out-of-distribution reasoning.

3. Feature-Level Recurrent Predictors (EAGLE / EAGLE-2)
   - Mechanism: Single-layer transformer operating on top-layer target embeddings.
   - Advantages: High acceptance rates (80-90%); robust across diverse temperatures.
   - Disadvantages: Requires training an auxiliary lightweight feature module.

4. N-Gram & Prompt Lookup (Prompt-Lookup / REST)
   - Mechanism: Matches context n-grams against prompt history to fetch candidate continuations.
   - Advantages: Zero compute overhead, zero parameter training.
   - Disadvantages: Effective only on high-redundancy tasks (retrieval-augmented generation, document summarization, code editing).

Production Serving Constraints and System Dynamics

While speculative decoding consistently accelerates single-stream generation ($B = 1$), its economic viability shifts under heavy batch serving conditions.

The Batch Size Throughput Trade-off

As concurrent serving batch size $B$ increases on a GPU:

1. High batch sizes ($B \ge 64$) naturally elevate arithmetic intensity $I = B \times 1 \text{ FLOP/byte}$, pushing the GPU out of the memory-bandwidth bound regime and into the compute-bound regime.
2. In the compute-bound regime, verifying $K$ speculative tokens requires $K\times$ additional FLOPs that compete directly with other concurrent requests for Tensor Core execution slots.
3. If the acceptance rate $\beta$ is insufficient, the compute cost of rejected speculative tokens reduces overall token throughput (tokens/second across all users) compared to standard continuous batching without speculation.

Consequently, modern LLM inference systems deploy speculative decoding strategically:

• Low concurrency / Strict SLA regimes: Speculative decoding is activated to minimize Time-to-First-Token (TTFT) and Inter-Token Latency (ITL).
• High throughput / Saturated queue regimes: Speculative decoding is dynamically throttled or disabled to maximize global cluster throughput.

KV Cache Management in Speculative Verification

In standard autoregressive decoding, exactly one KV entry is appended to the KV cache per sequence per step. In speculative decoding:

1. Speculative Allocation: Space for $\gamma$ candidate tokens must be reserved in the paged KV cache during draft generation.
2. Selective Pruning and Rollback: Upon target verification, only the accepted prefix of length $k \le \gamma$ is retained. The memory manager must immediately reclaim the unaccepted slots ($k+1$ through $\gamma$) without causing memory fragmentation or incurring memory copy overhead.

Engines such as vLLM implement paged tree-KV tables, using logical-to-physical block mapping tables to prune rejected branches with $O(1)$ pointer reassignments.

Sources

• Leviathan, Y., Kalman, M., & Matias, Y. (2022). Fast Inference from Transformers via Speculative Decoding. 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. arXiv:2302.01318.
• Miao, X., Oliaro, M., Zhi, Z., Le, T. V., Wang, X., Xu, Z., Lin, Z., Catalyurek, U., & Jia, Z. (2023). SpecInfer: Accelerating Large Language Model Serving with Tree-based Speculative Inference and Verification. arXiv:2305.09781.
• Cai, T., Li, Y., Geng, Z., Peng, B., Lee, J. D., Chen, D., & Dao, T. (2024). Medusa: Simple LLM Inference Acceleration with Multiple Decoding Heads. arXiv:2401.10774.
• Li, Y., Wei, F., Zhang, C., & Zhang, H. (2024). EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty. arXiv:2401.15077.
• Li, Y., Wei, F., Zhang, C., & Zhang, H. (2024). EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees. arXiv:2406.16858.
• Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., & Stoica, I. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.
• Zheng, L., Yin, L., Xie, Z., Sun, C., Huang, C. H., Yu, C. H., Cao, S., Christakis, C., Stoica, I., Gonzalez, J. E., & Sheng, Y. (2023). SGLang: Efficient Execution of Structured Language Model Programs. arXiv:2312.07104.