CLJun 25, 2024

OPT-Tree: Speculative Decoding with Adaptive Draft Tree Structure

arXiv:2406.17276v454 citationsHas Code
Originality Highly original
AI Analysis

This addresses the pressing problem of slow inference in large language models, offering a lossless acceleration method that is incremental over existing speculative decoding techniques.

The paper tackles the inference efficiency bottleneck in autoregressive language models by proposing OPT-Tree, an algorithm that constructs adaptive draft trees for speculative decoding, achieving a speed-up ratio of up to 3.2 compared to autoregressive decoding.

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.

Code Implementations1 repo
Foundations

The foundational work for this paper's niche, ranked by how specifically the neighbourhood builds on it — not by global fame.

Your Notes