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

Abstract: 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 of up to **5x**, which is 1.3x that of EAGLE. EAGLE-2 also ensures that the distribution of the generated text remains unchanged, making it a **lossless** acceleration algorithm.

## introduction modern large language models (llms) @xcite exhibit impressive capabilities and are widely applied across various domains. however, their parameter sizes have grown substantially, even exceeding hundreds of billions. during autoregressive generation, each token generation requires accessing all model parameters. in a single dialogue, hundreds to thousands of tokens might be generated, making llm inference slow and expensive. speculative sampling @xcite ) methods aim to address this issue by rapidly generating draft tokens v ic u n a 7 b v ic u n a 1 3 b l l a m a 2 -c h a t 7 b l l a m a 2 -c h a t 1 3 b models and then verifying them in parallel. these methods generate multiple tokens in a single forward pass, significantly reducing inference latency. standard speculative sampling @xcite uses a chain-structured draft. to improve acceptance length, recent work in speculative sampling has employed tree-structured drafts. sequoia @xcite explicitly assumes that the acceptance rate of a draft token depends only on its position in the tree. eagle @xcite and medusa @xcite use the same static draft tree structure in all contexts: at the i-th step of the draft phase, k candidates are added, with k being fixed. this implicitly assumes the aforementioned hypothesis. however, this assumption appears to contradict the insight of speculative sampling that some tokens are simpler and can be predicted by smaller models. our experiments (see section 3.1) reveal that the acceptance rate of draft tokens is not only position-dependent but also highly context-dependent. therefore, the static structure of draft trees has inherent limitations. dynamically adjusting the draft tree structure based v ic u n a 7 b v ic u n a 1 3 b ll a m a 2 -c h a t 7 b ll a m a 2 -c h a t 1 3 b ll a m a 2 -c h a t 7 0 b ll a m a 3 -i n st ru ct 8 b ll a m a 3 -i n st ru ct 7 0 b models on the acceptance rates of draft tokens in different contexts can yield better results. however, obtaining the acceptance rate of draft tokens requires the forward results from the original llm, which conflicts with the goal of speculative sampling to reduce the number of forwards for the original llm. fortunately, we find that eagle is well-calibrated: the confidence score (probability) of the draft model is a good approximation of the acceptance rate of draft tokens (see section 3.2). this makes it feasible to use a context-dependent dynamic draft tree structure. we propose eagle-2, which leverages the confidence scores from the draft model to approximate acceptance rates. based on this, it dynamically adjusts the draft tree structure, increasing the number of accepted tokens. we conducted comprehensive and extensive tests on six tasks: multi-turn conversation, code generation, mathematical reasoning, instruction following, summarization, and question answering. the datasets used were mtbench @xcite , humaneval @xcite , gsm8k @xcite , alpaca @xcite , cnn/daily mail @xcite , and natural questions @xcite . our comparisons included six advanced speculative sampling methods: standard speculative sampling @xcite @xcite , pld @xcite , medusa @xcite , lookahead @xcite , hydra @xcite , and eagle @xcite . we conducted experiments on three series of llms: vicuna, llama2-chat, and llama3-instruct. in all experiments, eagle-2 demonstrated the best performance, achieving a speedup of 2.5x-5x. figures 1 and 2 show the speedup ratios of eagle-2 and other speculative sampling methods on mt-bench. mt-bench is a multi-turn conversation dataset that closely resembles real-world scenarios for models like chatgpt and is frequently used to evaluate state-of-the-art open-source and closed-source models. on the mtbench dataset, eagle-2 is approximately 2x faster than medusa and about 2.3x faster than lookahead, while ensuring the output distribution remains unchanged. besides performance, eagle-2 offers the following advantages: ## context-aware dynamic draft tree building on the aforementioned observations, we introduce eagle-2, an acceleration algorithm for llm inference that dynamically adjusts the draft tree. eagle-2 does not alter the training and inference of the draft model, nor does it affect the verification stage. its improvements focus on two aspects: how to expand the draft tree (section 4.1) and how to rerank draft tokens (section 4.2). during the expansion phase, we input the most promising nodes from the latest layer of the draft tree into the draft model to form the next layer. during the reranking phase, we select the tokens with higher acceptance probabilities to form the input for the original llm during the verification phase. in the draft tree, a node represents a token. in the following text, we use "node" and "token" interchangeably. ## conclusion in this paper, we introduce eagle-2, an efficient and lossless speculative sampling method. we found that eagle's draft model confidence is a good approximation of the acceptance rate for draft tokens. based on this, eagle-2 employs a context-dependent draft tree structure, significantly increasing the number of accepted draft tokens and resulting in better speedup ratios. eagle-2 ensures that the generated results are consistent with the original llms and does not require additional training. we conducted extensive evaluations using various llms across multiple datasets and compared eagle-2 with several state-of-the-art speculative sampling methods. in all our experiments, eagle-2 achieved the highest speedup ratios.