Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

OPT-Tree tackles decoding inefficiency in autoregressive LLMs by learning adaptive draft-tree structures that maximize the expected acceptance length $E(A)$ in speculative decoding. It defines $T_{opt}$ via a greedy layer-wise construction controlled by a node budget and a termination threshold $\\delta$, leveraging per-node values $\\hat{p}$ and the subtree objective $E_{sub}(T,n)$. Empirically, OPT-Tree achieves up to $3.2\\times$ speedup over baseline methods across diverse models and datasets, with stronger drafts yielding larger $E(A)$ and robust performance across temperatures. The approach scales with larger draft trees and stronger draft models, indicating strong practical potential for accelerating inference on large language models.

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 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.

OPT-Tree: Speculative Decoding with Adaptive Draft Tree Structure

TL;DR

OPT-Tree tackles decoding inefficiency in autoregressive LLMs by learning adaptive draft-tree structures that maximize the expected acceptance length in speculative decoding. It defines via a greedy layer-wise construction controlled by a node budget and a termination threshold , leveraging per-node values and the subtree objective . Empirically, OPT-Tree achieves up to speedup over baseline methods across diverse models and datasets, with stronger drafts yielding larger and robust performance across temperatures. The approach scales with larger draft trees and stronger draft models, indicating strong practical potential for accelerating inference on large language models.

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 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.
Paper Structure (12 sections, 2 theorems, 3 equations, 8 figures, 3 tables, 2 algorithms)

This paper contains 12 sections, 2 theorems, 3 equations, 8 figures, 3 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. OPT-Tree
  4. Experiments
  5. Main Results
  6. Correlation between $E(A)$ and $A$
  7. Scaling the Draft Tree Size
  8. Impact of the Threshold
  9. Performance on Non-greedy Settings
  10. Case Study
  11. Related Work
  12. Conclusion

Key Result

Theorem 3.1

For any two nodes $v_{i}$ and $v_{j}$ in the tree, if $v_{i}$ is a node in the subtree of $v_{j}$, then $\hat{p}$ of $v_{i}$ is less than $\hat{p}$ of $v_{j}$.

Figures (8)

  • Figure 1: Draft structures used in speculative decoding. Nodes in the same layer share the same position index. OPT-Tree various in each decoding step to achieve a larger acceptance length.
  • Figure 2: An example of a draft tree containing $\hat{p}$ in each node. The value of $E(A)$ is 2.07.
  • Figure 3: The relationship between input length and the wall time for inference for models of different sizes on various GPUs.
  • Figure 4: Correlation between $E(A)$ and $A$. The horizontal axis represents $E(A)$, and the vertical axis represents $A$. Each square shows the number of times the corresponding situation occurs. The darker the color, the more times it indicates.
  • Figure 5: Mean acceptance length under different tree sizes under two sets of experiments.
  • ...and 3 more figures

Theorems & Definitions (3)

  • Theorem 3.1
  • proof
  • Theorem 3.2