Source-linked AI summary
Towards Optimal Multi-draft Speculative Decoding
Zhengmian Hu, Tong Zheng, Vignesh Viswanathan, Ziyi Chen, Ryan A. Rossi, Yihan Wu, Dinesh Manocha, Heng Huang
TL;DR
Autoregressive decoding is costly, and MDSD’s optimal acceptance rate is difficult to compute because its optimal-transport formulation grows rapidly. This paper uses the dual to derive a subset-selection approach and compute the bound for special draft distributions. On real text, sampling without replacement outperforms sampling with replacement, while existing verification algorithms remain below the bound.
Problem
The optimal-transport problem for MDSD grows exponentially with the number of draft tokens, leaving the theoretical acceptance rate and the gap of existing verification algorithms difficult to compute at practical vocabulary sizes.
Source: §1 · p. 2Method
The paper uses the dual optimal-transport problem and total unimodularity to obtain a subset-selection formulation, with efficient computation methods for special draft distributions.
Source: §1 · p. 2Results
On real text, sampling without replacement outperforms sampling with replacement, and existing verification algorithms remain significantly below the theoretical upper bound.
Source: §1 · p. 2Takeaways & Limitations
A greedy high-probability draft construction can in some cases achieve a higher acceptance rate than sampling without replacement, and its corresponding verifier can reach the theoretical upper bound.
Source: §1 · p. 2The paper considers only the single-step case of speculative decoding.
Source: §2.1 · p. 3
Abstract
from arXiv · showhide
Large Language Models (LLMs) have become an indispensable part of natural language processing tasks. However, autoregressive sampling has become an efficiency bottleneck. Multi-Draft Speculative Decoding (MDSD) is a recent approach where, when generating each token, a small draft model generates multiple drafts, and the target LLM verifies them in parallel, ensuring that the final output conforms to the target model distribution. The two main design choices in MDSD are the draft sampling method and the verification algorithm. For a fixed draft sampling method, the optimal acceptance rate is a solution to an optimal transport problem, but the complexity of this problem makes it difficult to solve for the optimal acceptance rate and measure the gap between existing verification algorithms and the theoretical upper bound. This paper discusses the dual of the optimal transport problem, providing a way to efficiently compute the optimal acceptance rate. For the first time, we measure the theoretical upper bound of MDSD efficiency for vocabulary sizes in the thousands and quantify the gap between existing verification algorithms and this bound. We also compare different draft sampling methods based on their optimal acceptance rates. Our results show that the draft sampling method strongly influences the optimal acceptance rate, with sampling without replacement outperforming sampling with replacement. Additionally, existing verification algorithms do not reach the theoretical upper bound for both without replacement and with replacement sampling. Our findings suggest that carefully designed draft sampling methods can potentially improve the optimal acceptance rate and enable the development of verification algorithms that closely match the theoretical upper bound.
1 INTRODUCTION
MDSD uses multiple draft tokens and parallel target-model verification, but the optimal acceptance rate is difficult to compute at practical vocabulary sizes. The paper develops efficient ways to compute it and compares existing methods and draft-sampling choices against the resulting upper bound.
1 INTRODUCTION
- The optimal verification algorithm for a fixed draft-sampling method corresponds to an optimal-transport solution, but its exponentially growing variables and constraints make practical computation difficult.For vocabulary sizes in the thousands, the optimal rate had not been computed, and existing algorithms’ gaps from the bound had not been quantified on real text.Source: §1 · p. 2
- The paper transforms the optimal-transport problem into a subset-selection problem using the dual and total unimodularity.This provides a different perspective for understanding MDSD efficiency.Source: §1 · p. 2
- For certain draft distributions, convexity-like structure enables efficient subset-selection methods and practical computation of the theoretical acceptance-rate upper bound.The special cases include sampling with and without replacement.Source: §1 · p. 2
- On real text, sampling without replacement outperforms sampling with replacement, while existing verification algorithms remain significantly below the theoretical upper bound.The paper measures the upper bound of MDSD efficiency and quantifies the gap for K-SEQ and RRS.Source: §1 · p. 2
- A greedy method that selects high-probability drafts and randomizes only the last draft can sometimes exceed the optimal acceptance rate of sampling without replacement.The paper also proposes a corresponding verification algorithm that reaches the theoretical acceptance-rate upper bound.Source: §1 · p. 2
2 PRELIMINARIES
The preliminaries describe single-step MDSD, formulate speculative decoding as matching draft and target outcomes, and introduce the multi-draft transport problem. They also explain why the multi-draft formulation is computationally difficult.
2.1 SPECULATIVE DECODING FOR ACCELERATING LLM INFERENCE
- In single-step MDSD, a draft model samples multiple tokens, the target model evaluates them in parallel, and a verification step determines the next token; this paper considers only the single-step case.If a draft is accepted, another target-model sample can generate a second token in the same step.Source: §2.1 · p. 3
2.2 SPECULATIVE DECODING WITH A SINGLE DRAFT TOKEN
- With one draft token, verification seeks a joint distribution with the target and draft distributions as marginals that maximizes the probability that their sampled tokens match.This is formulated as an optimal-transport problem.Source: §2.2 · p. 3
2.3 MULTI-DRAFT SPECULATIVE DECODING
- With multiple drafts, the objective is to maximize the chance that the target output matches any draft, but the transport formulation has exponentially many variables and constraints.Common draft constructions are sampling with replacement and sampling without replacement; RRS and K-SEQ are cited as existing approximations.Source: §2.3 · p. 4
3 OPTIMAL ACCEPTANCE RATE AS SUBSET SELECTION PROBLEM
This section expresses the optimal acceptance rate as a subset-selection problem. It derives that formulation through a dual problem, uses total unimodularity to obtain integer solutions, and then simplifies the integer formulation.
3 OPTIMAL ACCEPTANCE RATE AS SUBSET SELECTION PROBLEM
- The optimal acceptance rate can be expressed as a subset-selection problem.This reframes the problem of computing the rate as choosing a subset.Source: §3 · p. 4
3.1 DUAL PROBLEM
- An equivalent formulation turns the transportation problem into a b-matching problem, whose dual is a w-vertex cover problem.The paper derives this equivalent formulation from the linear program.Source: §3.1 · p. 5
3.2 TOTAL UNIMODULARITY
- The dual constraint matrix is totally unimodular, so integer right-hand sides guarantee an integer optimal solution.The paper establishes total unimodularity from a bipartite-graph incidence matrix and identity-matrix constraints.Source: §3.2 · p. 5
3.3 SUBSET SELECTION FORMULATION
- In the optimal solution, yi and zbar_i do not exceed 1 and therefore take values 0 or 1; substituting this solution yields the subset-selection formulation.Source: §3.3 · pp. 5–6
4 COMPUTING OPTIMAL ACCEPTANCE RATE IN SPECIAL CASES
The section expresses the optimal acceptance rate as a set-function minimization and uses q-convex structure to compute it efficiently in special cases. It develops a sorted-prefix algorithm for sampling with and without replacement.
4 COMPUTING OPTIMAL ACCEPTANCE RATE IN SPECIAL CASES
- The optimal acceptance rate is expressed as 1 + min_H f(H), with the section defining marginal value as the change in f when an element is added to a set.The marginal value is f(x|H) = f(H ∪ {x}) − f(H).Source: §4 · p. 6
4.1 q-CONVEX FUNCTIONS
- For both sampling methods, the rate minimization is an unconstrained submodular minimization problem with polynomial-time algorithms; exploiting q-convexity permits faster computation.The section establishes q-convexity for both sampling with and without replacement, and that every q-convex function is supermodular.Source: §4.1 · p. 6
4.2 EFFICIENT COMPUTATION
- Theorem 6 assumes Q is q-convex and monotone increasing, with p(x)>0 for all x∈Σ; when p(x)=0, there exists an optimal set H* for (8) containing x.Source: §4.2 · p. 6
4.2.1 ALGORITHM
- Sorting by q/p in non-increasing order and scanning prefixes computes the rate in O(|Σ| log |Σ|)+O(|Σ|) with replacement, or O(|Σ| log |Σ|)+O(n|Σ|) without replacement.For without-replacement sampling, the coefficient of a generating function is computed by dynamic programming.Source: §4.2.1 · pp. 6–7
5 A GREEDY APPROACH FOR SELECTING DRAFT TOKENS
The section introduces a greedy draft distribution that selects the highest-probability tokens first and randomizes only the last token. It pairs this construction with a verifier that reaches its optimal acceptance rate, and relates the two-draft case to SpecHub.
5 A GREEDY APPROACH FOR SELECTING DRAFT TOKENS
- The section proposes a greedy draft construction and a corresponding verification algorithm that achieves the distribution’s optimal acceptance rate.The subset-selection formulation is used to prove the verification result.Source: §5 · p. 7, §5.2 · p. 7
5.1 DRAFT CONSTRUCTION
- The greedy construction sets the first n−1 drafts to the top-probability tokens and samples only the last draft without replacement.The last token is sampled from q while excluding the already selected top n−1 tokens.Source: §5.1 · p. 7
5.2 VERIFICATION ALGORITHM
- A corresponding verification algorithm strictly achieves the optimal acceptance rate for the greedy draft distribution.The transport problem is simple because only one draft token is random.Source: §5.2 · p. 7
5.3 CONNECTION TO SPECHUB
- The greedy method for n=2 and SpecHub both ensure at least one draft is the highest-probability token, although their draft distributions differ.The difference yields a simpler verification algorithm for the greedy method.Source: §5.3 · p. 8
6 EXPERIMENTS
The experiments compare MDSD acceptance rates and theoretical bounds across real-text tasks, then examine how temperature, draft count, and greedy sampling affect those rates. Existing verifiers remain below the bound, while sampling strategy and experimental conditions shape the comparisons.
6.1 MAIN EXPERIMENT
- Existing RRS and K-SEQ verifiers remain below the theoretical acceptance-rate bound, while sampling without replacement beats sampling with replacement and greedy has the highest rate in the main experiment.The comparison covers multiple models and tasks; the authors attribute the replacement gap to duplicate drafts, and note that greedy is not always best at temperature 1.Source: Table 1 · p. 8, §6.1 · p. 9
6 EXPERIMENTS
- Experiments use Alpaca, WMT’14 De-En, and CNN-DailyMail distributions from generated responses, evaluating four LLM families with paired draft and target models.For each task, responses came from 1,024 samples and were capped at 128 tokens.Source: §6 · pp. 8–9
6.2 ABLATION STUDY I: IMPACT OF TEMPERATURE
- Temperature effects are non-monotonic and differ by method; as temperature rises, the gap between optimal and existing verifier rates grows, while the with-versus-without-replacement gap narrows.At T = 0, replacement methods effectively have one draft, whereas methods preventing duplicates select the top n draft tokens.Source: Figure 1 · p. 9, §6.2 · p. 10
6.3 ABLATION STUDY II: IMPACT OF NUMBER OF DRAFTS
- Increasing the number of drafts improves acceptance across methods, but sampling without replacement generally benefits more because replacement can produce redundant drafts.The gap between the optimal rate and existing RRS and K-SEQ rates also grows as draft count rises.Source: §6.3 · pp. 10–11
6.4 EVALUATING THE GREEDY SAMPLING METHOD ON GENERATION TASKS
- For two drafts, Greedy and SpecHub show no statistically significant acceptance-rate difference at any temperature; Greedy beats RRS without replacement at T=0.1, but its gain diminishes as temperature rises.The generation-task evaluation uses the EAGLE framework and MT-Bench; the temperature pattern agrees with the ablation.Source: §6.4 · p. 11
7 CONCLUSION
The paper presents efficient ways to compute optimal MDSD acceptance rates and reports measurements under real-text distributions. It also discusses verification methods and draft-model design trade-offs relevant to those results.
7 CONCLUSION
- The paper connects optimal acceptance rates to subset selection, provides efficient computation methods for common draft distributions, measures real-text bounds and verifier gaps, and proposes greedy draft construction that sometimes beats sampling without replacement.The conclusion frames these as theoretical and practical contributions, respectively.Source: §7 · p. 11
A.1 RECURSIVE REJECTION SAMPLING (RRS)
- Appendix A.1 defines the residual distribution Resp−q and gives recursive RRS formulations for drafts sampled with and without replacement.Source: §A.1 · p. 15
A.2 K-SEQ
- K-SEQ’s acceptance rate is theoretically guaranteed to achieve a (1 − e^-1)-approximation of the optimal rate.Source: §A.2 · p. 15
B.1 DRAFT MODEL DESIGN
- Draft-model design trades similarity to the target model, which can raise acceptance, against computational overhead; some approaches reuse target-model computations.Examples include look-ahead tokens, added prediction heads, reuse of the last-layer hidden state, and partial key-value caches.Source: §B.1 · p. 16
B.2 MULTI-DRAFT SPECULATIVE DECODING
The discussion distinguishes single-step MDSD from multi-step tree structures and situates the paper’s general theory alongside prior work on a special case.
B.2 MULTI-DRAFT SPECULATIVE DECODING
- This paper focuses on single-step MDSD, whereas multiple steps with multiple drafts form a tree structure for which Sequoia searches for an optimal topology by dynamic programming.Source: §B.2 · p. 16
- As trees deepen, some branch acceptance probabilities fall; Cascade Speculative Drafting assigns larger draft models to shallower, more likely-to-be-accepted branches.Source: §B.2 · p. 16
- Prior work’s optimal-rate result for two replacement-sampled drafts matches the paper’s result in that special case, while the paper’s theory covers arbitrary draft-sampling methods and draft counts.Source: §B.2 · p. 16
B.3 MULTI-STEP SPECULATIVE DECODING
Speculative decoding can be improved through multi-step or multi-draft methods. This work focuses on multi-draft decoding, while combining both directions is suggested as future work.
B.3 MULTI-STEP SPECULATIVE DECODING
- Multi-step methods draft a sequence, whereas this paper focuses on generating multiple draft tokens at each step.The two approaches are distinct directions for improving speculative decoding.Source: §B.3 · p. 16
- Multi-draft methods reduce to vanilla speculative decoding with one draft, while multi-step methods do so with one step.The cited multi-draft and multi-step works are Sun et al. (2024e); Khisti et al. (2024) and Sun et al. (2024d); Hu & Huang (2024); Sun et al. (2024c), respectively.Source: §B.3 · pp. 16–17
- Combining multi-draft and multi-step improvements is proposed as a direction for future research.Source: §B.3 · p. 17
C PROOFS
The proofs establish equivalence between two optimization formulations and derive acceptance-rate results for sampling and greedy drafting. They also use generating functions and convexity in the without-replacement analysis.
C PROOFS
- The two optimization formulations have equal optimal objective values, establishing their equivalence.The proof constructs feasible solutions in both directions while preserving or bounding objective values.Source: C PROOFS · p. 17
- For with-replacement sampling, Q(H) is derived from the probability that all n sampled tokens lie in H, using g(x) = x^n in the q-convexity proof.Source: C PROOFS · p. 17
- For without-replacement sampling, generating functions encode the probability that all n samples lie in H, with coefficients computed by recurrence.The proof uses non-negativity of the generating-function coefficients.Source: C PROOFS · p. 18
- The proof of Theorem 5 invokes q-convexity of Q and symmetry to establish its claim.Source: C PROOFS · p. 18
- The greedy method’s acceptance rate is derived, and its optimal set includes Top_n−1(q) together with tokens meeting a residual-probability condition.The stated set is H* = {i | q¬Top_n−1(q)(i) ≥ p(i)} ∪ Top_n−1(q).Source: C PROOFS · p. 19
C.1 DERIVATION OF THE DUAL PROBLEM
This merged block derives a dual problem, illustrates decoding procedures, and lists notation used throughout the paper.
C.1 DERIVATION OF THE DUAL PROBLEM
- The dual-problem derivation starts from the primal, introduces dual variables, forms and rearranges the dual function, and imposes boundedness before stating the dual.Source: §C.1 · p. 20
D ADDITIONAL ILLUSTRATION
- The additional illustration includes single-step draft-token generation and verification, plus pseudocode for multi-step speculative sampling with arbitrary tree topology.Source: D ADDITIONAL ILLUSTRATION · p. 21
E SUMMARY OF NOTATIONS
- The notation defines Ptarget(·|x1, x2, ..., xm) as the context-conditioned next-word model, Pdraft(·|x1, x2, ..., xm) as candidate-token generator, αK-SEQ as K-SEQ acceptance rate, P(H) and Q(H) as probability sums, and S, y, z as LP variables.Source: E SUMMARY OF NOTATIONS · pp. 23–24
F ADDITIONAL EXPERIMENTS
The additional experiments report average generation lengths for MDSD methods on MT-Bench using the Eagle framework. At one setting, three methods have close but distinct reported values.
F ADDITIONAL EXPERIMENTS
- Table 3 reports average generation length τ and ratio Δ on MT-Bench using the Eagle framework.Source: Table 3 · p. 24
- 3.51781, 3.51944, and 3.51975 are the average generation lengths for RRS without replacement, SpecHub, and Greedy, respectively, at 2 drafts, 4 steps, and T = 0.6.They appear equal after truncation to two decimal places, but the reported values differ.Source: F ADDITIONAL EXPERIMENTS · p. 24, Table 3 · p. 24