OPT-Tree: Speculative Decoding with Adaptive Draft Tree Structure

Large language models (LLMs) (Black et al., 2022; Touvron et al., 2023; Jiang et al., 2024; Zheng et al., 2024) have achieved remarkable performance in various NLP scenarios. As models increase in size and complexity, the computational demands for inference rise significantly. Consequently, accelerating decoding is becoming increasingly important to save computing resources and reduce response time.

Autoregressive models (Black et al., 2022; Zhang et al., 2022; Touvron et al., 2023) usually generate one token in one decoding step, leading to limited decoding efficiency. In recent work, speculative decoding (Stern et al., 2018; He et al., 2023; Yang et al., 2023; Fu et al., 2024; Cai et al., 2024; Li et al., 2024) has shown great potential for lossless accelerated decoding. It applies a “draft and then verify” mechanism to maintain the original output distribution of the target model to be accelerated. Drafting is performed by a less-overhead drafting model. The generated draft is verified in parallel by the target model to generate multiple tokens in one decoding step, bringing promising acceleration.

Existing work like EAGLE (Li et al., 2024) has proposed methods for training small but effective draft models. Previous work mainly adopts drafts with structures of sequences or fixed trees. However, we argue that neither of them is the optimal draft structure under a limited node budget. Sequence-structured drafts (Stern et al., 2018; Leviathan et al., 2023; Xia et al., 2023; Yang et al., 2023; Zhang et al., 2023; Fu et al., 2024) contain redundant nodes. For example, “A-B-C-D-E” and “A-B-C-F-G” have the same prefix “A-B-C”, which is calculated twice during verification. Therefore, there are only 7 valid tokens among the 10 nodes of these two sequences. Drafts with tree structure (He et al., 2023; Cai et al., 2024; Li et al., 2024; Jeon et al., 2024; Chen et al., 2024) solved this problem. The same token can appear only once in the same tree layer. A corresponding tree attention mask is designed for parallel verification. The specific structure of the tree is usually heuristic and remains constant. However, given a node budget, the best structure that maximizes the acceptance length during verification would change according to different inputs in each decoding step.

This paper proposes an adaptive and scalable tree structure called OPT-Tree. It can be applied to any autoregressive draft model. As is shown in Figure 1, the tree structure adaptively changes in each decoding step to maximize the mathematical expectation of the acceptance length. We apply a greedy algorithm to construct an OPT-Tree in each step. Details are elaborated in Section 3. We conduct comprehensive experiments in Section 4 to evaluate the effectiveness of OPT-Tree. Experimental results demonstrate that OPT-Tree outperforms the baselines and can be up to 3.2 times faster than vanilla autoregressive decoding. The mathematical expectation of the acceptance length is generally positively correlated with the actual acceptance length in practice. Moreover, OPT-Tree performs well when the tree size scales up. Using LLaMA-2-7B as the draft model, LLaMA-2-70B can generate 10 tokens in a single decoding step with OPT-Tree when the number of nodes is over 500, which indicates its great potential for adapting to more powerful computation resources and more effective draft models in the future.

In this section, we provide the necessary definitions to establish a clear and precise foundation for the concepts discussed in this paper.

After inputting x = (x₁, x₂,…, x_l), where l is the current sequence length, the target model M and the drafting model M_d return the next word distribution p(y^{l +1}|x₁, x₂,…, x_l) and $pd(ŷl+1∣x1,x2,…,xl)$⁠, respectively, where y^{l +1} and $ŷl+1$ are the sampled next words.

In speculative decoding with tree-structured draft, M_d first infers d steps to generate a draft tree T of depth d and then M verifies the draft. The verification depends on the sampling method. For greedy sampling, the ground truth is the sequence of tokens with the highest probability for each position output by M.

As is shown in Figure 2, the model can process the tree structure input in parallel by constructing the corresponding tree attention mask. For all branches in the tree that contain the root node, the longest branch with the same prefix as the ground truth is accepted. Therefore, multiple tokens can be generated in one decoding step while ensuring that the generated sequences are consistent with the original ones. Due to the parallel computing mechanism, with given computing resources, the time cost of verification can be considered constant when the length of the verification sequence is within a certain range. As is shown in Figure 3, on a 4090 GPU with a model containing 7B parameters, the time required to verify a sequence of length 10 is similar to that for a sequence of length 200. Consequently, when discussing the construction of a draft tree in Section 3, we consider the node budget of the draft tree as a given condition.

This section introduces OPT-Tree, an algorithm for constructing our defined optimal draft tree structure for any input sequence in speculative decoding with autoregressive draft models.

The process of solving T_opt is to find a subtree T of n nodes with the largest E(A) from a complete n-ary tree. If there are no other conditions, the time complexity of solving this tree is typically Ω(n²), which is unacceptable. However, by leveraging certain properties of the draft tree, we can design a more efficient algorithm to solve this problem.

Figure 4 shows a simple example of calculating $p^$ and E(A). E(A) should positively correlate with the acceptance length. We discuss their correlation in Section 4.2.

We use E_sub(T, n) to represent the maximum value of E(A) for all subtrees of T that contain the root node and have n nodes. Note that the root node is not considered when calculating node trees and mathematical expectations.

Then, we propose Algorithm 1 to construct T_opt during the drafting phase for each decoding step. We initialize T with a root node. At each drafting step, we greedily sample n tokens with the largest $p^$ in the next token distributions of nodes in the last layer of T to construct the next layer. T has d * n nodes at this time. Finally, we select the n nodes in T with the largest $p^$⁠. It is easy to prove that these n nodes are a subtree of T, which contains the root node:

We summarize this as Theorem 3.1.

According to Theorem 3.2, we can get the desired T_opt in theory by stopping drafting when E(T) no longer increases. However, the draft model brings additional overhead to the practice. For autoregressive draft models, the drafting overhead is proportional to the depth of the draft tree.

Taking this into consideration, we introduce a threshold δ when setting the conditions for terminating drafting. The value of δ should be controlled between μ and 1, where μ is the time taken for one drafting step divided by the time taken for one decoding step.

A complete decoding step of M is detailed in Algorithm 2. In practice, both M and M_d utilize key and value cache (Pope et al., 2023) to calculate attention. This approach ensures that the actual input length of each drafting step is n, effectively preventing computational bottlenecks during the inference process of the draft model, even when operating under larger tree size budgets.

Applying our proposed algorithm to solve T_opt incurs an acceptable time cost. Detailed time costs for each operation are presented in Section 4.4.

We adopt LLaMA-2-7B, LLaMA-2- 13B, LLaMA-2-70B (Touvron et al., 2023), and Vicuna-33B (Zheng et al., 2024) as target models to verify the effectiveness of OPT-Tree. We use a single GeForce RTX 4090 GPU for LLaMA-2- 7B, a single L20 GPU for LLaMA-2-13B and 4 A100-PCIE-40GB GPUs for LLaMA-2-70B and Vicuna-33B. We choose one or two smaller models in the same version as the draft model for each target model. Moreover, we adopt a corresponding EAGLE draft model for each target model. EAGLE (Li et al., 2024) is an effective speculation decoding method that trains additional autoregressive heads as draft models. It uses a well-designed heuristic draft tree structure with 25 nodes. In our experiments, we regard it as the EAGLE draft tree. EAGLE is certified by Xia et al. (2024) as the fastest speculative method in their experiments. For each target and draft model group, we perform speculative decoding with greedy sampling and compare OPT-Tree with the Binary tree and EAGLE tree. The temperature is set to zero.

We compare the average acceptance length and number of tokens generated per second decoding with different tree structures. The speedup ratio is calculated according to generation speed. The number of nodes needs to be controlled within a certain range to avoid excessive time consumption in the verification phase. It is treated as a hyperparameter chosen in [25,50,60] to maximize the speedup ratio according to different target models and GPU resources except for the EAGLE tree. We conduct evaluation on MT-Bench (Zheng et al., 2024) and GSM8K (Cobbe et al., 2021).

Experimental results are shown in Table 1 and Table 2. Note that using LLaMA-2-1B as the draft model can hardly speed up decoding when the target model is LLaMA-2-7B because the difference in inference time between the two models is too small. EAGLE draft models achieve strong performance with fewer parameters, thus providing better acceleration than the small models in the same series with the target models. OPT-Tree outperforms other tree structures in terms of mean acceptance length in each group of experiments, especially when the performance of the draft model is close to the target model (e.g., LLaMA-2-70B combined with L-7B and Vicuna-33B combined with Vicuna-7B), indicating its high upper limit. Since OPT-Trees are usually deeper than binary trees and EAGLE trees, they incur more overhead when drafting. Therefore, from the perspective of tokens per second, the improvement is not as significant as that from the mean acceptance length. Tokens per second are affected by different hardware resources and random errors. In addition, some method-independent techniques can also be used to reduce computation time. For example, the unchanged part of the attention mask in the drafting phase can be initialized only once and called multiple times, thus saving the time of multiple initializations. In order to make a fairer comparison in our experiments, we avoid these tricks to be consistent with EAGLE’s practice. Overall, OPT-Tree outperforms the baselines. It can be up to about 3.2 times faster than vanilla autoregressive decoding. The similar performance on both datasets verifies the robustness of the proposed method.

The theory of OPT-Tree is based on the premise that E(A) is positively correlated with actual A. We record the values of E(A) and A of OPT-Tree in about 8000 decoding steps for 4 groups of M and M_d. Figure 5 shows the results. The value of E(A) is rounded. The darker areas in the four images are basically distributed along the main diagonal line. When E(A) of the tree is larger, it also tends to get a more considerable acceptance length after verification. A stronger draft model shifts the distribution to the lower right corner. These phenomena corroborate our theoretical analysis. In addition, in the LLaMA-2-70B+LLaMA-2-7B group, high values of E(A) and A (e.g., E(A) = 14, A = 15) are generally found, which demonstrates the potential of OPT-Tree to adapt to stronger draft models and larger draft tree sizes.

We conduct experiments to explore the changes in mean acceptance length with larger tree sizes. We compare OPT-Tree with Sequoia (Chen et al., 2024) using LLaMA-2-7B and LLaMA-2-70B as target models. Sequoia is a scalable draft tree that uses dynamic programming to solve for the tree structure. It requires the target and draft models to be used in advance to infer some samples to determine the best structure. The tree structure is fixed when doing speculative decoding. We use 200 samples in C4 (Raffel et al., 2020) to construct the Sequoia trees. Temperature is set to 0 in the experiments.

The results are shown in Figure 6. OPT-Tree outperforms Sequoia under various tree sizes. For LLaMA-2-7B+LLaMA-2-68M, the mean acceptance length increases when the number of nodes is smaller than 130 for both OPT-Tree and Sequoia. When the number of nodes exceeds 140, the mean acceptance length increases slowly. For LLaMA-2-70B+LLaMA-2-7B, the growth of mean acceptance length with Sequoia stabilizes when the number of nodes exceeds 150. In contrast, OPT-Tree can consistently improve the mean acceptance length even with more than 500 nodes. Since LLaMA-2-7B is a strong draft model for LLaMA-2-70B, the mean acceptance length can achieve 10 with an OPT-Tree of 500 nodes. A tree with 500 nodes costs a large amount of computation time for LLaMA-2-70B with A100-PCIE-40GB GPUs, thus being unable to speed up decoding in our practice. However, this cost may be acceptable if more powerful computational resources are equipped in the future.

In this section, we discuss the overhead of applying OPT-Tree. We conduct experiments with 4 groups of models on A100 GPUs. The number of nodes is set to 50. The temperature is set to 0.

Figure 7 displays the time cost of the operations in speculative decoding with OPT-Tree. Each pie in the figure represents the average wall time of one speculative decoding step, composed of tree-related operations time, drafting time, and verification time. Since the cost of tree-related operations is independent of the size of the draft and target models, these costs will constitute a smaller proportion when the models are larger. In the model setting with the least number of parameters, the time overhead of tree-related operations accounts for 6.3%. However, given the improvements in draft quality achieved by OPT-Tree, these costs are justified. When the model becomes sufficiently large (e.g., LLaMA-2-70B+ LLaMA-2-7B), the time overhead associated with tree operations is negligible.

Considering that the overhead of the draft model is proportional to the depth of the tree, the tree that maximizes the acceptance length does not necessarily have the highest speed-up ratio. Therefore, we experiment to study the mean acceptance length and tokens/s under different thresholds.

Figure 8 shows the experimental results on LLaMA-2-7B. The mean acceptance length drops as the threshold grows when using LLaMA-2-68M as the draft model. However, there is a slight fluctuation for the EAGLE draft model. This is because E(A) and A are not completely equivalent. We calculate μ for each group of models, which is the time of one drafting step divided by the time of one decoding step. A threshold that is too large will reduce the tree’s depth, thus reducing the value of A. On the other hand, a threshold that is too small may make the tree too deep and increase the cost of drafting. When the depth of the tree increases by one but the increment of the E(A) does not exceed μ, it is not worth increasing the depth. Therefore, we set a threshold between μ and 1 in practice. LLaMA-2-68M and EAGLE achieve the highest acceleration when δ = 0.2 and δ = 0.8, respectively.

In the decoding setting of non-greedy sampling (random sampling), we only modify the acceptable tokens during the verification phase. We evaluate OPT-Tree on these non-greedy settings, where the temperature exceeds 0.

We perform speculative decoding with OPT-Tree on the MT-Bench dataset for all groups of models in 4.1 with the temperature set to 1. Table 3 displays the experimental results. The mean acceptance length and the speedup ratio of speculative decoding with OPT-Tree are slightly lower when the temperature is set to 1 than when the temperature is set to 0. Since the draft tree greedily samples tokens with higher probability, the positive correlation between E(A) and A will be weakened in the decoding of random sampling. Therefore, it is typical for the acceleration of speculative decoding to drop when the temperature is greater than 0. The improvement of OPT-Tree for EAGLE draft model is more significant in this setting. Figure 9 shows specific changes in mean acceptance length and tokens/s with different temperature values. Both metrics drop as the temperature rises in general. But even when the temperature is set to 1, OPT-Tree can still provide high speedup compared to vanilla autoregressive decoding.

We show an example of speculative decoding with an OPT-Tree of 50 nodes on LLaMA-2-70B with LLaMA-2-7B as the draft model in Figure 10. The threshold is 0.7, and the temperature is 0. The mean acceptance length is 9.34, and the generation speed is 12.07 tokens per second. Most words (blue text) are generated by the draft model and then verified by the target model. Each couple of red words and the continuous blue text in front of it is generated in a single decoding step of the target model. The appearance of red words is either because the depth of the draft tree is limited or because none of the candidates for this position hits the target. Prepositions (e.g., in, for, and with), conjunctions (e.g., and and or), articles (e.g., a and the), punctuation, and other words which have no apparent practical meanings in the drafts are prone to be rejected during the verification phase. In addition, the beginning of new sentences in drafts tends to be rejected because it has no solid sequential association with the previous word.

Speculative decoding (Stern et al., 2018; Xia et al., 2023; Leviathan et al., 2023; Chen et al., 2023a) accelerates autoregressive decoding by drafting and then verifying while ensuring consistent output. Drafting methods are mainly divided into independent drafting and self-drafting. Independent drafting leverages an external low-cost model. SpecDec (Xia et al., 2023) trains a non-autoregressive model for drafting while others (Leviathan et al., 2023; Chen et al., 2023a; Spector and Re, 2023; Chen et al., 2023b, 2024) directly utilize a smaller version of the target model. In addition, REST (He et al., 2023) proposed a retrieval-based drafting method. Self-drafting uses the original information of the target model to draft. Yang et al. (2023) adopt an early-exiting mechanism for drafting. Similarly, Zhang et al. (2023) performs adaptive layer skipping in the drafting phase. Lookahead Decoding (Fu et al., 2024) designed an algorithm for parallel drafting and verification. MEDUSA (Cai et al., 2024) trains multiple decoding heads to obtain candidates for multiple steps from original features in parallel. Considering that different sampling results at each step in drafting will affect the distribution of subsequent output, EAGLE (Li et al., 2024) designed an autoregressive head, which introduced the embedding of each word in the drafting stage.

The verification method has evolved from sequence-structured verification to tree-structured verification. Early work (Stern et al., 2018; Leviathan et al., 2023; Xia et al., 2023; Yang et al., 2023; Zhang et al., 2023; Fu et al., 2024) verifies drafts in the form of one or several sequences. However, as the number of verification tokens increases, there are a large number of prefix duplications between sequences, resulting in redundant calculations. To alleviate this problem, recent work (He et al., 2023; Cai et al., 2024; Li et al., 2024; Jeon et al., 2024) uses tree-structured heuristic drafts and designs the corresponding attention masks for parallel verification. Chen et al. (2024) proposed Sequoia, an algorithm for building draft trees, which performs well as the tree size scales up.

In this paper, we propose a novel and effective method called OPT-Tree to construct adaptive draft tree structures for speculative decoding, which is applicable to any autoregressive draft model. OPT-Tree maximizes the mathematical expectation of the acceptance length under any limited draft tree size. Experimental results across ten sets of target and draft models using two datasets demonstrate that OPT-Tree outperforms existing draft structures, achieving lossless acceleration of up to 3.2 times compared to vanilla autoregressive decoding. Furthermore, when paired with a robust draft model, OPT-Tree consistently increases the mean acceptance length even with over 500 nodes, showcasing its potential in scenarios with ample computational resources.

We want to thank all the anonymous reviewers and Action Editor Ivan Titov for their valuable comments. This work was supported by the National Science Foundation of China (NSFC No. 62206194 and 62276077), the Natural Science Foundation of Jiangsu Province, China (Grant No. BK20220488), Young Elite Scientists Sponsorship Program by CAST (2023QNRC001), and Tecorigin.

We counted the frequency of token acceptance under different draft model output probabilities when using LLaMA-2-7B and the EAGLE draft model for speculative decoding. Figure 11 shows the experimental results. The acceptance frequency is positively correlated with the output probability.