Title: A Theoretical Perspective for Speculative Decoding Algorithm
Venue: NeurIPS 2024
Accepted: yes
Keywords: probabilistic analysis

theory for large language models
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.

A Theoretical Perspective for Speculative Decoding Algorithm

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.

1 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.

Arecent 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?

His first Final was in Rio, Brazil, 2014.

Q: When and

where was

Final?

His first Final was in Lusail, Qatar, 2024.

Q: When and

where was

World Cup

Final?

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 E r N rej s for Speculative Decoding (Theorem 1). This renders a theoretical reference for understanding the acceleration rate T { E r N 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.

1.1 Related works

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

'

min

t

1

,

q

t

p

t

u

and the distribution

P

t

' r

q

t

´

p

t

s

`

in

2:

Require

:

p

t

p¨q

:

'

p

t

p¨|

x

1:

n

´

1

,

˜

x

n

:

t

´

1

q

,

q

t

p¨q

:

'

q

t

p¨|

x

1:

n

´

1

,

˜

x

n

:

t

´

1

q

.

@

t

ě

n

.

3:

for

t

'

n

:

T

do

4:

Sample

r

'

Uniform

r

0

,

1

s

.

5: if r ď min 1 , q t p ˜ x t q p t p ˜ x t q then

!

)

6:

Accept with

x

n

'

˜

x

t

.

n

Ð

n

`

1

.

7:

else

8:

Sample

x

n

' r

q

t

´

p

t

s

`

p¨q

.

9:

n

Ð

n

`

1

. Break.

10:

r

q

t

´

p

t

s

`

in Section 2.1

11:

end if

12:

end for

Algorithm 2

Framework for Rejection-based Decoding

1: Init : Horizon T , models q t , p t . Lookahead K ' T . Prompt x 0 . n ' 0 .

2:

Require

: Probability

b

t

, distribution

P

t

.

3:

while

n

ă

T

do

4:

for

t

'

n

:

T

do

5:

Sample drafts

˜

x

t

'

p

t

p¨|

x

1:

n

´

1

,

˜

x

n

:

t

´

1

q

.

6:

end for

7: Obtain the target logits in parallel for ˜ x n : T as q n p¨| x 1: n ´ 1 q , . . . , q T p¨| x 1: n ´ 1 , ˜ x n : T q .

8: for t ' n : T do

9:

Accept

x

n

'

˜

x

t

with prob.

b

t

.

n

Ð

n

`

1

.

10:

Else REJECTION:

11:

Sample

x

n

'

distribution

P

t

.

n

Ð

n

`

1

.

Break.

12:

end for

13:

end while

2 Preliminaries

2.1 Background for decoding problems

In 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 p x 0: t ` 1 | x 0: t q : ' p p x 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 r q ´ p s ` to denote the normalized distribution for max t 0 , q p x q ´ p p x qu , @ x P V with V being the token vocabulary. We also denote ¯ E x ' f r g s : ' ř x f p x q g p x q .

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 ˜ x n : n ` K ´ 1 via sampling from p p¨| x 1: n ´ 1 , ˜ x n : n ` i q , i P r K ´ 1 s ;

Conditional score computation : given ˜ x n : n ` K ´ 1 , computing the logits of the K tokens q p¨| x 1: n ´ 1 , ˜ x n : n ` i q , i P r K ´ 1 s in parallel ;

Token validation : Accept the candidate token ˜ x t (as x n ) with some probability 0 ď b t ď 1 . If accepted, continue to validate the next candidate ˜ x t ` 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 p ˜ x t q : ' min t 1 , q t p ˜ x t q p t p ˜ x t q u and distribution P t p¨q : ' r q t ´ p t s ` p¨q (see Algorithm 1 and Figure 1).

2.2 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 q p¨| x 1: n ´ 1 , ˜ x n : n ` i q , i P r T ´ n s . 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 ˜ x n : 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 O p 1 q .

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 ˜ x n : n ` K ´ 1 , we are allowed to sample until the end, i.e. ˜ x n : 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.

3 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 t 0 , 1 u 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):

(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 p x 1: T q ' q p x 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 E r TV p p n , q n qp¨| x 1: n ´ 1 qs ' 0 for all n , all tokens are accepted and the accelerate rate ' T ;

If E r TV p p 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 E r TV p p 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 { E r TV p p, q qs for K ' 8 . 4 Their result equals T { ř T n E r TV p p n , q n qp¨| x 1: n ´ 1 qs when E r TV p p 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. Wealso 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.

3.1 Optimality of Speculative Decoding

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 : ' p p, q q , where p, q stand for the distributions of the draft model and the target model respectively, defining the family of algorithms as F : ' t A : A is a specification of Algorithm 2 that satisfies P A t ' q t @ t p i.e., unbiased qu . For an algorithm A , denote N rej as the number of rejections. Then we have the lower bound

16.8

14.4

Average Rejections

16.7

Theoretical E[TV(pn，qn)]

n =1

Expected Rejections

Expected Rejections

16

14.2

16.6

15

16.5

14.0

14

16.4

13.8

Batch SpecDecode M = 4

True Value, batch M = 4

Batch SpecDecode M = 5

13

16.3

True Value, batch M = 5

13.6

0

50

100

150

200

0

50

100

150

200

0

500

1000

1500

2000

Number of runs N × 100

Number of runs N × 100

Number of Batches M

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 min p p p x q , q p x qq , which corresponds to the optimal expected rejection TV p p, q q . 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 .

3.2 Analysis for Batch Speculative Decoding

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 O p 1 q ; 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.

What are important components of AI?

…

-

&'(

ine

ing

learn

deep

learn

ing

rein

force

ment

learn

atten

tion

nism

What are important components of AI?

ine

learn

❌

❌

❌

❌

-

"

❌

Large Model

Decoder

Verify

Large Model

Decoder

Verify

✅

✅

❌

ing

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 ' r q m ´ p s ` , @ m P r M s with q 1 ' q being the target distribution. Then, for Batch Speculative Decoding 4, P Batch p x 1: T q ' q p x 1: T q . @ x 1: T P V T . Moreover,

where f p x 1: n q : ' P p x 1: n Xt n -th draft token rejected uq . f can be iteratively computed via p, q .

Takeaways. The expected rejection of Batch Decoding composes of two parts: (i) ř T n E q r TV r q, p ss 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 BI p q, p q : ' TV r q, p s ´ ś M m ' 1 TV r q m , p s 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 BI p Unif p V 1 q , Unif p V qq ' p 1 ´ 1 r q ´ p 1 ´ 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 ' Ber p u q , q ' Ber p v q . Then BI p Ber p v q , Ber p u qq ' | u ´ v | ¨ p 1 ´ u M ´ 1 q . Its pattern is exhibited in the upper right panel of Figure 3. Notice for both cases, the limiting batch improvement is TV p p, q q , 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.

Batch Simulation. The Middle panel of Figure 2(b) shows Batch Speculative Decoding simulations with batch M ' 4 , 5 . The green/purple lines are the empirical average rejections simulated via Algorithm 4 and the orange/pink lines are the theoretical values computed via Theorem 3. The right panel of Figure 2(c) plots the rejection vs. batch size. In particular, the the black dot with coordinate p 1 , 16 . 41 q represents the Speculative Decoding 1. The numeric example shows when M Ñ 8 , E r N rej s Û 0 . Intuitively, this is partial due to, if the distribution mismatch between p and q is large,

Rejection Probability

Distribution Distance TV

[ℙ

.

, *]

0

TV

[,, *]

TV

[,, *]

Attainable Region

Unattainable Region

0.4

probability

0.4

RejectionProb

0.3

TV[pAlg, q]

0.3

jection

0.2

0.2

0.1

0.1

0.0

0.0

0.0

0.1

0.2

0.3

0.4

0.0

0.1

0.2

0.3

0.4

Quality decrease measured by TV(pAlg,q)

OverAcceptanceThresholdε

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.

4 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?

4.1 An Optimization Formulation and Pareto-optimal Characterization

We measure the distribution bias between P A and q via TV distance. 6 Concretely, for any algorithm A : ' p b, P q in 2 with acceptance probability b and rejection distribution P , we fix b and aim to optimize the following objective

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

Suppose ř x r 1 ´ b p x qs p p x q ą 0 . Define A p x q : ' q p x q´ b p x q p p x q ř ˜ x r 1 ´ b p ˜ x qs p p ˜ x q , and the sets A ` ' t x P V : A p x q ě 0 u , A ´ ' t x P V : A p x q ă 0 u . Then the set of optimal distributions of objective (1) is characterized as t P ˚ : P ˚ | A ´ p¨q ' 0; 0 ď P ˚ | A ` p¨q ď A p¨qu , and the optimal value is Loss ˚ TV p b q ' 1 2 ř x | q p x q ´ b p x q p p x q| ´ 1 2 ř x p 1 ´ b p x qq p p x q ě 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 ď min t 1 , q { p u then A becomes a probability distribution and the optimal distribution P ˚ equals A which is also unbiased (Loss ˚ TV p b q ' 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 p b q ą 0 for Algorithm 2.

Main Takeaways. Via Theorem 4, we derive the pareto front (the optimal tradeoff) between rejection probability 7 P p reject q vs. distribution distance TV r P A , q s (Left panel of Figure 4). The blue region can be realized by some algorithm 2, and the red region cannot. p 0 , 0 q is the 'perfect algorithm' (no rejection, no bias) which does not exists, and, in particular, p 0 , TV r p, q sq stands for Speculative Decoding. Surprisingly, the pareto front is a straight line connecting p 0 , TV r p, q sq and p TV r p, q s , 0 q , 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 min t 1 , q p x q p p x q u ď b p x q ď 1 , @ x P V . Then

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 b p x q ' min t 1 , q p x q` ϵ p p x q u . The orange line in p c q and the green boundary in p b q are computed via Loss ˚ TV p b q 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.

4.2 Experiment

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 ą min t 1 , q { p u , 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.

Table 1: WinRate for Decoding-OPT vs Decoding-UNO with different over-acceptance threshold ϵ . The acceptance probability b p x q ' min t 1 , q p x q` ϵ p p x q u .

5 Discussions

On 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.

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.

Acknowledgments

The authors would like to thank anonymous reviewers for their valuable feedback. Mengdi Wang acknowledges the support by NSF IIS-2107304, NSF CPS-2312093, and ONR 1006977.

References

Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 , 2023.

Kwangjun Ahn, Ahmad Beirami, Ziteng Sun, and Ananda Theertha Suresh. Spectr++: Improved transport plans for speculative decoding of large language models. In NeurIPS 2023 Workshop Optimal Transport and Machine Learning , 2023.

Anurag Arnab, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Luˇ ci´ c, and Cordelia Schmid. Vivit: A video vision transformer. In Proceedings of the IEEE/CVF international conference on computer vision , pages 6836-6846, 2021.

Sangmin Bae, Jongwoo Ko, Hwanjun Song, and Se-Young Yun. Fast and robust early-exiting framework for autoregressive language models with synchronized parallel decoding. arXiv preprint arXiv:2310.05424 , 2023.

Benjamin Bergner, Andrii Skliar, Amelie Royer, Tijmen Blankevoort, Yuki Asano, and Babak Ehteshami Bejnordi. Think big, generate quick: Llm-to-slm for fast autoregressive decoding. arXiv preprint arXiv:2402.16844 , 2024.

Nikhil Bhendawade, Irina Belousova, Qichen Fu, Henry Mason, Mohammad Rastegari, and Mahyar Najibi. Speculative streaming: Fast llm inference without auxiliary models. arXiv preprint arXiv:2402.11131 , 2024.

Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning , pages 2397-2430. PMLR, 2023.

Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Jasmine Hsu, et al. Rt-1: Robotics transformer for real-world control at scale. arXiv preprint arXiv:2212.06817 , 2022.

F Warren Burton. Speculative computation, parallelism, and functional programming. IEEE Transactions on Computers , 100(12):1190-1193, 1985.

Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean-Baptiste Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.01318 , 2023.

Ziyi Chen, Xiaocong Yang, Jiacheng Lin, Chenkai Sun, Jie Huang, and Kevin ChenChuan Chang. Cascade speculative drafting for even faster llm inference. arXiv preprint arXiv:2312.11462 , 2023.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 , 2020.

Yann Dubois, Chen Xuechen Li, Rohan Taori, Tianyi Zhang, Ishaan Gulrajani, Jimmy Ba, Carlos Guestrin, Percy S Liang, and Tatsunori B Hashimoto. Alpacafarm: A simulation framework for methods that learn from human feedback. Advances in Neural Information Processing Systems , 36, 2023.

Yichao Fu, Peter Bailis, Ion Stoica, and Hao Zhang. Break the sequential dependency of llm inference using lookahead decoding. arXiv preprint arXiv:2402.02057 , 2024.

Freddy Gabbay and Avi Mendelson. Speculative execution based on value prediction . TechnionIIT, Department of Electrical Engineering, 1996.

Mukul Gagrani, Raghavv Goel, Wonseok Jeon, Junyoung Park, Mingu Lee, and Christopher Lott. On speculative decoding for multimodal large language models. arXiv preprint arXiv:2404.08856 , 2024.

Kai Han, Yunhe Wang, Hanting Chen, Xinghao Chen, Jianyuan Guo, Zhenhua Liu, Yehui Tang, An Xiao, Chunjing Xu, Yixing Xu, et al. A survey on vision transformer. IEEE transactions on pattern analysis and machine intelligence , 45(1):87-110, 2022.

Zhenyu He, Zexuan Zhong, Tianle Cai, Jason D Lee, and Di He. Rest: Retrieval-based speculative decoding. arXiv preprint arXiv:2311.08252 , 2023.

Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, et al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210.02303 , 2022.

Kaixuan Huang, Xudong Guo, and Mengdi Wang. Specdec++: Boosting speculative decoding via adaptive candidate lengths. arXiv preprint arXiv:2405.19715 , 2024.

Wonseok Jeon, Mukul Gagrani, Raghavv Goel, Junyoung Park, Mingu Lee, and Christopher Lott. Recursive speculative decoding: Accelerating llm inference via sampling without replacement. arXiv preprint arXiv:2402.14160 , 2024.

Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361 , 2020.

Sehoon Kim, Karttikeya Mangalam, Suhong Moon, Jitendra Malik, Michael W Mahoney, Amir Gholami, and Kurt Keutzer. Speculative decoding with big little decoder. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. In International Conference on Machine Learning , pages 19274-19286. PMLR, 2023.

Xiaoxuan Liu, Lanxiang Hu, Peter Bailis, Ion Stoica, Zhijie Deng, Alvin Cheung, and Hao Zhang. Online speculative decoding. arXiv preprint arXiv:2310.07177 , 2023.

Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang, Rae Ying Yee Wong, Zhuoming Chen, Daiyaan Arfeen, Reyna Abhyankar, and Zhihao Jia. Specinfer: Accelerating generative llm serving with speculative inference and token tree verification. arXiv preprint arXiv:2305.09781 , 2023.

Eric Mitchell, Rafael Rafailov, Archit Sharma, Chelsea Finn, and Christopher D Manning. An emulator for fine-tuning large language models using small language models. arXiv preprint arXiv:2310.12962 , 2023.

Jie Ou, Yueming Chen, and Wenhong Tian. Lossless acceleration of large language model via adaptive n-gram parallel decoding. arXiv preprint arXiv:2404.08698 , 2024.

Haifeng Qian, Sujan Kumar Gonugondla, Sungsoo Ha, Mingyue Shang, Sanjay Krishna Gouda, Ramesh Nallapati, Sudipta Sengupta, Xiaofei Ma, and Anoop Deoras. Bass: Batched attentionoptimized speculative sampling. arXiv preprint arXiv:2404.15778 , 2024.

Andrea Santilli, Silvio Severino, Emilian Postolache, Valentino Maiorca, Michele Mancusi, Riccardo Marin, and Emanuele Rodol` a. Accelerating transformer inference for translation via parallel decoding. arXiv preprint arXiv:2305.10427 , 2023.

Mohit Shridhar, Lucas Manuelli, and Dieter Fox. Perceiver-actor: A multi-task transformer for robotic manipulation. In Conference on Robot Learning , pages 785-799. PMLR, 2023.

Benjamin Spector and Chris Re. Accelerating llm inference with staged speculative decoding. arXiv preprint arXiv:2308.04623 , 2023.

Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit. Blockwise parallel decoding for deep autoregressive models. Advances in Neural Information Processing Systems , 31, 2018.

Qidong Su, Christina Giannoula, and Gennady Pekhimenko. The synergy of speculative decoding and batching in serving large language models. arXiv preprint arXiv:2310.18813 , 2023.

Hanshi Sun, Zhuoming Chen, Xinyu Yang, Yuandong Tian, and Beidi Chen. Triforce: Lossless acceleration of long sequence generation with hierarchical speculative decoding. arXiv preprint arXiv:2404.11912 , 2024.

Ziteng Sun, Jae Hun Ro, Ahmad Beirami, and Ananda Theertha Suresh. Optimal block-level draft verification for accelerating speculative decoding. arXiv preprint arXiv:2403.10444 , 2024.

Ziteng Sun, Ananda Theertha Suresh, Jae Hun Ro, Ahmad Beirami, Himanshu Jain, and Felix Yu. Spectr: Fast speculative decoding via optimal transport. arXiv preprint arXiv:2310.15141 , 2023.

Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11805 , 2023.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timoth´ ee Lacroix, Baptiste Rozi` ere, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023.

Siqi Wang, Hailong Yang, Xuezhu Wang, Tongxuan Liu, Pengbo Wang, Xuning Liang, Kejie Ma, Tianyu Feng, Xin You, Yongjun Bao, et al. Minions: Accelerating large language model inference with adaptive and collective speculative decoding. arXiv preprint arXiv:2402.15678 , 2024.

Heming Xia, Tao Ge, Peiyi Wang, Si-Qing Chen, Furu Wei, and Zhifang Sui. Speculative decoding: Exploiting speculative execution for accelerating seq2seq generation. In Findings of the Association for Computational Linguistics: EMNLP 2023 , pages 3909-3925, 2023.

Heming Xia, Zhe Yang, Qingxiu Dong, Peiyi Wang, Yongqi Li, Tao Ge, Tianyu Liu, Wenjie Li, and Zhifang Sui. Unlocking efficiency in large language model inference: A comprehensive survey of speculative decoding. arXiv preprint arXiv:2401.07851 , 2024.

Daliang Xu, Wangsong Yin, Xin Jin, Ying Zhang, Shiyun Wei, Mengwei Xu, and Xuanzhe Liu. Llmcad: Fast and scalable on-device large language model inference. arXiv preprint arXiv:2309.04255 , 2023.

Minghao Yan, Saurabh Agarwal, and Shivaram Venkataraman. Decoding speculative decoding. arXiv preprint arXiv:2402.01528 , 2024.

Sen Yang, Shujian Huang, Xinyu Dai, and Jiajun Chen. Multi-candidate speculative decoding. arXiv preprint arXiv:2401.06706 , 2024.

Yongchao Zhou, Kaifeng Lyu, Ankit Singh Rawat, Aditya Krishna Menon, Afshin Rostamizadeh, Sanjiv Kumar, Jean-Franc ¸ois Kagy, and Rishabh Agarwal. Distillspec: Improving speculative decoding via knowledge distillation. arXiv preprint arXiv:2310.08461 , 2023.

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 , ˜ x n with 0 ď b n p ˜ x n , x 1: n ´ 1 q ď 1 . 8 Based on this, we can further define the new function ϵ n : V ˆ V n ´ 1 ÞÑ R according to the following equation:

Indeed, we can choose ϵ n : ' b n ¨ p n ´ q n , and the validity of the definition is guaranteed by the Lemma 1. Then, the acceptance probability b n ą min t 1 , q n p n u implies ϵ n ą 0 .

Next, we show for any A P F , it must satisfy b n ď min t 1 , q n p n u for all n . To this end, we design the two token sets V ` ' t x : D n s.t. ϵ n p x q ą 0 u and V ´ ' t x : D n s.t. ϵ n p x q ď 0 u and prove V ` ' H , V ´ ' V .

Finally, by Lemma 2 in Appendix, any algorithm satisfies b n ď min t 1 , q n p n u , @ n P r T s must have E A P r N rej s ě ř T n ' 1 E q r TV p p 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.

A.2 High Level Proof Sketch for the second part of Theorem 3

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 p ˜ x n acc , ˜ x n ' x n | x 1: n ´ 1 q . Let ˜ x n ´ 1 ' p p¨| x 1: n ´ 2 q , then there are two cases: ˜ x n ´ 1 accepted or rejected. We have

In the process of finding p b , we need to compute the the quantity f p x 1: n q : ' P p x 1: n X t n -th draft token rejected uq and it can be recursively computed via (22) using p, q .

A.3 Proof sketch for Theorem 4

Due to space constraint, we only summarize the high-level proof ideas for Theorem 4. Since ř x A p x q ' 1 , the original objective (1) can be equivalently rewritten as

We now find the solution of objective (4) in two steps.

Step1: Recall A ` ' t x P V : A p x q ě 0 u and A ´ ' t x P V : A p x q ă 0 u , then any optimal P ˚ must satisfy P ˚ p x q ' 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 | A p x q ´ P ˚ p x q| ě ∥ 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 ˚ .

Appendix

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 ' q p¨| 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 t 0 , 1 u 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,

Here TV denote the TV distance between two distributions.

Proof. Given the verified the tokens x 1: n ´ 1 , we first compute P p Reject 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

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).

Lastly, by law of total expectation and above

Here the fourth equal sign comes from Speculative Decoding keep the distribution q (Proposition ?? ), the fifth equal sign comes from the event t R n ' 1 u ' t Reject at n u .

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 p x 1: T q ' q p x 1: T q .

Proof. In particular, we use induction to show the stronger result that @ t P r T s , @ x 1 , x 2 , . . . , x t P V , it holds P A t p x 1 , x 2 , . . . , x t q ' q t p x 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 p x 1 | x 0 q ' q 1 p x 1 | x 0 q , which further implies P A 1 p x 1 q ' q 1 p x 1 q (since the distribution of x 0 is independent of p, q ).

Step2: assume P A t p x 1 , x 2 , . . . , x t q ' q t p x 1 , x 2 , . . . , x t q , we first show P A p x t ` 1 | x 1: t q ' q p x t ` 1 | x 1: t q . Indeed, let ˜ x t ` 1 ' p p¨| x 1: t q , then by law of total probability

By Algorithm 1,

Next, the probability of rejection is:

where the second equal sign comes from (6). Lastly, by the construction of the algorithm,

Combining (8) and (7) yields

Plugging (6) and the above equation into (5) to obtain

Finally, applying the above we obtain

where the last equal sign uses the induction hypothesis.

C Lower Bound

Theorem 8 (Restatement of Theorem 2) . Define the arbitrary instance P : ' p p, q q , and define the family of algorithms as

F : ' t A : A is a specification of Algorithm 2 that satisfies P A t ' q t @ t p i.e., distribution unbiased qu . For an algorithm A , denote N rej as the number of rejections. Then we have

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 , ˜ x n with 0 ď b n p ˜ x n , x 1: n ´ 1 q ď 1 . 9 Based on this, we define the new function ϵ n : V ˆ V n ´ 1 ÞÑ R according to the following equation:

Indeed, we can choose ϵ n : ' b n ¨ p n ´ q n , 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 have

On the other hand, let ˜ x n ' p n p¨| x 1: n ´ 1 q , the decomposition holds

where the third equal sign uses P A n p ˜ x n ' x n | x 1: n ´ 1 q ' p n p x n | x 1: n ´ 1 q and the last equal sign is due to

Let's do some further simplifications. First off,

Plug (12) and (13) into (11), and then plug (11) into (10) to obtain

Now, we define the token set V ` ' t x : D n, x 1: n ´ 1 s.t. ϵ n p x, x 1: n ´ 1 q ą 0 u and similarly the token set V ´ ' t x : D n, x 1: n ´ 1 s.t. ϵ n p x, x 1: n ´ 1 q ď 0 u . Next, we show V ` ' H , and V ´ ' V .

and this implies 0 ě ϵ n p x n , x 1: n ´ 1 q , which means x n P V ´ ;

Case2.

q

n

If

p

p

x

n

n

|

p

x

x

n

1:

n

|

x

1:

1

´

n

´

q '

1

q ă

p

n

p

q

x

n

n

|

p

x

x

n

1:

|

n

x

´

1:

1

n

q

ě p n p x n | x 1: n ´ 1 q , and this implies

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 ď min t 1 , q n p n u . Finally, by Lemma 2, this implies E A P r N rej s ě C p P q . Since A P F is arbitrary, this concludes the proof.

Corollary 1. For any algorithm A P F , it follows @ n P r T s , x P V and x 1: n ´ 1 , b n p x, x 1: n ´ 1 q ď min ! 1 , q n p x | x 1: n ´ 1 q p n p x | x 1: n ´ 1 q ) . In this case, the distribution P n is defined as:

Applying ϵ n ' b n p n ´ q n , this is equivalent to

Proof of Corollary 1. We reutilize (14) here and call it (15).

By two cases discussion as in the proof of Theorem 8, we have @ A P F , it follows @ n P r T s , x P V and x 1: n ´ 1 , b n p x, x 1: n ´ 1 q ď min ! 1 , q n p x | x 1: n ´ 1 q p n p x | x 1: n ´ 1 q ) . Then we can directly solve (15) to obtain

Lastly, we verify such a P A n p x n | ˜ x n reject , x 1: n ´ 1 q is a valid distribution. First of all, since ϵ n ' b n ¨ p n ´ q n , then b n p x, x 1: n ´ 1 q ď min ! 1 , q n p x | x 1: n ´ 1 q p n p x | x 1: n ´ 1 q ) implies ϵ n p x, x 1: n ´ 1 q ď 0 , and this further implies

1

´

q `

ϵ

n

p

x

n

, x

1:

n

1

´

q

, then by (14) we have

which implies P A n p x n | ˜ x n reject , x 1: n ´ 1 q ě 0 . Second,

This concludes the proof.

Lemma 1. For any 0 ď b n p ˜ x n , x 1: n ´ 1 q ď 1 , there exists ϵ n p ˜ x n , x 1: n ´ 1 q P R such that (9) holds true.

Proof of Lemma 1. Indeed, set ϵ n ' b n ¨ p n ´ q n , then

where the first equal sign uses that ˜ x n is sampled from p n p¨| x 1: n ´ 1 q so p n p ˜ x n | x 1: n ´ 1 q ą 0 , and the second equal sign uses 0 ď b n ď 1 .

Lemma 2. For any instance P ' p p, q q , let F : ' t A : A is a realization of Algorithm 2 s.t. P A t ' q t @ t p i.e., unbiased qu . Suppose there is an A P F such that E A P r N rej s ă C p P q , then there exists a b n in Line 8 of Template 2 such that D x, x 1: n ´ 1

Proof of Lemma 2. Suppose for all n, x, x 1: n ´ 1 , b n p x | x 1: n ´ 1 q ď min ! 1 , q n p x | x 1: n ´ 1 q p n p x | x 1: n ´ 1 q ) . We define random variables R n P t 0 , 1 u indicating whether the n -th token is rejected ( 1 is rejected). Therefore, the expected number of rejection is

Here the second to last equality comes from the tower property and the last equal signs is due to A is unbiased. Next, denote ˜ x n ' p n p¨| x 1: n ´ 1 q to be the candidate token, we have

and this implies

contradicts E A P r N rej s ă C p P q ! This concludes the proof.

Algorithm 4 Batch Speculative Sampling

D Batch Speculative Decoding

We split the proofs for unbiasedness and rejections in two parts.

Theorem 9 (Unbiasedness of Theorem 3) . Denote the Algorithm 4 as A B . For any sequence x 1: T (where x i P V ), we have

Proof. Step1: Let x 1: n ´ 1 be the accepted tokens up to n ´ 1 . We first show P A B n p x n | x 1: n ´ 1 q ' P LLM n p x n | x 1: n ´ 1 q @ x n P V .

We partition the generation of x n into two cases: (i). accepted as the first token of the m -th responses ( m ' 1 , . . . , M ), or rejected by all M responses and sampled by Line 28 of Algorithm 4; (ii). accepted/rejected as the t -th token of the m -th responses ( m ' 1 , . . . , M ) for t ě 2 .

For the second case. Similar to the standard speculative decoding, let ˜ x m n ' p n p¨| x 1: n ´ 1 q , then

By the construction of Algorithm 4 (Line 3), when x n is accepted/rejected as the t pě 2 q -th draft token of certain response, we have q n p x n | x 1: n ´ 1 q ' P LLM n p x n | x 1: n ´ 1 q . This gives P A B n p x n | x 1: n ´ 1 q ' P LLM n p x n | x 1: n ´ 1 q .

For the first case. This part of the proof largely follows Theorem 4.2 of [26]. In this case, the n -th generated token x n has M ` 1 possibilities: accepted at the m -th response or rejected by all M responses and sample at Line 28. Since the algorithm will iterate all M responses if not accepted, we denote q m n as the q n for the m -th response. Then we have the recursion

where q 1 n ' P LLM n and r m is the rejection probability satisfies

Denote E m ' t m -th response rejected u and E 1: m ' t 1 : m -th responses all rejected u , then

Denote P A B n p x n | x 1: n ´ 1 q : ' P A B n p x n | E 0 , x 1: n ´ 1 q , by similar calculation we have in general P A B n p x n | E 0: m ´ 1 , x 1: n ´ 1 q ' min t p n p x n | x 1: n ´ 1 q , q m n | p x n | x 1: n ´ 1 qu ` r m ¨ P A B n p x n | E 0: m , x 1: n ´ 1 q .

Next we prove P A B n p¨| E 0: m , x 1: n ´ 1 q ' q m ` 1 n p¨| x 1: n ´ 1 q @ m P t 0 , 1 , . . . , M u by backward induction. First of all, P A B n p¨| E 0: M , x 1: n ´ 1 q ' ' q M n ` 1 p¨| x 1: n ´ 1 q ´ p n ` 1 p¨| x 1: n ´ 1 q ‰ ` '

In particular, we take m ' 0 to obtain

Combining both cases we finish the proof of Step1.

Step2: For any n , first of all we have P A B p x 0 q ' P LLM p x 0 q since x 0 is the prompt. Suppose P A B n ´ 1 p x 1: n ´ 1 q ' P LLM n ´ 1 p x 1: n ´ 1 q , @ x 1: n ´ 1 , then by Step1

where the second equal sign is by induction and the third equal sign is by Step1. This finish the proof.

D.1 Expected Rejections for Batch Speculative Decoding (Proof for the second part of Theorem 3)

In this section, we derive the number of expected rejections using the batch speculative decoding. However, the analysis is more involved than the Algorithm 4 due to the parallel response structure, since, given the verified token x 1: n ´ 1 , the probability of n -th token being rejected does not possess a unified expression and depends on the location of x n ´ 1 . We detail this below.

We recall the notion in Theorem 1 that random variables R n P t 0 , 1 u indicates whether the n -th token is rejected ( 1 is rejected). Therefore, the expected number of rejection is

where the last equality comes from the tower property of expectation and we assume x 0 is a given initial token and x 1:0 ' t x 0 u . Then

In this scenario, we cannot obtain P A p n -th token rej | x 1: n ´ 1 q ' TV p p n p¨| x 1: n ´ 1 q , P LLM n p¨| x 1: n ´ 1 qq since the conditional rejection probability implicitly encodes the location of p n ´ 1 q -th token: whether ˜ x n ´ 1 ' p p¨| x 1: n ´ 1 q is rejected (at the root of the tree) or ˜ x n ´ 1 ' p p¨| x 1: n ´ 1 q is accepted (at the branch of the tree). To formalize this, given validated token x 1: n ´ 1 , we denote q m n p¨| x 1: n ´ 1 q to be the m -th rejection distribution, then by the construction of Algorithm 4 (Line 19),

Here r m ' ř x 1 max t 0 , q m n p x 1 | x 1: n ´ 1 q ´ p n p x 1 | x 1: n ´ 1 qu is normalizing factor and q 1 n ' P LLM . Let ˜ x n ' p p¨| x 1: n ´ 1 q , then P A p n -th token acc | x 1: n ´ 1 q ' P A p ˜ x n acc | x 1: n ´ 1 q . Next, we compute the quantity P A p ˜ x n acc | x 1: n ´ 1 q .

We begin by first considering P A p ˜ x n acc , ˜ x n ' x n | x 1: n ´ 1 q . Let ˜ x n ´ 1 ' p p¨| x 1: n ´ 2 q , then there are two cases: ˜ x n ´ 1 accepted or rejected. We have

For p a . Given that ˜ x n ´ 1 is accepted, ˜ x n ' x n can only be accepted as the t -th token within the certain response for t ě 2 . This is because ˜ x n ´ 1 is accepted and has to be at least the first token in the response. In this case, q n p¨| x 1: n ´ 1 q ' P LLM n p¨| x 1: n ´ 1 q ' q 1 n p¨| x 1: n ´ 1 q , then

For p b . Given that ˜ x n ´ 1 is rejected, ˜ x n ' x n can only be accepted as the first token within the certain response. This is because ˜ x n ´ 1 is rejected will restart the parallel tree. Then

where the first equal sign comes from: since ˜ x n ´ 1 rej implies x n represents the first token of the parallel tree, then it is identical to the proof of the first case of Step1 in Theorem 9. The second equal sign is from: ˜ x n is rejected means all M responses since x n is the first token of the tree. The conditional rejection probability

so by chain rule, the total rejection probability is ś M m ' 1 r m . Plug (18), (19) into (17) to obtain

which is equivalent to

where the first line uses P A p x n | x 1: n ´ 1 q ' q 1 n p x n | x 1: n ´ 1 q , the second equivalence uses Bayes rule, the third equivalence uses P A p x 1: n q ' q 1 p x 1: n q ' q 1 n p x 1: n q again. Now denote f p x 1: n q : ' P A p ˜ x n rej , x n | x 1: n ´ 1 q q 1 n p x 1: n ´ 1 q , and sum over x 1: n for the above to obtain

where by (21) pseudo-measure f satisfies @ x 1: n

Here we used r m ' ř x max t 0 , q m n p x | x 1: n ´ 1 q ´ p n p x | x 1: n ´ 1 qu ' TV p q m , p qp x 1: n ´ 1 q . Plug the above back to (16), we finally have

The formulation for f is iteratively obtained in (22).

D.2 Increasing batch size to inf doesn't help

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:

M Ñ8 , E r N 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

where the second equality is due to Batch algorithm is unbiased (Theorem 3).

Second item: by

since ś M m ' 1 TV p q m , p qp x ă n q Ñ 0 as M Ñ 0 (note TV p q m , p qp x ă n q ' 1 iff there is no overlap between q m and p ), it implies f 8 p x 1: n q ' h p x n | x ă n qr q 1 p x 1: n ´ 1 q ´ f 8 p x 1: n ´ 1 qs .

Third item: Similar to the second item, it holds true via taking M Ñ8 . For proving E r N 8 rej s ą 0 , suppose E r N 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 E r N 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)!

E Proofs of Section 4

E.1 Proof of Theorem 4

For the sampling scheme where the acceptance probability b n goes beyond min t 1 , q n p x | x 1: n ´ 1 q p n p x | x 1: n ´ 1 q 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

Under the above, via equation (14) the objective is equivalent to the following (note that according to Algorithm 2 P A n p x n | draft token rejected , x 1: n ´ 1 q ' P n p x n | x 1: n ´ 1 q is an algorithmic choice)

where we removed x 1: n ´ 1 for notation simplicity, and recall again ϵ n : ' b n p n ´ q n . When ř ˜ x r p n p ˜ x q ´ q n p ˜ x q ´ ϵ n p ˜ x qs ` ' 0 , the objective degenerates to the constant (in P n ) TV p q n , p n q . Therefore, for the rest of the section, we focus on the case where ř ˜ x r p n p ˜ x q ´ q n p ˜ x q ´ ϵ n p ˜ x qs ` ą 0 . We have the following Theorem that characterizes the solution of (23).

Theorem 10 (Restatement of Theorem 4) . Suppose ř x r p n ´ q n ´ ϵ n s ` p x q ą 0 . Define

and define the positive token set A ` ' t x P V : A n p x q ě 0 u and the negative token set A ´ ' t x P V : A n p x q ă 0 u . Then the set of optimal distributions of objective (23) is characterized as

and the optimal value is

Remark 5. In the main context (Theorem 4) we define A n p x q : ' q n p x q´ b n p x q p n p x q ř ˜ x r 1 ´ b n p ˜ x qs p n p ˜ x q and Loss ˚ TV p b q ' 1 2 ř x | q n p x q ´ b n p x q p n p x q| ´ 1 2 ř x p 1 ´ b n p x qq p n p x q . This is equivalent to the above since ϵ n : ' b n p n ´ q n .

Proof of Theorem 10. In this case, note min t p n p x q , q n p x q ` ϵ n p x qu ` r p n p x q ´ q n p x q ´ ϵ n p x qs ` ' p n p x q , we have

By definition of A n then ř x A n p x q ' 1 , and the original objective (23) can be equivalently written as

We now find the solution of objective (25) in two steps.

Step1: Let the positive token set A ` ' t x P V : A n p x q ě 0 u and the negative token set A ´ ' t x P V : A n p x q ă 0 u , then any optimal P ˚ n must satisfy P ˚ n p x q ' 0 for all x P A ´ .

First, since ř x A n p x q ' 1 , it implies A ` ‰ H . Suppose for some optimal P ˚ n , there exists ¯ x P A ´ such that P ˚ n p ¯ x q ą 0 , then we show there exists ˇ x P A ` such that A n p ˇ x q ą P ˚ n p ˇ x q . Suppose this is

not the case, i.e. A n p x q ď P ˚ n p x q @ x P A ` , then

Contradiction! Hence, there exists ˇ x P A ` such that A n p ˇ x q ą P ˚ n p ˇ x q .

Second, by A n p ˇ x q ą P ˚ n p ˇ x q , ´ P ˚ n p ¯ x q ă 0 , and triangular inequality, we have

Note A n p ¯ x q ă 0 , the above is equivalent to

Now set then apply (26) we have

This contradicts P ˚ n is the optimal solution! Therefore, for any optimal P ˚ n , it holds P ˚ n p x q ' 0 @ x P A ´ .

Step2: We characterize the optimal solutions of the objective. Indeed, by Step1, any optimal solution P ˚ n satisfies

The inequality becomes equality if and only if A n p x q ´ P ˚ n p x q ě 0 . Finally, recall the definition of A n we receive the optimal value for the original objective is

Replace ϵ n by b n gives the desired result.

Remark 6. The general optimization should follow

Meanwhile, our current analysis only considers the single token setting. We mention that solving the (27) is challenging as it corresponds to a high-dimensional discrete optimization with dimension T and it might not have closed-form solutions in the general cases.

E.2 Proof of Theorem 5

Proof. For an algorithm A with the rejection probability b p¨q . Let ˜ x ' p p¨q , then the rejection probability is computed as

Also, from Theorem 4

Next, we show

Indeed, since min t 1 , q p x q p p x q u ď b p x q ď 1 , @ x P V , then b p x q p p x q ě min t p p x q , q p x qu . Then we prove the following stronger claim

If q p x q ě p p x q , then 1 ' min t 1 , q p x q p p x q u ď b p x q ď 1 implies b p x q ' 1 , so the above is equivalent to | q p x q ´ p p x q| ' | p p x q ´ q p x q| is always true;

If q p x q ă p p x q , then b p x q p p x q ě min t p p x q , q p x qu ' q p x q . In this case | q p x q´ b p x q p p x q|`p 1 ´ b p x qq p p x q ' b p x q p p x q´ q p x q`p 1 ´ b p x qq p p x q ' p p x q´ q p x q ' | p p x q´ q p x q| .

This concludes the proof.

F Numerical Simulation Details

To 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 p p x n | x 1: n ´ 1 q , for Markov Chain, the one step transition is Markovian with p p x n | x 1: n ´ 1 q ' p p x 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

and the shaded regions are error bars.

Below presents the simulation for horizon T ' 100 . Left panel of Figure 5 is the standard speculative decoding, the middle panel is batch speculative decoding with batch size M ' 2 , and Right panel shows the expected rejections with varying batch sizes M computed from Theorem 1.

31.6

33.8

Theoretical ≥E[TV(pn,qn)]

Average Rejections

n =1

Expected Rejections

31.4

33.6

31.2

33.4

31.0

33.2

30.8

Batch SpecDecode M = 2

True Value,batch M = 2

0

50

100

150

200

0

50

100

150

200

Number of runs N × 100

Number of runs N × 100

32

Expected Rejections

M =1

M= 2

30

M=5

M =10

28

M=50

M=150

M=200

M=400

26

M=2000

0

500

1000

1500

2000

G Details for Experiment Section 4.2

Consider objective (1)

For any biased algorithm with the given acceptance probability b ą min t 1 , q { p u , we can rewrite:

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 the 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.

Implementation detail for Table 1. Concretely, we leverage EleutherAI/pythia-2.8b and EleutherAI/pythia-70m from HuggingFace. To perform speculative decoding, we specify assistant model=EleutherAI/pythia-70m in the generation function for target model EleutherAI/pythia-2.8b .

Notice that for the biased speculative decoding with b p x q ' min t 1 , q p x q` ϵ p p x q u , A p x q in Theorem 4 can equivalently expressed as A p x q : ' max t q p x q´ p p x q , ´ ϵ u ř ˜ x max t q p ˜ x q´ p p ˜ x q , ´ ϵ u , and we can select P ˚ : ' r A s ` (recall r¨s ` in Section 2.1), which satisfies P ˚ | A ´ p¨q ' 0; 0 ď P ˚ | A ` p¨q ď A p¨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.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

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.

2. Limitations

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.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We provide the full proof in appendix.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The code is written down in the appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

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.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We have detailed section in Appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We have error bars in Figure 2.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer 'Yes' if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

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.

8. 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).

9. 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 ?

Answer: [Yes]

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).

10. 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).

11. 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.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: NO new assets intorduced.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: No crowdsourcing involved.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: 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.

STRENGTHS:
- The authors provide a robust theoretical foundation, illustrating the practical implications of speculative decoding, such as the improvement of rejection accuracy, which cannot be achieved by simply changing the acceptance probability.
 - The study explores the trade-offs between inference cost and quality degradation, supported by an optimization model. This analysis is valuable for practical applications.

WEAKNESSES:
- 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:
NA

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13

comments: 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.

STRENGTHS:
The writing is very clear, with takeaways provided under each theorem to explain the theory.

Theorems 1 and 2 are crucial for speculative sampling. In paper [23], the authors showed that speculative sampling is unbiased but did not prove its efficiency compared to other rejection sampling methods. The proof provided here is very important.

WEAKNESSES:
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:
see weakness

LIMITATIONS:
na
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1) The theoretical analysis presented by the authors provide several interesting insights about the inference efficiency observed by spectral decoding.\n\n2) All the results are backed up with simulation experiments, which strengthen the results presented in the paper.\n\nWEAKNESSES:\n1) 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.\n\n2) 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.\n\nQUESTIONS:\nI request the authors to kindly address the questions in the weaknesses section.\n\nLIMITATIONS:\nYes, the authors have addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. The paper provides comprehensive theoretical analysis.\n\n2. The findings in Theorem 4 and 5 are interesting.\n\n3. The paper is easy to understand.\n\nWEAKNESSES:\n1. 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. \n\n2. 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.\n\n3. 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.\n\nQUESTIONS:\nsee weakness above\n\nLIMITATIONS:\nsee weakness above", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-13"}