Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Fast Inference via Hierarchical Speculative Decoding

Clara Mohri, Haim Kaplan, Tal Schuster, Yishay Mansour, Amir Globerson

TL;DR

The paper tackles slow autoregressive inference in transformers by introducing Hierarchical Speculative Decoding (HSD), which stacks multiple drafters so only the smallest model generates tokens while larger models verify and draft for the next level. It derives a per-token latency expression $\mathbb{E}[\text{latency per token}] = \sum_{i=0}^K c_i \prod_{j=i}^K R(\alpha_{j-1,j}, j)$ using a function $R(\cdot)$ that encapsulates drafting rounds and acceptance rates, and shows this optimization can be solved in polynomial time via a reduction to the Generalized Shortest Path problem. The authors prove correctness, provide a reduction scheme, and establish a tight complexity bound $O(\overline{T}^4 K^4 \log(\overline{T}K))$ for finding the optimal hierarchy. Empirically, on open-source LLMs and benchmarks CNN-DM and XSUM, HSD achieves up to 1.17× speedups over the best single-draft setup and outperforms autoregressive decoding, validating the practicality of hierarchical speculative decoding for faster, high-quality generation.

Abstract

Transformer language models generate text autoregressively, making inference latency proportional to the number of tokens generated. Speculative decoding reduces this latency without sacrificing output quality, by leveraging a small draft model to propose tokens that the larger target model verifies in parallel. In practice, however, there may exist a set of potential draft models- ranging from faster but less inaccurate, to slower yet more reliable. We introduce Hierarchical Speculative Decoding (HSD), an algorithm that stacks these draft models into a hierarchy, where each model proposes tokens, and the next larger model verifies them in a single forward pass, until finally the target model verifies tokens. We derive an expression for the expected latency of any such hierarchy and show that selecting the latency-optimal hierarchy can be done in polynomial time. Empirically, HSD gives up to 1.2x speed-up over the best single-draft baseline, demonstrating the practicality of our algorithm in reducing generation latency beyond previous techniques.

Fast Inference via Hierarchical Speculative Decoding

TL;DR

The paper tackles slow autoregressive inference in transformers by introducing Hierarchical Speculative Decoding (HSD), which stacks multiple drafters so only the smallest model generates tokens while larger models verify and draft for the next level. It derives a per-token latency expression using a function that encapsulates drafting rounds and acceptance rates, and shows this optimization can be solved in polynomial time via a reduction to the Generalized Shortest Path problem. The authors prove correctness, provide a reduction scheme, and establish a tight complexity bound for finding the optimal hierarchy. Empirically, on open-source LLMs and benchmarks CNN-DM and XSUM, HSD achieves up to 1.17× speedups over the best single-draft setup and outperforms autoregressive decoding, validating the practicality of hierarchical speculative decoding for faster, high-quality generation.

Abstract

Transformer language models generate text autoregressively, making inference latency proportional to the number of tokens generated. Speculative decoding reduces this latency without sacrificing output quality, by leveraging a small draft model to propose tokens that the larger target model verifies in parallel. In practice, however, there may exist a set of potential draft models- ranging from faster but less inaccurate, to slower yet more reliable. We introduce Hierarchical Speculative Decoding (HSD), an algorithm that stacks these draft models into a hierarchy, where each model proposes tokens, and the next larger model verifies them in a single forward pass, until finally the target model verifies tokens. We derive an expression for the expected latency of any such hierarchy and show that selecting the latency-optimal hierarchy can be done in polynomial time. Empirically, HSD gives up to 1.2x speed-up over the best single-draft baseline, demonstrating the practicality of our algorithm in reducing generation latency beyond previous techniques.
Paper Structure (28 sections, 6 theorems, 13 equations, 1 figure, 9 tables, 1 algorithm)

This paper contains 28 sections, 6 theorems, 13 equations, 1 figure, 9 tables, 1 algorithm.

Table of Contents

  1. Speculative decoding.
  2. Hierarchical models.
  3. Background
  4. Hierarchical Speculative Decoding
  5. Motivation
  6. Preliminaries
  7. Main algorithm
  8. Latency analysis
  9. Motivating example
  10. Efficient Optimization of Hierarchies
  11. Latency of a subset of models
  12. Preliminaries
  13. Reduction from HSD to GSP
  14. Computational complexity
  15. Empirical validation
  16. ...and 13 more sections

Key Result

Theorem 2.1

For any set of models $\mathcal{M}_0, \ldots, \mathcal{M}_K$, where $\mathcal{M}_K$ is the target model and any parameters, the output distribution of Algorithm alg:seq_generation follows that of target model $\mathcal{M}_K$.

Figures (1)

  • Figure 1: An example stack trace of HSD for $T_0 = 3, T_1 = 6, T_2 = 12$. The color of a token represents which model generated that token. A token can be generated either auto-regressively by the base model $\mathcal{M}_0$, or by the verification rule which can either replace a token or generates an additional token when all draft tokens are accepted. A token is considered to be part of the context for a certain model if a model above it accepted this token.

Theorems & Definitions (14)

  • Theorem 2.1: Correctness of HSD
  • Theorem 2.2
  • proof
  • Corollary 3.1
  • Definition 3.2: HSD problem
  • Definition 3.3: Generalized Shortest Path (GSP) Problem
  • Definition 3.4
  • Definition 3.5
  • Lemma 3.6: oldham2001combinatorial
  • Theorem 3.7
  • ...and 4 more