Source-linked AI summary
EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty
Yuhui Li, Fangyun Wei, Chao Zhang, Hongyang Zhang
TL;DR
Autoregressive decoding is slow and costly, while existing speculative sampling faces draft-model and feature-prediction challenges. EAGLE drafts at the second-to-top-layer feature level, incorporates tokens from one time step ahead, and theoretically preserves output distributions while accelerating generation, including 2.7x-3.5x latency speedup for LLaMA2-Chat 70B.
Problem
Autoregressive decoding is slow and costly, while speculative sampling can require unsuitable or high-overhead draft models.
Method
EAGLE autoregressively drafts second-to-top-layer features and incorporates tokens from one time step ahead to mitigate feature-prediction uncertainty.
Results
2.7x-3.5x latency speedup and doubled throughput were achieved for LLaMA2-Chat 70B while preserving the generated text’s distribution.
Takeaways & Limitations
EAGLE provides an efficient speculative sampling framework with theoretically guaranteed output-distribution preservation in greedy and non-greedy settings.
Takeaways & Limitations
EAGLE’s distribution-preservation guarantee makes evaluating the quality of its generated results unnecessary and meaningless under the paper’s evaluation rationale.
Abstract
from arXiv · showhide
Autoregressive decoding makes the inference of Large Language Models (LLMs) time-consuming. In this paper, we reconsider speculative sampling and derive two key observations. Firstly, autoregression at the feature (second-to-top-layer) level is more straightforward than at the token level. Secondly, the inherent uncertainty in feature (second-to-top-layer) level autoregression constrains its performance. Based on these insights, we introduce EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency), a simple yet highly efficient speculative sampling framework. By incorporating a token sequence advanced by one time step, EAGLE effectively resolves the uncertainty, enabling precise second-to-top-layer feature prediction with minimal overhead. We conducted comprehensive evaluations of EAGLE, including all models from the Vicuna and LLaMA2-Chat series, the MoE model Mixtral 8x7B Instruct, and tasks in dialogue, code generation, mathematical reasoning, and instruction following. For LLaMA2-Chat 70B, EAGLE achieved a latency speedup ratio of 2.7x-3.5x, doubled throughput, while maintaining the distribution of the generated text.
1. Introduction
EAGLE addresses the latency and draft-model overhead of speculative sampling by predicting structured second-to-top-layer features and using a one-step-advanced token sequence to reduce feature uncertainty. Evaluations across multiple models and tasks report substantial acceleration while preserving the target distribution.
- Autoregressive decoding is slow because LLMs generate tokens sequentially, motivating draft-and-verify speculative sampling.
- Finding a low-overhead draft model is difficult, especially for the smallest models, because existing alternatives may be incompatible or reduce acceleration gains.
- EAGLE introduces speculative drafting at the more structured second-to-top-layer feature level and incorporates a token sequence advanced by one time step.
- Feature-level autoregression achieves a 1.9x speedup ratio compared with 1.5x for token-level autoregression in the reported comparison.
- EAGLE was evaluated on dialogue, code generation, mathematical reasoning, and instruction-following tasks across Vicuna, LLaMA2-Chat, and Mixtral models.
- EAGLE provides generality through zero-shot application and a lightweight single-transformer-decoder-layer plug-in, without fine-tuning the original LLM.
2. Preliminaries
Speculative sampling separates inexpensive drafting from parallel verification and accepts or resamples drafted tokens to match the target LLM's distribution. EAGLE adopts this procedure for both greedy and non-greedy decoding.
- Speculative sampling drafts γ tokens and their distributions with a smaller model, then verifies them in one target-LLM forward pass.
- A drafted token is accepted with probability min(1, p_j+i(ˆt_j+i)/ˆp_j+i(ˆt_j+i)); after rejection, later drafted tokens are discarded and the token is resampled.
- The procedure is theoretically equivalent to direct sampling from the target LLM, so EAGLE preserves the generated-text distribution in greedy and non-greedy settings.
3. EAGLE
EAGLE combines feature-level drafting with an advanced token sequence, then verifies the resulting tree using the target LLM. Its training jointly optimizes feature regression and token prediction while reusing target-model embedding and LM-head parameters.
- Drafting and verification: EAGLE uses both drafting and verification phases, with its main distinction from other methods occurring in the drafting phase.The target LLM verifies the tree-structured draft in a single forward pass using tree attention.
- Drafting and verification: EAGLE predicts the next feature from prior features and a token sequence advanced by one time step, then samples the next token through the target LLM’s LM Head.The predicted feature and sampled token are subsequently concatenated for continued drafting.
- Draft model: The draft model consists of an Embedding layer, LM Head, and Autoregression Head, with the first two reusing frozen target-LLM parameters.The feature and advanced-token inputs are embedded and fused before processing by the Autoregression Head.
- Training of the draft model: Feature prediction is trained with Smooth L1 regression loss, while classification loss directly optimizes the draft model’s token-prediction objective.The combined objective is L = Lreg + wclsLcls, with wcls set to 0.1 because classification loss is typically an order of magnitude larger.
- Training and verification: EAGLE uses a fixed dataset rather than target-LLM-generated text and applies tree attention to compute token probabilities for the draft tree in one target-model pass.Accepted tokens and their features are recorded for subsequent drafting steps.
4. Experiments
EAGLE was evaluated across multiple LLMs and tasks using latency, acceptance, and throughput metrics. It accelerates decoding while preserving the target LLM’s output distribution, with performance affected by temperature, tree attention, model architecture, and batch size.
- Evaluation Setup: EAGLE was evaluated on Vicuna, LLaMA2-Chat, and Mixtral models across dialogue, code, mathematics, and instruction-following tasks.
- Evaluation Setup: EAGLE preserves the target LLM’s output distribution, making generated-text quality evaluation unnecessary under its guarantee.
- Effectiveness: At temperature=0, LLaMA2-Chat 13B achieved 3.01x-3.76x speedup, compared with 2.66x-2.89x at temperature=1.
- Effectiveness: EAGLE generated 3.2-4.5 tokens per target-LLM forward pass, versus one token for vanilla decoding.
- Effectiveness: 1.5x speedup was achieved on Mixtral 8x7B Instruct, with lower acceleration attributed to shorter acceptance length and MoE verification complexity.
- Ablations: Tree attention increased average acceptance length by approximately 0.6-0.8 and speedup ratio by about 0.3-0.5 over chain drafting.
- Ablations: The feature&shifted-token input markedly outperformed alternatives without added complexity by accounting for sampling randomness through a one-step token shift.
- Training and Systems: EAGLE’s fixed training dataset marginally trails target-generated data while reducing training-data generation overhead, and it requires slightly more CUDA memory than vanilla decoding.
5. Related Work
Related work accelerates language-model inference through model compression, architectural changes, and speculative sampling. EAGLE belongs to speculative sampling methods that preserve the original output distribution without relaxation, while differing mainly in draft generation.
- General Acceleration: Language-model acceleration research includes distillation, quantization, pruning, and alternative network architectures.
- Speculative Sampling: Speculative sampling divides decoding into draft and verification stages and extends acceleration to non-greedy sampling while preserving the original output distribution.
- Speculative Sampling: DistillSpec, BiLD, and Medusa trade exact distribution preservation for relaxed or threshold-based acceptance decisions, unlike EAGLE.
- Drafting Methods: Speculative methods differ chiefly in drafting: they use smaller models, skipped target layers, parallel small models, or other specialized draft mechanisms.
6. Conclusion
EAGLE accelerates speculative sampling by autoregressing structured second-to-top-layer features and using one-step-ahead tokens to mitigate sampling uncertainty. It preserves the target distribution and reports substantial speedups over vanilla decoding and competing methods.
- EAGLE drafts autoregressively at the second-to-top-layer feature level and incorporates one-step-ahead tokens to mitigate sampling uncertainty.
- On MT-bench, EAGLE was 2.1x-3.8x faster than vanilla autoregressive decoding, 1.7x-2.1x faster than Lookahead, and 1.5x-1.6x faster than Medusa.
- EAGLE preserves the output distribution of the target LLM while enhancing generation speed.
A.1. Tree Structure
EAGLE uses tree attention to organize draft generation, with higher-probability branches intended to become deeper and wider. The fixed structure is intuitive rather than rigorously optimized and may not suit every context.
- Tree Structure: Tree attention produces a tree-structured draft, whereas disabling it produces a chain-structured draft.
- Tree Structure: In greedy decoding, child nodes are the top k tokens by probability; in non-greedy decoding, k tokens are sampled.
- Tree Structure: The chosen tree structure was based on intuition rather than rigorous optimization, with higher-probability branches intended to be deeper and wider.
- Limitations: The optimal tree may depend on context, and larger batch sizes may favor smaller trees because redundant computational resources decrease.
A.2. Multi-Round Speculative Sampling
Multi-round speculative sampling adapts speculative sampling to tree-structured drafts by recursively applying its sampling algorithm across candidate tokens.
- Multi-round speculative sampling recursively invokes the standard sampling algorithm for tree-structured drafts with k candidate tokens.Accepted tokens are returned; otherwise, sampling proceeds from an adjusted distribution.
B. Detailed experimental results
The experiments report EAGLE’s speedup ratio, average acceptance length, and acceptance rate across code generation, mathematical reasoning, and instruction-following datasets at temperature 0.
- Table 8 reports EAGLE’s speedup ratio, average acceptance length τ, and acceptance rate α on HumanEval, GSM8K, and Alpaca.The evaluations use temperature = 0.
- Algorithm 1 specifies the multi-round speculative sampling procedure.It takes the target distribution, candidate samples, and their distributions as input and returns a sample from the target distribution.