Source-linked AI summary
EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees
Yuhui Li, Fangyun Wei, Chao Zhang, Hongyang Zhang
TL;DR
LLM inference is expensive, while static draft trees assume token acceptance depends only on position. EAGLE-2 uses draft confidence to build context-dependent dynamic trees, achieving the highest speedups across the reported evaluations while preserving the original output distribution. Its confidence estimate remains a local approximation rather than the actual acceptance probability.
Problem
Static draft trees implicitly assume draft-token acceptance depends only on position, despite acceptance varying with context.
Method
EAGLE-2 uses EAGLE draft-model confidence scores to approximate acceptance rates and dynamically adjust the draft tree without additional training.
Results
EAGLE-2 achieved the highest speedup ratios in experiments across three LLM series and six tasks while preserving the original generated-text distribution.
Takeaways & Limitations
EAGLE-2 provides efficient, lossless speculative sampling without fine-tuning the original LLM or relaxing acceptance conditions.
Takeaways & Limitations
Draft confidence is a local approximation and cannot reflect the actual probability of a draft token being accepted.
Abstract
from arXiv · showhide
Inference with modern Large Language Models (LLMs) is expensive and time-consuming, and speculative sampling has proven to be an effective solution. Most speculative sampling methods such as EAGLE use a static draft tree, implicitly assuming that the acceptance rate of draft tokens depends only on their position. Interestingly, we found that the acceptance rate of draft tokens is also context-dependent. In this paper, building upon EAGLE, we propose EAGLE-2, which introduces a new technique of context-aware dynamic draft tree into drafting modeling. This improvement leverages the fact that the draft model of EAGLE is well-calibrated: the confidence scores from the draft model approximate acceptance rates with small errors. We conducted extensive evaluations on three series of LLMs and six tasks, with EAGLE-2 achieving speedup ratios 3.05x-4.26x, which is 20%-40% faster than EAGLE-1. EAGLE-2 also ensures that the distribution of the generated text remains unchanged, making it a lossless acceleration algorithm.
1. Introduction
LLM inference is costly, and speculative sampling reduces latency by drafting and verifying multiple tokens per forward pass. EAGLE-2 addresses context-dependent acceptance rates with a confidence-guided dynamic draft tree while preserving output distributions and requiring no additional training.
- Motivation: Speculative sampling reduces inference latency by generating draft tokens and verifying them in parallel.Autoregressive LLM generation repeatedly accesses all model parameters, making long outputs expensive.
- Motivation: Static draft trees assume acceptance rates depend only on tree position, but acceptance also varies across contexts.EAGLE and Medusa use fixed branching at each draft step, implicitly adopting the position-only assumption.
- EAGLE-2: EAGLE-2 uses draft-model confidence scores to approximate acceptance rates and dynamically adjust the draft tree.This avoids obtaining acceptance rates directly from the original LLM, which would require additional forward results.
- Results: 2.5x-5x speedup: EAGLE-2 achieved the best performance across six tasks and three LLM series in the reported experiments.Evaluations covered conversation, code, mathematics, instruction following, summarization, and question answering.
- Results: Approximately 2x faster than Medusa and about 2.3x faster than Lookahead on MT-bench while preserving the output distribution.The comparison concerns a turn-conversation dataset designed to resemble real-world model interactions.
- Advantages: EAGLE-2 requires no additional training and keeps the generated-text distribution exactly the same as the original LLM.It adjusts the tree from draft confidence rather than training a separate tree-prediction model, and does not relax acceptance conditions.
2. Preliminaries
Speculative sampling drafts tokens with a smaller model and verifies them with the original LLM, while EAGLE improves this process using feature-level drafting and tree-structured verification. EAGLE-2 replaces EAGLE’s fixed draft shape with a context-dependent structure.
- Speculative Sampling: Speculative sampling alternates between drafting candidate tokens with a smaller model and verifying them with the original LLM.The original model computes verification probabilities, and tokens are accepted sequentially according to the probability ratio.
- EAGLE: EAGLE performs autoregression at the feature level before using the original LLM’s head to obtain draft tokens.It also inputs a token sequence advanced by one timestep to address uncertainty in the feature sequence.
- EAGLE: EAGLE uses a tree-structured draft so alternative branches can be attempted when a draft token is rejected.This avoids discarding every subsequent token as in a chain-structured draft.
- EAGLE-2: EAGLE uses the same fixed draft-tree shape across contexts, whereas EAGLE-2 dynamically adjusts the shape.For “10+2=”, EAGLE-2 adds only one candidate because the next token is easier to predict.
3. Observations
Acceptance rates vary both by position in the draft tree and by context. This supports using a context-aware dynamic tree rather than relying only on a fixed positional structure.
- Position Dependence: Acceptance rates are highest at position P1 and lowest at position P6 in the evaluated draft tree.The figure maps positions P1-P6 in the tree to positions 1-6 on the horizontal axis.
- Position Dependence: Tokens on the upper-left side of the tree have higher acceptance rates than tokens on the lower-right side.This pattern supports placing more nodes in the upper left and fewer nodes in the lower right.
- Context Dependence: Significant variance at the same tree position shows that acceptance depends on context as well as position.The observation motivates a context-aware dynamic draft tree.
- Confidence Calibration: Draft-model confidence is strongly positively correlated with acceptance rate, ranging from approximately 0.04 below confidence 0.05 to about 0.98 above confidence 0.95.This relationship provides a low-cost way to estimate acceptance without invoking the original LLM.
4. Context-Aware Dynamic Draft Tree
EAGLE-2 dynamically adjusts and reranks a draft tree using draft-model confidence as an approximation of token acceptance. Its expansion strategy selectively grows promising branches while verification remains consistent with the existing process.
- Overview: EAGLE-2 dynamically adjusts the draft tree without altering draft-model training or the verification stage.Its improvements concern tree expansion and token reranking.
- Expansion Phase: Tree attention expands all tokens in a layer in one draft-model forward pass, but selective expansion limits the cost of exponentially growing layers.Too many simultaneous inputs can slow the draft model.
- Expansion Phase: EAGLE-2 expands the top-k current-layer nodes with the highest global acceptance values.A node’s global acceptance value is the product of acceptance rates along its root-to-node path.
- Expansion Phase: Draft-model confidence scores approximate node acceptance rates when computing global values.The confidence score c_j is used as an approximation for the acceptance rate p_j.
- Reranking Phase: After expansion, EAGLE-2 reranks all draft tokens and selects the top m by value for verification.The selected tokens are flattened into a one-dimensional sequence, with an adjusted attention mask preserving branch separation.
5. Experiments
EAGLE-2 was evaluated across multiple LLMs, tasks, and lossless speculative-sampling baselines using speedup and average acceptance length. It achieved the strongest reported acceleration, while ablations supported value-based expansion and reranking.
- Experimental Setup: EAGLE-2 was evaluated on Vicuna, LLaMA2-Chat, and LLaMA3-Instruct models across six generation tasks.The tasks included conversation, code generation, mathematical reasoning, instruction following, summarization, and question answering.
- Metrics and Comparisons: Speedup ratio measures acceleration against vanilla autoregressive decoding, whose baseline is 1.00x.Average acceptance length τ measures accepted tokens per drafting-verification cycle and is hardware-independent, but excludes draft-model overhead.
- Effectiveness: Across all tested datasets and LLMs, EAGLE-2 achieved the highest speedup ratios.The evaluation compared it with standard speculative sampling, PLD, Lookahead, Medusa, Hydra, and EAGLE.
- Effectiveness: Approximately 4-5.5 tokens were accepted per EAGLE-2 drafting-verification cycle, roughly twice the average for standard speculative sampling and Medusa.PLD and Lookahead had shorter acceptance lengths but low drafting overhead, producing speedup ratios close to those lengths.
- Effectiveness: EAGLE-2 showed lower acceptance lengths and speedup ratios on Natural Questions and CNN/Daily Mail than on other tasks, yet still outperformed standard speculative sampling.The paper attributes this pattern to draft-model training data containing less world knowledge and summarization coverage.
- Ablation Studies: Value-based expansion produced higher speedup ratios and average acceptance lengths than confidence-based expansion, while reranking improved both metrics.The value combines a token’s confidence with the confidence scores of its ancestor nodes.
6. Related Work
Related work accelerates LLM generation through model compression, speculative sampling, retrieval, dynamic drafting, and other draft-model strategies. EAGLE-2 differs from partially dynamic methods by using a context-dependent draft tree without the stated structural limitations.
- Inference Acceleration: Quantization, pruning, and knowledge distillation reduce forward-pass cost but often trade generation quality for lower computational overhead.These approaches differ from speculative sampling, which uses the original LLM for verification.
- Speculative Sampling: Speculative sampling extends draft-verification decoding to non-greedy generation and has pursued lower draft overhead and closer draft-target consistency.Related methods include tree aggregation, layer skipping, early exit, retrieval, additional alignment training, and cascaded draft models.
- Lossless Acceleration: Some speculative methods increase speedup by relaxing acceptance conditions or fine-tuning the original LLM, which can compromise losslessness.BiLD relaxes acceptance conditions, while Medusa-2, CLLMs, and SPACE fine-tune the original model.
- Dynamic Draft Trees: Partially dynamic draft-tree methods use confidence-based early stopping or add candidates without fully expanding them.BiLD and Kangaroo control depth through early stopping, while GLIDE and CAPE add candidates that are not further expanded.
- Dynamic Draft Trees: EAGLE-2 uses a context-dependent draft tree and is presented as avoiding the structural limitations described for these partially dynamic methods.The supplied passage truncates the final comparison.
7. Conclusion
EAGLE-2 is presented as an efficient, lossless speculative-sampling method that uses calibrated draft confidence to build context-dependent trees. Across the reported experiments, it achieved the highest speedup ratios without additional training or changing the original output distribution.
- Conclusion: EAGLE-2 approximates draft-token acceptance rates with EAGLE confidence scores and dynamically adjusts the draft tree to increase accepted tokens.The method is described as context-dependent and built upon EAGLE.
- Conclusion: EAGLE-2 achieved the highest speedup ratios in all reported experiments across various LLMs and datasets.The conclusion describes comparisons with several state-of-the-art speculative-sampling methods.
- Conclusion: EAGLE-2 preserves consistency with the original LLM’s generated results and requires no additional training.These properties are stated as part of the method’s lossless acceleration design.
A. Implementation Details
The implementation evaluates EAGLE-2 and comparison methods using Hugging Face and PyTorch-based setups with released or specified draft-model weights. EAGLE-2 uses fixed token-count and tree-depth settings by original-model size.
- Baselines: Vanilla decoding uses Hugging Face Transformers models with PyTorch and a pre-allocated KV cache.The same base models are used for other methods.
- Baselines: Standard speculative sampling uses Hugging Face Transformers’ assisted-generation feature.
- Baselines: PLD, Lookahead, Medusa, and Hydra use default settings and officially released weights.
- EAGLE: EAGLE uses officially released Vicuna and LLaMA2-Chat draft weights, while LLaMA3-Instruct is trained on ShareGPT.The LLaMA3-Instruct setup is consistent with Medusa and Hydra.
- EAGLE-2: For 7B/8B, 13B, and 70B original LLMs, EAGLE-2 uses 60, 50, and 48 draft tokens, respectively, with depth 6 and 10 expansion nodes.