OPT -Tree: Speculative Decoding with Adaptive Draft Tree Structure (2025)

Abstract: Autoregressive language models demonstrate excellent performance in various scenarios. However, the inference efficiency is limited by its one-step-one-word generation mode, which has become a pressing problem recently as the models become increasingly larger. Speculative decoding employs a “draft and then verify” mechanism to allow multiple tokens to be generated in one step, realizing lossless acceleration. Existing methods mainly adopt fixed heuristic draft structures, which do not adapt to different situations to maximize the acceptance length during verification. To alleviate this dilemma, we propose OPT-Tree, an algorithm to construct adaptive and scalable draft trees, which can be applied to any autoregressive draft model. It searches the optimal tree structure that maximizes the mathematical expectation of the acceptance length in each decoding step. Experimental results reveal that OPT-Tree outperforms the existing draft structures and achieves a speed-up ratio of up to 3.2 compared with autoregressive decoding. If the draft model is powerful enough and the node budget is sufficient, it can generate more than ten tokens in a single step. Our code is available at https://github.com/Jikai0Wang/OPT-Tree.

## introduction large language models (llms) @xcite @xcite 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 @xcite @xcite usually generate one token in one decoding step, leading to limited decoding efficiency. in recent work, speculative decoding @xcite @xcite @xcite 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 lessoverhead 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 @xcite 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 @xcite @xcite @xcite 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 @xcite @xcite @xcite 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. ## preliminaries in this section, we provide the necessary definitions to establish a clear and precise foundation for the concepts discussed in this paper. inference. after inputting @xmath0, x 2 , . . . , 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 1 , x 2 , . . . , x l ) and p d (ŷ l+1 |x 1 , x 2 , . . . , x l ), respectively, where y l+1 and ŷl+1 are the sampled next words. speculative decoding. 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. ## opt-tree 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. considering a certain step in speculative decoding whose input is x, the draft model m d generates a draft tree based on x and the given tree structure t . draft tree t is defined as follows: given m , m d and n, for input x, an optimal tree structure t opt should maximize the mathematical expectation of the acceptance length. note that t opt changes as the input changes. since the optimization goal of the draft model is to make its output distribution close to the target model distribution, it is intuitive that the token with a larger output probability in draft model is more likely to be accepted. see appendix a for verification of this property. based on this, we can use p d as the probability that it will be accepted to approximately calculate the mathematical expectation of the acceptance length. we use e(a) to denote the approximation of the expected average acceptance length, which can be calculated by: 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 2 ), which is unacceptable. however, by leveraging certain properties of the draft tree, we can design a more efficient algorithm to solve this problem. for simplicity, we define the probability of the prefix as p: 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. algorithm 1 construct an opt-tree t opt input: 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: proof. (1) if these nodes can not form a tree with the root, there is at least one node v i whose parent node v j is not among these nodes. (2) since v j is the parent node of v i , p of v j is larger than p of algorithm 2 speculative decoding with adaptive draft tree structure input: input sequence @xmath1, x 2 , . . . , x l ), target model m , draft model m d , number of nodes n, threshold δ. output: new input sequence @xmath2, x 2 , . . ., we summarize this as theorem 3.1. theorem 3.1. given the depth of the tree, the top n nodes of the complete n-ary draft tree with the largest p form the t opt . theorem 3.2. as the drafting step increases, 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 @xcite 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. ## related work speculative decoding @xcite @xcite accelerates autoregressive decoding by drafting and then verifying while ensuring consisthe verification method has evolved from sequence-structured verification to tree-structured verification. early work @xcite @xcite @xcite 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 @xcite @xcite uses tree-structured heuristic drafts and designs the corresponding attention masks for parallel verification. @xcite proposed sequoia, an algorithm for building draft trees, which performs well as the tree size scales up. ## conclusion 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.