Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Greedy Multi-Path Block Verification for Faster Decoding in Speculative Sampling

Rahul Thomas, Arka Pal

TL;DR

This work first shows that block verification is optimal even over verification algorithms that use off-path probabilities, by constructing an information-agnostic linear program (LP), and extends this LP to the setting where the draft model samples multiple candidate paths, and uses it to construct a natural class of multi-path block verification generalizations.

Abstract

The goal of $L$-step speculative decoding is to accelerate autoregressive decoding of a target model by using a cheaper draft model to generate a candidate path of $L$ tokens. Based on a verification algorithm involving target and draft model probabilities, a prefix of the candidate sequence is accepted, and an additional correction token is sampled from a residual distribution to ensure that the final output adheres to the target distribution. While standard speculative decoding uses a verification algorithm which is independent at each token on the path, a recent extension called block verification uses a joint condition involving all sampled on-path probabilities. Block verification (BV) was shown to be optimal over all verification algorithms which use only on-path probabilities, improving on standard speculative decoding. In this work, we first show that block verification is optimal even over verification algorithms that use off-path probabilities, by constructing an information-agnostic linear program (LP). Further, we can extend our LP to the setting where the draft model samples multiple candidate paths, and use it to construct a natural class of multi-path block verification generalizations. While computing the optimal algorithm in this class is not tractable, by considering a stricter class of greedy algorithms, we can formulate an efficient method called greedy multi-path block verification (GBV). Empirically, GBV can improve block efficiency by over 30% and reduce decoding walltimes by over 15% relative to BV. On Llama-3 70B, GBV can improve the end-to-end decoding throughput over SOTA multi-path verification methods by more than 15%.

Greedy Multi-Path Block Verification for Faster Decoding in Speculative Sampling

TL;DR

This work first shows that block verification is optimal even over verification algorithms that use off-path probabilities, by constructing an information-agnostic linear program (LP), and extends this LP to the setting where the draft model samples multiple candidate paths, and uses it to construct a natural class of multi-path block verification generalizations.

Abstract

The goal of -step speculative decoding is to accelerate autoregressive decoding of a target model by using a cheaper draft model to generate a candidate path of tokens. Based on a verification algorithm involving target and draft model probabilities, a prefix of the candidate sequence is accepted, and an additional correction token is sampled from a residual distribution to ensure that the final output adheres to the target distribution. While standard speculative decoding uses a verification algorithm which is independent at each token on the path, a recent extension called block verification uses a joint condition involving all sampled on-path probabilities. Block verification (BV) was shown to be optimal over all verification algorithms which use only on-path probabilities, improving on standard speculative decoding. In this work, we first show that block verification is optimal even over verification algorithms that use off-path probabilities, by constructing an information-agnostic linear program (LP). Further, we can extend our LP to the setting where the draft model samples multiple candidate paths, and use it to construct a natural class of multi-path block verification generalizations. While computing the optimal algorithm in this class is not tractable, by considering a stricter class of greedy algorithms, we can formulate an efficient method called greedy multi-path block verification (GBV). Empirically, GBV can improve block efficiency by over 30% and reduce decoding walltimes by over 15% relative to BV. On Llama-3 70B, GBV can improve the end-to-end decoding throughput over SOTA multi-path verification methods by more than 15%.
Paper Structure (31 sections, 11 theorems, 55 equations, 7 tables, 1 algorithm)

This paper contains 31 sections, 11 theorems, 55 equations, 7 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Single-Path
  4. Speculative sampling.
  5. Draft extensions.
  6. Tree verification.
  7. Block verification (BV).
  8. Multi-Path
  9. OTLP-based algorithms.
  10. Traversal verification.
  11. Single-Path Information-Agnostic LP
  12. Multi-Path Information Agnostic LP
  13. Greedy Multi-Path Block Verification
  14. Experiments
  15. Other Model Families and Datasets
  16. ...and 16 more sections

Key Result

Theorem 3.3

A single-path verification algorithm $\Phi$ is valid if and only if the node budgets $D_\Phi:\mathcal{V}^{\leq L+1} \to \mathbb{R}$ are feasible in the following LP: The objective value is the block efficiency $\mathbb{E}[\tau+1]$ for $\Phi$.

Theorems & Definitions (29)

  • Definition 3.1: Single-path draft verification algorithm
  • Definition 3.2
  • Theorem 3.3
  • Theorem 3.4
  • Definition 4.1: $K$-path draft verification algorithm
  • Definition 4.2
  • Theorem 4.3
  • Lemma 4.4
  • Lemma 4.5
  • Theorem 4.6
  • ...and 19 more