Speculative Decoding in Language Models

Authors: Sanjit Neelam, Daniel Heinlein, Vaclav Cvicek, Akshay Mishra, and Reiner Pope
Date: April 08, 2025

Speculative decoding (SD) has been shown to reduce the latency of autoregressive decoding (AD) by 2-3× for small batch sizes. However, increasing throughput and therefore reducing the cost per token requires decoding with large batch sizes. Recent work shows that SD can accelerate decoding with large batch sizes too if the context is sufficiently long and the draft model’s KV cache is sparse. We introduce SPIRe, a draft model that combines static sparse attention, pruned initialization, and feedback memory to increase the modeled throughput of speculative decoding by over 100% compared to speculation with a much smaller draft model and by over 35% compared to the strong baseline of sparse self-speculation. Our approach is particularly effective when context lengths vary significantly across requests. A PDF version of this article is available at arxiv.org.

Introduction

Speculative decoding uses a cheap draft model to generate candidate tokens that are verified in parallel by an expensive target model (Leviathan et al., 2023; Chen et al., 2023). Accepted draft tokens are guaranteed to be samples from the target model’s distribution over the next token by a modified rejection sampling scheme. If draft tokens are generated quickly enough and have a high enough acceptance rate, speculative decoding improves both the latency and throughput of decoding by reducing the number of memory accesses needed to generate each token 1.

Optimizing speculative decoding for higher throughput requires different strategies than optimizing for lower latency; the latter is well-documented in the literature (Xia et al., 2024; Miao et al., 2024; Yan et al., 2025). For example, low latency can often be achieved with small batch sizes (Pope et al., 2022) and small draft models are preferred in this setting since fetching weights is the bottleneck. On the other hand, high throughput requires large batch sizes, which enables the use of larger draft models since in this regime loading the KV cache is the bottleneck (Chen et al., 2024). See Figure 1 and Figure 4 for examples of how the speedup due to speculative decoding varies with draft model architecture, batch size, and context length according to our performance model.

Figure 1: Speedup of generating tokens in a batch of size B=64 given a context of length L. A round of speculation is the generation and subsequent verification of each block of k=4 draft tokens, the average generation length τ is the average number of tokens generated per round of speculation, and speedup is τ divided by how much longer a round of speculation takes compared to autoregressively decoding one token.

Prior work has shown that the acceptance rate of draft tokens can be increased using the target model’s intermediate activations and output logits (Du et al., 2024; Zhou et al., 2024). Since a target model’s state can be cached during training at negligible cost, and since inference rather than training accounts for most of the cost of serving an LLM, the investment in training a draft model aligned to a particular target model may be worthwhile. Indeed, we show below that this is the case for SPIRe under conservative assumptions.

Our work, along with Chen et al., 2024, demonstrates that speculative decoding is an effective technique for reducing the cost of decoding, and our work enables more efficient exploration of the design space of draft models based on the batch sizes and context lengths expected in production. Our main contributions are the following:

• An implementation-agnostic model for evaluating the throughput of speculative decoding with different draft models.
• An efficient method for training a Feedback Transformer draft model.
• A draft model, SPIRe, that increases the modeled throughput of speculative decoding by over 100% compared to vanilla speculative decoding and by over 35% compared to the strong baseline of MagicDec (Chen et al., 2024).

Evaluation Metric

It is problematic to compare different methods of training the draft model using the throughput achieved by the resulting models in a production environment. Measured throughput depends on myriad implementation details such as use of specialized kernels (Dao et al., 2022), how models are partitioned across chips (Pope et al., 2022), how queries are batched (Daniel et al., 2023), and which queries are used for evaluation (Zhou et al., 2024). Any of these variables could be confounders when comparing different draft model architectures.

To derive a proxy for measured throughput, we first note that the throughput of speculative decoding is equal to the throughput of vanilla decoding multiplied by:

ThroughputMultiplier=E[# Tokens Generated per Round of Speculation] / (k × t_draft + t_verify) / t_target

To obtain a more implementation-agnostic expression, we can break down the cost of a forward pass into compute cost and memory cost. By multiplying the memory cost in bytes with the hardware operational intensity (HOI) in FLOPs per byte, we obtain an estimate of the memory cost in "/FLOP-equivalents" and compare the compute and memory costs on equal footing. Given this, we obtain the following expression for the cost of a forward pass:

Forward Cost=max(Compute Cost,FLOPs),(Weight Cost + KV Cache Cost) × HOI.

Finally, substituting the latencies for the draft, verify, and target model with their corresponding forward costs yields an implementation-agnostic approximation of the throughput multiplier which serves as our evaluation metric.

SPIRe

We combine the following techniques when training our draft model SPIRe to increase the throughput multiplier of speculative decoding. The name SPIRe stands for Sparse KV cache, Pruned Initialization, and Recurrent Feedback Transformer.

1. Sliding window KV cache with attention sink

In the large-batch long-context regime, the memory cost of loading the KV cache is the dominant cost of decoding, and reducing the size of the KV cache can significantly reduce the iteration time multiplier (Δt). Following MagicDec, we use a StreamingLLM attention mask during training and inference to ensure the draft model’s KV cache is sparse and that the size of the KV cache is constant with respect to the decoding sequence length.

2. Initialize by pruning the target model

In our experiments, we initialized the draft model using the embedding layer, the last two transformer blocks, and the unembedding layer of the trained target model, producing higher τ than random initialization.

3. Feedback Transformer and attending to target model activations

Feedback memory was proposed to increase the representation capacity of Transformers, creating Feedback Transformers with stronger performance at any given inference cost. We require our Feedback Transformer draft model to be able to generate blocks of k tokens given contexts of various lengths.

4. Distillation loss

Early implementations of speculative decoding for LLMs optimize the draft model’s parameters by minimizing the standard cross-entropy loss. We minimize a MixedLoss to train the draft model to produce outputs that are more similar to the outputs of the target model.

Evaluation

We fix an 8-layer, 67-million parameter multi-head attention target model and vary the draft model, with all model architecture details provided in the Appendix. Each model with N total-parameters is trained on 20 × N tokens (Hoffmann et al., 2022) from the LongCrawl64 dataset using sequences of length 1024.

We measure average generation lengths τ by generating G=64 tokens given n=65536 contexts from the validation split of the LongCrawl64 dataset.

Table 1: Average generation lengths τ, with 95% confidence intervals, for different draft models and context lengths. Using a maximum speculation depth of k=4, each round of speculation generates between 1 and 5 tokens, with between 0 and 4 tokens accepted before the first rejection.

Results

Our draft model SPIRe produces the largest increase in the modeled throughput of generating tokens for most batch sizes and context lengths.

Figure 4: Speedup of generating tokens in a batch of size B given a context of length L.

Ablations

See Figure 5 for ablations of the techniques used to train our draft model SPIRe. The techniques are in order of importance: 1) sparse KV cache, 2) attending to target model activations, 3) distillation loss, 4) initializing by pruning the target model, 5) feedback memory, and 6) MixedLoss.

Extrapolation

What if we extrapolate the results above to drafting for Llama 3 70B while maintaining the ratio of draft to target model parameters? Our draft model SPIRe continues to produce the largest increase in the modeled throughput.

Figure 6: Estimated speedup of generating tokens from Llama 3 70B in a batch of size B given a context of length L.

Cost Analysis

Each draft model costs a different amount in training FLOPs. For SPIRe, the training cost is around 14 times lower than for the target model, which is significant when considering deployment costs.

Conclusion

Minimizing the cost per token requires maximizing throughput. We demonstrated that draft model architecture significantly impacts throughput, highlighting the need for careful design based on expected batch sizes and context lengths in production.