Source-linked AI summary
OPT-Tree: Speculative Decoding with Adaptive Draft Tree Structure
Jikai Wang, Yi Su, Juntao Li, Qingrong Xia, Zi Ye, Xinyu Duan, Zhefeng Wang, Min Zhang
TL;DR
Autoregressive decoding is inefficient because it generates one token per step, while existing speculative-decoding trees are fixed and may not maximize acceptance across inputs. OPT-Tree adaptively constructs a scalable tree that maximizes expected acceptance length, achieving up to 3.2× acceleration and more than ten tokens in one step under favorable conditions.
Problem
Existing speculative-decoding methods mainly use fixed heuristic draft structures that do not adapt to different inputs to maximize acceptance length under a limited node budget.
Method
OPT-Tree greedily constructs an adaptive draft tree for each decoding step to maximize the mathematical expectation of acceptance length under a limited tree size.
Results
OPT-Tree outperforms existing draft structures and achieves up to 3.2 times faster lossless decoding than vanilla autoregressive decoding.
Takeaways & Limitations
With a sufficiently powerful draft model and more than 500 nodes, OPT-Tree can generate 10 tokens in a single decoding step and continues scaling with tree size.
Takeaways & Limitations
Reported throughput depends on hardware and environment, and the paper does not explore other decoding optimizations that could further improve speed.
Abstract
from arXiv · showhide
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 fail to adapt to different situations to maximize the acceptance length during verification. To alleviate this dilemma, we proposed OPT-Tree, an algorithm to construct adaptive and scalable draft trees. 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.
1 Introduction
Speculative decoding accelerates autoregressive inference by drafting and verifying multiple tokens, but existing sequence and fixed-tree structures do not adapt their shape to each input. OPT-Tree adaptively changes its draft tree to maximize expected acceptance length and achieves up to 3.2× faster decoding than vanilla autoregressive decoding.
- Autoregressive models typically generate one token per decoding step, making inference efficiency increasingly important as model size and complexity grow.
- Speculative decoding drafts tokens with a lower-overhead model and verifies them in parallel with the target model while preserving its output distribution.
- Sequence drafts redundantly calculate shared prefixes, while fixed heuristic trees do not adapt their structure to different inputs under a limited node budget.
- OPT-Tree adaptively changes its tree structure each decoding step to maximize the mathematical expectation of acceptance length.
- Up to 3.2 times faster decoding than vanilla autoregressive decoding demonstrates OPT-Tree’s reported acceleration.
2 Preliminaries
Speculative decoding uses a draft model to construct a tree and a target model to verify its branches in parallel. Under greedy sampling, the longest branch matching the target model’s highest-probability sequence is accepted, allowing multiple tokens per step while preserving the original sequence.
- The target and draft models return next-token distributions conditioned on the current input sequence.
- A draft model generates a tree for several steps, and the target model verifies all tree tokens in parallel using a tree attention mask.
- Under greedy sampling, the longest branch sharing the target model’s highest-probability prefix is accepted.
- This verification procedure can generate multiple tokens in one decoding step while keeping the generated sequence consistent with the original model.
3 OPT-Tree
OPT-Tree constructs an adaptive draft tree for speculative decoding by maximizing expected acceptance length under a node budget. It uses greedy expansion, subtree selection, and a threshold to balance acceptance gains against drafting overhead.
- OPT-Tree represents a draft tree with node and edge sets, layer-wise sampled-token counts, and node scores computed from draft probabilities over parent paths.The node score simplifies subsequent tree-construction operations.
- OPT-Tree defines an optimal draft tree that changes with the input and maximizes the expected acceptance length for a given node budget.The expectation is approximated from draft-model probabilities, whose node scores correlate positively with acceptance likelihood for effective draft models.
- The tree is built greedily by expanding frontier nodes with the largest p-hat values, then retaining the n highest-scoring nodes as a root-containing subtree.The monotonicity theorem supports selecting a subtree through the parent-score ordering.
- Each decoding step drafts Topt, applies a corresponding tree attention mask in the target model, and accepts the longest candidate matching the target output.This verification procedure allows multiple tokens to be generated in one decoding step.
- Drafting stops when additional expansion yields insufficient expected-acceptance improvement, using threshold δ to account for draft-model overhead.For autoregressive draft models, overhead grows with tree depth, and δ is controlled between μ and 1, where μ is the relative cost of one drafting step.
- Table 1 reports MT-Bench results using Mean Acceptance Length, with vanilla autoregressive decoding represented by no draft model.The table distinguishes LLaMA-2 and Vicuna draft models through the Md column.
4 Experiments
Experiments evaluate OPT-Tree across model groups, tree sizes, thresholds, and sampling temperatures, comparing acceptance length and generation speed with established draft structures. OPT-Tree generally improves acceptance length and can achieve up to 3.2× the speed of vanilla autoregressive decoding, while deeper trees and higher temperatures introduce practical trade-offs.
- 4.1 Main Results: OPT-Tree outperforms other tree structures in mean acceptance length across experiment groups and reaches about 3.2× the speed of vanilla autoregressive decoding.Its advantage is especially pronounced when the draft model is close in capability to the target model, while deeper trees reduce the corresponding tokens-per-second gain.
- 4.2 Correlation between E(A) and A: OPT-Tree’s expected acceptance length generally tracks actual acceptance length, with high values such as E(A) = 14 and A = 15 observed for LLaMA-2-70B with LLaMA-2-7B.The distribution shifts toward higher values with a stronger draft model, supporting adaptation to stronger drafts and larger trees.
- 4.3 Scaling the Draft Tree Size: OPT-Tree continues improving mean acceptance length beyond 500 nodes for LLaMA-2-70B with LLaMA-2-7B, reaching 10, whereas Sequoia becomes flat beyond 150 nodes.The 500-node tree incurs substantial computation on the tested A100 GPUs and cannot speed up decoding in that practice setting.
- 4.4 Impact of the Threshold: Threshold choice trades off acceptance length against drafting cost: the highest acceleration occurs at δ = 0.2 with LLaMA-2-68M and δ = 0.8 with EAGLE.A threshold that is too large reduces tree depth and acceptance length, while one that is too small increases drafting cost; practical thresholds lie between µ and 1.
- 4.5 Performance on Non-greedy Settings: At temperature 1, OPT-Tree’s mean acceptance length and speedup are slightly lower than at temperature 0, but it still provides high speedup over vanilla autoregressive decoding.Both metrics generally decrease as temperature rises because random sampling weakens the positive correlation between E(A) and A.
5 Related Work
Speculative decoding accelerates autoregressive generation by drafting and verifying tokens while preserving the target model’s output distribution. Research has progressed from sequence-based verification with redundant prefixes toward heuristic tree structures that reduce duplicated computation.
- Speculative decoding uses a draft-and-verify mechanism to generate multiple tokens per decoding step while maintaining the target model’s original output distribution.Drafting may use an external low-cost model or a smaller version of the target model, and verification is performed in parallel.
- Early speculative-decoding methods verify one or several sequences, but increasing verification length creates substantial prefix duplication and redundant computation.
- Recent methods use heuristic tree-structured drafts and corresponding attention mechanisms to reduce redundancy during verification.
6 Conclusion
OPT-Tree constructs adaptive draft trees that maximize expected acceptance length under a limited tree size. Across ten target–draft model groups and two datasets, it outperforms existing structures, reaches 3.2× lossless acceleration, and continues scaling with strong draft models.
- 3.2× maximum lossless acceleration over vanilla autoregressive decoding is achieved by OPT-Tree across experiments with ten target–draft model groups and two datasets.The method also shows robustness across different datasets and temperatures.
- OPT-Tree constructs adaptive draft trees that maximize the mathematical expectation of acceptance length under any limited draft-tree size.
Limitations
Reported throughput depends on hardware resources and execution environments, and the experiments use the EAGLE decoding framework for fair comparison.
- Throughput speeds reported by the experiments vary with hardware resources and execution environments.
- The experiments use the same decoding framework as EAGLE for fair comparison, while alternative decoding optimizations are not explored.
- Further speed improvements from optimizing the decoding algorithm from other perspectives remain outside this paper’s scope.