PaperlayerView exact revision on arXiv

Source-linked AI summary

EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees

Yuhui Li, Fangyun Wei, Chao Zhang, Hongyang Zhang

arXiv:2406.16858v2cs.CLcs.LG

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.

    Source: §1 · pp. 1–2

  • Method

    EAGLE-2 uses EAGLE draft-model confidence scores to approximate acceptance rates and dynamically adjust the draft tree without additional training.

    Source: §1 · p. 2

  • Results

    EAGLE-2 achieved the highest speedup ratios in experiments across three LLM series and six tasks while preserving the original generated-text distribution.

    Source: §1 · p. 2, §7 · p. 8

  • Takeaways & Limitations

    EAGLE-2 provides efficient, lossless speculative sampling without fine-tuning the original LLM or relaxing acceptance conditions.

    Source: §1 · p. 2, §7 · p. 8

    Draft confidence is a local approximation and cannot reflect the actual probability of a draft token being accepted.

    Source: §4.1 · p. 5

Abstract

from arXiv · show

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.Source: §1 · p. 1
  • 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.Source: §1 · pp. 1–2

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.Source: §1 · p. 2

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.Source: §1 · p. 2
  • 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.Source: §1 · p. 2

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.Source: §1 · p. 2

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.Source: §2.1 · pp. 2–3

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.Source: §2.2 · p. 3
  • 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.Source: §2.2 · p. 3

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.Source: §2.2 · p. 3, Figure 4 · p. 3

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.Source: §3.1 · p. 3, Figure 5 · p. 4
  • 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.Source: §3.1 · pp. 3–4

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.Source: §3.1 · p. 4

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.Source: §3.2 · p. 4

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.Source: §4 · p. 4

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.Source: §4.1 · p. 4
  • 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.Source: §4.1 · pp. 4–5
  • 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.Source: §4.1 · pp. 4–5

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.Source: §4.2 · p. 5

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.Source: §5 · p. 5

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.Source: §5 · p. 6

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.Source: §5.1 · p. 6, §5 · p. 6
  • 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.Source: §5.1 · p. 6
  • 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.Source: §5.1 · p. 6

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.Source: §5.2.1 · p. 6, §5.2.2 · p. 7

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.Source: §6 · p. 7

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.Source: §6 · pp. 7–8

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.Source: §6 · p. 8

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.Source: §6 · p. 8
  • 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.Source: §6 · p. 8

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.Source: §7 · p. 8
  • 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.Source: §7 · p. 8
  • 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.Source: §7 · p. 8

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

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.Source: A. Implementation Details · p. 12

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.Source: A. Implementation Details · p. 12

Related papers

Loading 2406.16858v2…