A Theoretical Perspective for Speculative Decoding Algorithm

Authors: ming  yin
minshuo  chen
kaixuan  huang
mengdi  wang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 19:20:15

Abstract: Transformer-based autoregressive sampling has been the major bottleneck for slowing down large language model inferences. One effective way to accelerate inference is Speculative Decoding, which employs a small model to sample a sequence of draft tokens and a large model to validate. Given its empirical effectiveness, the theoretical understanding of Speculative Decoding is falling behind. This paper tackles this gap by conceptualizing the decoding problem via markov chain abstraction and studying the key properties, output quality and inference acceleration, from a theoretical perspective. Our analysis covers the theoretical limits of speculative decoding, batch algorithms, and output quality-inference acceleration tradeoffs. Our results reveal the fundamental connections between different components of LLMs via total variation distances and show how they jointly affect the efficiency of decoding algorithms.

Main paper: 
 introduction: 
 The recent surge of scaling Transformer models has led to the flourishing of AI, where success has been witnessed in wide areas such as natural language [39, 1] , computer vision [12, 17] , video generations [3, 19] , and robotics [8, 31] . In the meantime, the decoding process also becomes more and more time-consuming as the model size scales up. This is mainly due to the autoregressive nature of Transformers, where each generated token also serves as the input for future generations. As a result, decoding T tokens would take T forward passes of the full model.A recent effort to tackle this challenge is speculative decoding (SD) [10, 24] , where the autoregressive sampling is performed on a small draft model and the large language model verifies tokens generated by draft model to decide whether it should be accepted/rejected. Once a token is rejected, the generation process will start from the most recently accepted token, until a full response is completed. Speculative decoding achieves 2-2.5ˆLLM inference speedup empirically, while preserving the quality of generation.Subsequently, numerous studies [26, 25, 23, 46] have expanded this methodology, enabling further inference acceleration. Intuitively, for speculative decoding, when the generation distribution of small model p and large model q are close to each other, decoding is faster (since less rejection occurs), and when the distribution overlap between p and q is small, the opposite happens. However, a precise understanding of inference accelerating given the small model p and large model q remains elusive. This motivates us to ask the following question: What is the fundamental limit for inference acceleration via speculative decoding? In addition, what is the best trade-off between inference acceleration and output quality for speculative decoding? In this paper, we answer these questions from the theoretical lens. Our contributions are summarized as follows.We formalize the decoding problem through the Markov Chain abstraction that establishes the theoretical setup. We draw the connection runtime " # rejections and use it to measure efficiency. We derive the exact formula, fully characterized by distribution p and q, for the expected rejections ErN rej s for Speculative Decoding (Theorem 1). This renders a theoretical reference for understanding the acceleration rate T {ErN rej s.Next, to understand whether Speculative Decoding can be further improved, we generalize it to a class of rejection-based algorithms 2 where probability b t and distribution P t can be customized. We prove in Theorem 2 that any unbiased algorithm cannot have fewer rejections than Speculative Decoding. This indicates its optimality among the class, and having fewer rejections needs to suffer quality loss or requires extra information.Furthermore, we consider a batch version of Speculative Decoding (Algorithm 4) that utilizes multiple draft sequences. We show our batch algorithm is unbiased. We derive the expected rejections that is fully expressed by p and q and exhibit the improvement over non-batch version in Theorem 3. We provide examples and detailed discussion to explain how our theory characterize the improvement.In section 4, we shift from unbiased algorithms and study the tradeoff between inference cost and quality degradation. We formulate this into an optimization model (1) . Theorem 5 established a linear Pareto front that characterizes the tradeoff between inference cost and quality degradation (Figure 4 ). A simple experiment in Section 4.2 is also consistent with our theoretical finding.Last but not least, our technical results involve novel analysis, for instance, the design of V `, V ´in the lower bound proof C and the iterative computation for f in D.1. They are the first of its kind and consist of our technical contributions. We provide a proof sketch section in Appendix A.Speculative Decoding and its applications. Speculative execution, where performing speculative work can expedite computation on parallel machines by allowing certain tasks to commence before their necessity is confirmed, can date back to [9, 15] . Recently, [10, 24] formalize this idea with rejection sampling based design for LLM Decoding and achieve multiple-time inference acceleration compared to vanilla auto-regressive decoding. There are fruitful studies [41, 26, 25, 37, 23, 46, 32, 27, 18, 4, 34, 2, 30, 43, 11, 42, 45, 29, 35, 5, 44, 40, 20] since then, and they improve speculative decoding from different angles such as online updating [25] , multiple candidates [45] , retrieval technique [18] , Multimodality [16] or even decoding without draft models [14, 6] .Theoretical endeavor for Speculative Decoding. There are also works that study the theoretical properties for speculative decoding (SD). In particular, [37] considers speculative decoding from the optimal transport perspective and show it is optimal in the single token regime. It further extends to the k multiple draft token setting and formulates the optimal transport solution via linear programming with exponential in k computation time. An approximate sequential selection algorithm is also proposed with linear time. [2] further proposes the improved plan, and [36] extends [37] to the block-level optimal transport for SD. [34] investigates the synergy between draft length and batch size for Speculative Decoding and formulate the optimal speculation length as the root of a polynomial equation. [26] proposes the SpecInfer, a batch algorithm that uses small speculative models to form the token tree, and proves its output matches the distribution of the large model. [45] considers batch speculative decoding without replacement to avoid repeatedly sampling rejected tokens and proves it keeps the decoding quality. Nevertheless, the findings for inference acceleration of these works are mostly empirical, lacking theoretical guarantees.Algorithm 1 Speculative Decoding [10, 24] 1: Input: Set probability b t " mint1, qt pt u and the distribution P t " rq t ´pt s `in Algorithm 2. 2: Require: p t p¨q :" p t p¨|x 1:n´1 , xn:t´1 q, q t p¨q :" q t p¨|x 1:n´1 , xn:t´1 q. @t ě n. Sample r " Uniformr0, 1s.if r ď minIn this section, we provide the mathematical formulation for decoding problems using Markov Chains, and we explain auto-regressive models and speculative decoding algorithm based upon that.A Markov Chain Model for Decoding. We denote x 0 as the prompt. 2 x n is the n-th token output and x 1:T is the trajectory output by the algorithm with T to be the fixed decoding horizon. 3 . Then any decoding algorithm can be characterized by a Markov Chain: state at t is described by history x 0:t , the initial state is x 0 ; the state transition P t maps state x 0:t to state x 0:t`1 . In the context of decoding problem, the transition matrix is defined as P px 0:t`1 |x 0:t q :" ppx t`1 |x 1:t q. In particular, we use p to denote the (conditional) distribution for the draft model that resembles small speculative model, and q for the target model that represents large language model. We use rq ´ps `to denote the normalized distribution for maxt0, qpxq ´ppxqu, @x P V with V being the token vocabulary. We also denote Ēx"f rgs :" ř x f pxqgpxq. Auto-Regressive Decoding. Consider sampling a trajectory x 1:T from an auto-regressive model, where for the given x 1:n´1 , the next token x n follows the conditional distribution q. This decoding mechanism (Algorithm 3) is the prototype for Transformer-based LLMs (e.g. GPT-4). As mentioned in [29] , the accuracy performance of Transformer-based LLMs has been shown to scale with model size, with larger models demonstrating improved capabilities [22] . However, this improvement comes at the cost of higher latency during inference and increased computational requirements. Speculative Decoding. Different from large models, small models are usually much faster at the inference stage. Consequently, we can use a small model to perform auto-regressive sampling and assign large model as a verifier, where the goal of the large model is to check whether the token sampled by the small model should be accepted/rejected. Concretely, this procedure can be summarized into the following three steps:• Draft sampling: given the verified tokens x 1:n´1 , the draft model obtains K sequential candidate tokens xn:n`K´1 via sampling from pp¨|x 1:n´1 , xn:n`i q, i P rK ´1s; • Conditional score computation: given xn:n`K´1 , computing the logits of the K tokens qp¨|x 1:n´1 , xn:n`i q, i P rK ´1s in parallel;• Token validation: Accept the candidate token xt (as x n ) with some probability 0 ď b t ď 1.If accepted, continue to validate the next candidate xt`1 , otherwise reject and sample x n from some designed distribution P t .The above process repeats until x 1:T are generated, and the whole algorithm is summarized as the general rejection-based decoding 2. Specifically, Speculative Decoding [10, 24] designs b t px t q :" mint1, qtpxtq ptpxtq u and distribution P t p¨q :" rq t ´pt s `p¨q (see Algorithm 1 and Figure 1 ). 
problem setup: 
 Speculative Decoding has two key merits. First, it maintains output quality, meaning the output distribution by the algorithm is identical to the output distribution of the large model q, and we term it distribution unbiasedness. Second, it has fast inference property since the auto-regressive sampling is performed on the small model and the target model only verifies. With (possibly) multiple draft tokens being accepted very round, Speculative Decoding can be much faster than direct decoding on large models, and its inference is bottlenecked by the parallel score computation qp¨|x 1:n´1 , xn:n`i q, i P rT ´ns. To highlight this, we defined it as the oracle call and make an assumption based on that. Definition 1. We define one trigger of obtaining logits for xn:T in Step 7 of 2 as one Oracle call. Assumption 1. We assume that: (1) compared to the large model, the computational cost of the draft/small model is negligible. ( 2) each oracle call has runtime Op1q. Remark 1. We assume the above only for the theoretical cleanliness. With negligible draft model, the lookhead K " T in Algorithm 2. In other words, given x 1:n´1 , instead of sampling the next K draft tokens xn:n`K´1 , we are allowed to sample until the end, i.e. xn:T . In practice, Assumption 1(1) also holds true in many cases. One example of negligible-cost model c « 0 is n-gram models, and the empirical evidence in [24, 28] shows n-gram draft model speeds up inference pretty well. In addition, for summarization tasks where long sequences are likely to repeat, any draft model that reuses tokens from the context with a matching prefix, is also cost negligible. Assumption 1(2) is also a standard abstraction, since parallelization consumes (roughly) equal computation as for a single logit. It will consume more memory, but such aspect is beyond the scope of our theoretical study.Metric for measuring efficiency. Desired algorithms should preserve the output quality and be efficient for decoding. To formalize, our theory aims to study the following two quantities:• Inference acceleration: the ratio between the inference time of decoding large model q and the inference time of decoding the algorithm; • Output quality: the distribution bias between the algorithm and the large model distribution q. An algorithm maintains the output quality if it is distribution unbiased.Rejections, and why consider it? A third metric for decoding is the number of draft tokens being rejected. With more draft tokens being accepted, the fewer rejections would occur. Therefore, algorithms with large number of rejections are slower than those with fewer rejections. This means Rejections serves as an alternative metric for the inference speedups. Throughout the paper, we use number of rejections for measuring inference acceleration.To further motivate why choosing Rejections is appropriate, we draw its connection to inference runtime. For Speculative Decoding, the runtime is dominated by the number of oracle calls (Definition 1).After each rejection (Step 10 of 2 or Step 7 of 1), there is one oracle call, thus we obtain inference time " # oracle calls " # rejections.Notice the runtime for auto-regressive decoding (3) is T , therefore we have the relation: inference acceleration " T {# rejections. Based on this setup, we present our main results in next sections.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: In the discussion section, we mentioned that we don't have a lower bound the batch SD algorithms.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
analysis on efficiency and optimality for speculative decoding: 
 We start with the following theorem at the beginning of the section. It covers output quality, which is measured by distribution bias, and expected number of rejections for speculative decoding. Its proof is deferred to Appendix B.Theorem 1. We have the following two results for Speculative Decoding.(1) We define random variables R n P t0, 1u that indicates whether the n-th token is rejected (with 1 being rejected). Here rejection means Line 7 of Algorithm 1 is executed. Then, the total number of rejections N rej " ř T n"1 R n . For Speculative Decoding (here TV denote the TV distance):ErN rej s " T ÿ n"1E x1:n´1"q rTVpp n p¨|x 1:n´1 q, q n p¨|x 1:n´1 qqs.(2) The output distributions of Algorithm 1 and the target model q are identical, i.e. for any output sequence x 1:T P V T , the joint the distributions over x 1:T satisfies: P SD px 1:T q " qpx 1:T q.The first part of Theorem 1, to our knowledge, is the first result that characterizes the expected rejection for speculative decoding a sequence of length T using p and q. The second part of Theorem 1 shows the distribution unbiasedness for SD, which has been presented in [10, 24] . There are three interesting indications:• If ErTVpp n , q n qp¨|x 1:n´1 qs " 0 for all n, all tokens are accepted and the accelerate rate " T ;• If ErTVpp n , q n qp¨|x 1:n´1 qs " 1 for all n, all tokens are rejected and the accelerate rate " 1, i.e. all T tokens are sampled from the large model q;• In general, the accelerate rate for SD is T { ř T n"1 ErTVpp n , q n qp¨|x 1:n´1 qs. Remark 2. [24] derive the expected number of token generated per run of Speculative Decoding as 1{ErTVpp, qqs for K " 8. 4 Their result equals T { ř T n ErTVppn, qnqp¨|x1:n´1qs when ErTVpp n , q n qp¨|x 1:n´1 qs is identical for all n, and this is due to their assumption that the acceptance rates β are i.i.d.. In contrast, our guarantee holds for the case that allows the sequential dependence between different decoding steps.Simulation. We also provide a simulation of Speculative Decoding and compare it with our Theorem 1 in the left panel of Figure 2 (a) with horizon T " 50, p n , q n , n " 1, . . . , 50 are nonstationary Markov Chains. The green line is the empirical average rejections among 100 ˆN runs of Algorithm 1 and the orange line the theoretical value computed via Theorem 1. From the simulation, after 5000 runs the empirical average rejections converge to our theoretical value 16.41. In this example, the acceleration rate is 50{16.41 " 3.05. The specifications of the simulation is included in Appendix F.Now we have analyzed the efficiency of speculated decoding and provided mathematical characterization of the number of expected rejections, which is depicted by the TV distance and scales linearly with the inference time. A natural next-step question to ask is: Is there any other rejection-based algorithm 2 that can do better? We answer this question in the next theorem.Theorem 2 (Instance-dependent Rejection Lower Bound). For an instance P :" pp, qq, where p, q stand for the distributions of the draft model and the target model respectively, defining the family of algorithms as F :" tA : A is a specification of Algorithm 2 that satisfies P A t " q t @t pi.e., unbiasedqu. For an algorithm A, denote N rej as the number of rejections. Then we have the lower boundinf APF E A P rN rej s ě T ÿ n"1E x1:n´1"q rTVpp n , q n qp¨|x 1:n´1 qs . Takeaways. Via Theorem 2, the answer to the key question is: No rejection improvement can be made in Algorithm 2 by changing the acceptance probability b t and the distribution P t if we want to keep the distribution unbiasedness. We point out that lower bound of Theorem 2 matches the complexity upper bound of Speculative Decoding given by Theorem 1. This result confirms that speculative decoding is optimal in the class of all rejection-based methods.The practical implication is that there is no need to tweak acceptance probability, as it will not make performance better. In the next Section 3.2, we will see that Speculative Decoding can be provably improved, as long as one can decode and verify multiple sequence of tokens in parallel.Connection to optimal transport. We mention [37] nicely consider maximizing the acceptance probability from the optimal transport perspective. For a single token, the optimal transport cost is ř x minpppxq, qpxqq, which corresponds to the optimal expected rejection TVpp, qq. However, for the sequential T tokens, their Section 5,6 does not provide a explicit formulation for optimal acceptance/rejections. In this sense, their optimality result can be cast as a special case of ours. However, we do emphasis there are differences in the settings, where [37] consider the optimal transport and we study the class F.To further improve the provable efficiency, we consider batch algorithms that extend the speculative decoding with multiple candidate draft sequences. Most of the existing works [32, 34, 26, 45, 21] formulates batch sequences via a speculation tree structure. In particular, the representative work [26] propose the merged token tree structure that combines sequences with the same prefix. A depth-first search is designed for speculation to ensure the unbiasedness of the algorithm. In addition, [33] devise the parallel decoding structure that speculate the token of each responses given its previous tokens are accepted. Motivated by these works, we consider a batch version of speculative decoding algorithm using a simpler parallel structure (Left of Figure 3 ). Our Algorithm 4 can be viewed as a simplified approximation to those batch algorithms. There are several differences that distinguish batch algorithm from the non-batch version. We highlighting them as follows.Difference1: Oracle call. At the beginning of batch speculation, M draft sequences are generated in parallel as well as the corresponding logits q. This corresponds to Step 4-9 of Alg 4 and is defined as one oracle call which we assume to have unit computation Op1q;5 Difference2: Speculation procedure. It follows the DFS principle: If the first token of a response is accepted, only that response will be speculated until rejection. For instance in the Left panel of Figure 3 , if the token 'deep' is accepted, the algorithm will keep verifying token 'learn', 'ing' until rejection and the rest of sequences won't be examined; if 'deep' is not verified, then the algorithm will keep examing 'rein'. In this case, rejection happens only if 'deep', 'rein' and 'atten' are all rejected. Once rejection happens, the process will restart. By this design, it still holds true that: inference time " # oracle calls " # rejections.Again, we measure the output quality and rejections for the batch algorithm. The following main result (Theorem 3) provides the explicit characterization for rejections and batch improvement. Detailed proofs and discussions can be found in D. Bernoulli distributions with q " Berp0.5q. Lower: p " UnifpV q, q " UnifpV 1 q with r " V {V 1 .Theorem 3 (Unbiasedness and efficiency of batch SD). Recall the definition of R n and N rej in Thm 1, and iteratively define: q m`1 " rq m ´ps `, @m P rM s with q 1 " q being the target distribution.Then, for Batch Speculative Decoding 4, P Batch px 1:T q " qpx 1:T q. @x 1:T P V T . Moreover,ErNrejs " T ÿ n"1 Ex 1:n´1 "q rTVrq, psp¨|x1:n´1qs ´T ÿ n"1 Ēx 1:n´1 "f rTVpq, pqpx1:n´1q ´r M ź m"1TVpq m , pqpx1:n´1qss looooooooooooooooooooooooooooooooooooooooomooooooooooooooooooooooooooooooooooooooooon Batch Improvement where f px 1:n q :" Ppx 1:n X tn-th draft token rejecteduq. f can be iteratively computed via p, q.Takeaways. The expected rejection of Batch Decoding composes of two parts: (i) ř T n E q rTVrq, pss which is the rejection that matches the non-batch speculative decoding; (ii) the batch improvement (BI) which comes from our design and is always non-negative. To better understand how the batch improvement scale with M , we instantiate the single token BIpq, pq :" TVrq, ps ´śM m"1 TVrq m , ps with two simple examples.Uniform Distribution. For the whole space V, let p be a uniform distribution with support V and q be a uniform distribution over a subset of V with size V 1 ă V . Let V {V 1 " r, in this case BIpUnifpV 1 q, UnifpV qq " p1 ´1 r q ´p1 ´1 r q M , r " V {V 1 . Its pattern is visualized in the lower right panel of Figure 3 . The improvement converges to 1 ´1{r and is always positive as long as batch size is at least 2. Also, when the vocabulary size V scales up, i.e. r goes up, the improvement is going to zero. This is not surprising, since the probability of draft sample to fall within in the support of target distribution is very small. Bernoulli Distribution. Suppose u ě v and let p " Berpuq, q " Berpvq. Then BIpBerpvq, Berpuqq " |u ´v| ¨p1 ´uM´1 q. Its pattern is exhibited in the upper right panel of Figure 3 . Notice for both cases, the limiting batch improvement is TVpp, qq, which is only significant when p deviates from q. This provides the practical design heuristic: batch design can significantly reduce the number of rejections when the draft distribution p and target distribution q are different from each other and won't make too much improvement when p and q are close. rejection is unavoidable no matter how many draft responses are sampled from p. We also formally prove this in Proposition 1. This theoretical discovery indicates that having a very large draft batches does not necessarily result in significant inference speedups compared to small batches.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: The code is in the appendix, the Alpaca data is also public. Guidelines:• The answer NA means that paper does not include experiments requiring code. • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
analysis on the optimal rejection-distribution bias tradeoff: 
 In earlier sections, we focus on analyzing SD or Batch SD which are distribution unbiased. To further reduce rejections, the algorithm may have to compromise on output quality. Nevertheless, it is usually acceptable for many real-world applications. For instance, for the family of Gemini models [38] , Google may deploy the small model Gemini Nano as the draft model p and Gemini Ultra as the target model q and tune the acceptance probability b (in Algorithm 2) higher such that Gemini Nano is applied more often for the purpose of less expenditure. Therefore, an intriguing question to ask is: For biased algorithms, what is the optimal tradeoff between rejection and distribution bias?We measure the distribution bias between P A and q via TV distance. 6 Concretely, for any algorithm A :" pb, Pq in 2 with acceptance probability b and rejection distribution P, we fix b and aim to optimize the following objective Loss TV pbq :" min P TVrP A , qs, where A :" pb, Pq.We wish to solve the above since it characterizes the minimal distribution bias for any design b. We present our solution in the next. Theorem 4 (Optimal solution to optimization (1)). We can show the objective (1) is equivalent to Loss TV pbq :" 1 2 min P ÿx ˇˇˇˇq pxq ´bpxqppxq ´Ppxq ÿx r1 ´bpxqsppxq ˇˇˇˇs .t. ÿx Ppxq " 1, Ppxq ě 0, @x.x r1 ´bpxqsppxq ą 0. Define Apxq :" qpxq´bpxqppxq ř x r1´bpxqsppxq , and the sets A `" tx P V : Apxq ě 0u, A ´" tx P V : Apxq ă 0u. Then the set of optimal distributions of objective (1) is characterized as tP ˚: P ˚|A´p ¨q " 0; 0 ď P ˚|A`p ¨q ď Ap¨qu, and the optimal value is Loss TV pbq "1 2 ř x |qpxq ´bpxqppxq| ´1 2 ř x p1 ´bpxqqppxq ě 0.Theorem 4 is a universal characterization that contains both biased and unbiased situations. 1. If b is less than Speculative Decoding threshold, i.e. b ď mint1, q{pu then A becomes a probability distribution and the optimal distribution P ˚equals A which is also unbiased (Loss TV pbq " 0); 2. If b exceeds the Speculative Decoding threshold, then A is no longer a distribution and there are multiple optimal distributions P ˚. In this case, the optimal distribution bias Loss TV pbq ą 0 for Algorithm 2.Main Takeaways. Via Theorem 4, we derive the pareto front (the optimal tradeoff) between rejection probability 7 Pprejectq vs. distribution distance TVrP A , qs (Left panel of Figure 4 ). The blue region can be realized by some algorithm 2, and the red region cannot. p0, 0q is the "perfect algorithm" (no rejection, no bias) which does not exists, and, in particular, p0, TVrp, qsq stands for Speculative Decoding. Surprisingly, the pareto front is a straight line connecting p0, TVrp, qsq and pTVrp, qs, 0q, which represents a linear relationship between the rejection probability and the optimal TV deviation. This is guaranteed by the following theorem.Theorem 5 (Pareto front). For any Algorithm A in the class 2 that satisfies mint1, qpxq ppxq u ď bpxq ď 1, @x P V. Then P A prejectq `Loss TV pbq " TVrp, qs. Here P A prejectq " 1 ´řx bpxqppxq and Loss TV pbq :" min P TVrP A , qs. We plot the numerical example using two Markov Chains in the middle and right panel of Figure 4 that coincides with our theoretical finding. In the figure, the acceptance probability is set to be bpxq " mint1, qpxq`ϵ ppxq u. The orange line in pcq and the green boundary in pbq are computed via Loss TV pbq from Theorem 4. The complete proofs for Theorem 5 and Theorem 4 are deferred to Appendix E. For the clearness of illustration, we focus on the pareto front between rejection vs. the minimal TV deviation for a single token.We provide an additional simple experiment to show the effectiveness of our Pareto-optimal solution in Theorem 4. Consider objective (1) . For the given acceptance probability b ą mint1, q{pu, we have two options for P: 1. Baseline: select P to be suboptimal to (1), i.e. simply set P :" q, which is the target distribution; 2. Set P :" P ˚be the optimal distribution of Theorem 4. We call the first method Decoding-UNO and the latter one Decoding-OPT. Instead of TV distance, we measure the quality via WinRate in our experiment. Concretely, for each prompt, we let Decoding-UNO and Decoding-OPT to generate responses independently, and use score models to compare whose generation has higher quality. We specify draft model p as pythia-70m and target model q as pythia-2.8b from EleutherAI [7] . We apply the score model to be RM-Mistral-7B or GPT-4. We test 200 prompts from Alpaca-Farm-Eval Dataset [13] with 500 responses/comparisons per prompt. Table 1 shows that Decoding-OPT achieves better performance than Decoding-UNO across different choice of ϵ's. Due to space constraint, the missing details are deffered to Appendix G.RM-Mistral-7B GPT-4 ϵ " 0.1 ϵ " 0.2 ϵ " 0.4 ϵ " 0.8 ϵ " 0.1 ϵ " 0.2 ϵ " 0. Extending Speculative Decoding to other studies. Speculative Decoding is a generic sampling approach that extends beyond mere decoding tasks. It holds potential for wider applications such as search engines and recommendation systems, where it can be employed to quickly generate and refine search outcomes or content suggestions, enhancing the overall efficiency and user experience of these systems. We leave these as future works.A Proof Sketch A.1 Proof sketch of Theorem 2.For any algorithm A P F, its design b n can be written as a function of any sequence x 1:n´1 , xn with 0 ď b n px n , x 1:n´1 q ď 1. 8 Based on this, we can further define the new function ϵ n : V ˆVn´1 Þ Ñ R according to the following equation: Indeed, we can choose ϵ n :" b n ¨pn ´qn , and the validity of the definition is guaranteed by the Lemma 1. Then, the acceptance probability b n ą mint1, qn pn u implies ϵ n ą 0. Next, we show for any A P F, it must satisfy b n ď mint1, qn pn u for all n. To this end, we design the two token sets V `" tx : Dn s.t. ϵ n pxq ą 0u and V ´" tx : Dn s.t. ϵ n pxq ď 0uand prove V `" H, V ´" V.Finally, by Lemma 2 in Appendix, any algorithm satisfies b n ď mint1, qn pn u, @n P rT s must have E A P rN rej s ě ř T n"1 E q rTVpp n , q n qp¨|x 1:n´1 qs. Since A P F is arbitrary, this concludes the proof. The full proof is included in C.The derivation for the number of expected rejections using the batch speculative decoding is more involved than the Algorithm 4 due to the parallel response structure. The key step is to compute the intermediate quantity P A px n acc, xn " x n |x 1:n´1 q. Let xn´1 " pp¨|x 1:n´2 q, then there are two cases: xn´1 accepted or rejected. We have P A px n acc, xn " x n |x 1:n´1 q " P A px n acc, xn " x n , xn´1 acc|x 1:n´1 q `PA px n acc, xn " x n , xn´1 rej|x 1:n´1 q " P A px n acc, xn " x n |x n´1 acc, x 1:n´1 q looooooooooooooooooooooomooooooooooooooooooooooon pa P A px n´1 acc|x 1:n´1 q `PA px n acc, xn " x n |x n´1 rej, x 1:n´1 q looooooooooooooooooooooomooooooooooooooooooooooonEQUATIONIn the process of finding p b , we need to compute the the quantity f px 1:n q :" Ppx 1:n X tn-th draft token rejecteduq and it can be recursively computed via (22) using p, q.A.3 Proof sketch for Theorem 4Due to space constraint, we only summarize the high-level proof ideas for Theorem 4. Since řx Apxq " 1, the original objective (1) can be equivalently rewritten asEQUATIONWe now find the solution of objective (4) in two steps.Step1: Recall A `" tx P V : Apxq ě 0u and A ´" tx P V : Apxq ă 0u, then any optimal P ˚must satisfy P ˚pxq " 0 for all x P A ´. This can be shown by contradiction via an alternative construction P 1 to reason P ˚is suboptimal to P 1 .Step2: We characterize the optimal solutions of the objective. By Step1, we can show any optimal solution P ˚satisfies ř x |Apxq ´P˚p xq| ě ∥A∥ 1 ´1 and the equal sign can be achieved. Then we can convert ∥A∥ 1 ´1 back to Loss TV . This also helps identify the optimal set of P ˚.Algorithm 3 Auto-Regressive Model Decoding 1: Init: Horizon T . Distribution q. Prompt x 0 . n " 0. 2: while n ă T do 3: Sample x n " qp¨|x 1:n´1 q. n Ð n `1. 4: end while B Proof of Theorem 1 Theorem 6 (Restatement of the first part of Theorem 1). We define random variables R n P t0, 1u indicating whether the n-th token is rejected with 1 being rejected (here rejection means Line 6 of Algorithm 1 is executed). Then, the total number of rejections N rej " ř T n"1 R n . For Speculative Decoding,ErN rej s " T ÿ n"1E x1:n´1"q rTVpp n p¨|x 1:n´1 q, q n p¨|x 1:n´1 qqs.Here TV denote the TV distance between two distributions.Proof. Given the verified the tokens x 1:n´1 , we first compute PpReject at n|x 1:n´1 q. Denote the candidate draft token x " p n p¨|x 1:n´1 q, then by law of total probability PpReject at n|x1:n´1q " where the third equal sign uses draft token is sampled from p n and the fourth equality is by design of Speculative Decoding (Algorithm 1, Line 4).ErN rej s "T ÿ t"1 ErR t s " T ÿ t"1 E rErR t |x 1:t´1 ss " T ÿ t"1 ÿ x1:n´1 ErR t |x 1:t´1 sPpx 1:n´1 q " T ÿ t"1 ÿ x1:n´1 ErR t |x 1:t´1 sqpx 1:n´1 q " T ÿ t"1 ÿ x1:n´1PpReject at n|x 1:n´1 qqpx 1:n´1 q " T ÿ n"1 E x1:n´1"q rTVpp n p¨|x 1:n´1 q, q n p¨|x 1:n´1 qqs.Here the fourth equal sign comes from Speculative Decoding keep the distribution q (Proposition ??), the fifth equal sign comes from the event tR n " 1u " tReject at nu.Theorem 7 (Restatement of the second part of Theorem 1). The output distributions of Speculative Decoding Algorithm 1 and the large model q are identical, i.e. for any output sequence x 1:T P V T , the joint the distributions over x 1:T satisfies: P SpecDecoding px 1:T q " qpx 1:T q.Remark 3. The proof of Theorem 7 is very similar to [10] except we have K " 8. In addition, [10] proves the distribution match for a single token, we complement the proof to show the distribution match holds for a sequence of tokens x 1:T .Proof. In particular, we use induction to show the stronger result that @ t P rT s, @x 1 , x 2 , . . . , x t P V, it holds P A t px 1 , x 2 , . . . , x t q " q t px 1 , x 2 , . . . , x t q. Step1: Since x 0 is the prompt, its distribution is independent to p, q. Then for t " 1, applying Theorem 1 of [10] (with K " 8) directly with p and q to be conditional distributions p 1 p¨|x 0 q and q 1 p¨|x 0 q gives P A 1 px 1 |x 0 q " q 1 px 1 |x 0 q, which further implies P A 1 px 1 q " q 1 px 1 q (since the distribution of x 0 is independent of p, q).Step2: assume P A t px 1 , x 2 , . . . , x t q " q t px 1 , x 2 , . . . , x t q, we first show P A px t`1 |x 1:t q " qpx t`1 |x 1:t q. Indeed, let xt`1 " pp¨|x 1:t q, then by law of total probabilityP A px t`1 |x 1:t q "P A px t`1 " x t`1 |x 1:t qPpx t`1 acc|x t`1 " x t`1 , x 1:t q `PA px t`1 rej|x 1:t qP A px t`1 |x t`1 rej, x 1:t q (5)By Algorithm 1,EQUATIONNext, the probability of rejection is:EQUATIONwhere the second equal sign comes from (6) . Lastly, by the construction of the algorithm, P A px t`1 |x t`1 rej, x 1:t q " maxt0, qpx t`1 |x 1:t q ´ppx t`1 |x 1:t qu ř x 1 maxt0, qpx 1 |x 1:t q ´ppx 1 |x 1:t qu .Combining ( 8) and (7) yields P A px t`1 rej|x 1:t qP A px t`1 |x t`1 rej, x 1:t q " maxt0, qpx t`1 |x 1:t q ´ppx t`1 |x 1:t qu.Plugging (6) and the above equation into (5) to obtain P A px t`1 |x 1:t q " mintppx t`1 |x 1:t q, qpx t`1 |x 1:t qu `maxt0, qpx t`1 |x 1:t q ´ppx t`1 |x 1:t qu " qpx t`1 |x 1:t q.Finally, applying the above we obtain P A t`1 px 1:t`1 q " P A t px 1:t q ¨PA px t`1 |x 1:t q " P A t px 1:t q ¨qpx t`1 |x 1:t q " q t`1 px 1:t`1 q, where the last equal sign uses the induction hypothesis.Theorem 8 (Restatement of Theorem 2). Define the arbitrary instance P :" pp, qq, and define the family of algorithms as F :" tA : A is a specification of Algorithm 2 that satisfies P A t " q t @t pi.e., distribution unbiasedqu. For an algorithm A, denote N rej as the number of rejections. Then we haveinf APF E A P rN rej s ě T ÿ n"1E x1:n´1"q rTVpp n , q n qp¨|x 1:n´1 qs :" CpPq.Remark 4. Theorem 8 shows the rejections of Algorithm 1 is tight at the instance level (over the family of algorithms F). Therefore, it attains the instance-optimality over sequential decoding algorithms family F.Proof. For any algorithm A P F, its b n can be written as a function of the sequence x 1:n´1 , xn with 0 ď b n px n , x 1:n´1 q ď 1. 9 Based on this, we define the new function ϵ n : V ˆVn´1 Þ Ñ R according to the following equation:EQUATIONIndeed, we can choose ϵ n :" b n ¨pn ´qn , and the validity of the definition is guaranteed by the Lemma 1. Recall A P F is a distribution unbiased algorithm w.r.t. q. Let x 1 , . . . , x n be the validated sequence of A, then we haveP A n px n |x 1:n´1 q " P A n px 1:n q P An px 1:n´1 q " q n px 1:n q q n px 1:n´1 q " q n px n |x 1:n´1 q.On the other hand, let xn " p n p¨|x 1:n´1 q, the decomposition holds P A n px n |x 1:n´1 q " P A n px n , xn accept|x 1:n´1 q `PA n px n , xn reject|x 1:n´1 q "P A n px n accept|x n " x n , x 1:n´1 qP A n px n " x n |x 1:n´1 q `PA n px n |x n reject, x 1:n´1 qP A n px n reject|x 1:n´1 q "b n px n , x 1:n´1 q ¨pn px n |x 1:n´1 q `PA n px n |x n reject, x 1:n´1 q ¨PA n px n reject|x 1:n´1 q "b n px n , x 1:n´1 q ¨pn px n |x 1:n´1 q `PA n px n |x n reject, x 1:n´1 q ¨p1 ´ÿ x b n px, x 1:n´1 qp n px|x 1:n´1 qq, (11) where the third equal sign uses P A n px n " x n |x 1:n´1 q " p n px n |x 1:n´1 q and the last equal sign is due toP A n px n reject|x 1:n´1 q " 1 ´PA n px n accept|x 1:n´1 q " 1 ´ÿ x P A n px n accept, xn " x|x 1:n´1 q "1 ´ÿ x P A n px n accept|x n " x, x 1:n´1 qP A n px n " x|x 1:n´1 q " 1 ´ÿ x b n px, x 1:n´1 qp n px|x 1:n´1 q.Let's do some further simplifications. First off, b n px n , x 1:n´1 q ¨pn px n |x 1:n´1 q "p n px n |x 1:n´1 q min " 1, q n px n |x 1:n´1 q `ϵn px n , x 1:n´1 q p n px n |x 1:n´1 q * " min tp n px n |x 1:n´1 q, q n px n |x 1:n´1 q `ϵn px n , x 1:n´1 qu , (12) and 1 ´ÿ x b n px, x 1:n´1 qp n px|x 1:n´1 q "1 ´ÿ x min tp n px|x 1:n´1 q, q n px n |x 1:n´1 q `ϵn px, x 1:n´1 qu " ÿ x max t0, p n px|x 1:n´1 q ´qn px|x 1:n´1 q ´ϵn px, x 1:n´1 qu (13) 9 For Speculative Decoding, its bnpxn, x1:n´1q " min ! 1,qnpxn|x 1:n´1 q pnpxn|x 1:n´1 q ) .Plug ( 12) and ( 13) into (11) , and then plug ( 11) into (10) to obtain q n px n |x 1:n´1 q " min tp n px n |x 1:n´1 q, q n px n |x 1:n´1 q `ϵn px n , x 1:n´1 qu `PA n px n |x n reject, x 1:n´1 q ¨ÿ x pp n px|x 1:n´1 q ´qn px|x 1:n´1 q ´ϵn px, x 1:n´1 qq `.Now, we define the token set V `" tx : D n, x 1:n´1 s.t. ϵ n px, x 1:n´1 q ą 0u and similarly the token set V ´" tx : D n, x 1:n´1 s.t. ϵ n px, x 1:n´1 q ď 0u. Next, we show V `" H, and V ´" V.Case1. If p n px n |x 1:n´1 q ě q n px n |x 1:n´1 q `ϵn px n , x 1:n´1 q, then by ( 14) we have q n px n |x 1:n´1 q " q n px n |x 1:n´1 q `ϵn px n , x 1:n´1 q `PA n px n |x n reject, xn " x n , x 1:n´1 q ¨ÿ x pp n px|x 1:n´1 q ´qn px|x 1:n´1 q ´ϵn px, x 1:n´1 qq ěq n px n |x 1:n´1 q `ϵn px n , x 1:n´1 q, and this implies 0 ě ϵ n px n , x 1:n´1 q, which means x n P V´;Case2. If p n px n |x 1:n´1 q ă q n px n |x 1:n´1 q `ϵn px n , x 1:n´1 q, then by ( 14) we have q n px n |x 1:n´1 q " p n px n |x 1:n´1 q `PA n px n |x n reject, xn " x n , x 1:n´1 q ¨ÿ x pp n px|x 1:n´1 q ´qn px|x 1:n´1 q ´ϵn px, x 1:n´1 qq ěp n px n |x 1:n´1 q, and this implies ϵ n px n , x 1:n´1 q "b n px n , x 1:n´1 qp n px n |x 1:n´1 q ´qn px n |x 1:n´1 q ďp n px n |x 1:n´1 q ´qn px n |x 1:n´1 q ď 0, which means x n P V´.Therefore, combining the two cases we always have x P V ´, which indicates V `" H. By (9) , this implies for all n, b n ď mint1, qn pn u. Finally, by Lemma 2, this implies E A P rN rej s ě CpPq. Since A P F is arbitrary, this concludes the proof.Corollary 1. For any algorithm A P F, it follows @ n P rT s, x P V and x 1:n´1 , b n px, x 1:n´1 q ď min ! 1, qnpx|x1:n´1q pnpx|x1:n´1q) . In this case, the distribution P n is defined as:P n px|x 1:n´1 q " q n px|x 1:n´1 q ´min tp n px|x 1:n´1 q, q n px|x 1:n´1 q `ϵn px, x 1:n´1 qu ř x pp n px|x 1:n´1 q ´qn px|x 1:n´1 q ´ϵn px|x 1:n´1 qq `. Applying ϵ n " b n p n ´qn , this is equivalent to P n px|x 1:n´1 q " q n px|x 1:n´1 q ´bn px, x 1:n´1 qp n px|x 1:n´1 q ř x p1 ´bn px, x 1:n´1 qqp n px|x 1:n´1 q .Proof of Corollary 1. We reutilize ( 14) here and call it (15) . q n px n |x 1:n´1 q " min tp n px n |x 1:n´1 q, q n px n |x 1:n´1 q `ϵn px n , x 1:n´1 qu `PA n px n |x n reject, x 1:n´1 q ¨ÿ x pp n px|x 1:n´1 q ´qn px|x 1:n´1 q ´ϵn px,EQUATION)By two cases discussion as in the proof of Theorem 8, we have @ A P F, it follows @ n P rT s, x P V and x 1:n´1 , b n px, x 1:n´1 q ď min ! 1, qnpx|x1:n´1q pnpx|x1:n´1q) . Then we can directly solve (15) to obtainP A n px|x n reject, x 1:n´1 q "q n px|x 1:n´1 q ´min tp n px|x 1:n´1 q, q n px|x 1:n´1 q `ϵn px, x 1:n´1 qu ř x pp n px|x 1:n´1 q ´qn px|x 1:n´1 q ´ϵn px|x 1:n´1 qq `.Lastly, we verify such a P A n px n |x n reject, x 1:n´1 q is a valid distribution. First of all, since ϵ n " b n ¨pn ´qn , then b n px, x 1:n´1 q ď min ! 1, qnpx|x1:n´1q pnpx|x1:n´1q) implies ϵ n px, x 1:n´1 q ď 0, and this further implies q n px|x 1:n´1 q ´min tp n px|x 1:n´1 q, q n px|x 1:n´1 q `ϵn px, x 1:n´1 qu ěq n px|x 1:n´1 q ´min tp n px|x 1:n´1 q, q n px|x 1:n´1 qu ě 0 which implies P A n px n |x n reject, x 1:n´1 q ě 0. Second,ÿ x P A n px|x n reject, x 1:n´1 q " ÿ xq n px|x 1:n´1 q ´min tp n px|x 1:n´1 q, q n px|x 1:n´1 q `ϵn px, x 1:n´1 qu ř x pp n px|x 1:n´1 q ´qn px|x 1:n´1 q ´ϵn px|x 1:n´1 qq " 1 ´řx min tp n px|x 1:n´1 q, q n px|x 1:n´1 q `ϵn px, x 1:n´1 qu ř x pp n px|x 1:n´1 q ´qn px|x 1:n´1 q ´ϵn px|x 1:n´1 qq " 1 `řx max t´p n px|x 1:n´1 q, ´qn px|x 1:n´1 q ´ϵn px, x 1:n´1 qu ř x pp n px|x 1:n´1 q ´qn px|x 1:n´1 q ´ϵn px|x 1:n´1 qq " řx p n px|x 1:n´1 q `řx max t´p n px|x 1:n´1 q, ´qn px|x 1:n´1 q ´ϵn px, x 1:n´1 qu ř x pp n px|x 1:n´1 q ´qn px|x 1:n´1 q ´ϵn px|x 1:n´1 qq " řx max t0, p n px|x 1:n´1 q ´qn px|x 1:n´1 q ´ϵn px, x 1:n´1 qu ř x pp n px|x 1:n´1 q ´qn px|x 1:n´1 q ´ϵn px|x 1:n´1 qq `" 1. This concludes the proof.Lemma 1. For any 0 ď b n px n , x 1:n´1 q ď 1, there exists ϵ n px n , x 1:n´1 q P R such that (9) holds true.Proof of Lemma 1. Indeed, set ϵ n " b n ¨pn ´qn , thenmin " 1, q n px n |x 1:n´1 q `ϵn px n , x 1:n´1 q p n px n |x 1:n´1 q * " mint1, b n px n , x 1:n´1 qu " b n px n , x 1:n´1 qwhere the first equal sign uses that xn is sampled from p n p¨|x 1:n´1 q so p n px n |x 1:n´1 q ą 0, and the second equal sign uses 0 ď b n ď 1.Lemma 2. For any instance P " pp, qq, let F :" tA : A is a realization of Algorithm 2 s.t. P A t " q t @t pi.e., unbiasedqu. Suppose there is an A P F such that E A P rN rej s ă CpPq, then there exists a b n in Line 8 of Template 2 such that Dx, x 1:n´1 b n px|x 1:n´1 q ą min " 1, q n px|x 1:n´1 q p n px|x 1:n´1 q * .Proof of Lemma 2. Suppose for all n, x, x 1:n´1 , b n px|x 1:n´1 q ď min ! 1, qnpx|x1:n´1q pnpx|x1:n´1q) . We define random variables R n P t0, 1u indicating whether the n-th token is rejected (1 is rejected). Therefore, the expected number of rejection isE A « T ÿ n"1 R n ff " T ÿ n"1 E A rR n s " T ÿ n"1 E x1:n´1"P A " E A rR n |x 1:n´1 s ‰ " T ÿ n"1 E x1:n´1"q " E A rR n |x 1:n´1 s ‰ .Here the second to last equality comes from the tower property and the last equal signs is due to A is unbiased. Next, denote xn " p n p¨|x 1:n´1 q to be the candidate token, we haveE A rR n |x 1:n´1 s " P A rR n " 1|x 1:n´1 s " ÿ xn P A rR n " 1, xn |x 1:n´1 s " ÿ xn P A rR n " 1|x n , x 1:n´1 sP A rx n |x 1:n´1 s " ÿ xn P A rR n " 1|x n , x 1:n´1 sp n rx n |x 1:n´1 s " ÿ xn p1 ´bn q ¨pn rx n |x 1:n´1 s ě ÿ xn p1 ´min " 1, q n px n |x 1:n´1 q p n px n |x 1:n´1 q * q ¨pn rx n |x 1:n´1 s " ÿ xnrp n px n |x 1:n´1 q ´qn px n |x 1:n´1 qs `" TVrpp n p¨|x 1:n´1 q ´qn p¨|x 1:n´1 qs and this impliesE A P rN rej s " E A « T ÿ n"1 r n ff ě T ÿ n"1E x1:n´1"q rTVrpp n p¨|x 1:n´1 q ´qn p¨|x 1:n´1 qss " CpPq contradicts E A P rN rej s ă CpPq! This concludes the proof.Algorithm 4 Batch Speculative Sampling 1: Init: Horizon T , Distributions q t and p t , with q " P LLM , p " P Draft . Lookahead K " 8.2: while n ă T do 3:Reset q t " P LLM t @t. n 0 " n.for m " 1 : M ˛Sample M draft responses in parallel˛do 5:for t " n : T do 6:Sample xm t " p t p¨|x 1:n´1 , xm n:t´1 q.7:end for 8:end for 9:Obtain logits q n p¨|x 1:n´1 q, . . . , q T p¨|x 1:n´1 , xm n:T q, @m P rM s in parallel for xm n:T .10:---˛Verification Begins ---11:Set Sample"False.12:for m " 1, . . . , M do 13:for t " n : T do 14:Sample r " Uniformr0, 1s.15:if r ď min ! 1, qtpx m t |x1:n´1,x m n:t´1 q ptpx m t |x1:n´1,x m n:t´1 q) then 16:Accept with x n " xm t . n Ð n `1. Sample"True. if t " n 0 then 19:Update q n Ð rq n p¨|x 1:n´1 q ´pn p¨|x 1:n´1 qs `. Break. //Here n 0 equals n. ˛Rejection Sample x n " rq n p¨|x 1:n´1 q ´pn p¨|x 1:n´1 qs `. n Ð n `1. Break.end ifend ifend for 25:if Sample" TRUE then˛Rejection Sample x n " rq n p¨|x 1:n´1 q ´pn p¨|x 1:n´1 qs `. n Ð n `1. Break. end for 31: end while where by (21) pseudo-measure f satisfies @x 1:n f px 1:n q " maxt0, q 1 n px n |x 1:n´1 q ´pn px n |x 1:n´1 qquq 1 n px 1:n´1 q ´˜maxt0, q 1 n px n |x 1:n´1 q ´pn px n |x 1:EQUATIONHere we used r m " ř x maxt0, q m n px|x 1:n´1 q ´pn px|x 1:n´1 qu " TVpq m , pqpx 1:n´1 q. Plug the above back to ( 16), we finally haveE A B r T ÿ n"1 R n s " T ÿ n"1 E x1:n´1"q 1 rTVpq 1 p¨|x 1:n´1 , pp¨|x 1:n´1 qqqs ´T ÿ n"1 Ēx1:n´1"f « TVpq 1 , pqpx 1:n´1 q ´r M ź m"1TVpq m , pqpx 1:n´1 qs ff .The formulation for f is iteratively obtained in (22) .Proposition 1. Let f M be the f in Theorem 3 with batch M , and let f 8 " lim M Ñ8 f M . Then we have:• f M p¨q ď q 1 p¨q, @M P N; f 8 p¨q ď q 1 p¨q.• f 8 px 1:n q " hpx n |x ăn qrq 1 px 1:n´1 q ´f 8 px 1:n´1 qs.• lim M Ñ8 ErN rej s " ř T n"1 pE q 1 rTVrq 1 , pss ´Ē f 8 rTVpq 1 , pqsq. • M Ñ 8, ErN rej s Û 0. This indicates increasing batch size to 8 doesn't help.Proof. First item: Recall in the proof of Theorem 3, f is defined as f px 1:n q :"P A px n rej, x n |x 1:n´1 qq 1 n px 1:n´1 q "P A px n rej, x n |x 1:n´1 qP A px 1:n´1 q "P A px n rej, x 1:n q ď P A px 1:n q " q 1 px 1:n q, where the second equality is due to Batch algorithm is unbiased (Theorem 3).f px 1:n q " hpx n |x ăn qq 1 n px 1:n´1 q´rhpx n |x ăn q´p M ź m"1TVpq m , pqpx ăn qqq M `1px n |x ăn qsf px 1:n´1 q, since ś M m"1 TVpq m , pqpx ăn q Ñ 0 as M Ñ 0 (note TVpq m , pqpx ăn q " 1 iff there is no overlap between q m and p), it implies f 8 px 1:n q " hpx n |x ăn qrq 1 px 1:n´1 q ´f 8 px 1:n´1 qs.Third item: Similar to the second item, it holds true via taking M Ñ 8. For proving ErN 8 rej s ą 0, suppose ErN 8 rej s " 0. Then it holds q 1 " f 8 . By the second item, this further implies q 1 " f 8 " 0, which is impossible.Fourth item: We prove by contradiction. Suppose ErN rej s Û 0, then by the first item and the third item this implies q 1 " f 8 . Plug this back to the second item, this further implies f 8 " 0, so q 1 " f 8 " 0 is a contradiction (q 1 is a probability distribution)!For the sampling scheme where the acceptance probability b n goes beyond mint1, qnpx|x1:n´1q pnpx|x1:n´1q u, there is a quality degradation for the output sequence as the algorithm is leaning towards the draft model (smaller model). In this case, the objective is to minimize the quality degradation via considering the TV distance minTVrP A n p¨|x 1:n´1 q, q n p¨|x 1:n´1 qs, @x 1:n´1 P V n´1 .Under the above, via equation ( 14) the objective is equivalent to the following (note that according to Algorithm 2 P A n px n |draft token rejected, x 1:n´1 q " P n px n |x 1:n´1 q is an algorithmic choice)min Pn 1 2 ÿ x ˇˇˇˇqn pxq ´mintp n pxq, q n pxq `ϵn pxqu ´Pn pxq ÿx rp n pxq ´qn pxq ´ϵn pxqs `š .t.EQUATIONwhere we removed x 1:n´1 for notation simplicity, and recall again ϵ n :" b n p n ´qn . When ř xrp n pxq ´qn pxq ´ϵn pxqs `" 0, the objective degenerates to the constant (in P n ) TVpq n , p n q. Therefore, for the rest of the section, we focus on the case where ř xrp n pxq ´qn pxq ´ϵn pxqs `ą 0. We have the following Theorem that characterizes the solution of (23) . Theorem 10 (Restatement of Theorem 4). Suppose ř x rp n ´qn ´ϵn s `pxq ą 0. Define A n pxq :" q n pxq ´mintp n pxq, q n pxq `ϵn pxqu ř xrp n pxq ´qn pxq ´ϵn pxqs`,and define the positive token set A `" tx P V : A n pxq ě 0u and the negative token set A ´" tx P V : A n pxq ă 0u. Then the set of optimal distributions of objective ( 23) is characterized as tP n : P n pxq " 0, @x P A ´; 0 ď P n pxq ď A n pxq, @x P A `; ÿ x P n pxq " 1u, and the optimal value is Loss TV pbq " 1 2 ÿ x |q n pxq ´mintp n pxq, q n pxq `ϵn pxqu| ´1 2 ÿ x rp n pxq ´qn pxq ´ϵn pxqs `.Remark 5. In the main context (Theorem 4) we define A n pxq :" qnpxq´bnpxqpnpxq ř x r1´bnpxqspnpxq and Loss TV pbq "1 2 ř x |q n pxq ´bn pxqp n pxq| ´1 2 řx p1 ´bn pxqqp n pxq. This is equivalent to the above since ϵ n :" b n p n ´qn .Proof of Theorem 10. In this case, note mintp n pxq, q n pxq `ϵn pxqu `rp n pxq ´qn pxq ´ϵn pxqs `" p n pxq, we have ÿ x q n pxq ´mintp n pxq, q n pxq `ϵn pxqu " ÿ x rp n pxq ´qn pxq ´ϵn pxqs `. We now find the solution of objective (25) in two steps.Step1: Let the positive token set A `" tx P V : A n pxq ě 0u and the negative token set A ´" tx P V : A n pxq ă 0u, then any optimal P n must satisfy P n pxq " 0 for all x P A ´.x A n pxq " 1, it implies A `‰ H. Suppose for some optimal P n , there exists x P A śuch that P n pxq ą 0, then we show there exists x P A `such that A n pxq ą P n pxq. Suppose this isTo validate the correctness of our theory, we provide the numeric simulations that are displayed in Figure 4 ,2. We model the distribution p 1:T and q 1:T to be two non-stationary Markov Chains with 7 states/tokens. Instead of being ppx n |x 1:n´1 q, for Markov Chain, the one step transition is Markovian with ppx n |x 1:n´1 q " ppx n |x n´1 q. We set the random seed to be 10. The prompt distribution p 0 is set to be Uniform distribution. For Figure 2 , the true value is computed via Theorem 1,3 respectively, and solid line is computed by G Details for Experiment Section 4.2Consider objective (1)Loss TV pbq :" min P TVrP A , qs, where A :" pb, Pq.For any biased algorithm with the given acceptance probability b ą mint1, q{pu, we can rewrite:b " mint1, q `ϵ p u, for some ϵ ą 0.We consider the following two decoding options for P:• Baseline: select P to be suboptimal to (28) , i.e. simply set P :" q, which is target distribution; 10 We call this method Decoding-UNO. • Set P :" P ˚to be the optimal solution of ( 28), whose solution is presented in Theorem 4.We call this method Decoding-OPT.Measuring performance. Instead of TV distance, we measure the quality via WinRate in our experiment. Concretely, for each prompt, we let Decoding-UNO and Decoding-OPT to generate responses independently, and use score models to compare whose generation has higher quality. We specify draft model p as pythia-70m and target model q as pythia-2.8b from EleutherAI [7] . We apply the score model to be RM-Mistral-7B or GPT-4. We test 200 prompts from Alpaca-Farm-Eval Dataset [13] with 500 responses/comparisons per prompt. For a given prompt, Decoding-OPT wins if for more than 250 comparisons, score model prefer its response 11 over Decoding-UNO. Then the WinRate for each method is computed as #wins{200.Sanity Check for the experiment. To validate having smaller distance w.r.t. the large model q (pythia-2.8b) indicates higher performance, we preform the WinRate test for decoding via pythia-70m only against decoding via pythia-2.8b only. Table 2 shows that there is a significant performance gap between large model and small model, therefore validate the legitimacy of the experiment in Table 1 .Method RM-Mistral-7B GPT-4 pythia-2.8b 64.5% 69% pythia-70m35.5% 31% Notice that for the biased speculative decoding with bpxq " mint1, qpxq`ϵ ppxq u, Apxq in Theorem 4 can equivalently expressed as Apxq :" maxtqpxq´ppxq,´ϵu ř x maxtqpxq´ppxq,´ϵu , and we can select P ˚:" rAs `(recall r¨s ìn Section 2.1), which satisfies P ˚|A´p ¨q " 0; 0 ď P ˚|A`p ¨q ď Ap¨q.To implement Decoding-UNO, we modify the speculative sampling function in HuggingFace transformers/generation/utils.py file as follows 12 (where variable eps is ϵ in Table 1 ). This is conducted in a single A100 GPU. Justification: We study the theoretical properties for Speculative decoding, which is identical what is stated in abstract.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: See Appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Justification: Yes, we followed that.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: Our work is most theoretical.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: This paper has no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [NA] Justification: All the materials are open-sourced.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.˚Correspondence to: my0049@princeton.edu, mengdiw@princeton.edu.We use the abstraction x0 :" py1, y2, . . . , yn prompt q with yi's being the tokens in the prompt.In general, T is a random variable. However, we can append [EOS] tokens to keep the output length fixed.This is from their equation (1) that has 1{p1 ´αq with α " Epminpp, qqq and 1 ´Epminpp, qqq " ErTVpp, qqs.We mention batch drafting in parallel will cause more memory and arithmetic operations, but not computation time.We mention TV distance is only one metric for measuring distribution bias. In general, any distribution distance can be considered. We leave a thorough study for all other distances as future work.DiscussionsOn the optimality for batch algorithms. Unlike their non-batch counterparts, our batch algorithm studies do not come with optimality guarantees. This is largely due to the diverse and arbitrary nature of batch algorithm designs, making it challenging to define a comprehensive class that encompasses a wide range of batch algorithms. While[37] investigate optimal batch algorithms through an optimal transport lens, their work does not extend to calculating optimal rejection rates or developing an efficient algorithm to achieve this (they only propose an approximate solution). Consequently, the pursuit of batch optimality remains an open field. Identifying the optimal batch algorithm could yield valuable insights for enhancing practical applications in real-world scenarios.7 Here the rejection probability is computed as Pprejectq " ř x Ppreject, xq "It needs to follow bn P r0, 1s since bn is a probability.To be rigorous, we mention we didn't choose P :" rq ´ps`as the baseline since P :" rq ´ps`might fall in the optimal solution sets of P ˚defined in Theorem 4.We mention for RM-Mistral-7B it output values, then the response with higer value wins. For GPT-4, it outputs preference.Note the HuggingFace code uses p as target model and q as the draft model, which is different from us. 
 
 Appendix:  


Decision: This paper provides detailed analysis of speculative decoding and batch speculative decoding. The conclusions of the paper are: (1) speculative decoding is unbiased and it shows the expected rejection rate; (2) speculative decoding has the lowest rejection rate in all the unbiased algorithm that belongs to the family defined in Algorithm 2; (3) batch speculative decoding has lower rejection rate than speculative decoding; (4) it analyzes the trade-off between efficiency and effectiveness of the family of algorithm defined in Algorithm 2.
The paper is well written and easy to understand. The theoretical analysis presented by the authors provide several interesting insights about the inference efficiency observed by speculative decoding. All the results are backed up with simulation experiments, which strengthen the results presented in the paper. After the author rebuttal, the only major remaining ask would be to see improvements in batch speculative sampling in real-world scenarios which is currently impossible.

Reviews: review: paper_summary: paper_summary: This paper provides detailed analysis to speculative decoding and batch speculative decoding. The conclusions of the paper are: (1) speculative decoding is unbiased and it shows the expected rejection rate; (2) speculative decoding has the lowest rejection rate in all the unbiased algorithm that belongs to the familty defined in Algorithm 2; (3) batch speculative decoding has lower rejection rate than speculative decoding; (4) it analyzes the trade-off between efficiency and effectiveness of the family of algorithm defined in Algorithm 2.
strength_weakness: strength_weakness: 1. Although the paper provides lots of theoretical analysis. But I find only Theorem 4 and 5 are somewhat interesting. Theorem 1 is already derived in the original speculative decoding paper. For Theorem 2, although speculative decoding is proven to be optimal in the family of algorithms defined in Algorithm 2, but I don't think there are a lot of existing algorithms can be formulated in Algorithm 2. In fact, is there any algorithm that belongs to Algorithm 2 and is unbiad and it not speculative decoding? For Theorem 3, the finding that batch speculative decoding has lower rejection rate than vanilla speculative decoding is not surprising. 

2. Although Theorem 4 and Theorem 5 are interesting, it only solves half of the problem: given b, what should P be. It would be better if the authors could also discuss the design of b.

3. I think the paper can also be improved if the authors could summarize a new speculative algorithm from Theorem 4 and 5 and running experiments to compare with vanilla speculative decoding.
questions: questions: see weakness above
limitations: limitations: see weakness above
paper_summary: paper_summary: The author aim to develop theoretical understanding of speculative decoding. The authors assume that given a large and small model participating in speculative decoding, the computation complexity of the small model is negligible. Under this assumption, they characterize the expected rejection rate of speculative decoding. They show that this bound depends on the total variation distance between the generations from small and large model. Next, the authors show that spectral decoding gives optimal rejection bounds in class of all rejection based methods. Motivated by recent works analyzing batch speculative decoding, where the rejection is done only if M tokens are rejected from a given sample. Finally, given an acceptance probability, the authors show an optimal solution solution to the total variation loss between the distributions of large model and the one found by speculative decoding. This objective changes linearly with the rejection probability. This provides insights on selecting the optimal value of rejection threshold as per requirement. The presented theoretical results are backed up with appropriate experiments validating them.
strength_weakness: strength_weakness: 1) It is not completely clear, why making the assumption about negligible compute of the small model is not a strong assumption. Since the small model needs to generate the tokens autoregressively therefore even though its single pass could be small as compared to the larger model but it the context length is high i.e. several autoregressive passes are made, the compute of small model might not be negligible. It would be great if the authors can provide some empirical evidence to justify this assumption.

2) It would have been great if the authors provided evidence using real world models in support of their theory. Although, this is not a major weakness, but authors should consider it in the camera ready version.
questions: questions: I request the authors to kindly address the questions in the weaknesses section.
limitations: limitations: Yes, the authors have addressed the limitations.
paper_summary: paper_summary: The paper presents a theoretical perspective on speculative sampling. Through Theorems 1 and 2, the authors demonstrate that the sampling method employed by speculative sampling is optimal and unbiased. Subsequently, Theorem 3 introduces a multi-candidate approach to enhance the acceptance rate of speculative sampling.
strength_weakness: strength_weakness: The experiments are not sufficient. I would like to see improvements in batch speculative sampling in real-world scenarios.

I am curious if batch speculative sampling can be combined with tree-style methods, e.g., [1] CaPE and [2] Medusa?

[1] Du, C., Jiang, J., Yuanchen, X., Wu, J., Yu, S., Li, Y., ... & You, Y. GliDe with a CaPE: A Low-Hassle Method to Accelerate Speculative Decoding. In Forty-first International Conference on Machine Learning.
[2] Cai, T., Li, Y., Geng, Z., Peng, H., Lee, J. D., Chen, D., & Dao, T. (2024). Medusa: Simple LLM inference acceleration framework with multiple decoding heads. In Forty-first International Conference on Machine Learning.
questions: questions: see weakness
limitations: limitations: na
paper_summary: paper_summary: This paper presents a theoretical study on speculative decoding, an efficient inference method for large autoregressive models. It highlights practical implications, proposing a Pareto-optimal solution for the rejection-distribution bias tradeoff.
strength_weakness: strength_weakness: - The main figure does not clearly communicate the core concept of speculative decoding. It might lead readers to believe that speculative decoding primarily addresses hallucination, which is not its main advantage.
 - The experimental results are not distinctly highlighted, and the authors do not explain how these results support their theoretical analysis. While the theoretical contributions are significant, the paper would benefit from more extensive empirical validation.
questions: questions: NA
limitations: limitations: The results may not guarantee optimality in practical situations because real-world circumstances are more complex and varied than those considered in the theoretical analysis.

Comments: title: Reply to Reviewer GGsr
Reply to Reviewer qBZ2
Reply to reviewer Ceup
comment: I thank the authors for their rebuttal. I acknowledge that I have read their reply. This increases my confidence in this work.
Dear reviewer GGsr,

Thank you for checking our rebuttal. To be academically rigorous, we do agree with the reviewer that our proposed batch algorithm 4 is different from the existing tree spec such as SpecInfer [Miao et al.], CAPE and Medusa since our setup (alg4) is suitable for the theoretical understanding (batch improvement in Theorem 3 and Figure 2 Right) for the efficiency of spec decoding, which is the main focus of the paper and this theory is lacking in the existing batch tree spec algorithms. We promise to discuss it carefully in the revision and will mention the limitation that it remains unclear on how to combine batch algorithm 4 with other tree-spec algorithms as you suggested. Thank you again for the great point.  

Best, Authors
Dear reviewer qBZ2,

Thank you for checking our rebuttal. Please let us know if you have any final questions.

Best, Authors
Due to this method can not be used in the real setting like tree spec, I still maintain my initial score.
Thank you for your response. I would like to remain the score.
Dear reviewer Ceup,

Thank you for spending time reading our rebuttal. We authors are still available here to answer your questions in the next two days in case you have any last-minute questions.

Best, Authors
I appreciate the author’s efforts in providing the rebuttal. After reading the rebuttal and other reviewer’s opinions, I decided to raise my score.