S pec H ub: Provable Acceleration to Multi-Draft Speculative Decoding (2024)

Abstract: Large Language Models (LLMs) have become essential in advancing natural language processing (NLP) tasks, but their sequential token generation limits inference speed. Multi-Draft Speculative Decoding (MDSD) offers a promising solution by using a smaller draft model to generate multiple token sequences, which the target LLM verifies in parallel.However, current heuristic approaches, such as Recursive Rejection Sampling (RRS), suffer from low acceptance rates in subsequent drafts, limiting the advantages of using multiple drafts. Meanwhile, Optimal Transport with Membership Cost (OTM) can theoretically improve acceptance rates, but its computational cost is too high for real-time use.We present SpecHub, a novel, efficient sampling-verification method for MDSD that improves acceptance rates with only linear computational overhead. By simplifying the OTM problem into a compact Linear Programming model, SpecHub significantly reduces computational complexity. It further accelerates sampling by leveraging a sparse joint distribution, focusing computation on high-probability token sequences.%It integrates seamlessly into existing MDSD frameworks.In extensive experiments, Spechub consistently generates 0.05-0.27 and 0.02-0.16 more tokens per step than RRS and RRS without replacement. We attach our code at https://github.com/MasterGodzilla/Speculative_decoding_OT.

## introduction with the growing adoption of large language models (llms) in diverse applications, there is a significant demand for faster inference and lower latency in both local computing and online api services. however, the sequential generation process of autoregressive language models complicates parallel computation. this challenge is exacerbated by the memory limitations of current hardware architectures, where ram and cache communication latencies often constrain performance, resulting in underutilized computing capacity. speculative decoding @xcite accelerates llm inference while preserving the model's output distribution. by generating a sequence of draft tokens in advance using a smaller model, it leverages gpus to verify tokens simultaneously through rejection sampling. recent advancements @xcite @xcite have further enhanced this approach by introducing tree-structured multi-drafts, where each path represents a draft. these tokens are verified in parallel during a single forward pass of the llm. using a token tree increases the number of accepted tokens by providing multiple options for each token position, thus increasing the overall acceptance rate of the algorithm and generation efficiency. despite having various tree constructions, draft model designs, and hardware optimizations, existing multi-draft speculative decoding (mdsd) methods depend on recursive rejection sampling (rrs) for acceptance, which is far from optimal. while rrs greedily accepts the token from the first draft, it does not consider the subsequent drafts and misses the opportunity to dynamically adjust the current token's acceptance strategy to improve the acceptance rates of the later drafts. rrs prioritizes the first draft's tokens but fails to dynamically adjust acceptance strategies for subsequent drafts. as a result, later iterations use a residual distribution modified by previous acceptances, leading to misalignment with the original draft distribution and lower acceptance rates @xcite . @xcite shows that the acceptance rule could be optimized through an optimal transport problem with membership cost (otm), which maximizes acceptance rates by better aligning draft tokens with the accepted token. however, otm requires tremendous computation overhead and is not practically feasible. building on the new lp formulation and insights, we introduce spechub, a faster samplingverification paradigm with only linear computational overhead. instead of constructing a dense distribution of k-draft and the accepted token, spechub strategically selects drafts containing the highest probability token sampled from the draft model. the top draft token serves as a transport hub for an oversampled token 1 to transfer its excessive probability mass to an undersampled token. this sparse structure simplifies and accelerates the underlying linear programming. spechub performs particularly well on llms whose output distributions are concentrated on the top token, resulting in higher acceptance rates than rrs. it even provably outperforms otm under certain situations. the algorithm is widely applicable and can seamlessly integrate into various mdsd algorithms, enhancing their efficiency and overall decoding speed. 1 draft model probability exceeds that of the target model. we empirically test spechub by implementing it to various mdsd frameworks @xcite @xcite . we observe a 1 -5% increase in the second draft acceptance rate, which yields a consistent 0.02 -0.16 improvement in batch efficiency over current methods. more impressively, spechub uses a tree with only half the nodes of other methods to reach the same level of batch efficiency. in our ablation study, spechub brings consistent acceleration to llm decoding under different temperatures. our toy experiments further show that spechub sometimes outperforms otm in high-entropy regions. ## background and related work here, we review the sampling-verification schema of speculative decoding. we discuss the theory behind rejection sampling and explain why naively extending it to multi-draft speculative decoding (mdsd) becomes inefficient. speculative sampling language model decoding is intrinsically serial. let v denote the vocabulary, a discrete set of tokens that the language model may generate. let x 1:@xmath0, . . . , x t ) ∈ v ⊗t denote a sequence of tokens. then, the target language model produces a conditional probability p(•|x 1:t ), from which we sample the next token x t+1 ∼ p(•|x 1:t ). however, this process is slow for its serial execution. speculative decoding @xcite addresses the issue by parallelizing the decoding process with a draft and verify phase. it first uses a smaller draft model q(•|x 1:t ) to generate a draft (x t+1 , . . . , x t+d ) sequentially. the depth of the draft, d, is usually around 5. this draft allows us to compute the target distributions p(x t+τ |x 1:t+τ -1 ) in parallel for τ ≤ d. then, we iteratively accept each draft token using rejection sampling with acceptance probability min 1, p(x t+τ |x 1:t+τ -1 ) q(x t+τ |x 1:t+τ -1 ) . in this single draft setting, speculative decoding equates to sampling directly from the target distribution. after rejection, we sample from the residual distribution norm(max(0, p(•|x 1:t+τ -1 )q(•|x 1:t+τ -1 ))). with only a single draft, the expected number of tokens generated at each iteration is upper-bounded. assume the average acceptance rate for each token is α, the maximum acceleration is 1/(1α) @xcite . multi-draft speculative decoding solves this issue @xcite . instead of verifying one sequence per time, mdsd generates a tree of tokens and calculates their target probability in parallel. thus, when the first draft gets rejected, the other drafts can be picked up, and their offspring get verified in the current step. by doing so, we trade more parallel inference for more tokens generated in each step. in the rest of the paper, we ignore any temporal relationship and only focus on a single temporal step in the decoding process. in particular, given q(•|x 1:t-1 ) and p(•|x 1:t-1 ), we discuss the sampling and verification algorithm for generating the offspring drafts and accepting one. we simplify the notation and use @xmath1:t-1 ) ∈ ∆ |v|-1 to denote the target model's probability distribution and @xmath2:t-1 ) ∈ ∆ |v|-1 to denote the draft model's distribution. here, we also notate the probability simplex of joint distributions over a group of k drafts at a single temporal step as x 1:@xmath3, . . . , x k ) as: ## mathematical formulation of multi-draft speculative decoding in this section, we lay out the mathematical formulation of the sampling and verification paradigm of mdsd. we start by reviewing the optimal transport with membership cost framework by @xcite in section 3.1. we show that it can simplified and propose an equivalent lp formulation that greatly reduces computation complexity in section 3.2. lastly, we point out that changing the design of sampling can make the lp feasible for real-world calculation in section 3.3 while preserving the acceleration. we also discuss some considerations for a real-world algorithm. ## spechub building on the aforementioned insights, we introduce spechub, a faster sampling-and-verifying paradigm with only linear computational overhead. it effectively captures the transport features of otm solutions to enhance the acceptance rate and can be applied to various multi-draft speculative sampling algorithms. since using more than two drafts offers little gains in efficiency, spechub uses two drafts (i.e., @xmath4. we thoroughly discuss expanding the algorithm to more drafts in appendix c. first, we identify the token with the highest draft probability, denoted as a, and sample it alongside other tokens. we only populate the first column and the first row in the joint draft distribution q. in particular, we define the joint draft distribution q(x 1 , x 2 ) as follows: this specific design of q makes the solution to the simplified lp formulation straightforward. ∀@xmath5, respectively. this method ensures efficient allocation of acceptance probabilities. after transporting draft probabilities to target probabilities of non-top tokens, the remaining draft accepts the top token a evenly out of p(a) the remaining entries in π can be reconstructed as described in the previous section. this solution effectively allocates as much probability mass as possible to the non-hub draft tokens while ensuring that the hub token a is never undersampled. this strategy maximizes the utilization of the draft distribution and leads to a higher acceptance rate compared to traditional methods like rrs. the pseudocode implementation is in appendix b.1. analysis spechub offers several theoretical advantages. first, since all drafts contain the top token a, it is accepted with a probability of p(a) and is never undersampled (see corollary 1). additionally, let α be the first draft acceptance rate of rejection sampling, defined as α = while spechub might theoretically decrease the first draft acceptance rate for the top token a in rare cases, our empirical results, detailed in appendix b.4, show that this effect is negligible. ## experiments in this section, we empirically show that spechub improves batch efficiency in speculative multi-draft decoding. we first show that spechub gives a significantly higher acceptance rate for its better coupling properties in the second draft acceptance rate. we then illustrate how the improvement transfers to higher batch efficiency. ## related work speculative decoding speculative decoding aims to execute multiple decoding steps in parallel. early work @xcite predicts future tokens to accelerate greedy decoding. speculative sampling @xcite extends to non-greedy decoding and uses rejection sampling to recover target distribu-tion optimally. recent works focus on reducing the running time of the draft model and increasing the acceptance rate. osd @xcite and @xcite train draft models on text generated by the target model. rest @xcite constructs drafts through retrieval. lookahead decoding @xcite breaks the sequential dependency with jacobi iterations. self-speculative decoding @xcite avoids additional models and generates draft tokens by skipping intermediate layers. several works, such as medusa @xcite and eagle @xcite , reuse the feature embedding of llms' last attention layer to predict multiple future tokens in a non-causal or autoregressive manner. ## conclusion we presented spechub, a versatile and provably faster sampling-verification paradigm for multi-draft speculative decoding. using an optimal transport map between a sparse draft and target distributions, spechub increases the acceptance rate of the second draft by 1 -5%, which leads to higher batch efficiency of llm inference by up to 0.27 tokens per iteration. in addition to practical speedups, spechub also provides insight into the underlying mathematical structure in mdsd. we hope to promotes future research in this area.