Multi-Draft Speculative Sampling: Canonical Decomposition and Theoretical Limits

TL;DR

This paper tackles accelerating LLM decoding by leveraging multiple draft models through multi-draft speculative sampling. It proves a canonical two-step decomposition: an importance-weighted selection of an intermediate token from the drafts, followed by a single-draft speculative sampling step to generate the final token, achieving the optimal acceptance probability P^⋆(acc). For the case of two identical drafts, it derives a necessary and sufficient condition for unit acceptance probability and provides an explicit closed-form expression, along with a weighted importance sampling framework and fast approximations. Empirical results on OPT across Dolly, XSum, and WMT show consistent gains in block efficiency and token rates over baselines, highlighting practical impact for memory-bound LLM inference.

Abstract

We consider multi-draft speculative sampling, where the proposal sequences are sampled independently from different draft models. At each step, a token-level draft selection scheme takes a list of valid tokens as input and produces an output token whose distribution matches that of the target model. Previous works have demonstrated that the optimal scheme (which maximizes the probability of accepting one of the input tokens) can be cast as a solution to a linear program. In this work we show that the optimal scheme can be decomposed into a two-step solution: in the first step an importance sampling (IS) type scheme is used to select one intermediate token; in the second step (single-draft) speculative sampling is applied to generate the output token. For the case of two identical draft models we further 1) establish a necessary and sufficient condition on the distributions of the target and draft models for the acceptance probability to equal one and 2) provide an explicit expression for the optimal acceptance probability. Our theoretical analysis also motives a new class of token-level selection schemes based on weighted importance sampling. Our experimental results demonstrate consistent improvements in the achievable block efficiency and token rates over baseline schemes in a number of scenarios.

Figures (9)

• Figure 1: Optimal Approach for Multi-Draft Speculative Sampling
• Figure 2: Numerical evaluation of $Pr(\mathrm{accept})$ for the optimal scheme (Theorem \ref{['thm:opt']}) as well as two baseline schemes -- SpecTr sun2024spectr and SpecInfer miao2024specinfer. For sake of illustration we select alphabet $\Omega = \{1,2,3\}$ and ${\mathbf{p}} = [1/3, 1/3, 1/3]$. The left plot sets ${\mathbf{q}} = [1/3, q_2, 2/3-q_2]$ while the right plot sets ${\mathbf{q}} = [1/6, q_2, 5/6-q_2]$ where $q_2$ is varied on the x-axis.
• Figure 3: Performance comparison of different multi-draft schemes. The temperature of the first draft models is set to $1.2$, while we vary the temperature of the other draft.
• Figure 4: Performance comparison of different schemes on the Dolly task, while we vary the temperature of the target model and keeping the temperature of the two drafts to $1.0$.
• Figure 5: Truncated Histogram for OPT Draft and Target Models for the effective alphabet size after top-p sampling with p=0.95 on XSum dataset

Theorems & Definitions (15)

• Definition 1: Importance Weighted Sampling
• Theorem 1: Optimal Acceptance Probability and Canonical Decomposition
• Remark 1
• Theorem 2
• Example 1
• Theorem 3
• Remark 2
• Remark 3
• Theorem 4
• Example 2