Title: SpecTr: Fast Speculative Decoding via Optimal Transport
Venue: NeurIPS 2023
Accepted: yes
Keywords: autoregressive sampling; computation efficiency; optimal transport
Abstract: Autoregressive sampling from large language models has led to state-of-the-art results in several natural language tasks.
However, autoregressive sampling generates tokens one at a time making it slow, and even prohibitive in certain tasks. One way to speed up sampling is *speculative decoding*: use a small model to sample a *draft* (block or sequence of tokens), and then score all tokens in the draft by the large language model in parallel. A subset of the tokens in the draft are accepted (and the rest rejected) based on a statistical method to guarantee that the final output follows the distribution of the large model. 
In this work, we provide a principled understanding of speculative decoding through the lens of optimal transport (OT) with *membership cost*. This framework can be viewed as an extension of the well-known *maximal-coupling* problem. This new formulation enables us to generalize the speculative decoding method to allow for a set of $k$ candidates at the token-level, which leads to an improved optimal membership cost. We show that the optimal draft selection algorithm (transport plan) can be computed via linear programming, whose best-known runtime is exponential in $k$. We then propose a valid draft selection algorithm whose acceptance probability is $(1-1/e)$-optimal multiplicatively. Moreover, it can be computed in time almost linear with size of domain of a single token.
Using this new draft selection algorithm, we develop a new autoregressive sampling algorithm called *SpecTr*, which provides speedup in decoding while ensuring that there is no quality degradation in the decoded output.
We experimentally demonstrate that for state-of-the-art large language models, the proposed approach achieves a wall clock speedup of 2.13X, a further 1.37X speedup over speculative decoding on standard benchmarks.

SpecTr: Fast Speculative Decoding via Optimal Transport

Abstract

Autoregressive sampling from large language models has led to state-of-the-art results in several natural language tasks. However, autoregressive sampling generates tokens one at a time making it slow, and even prohibitive in certain tasks. One way to speed up sampling is speculative decoding : use a small model to sample a draft (block or sequence of tokens), and then score all tokens in the draft by the large language model in parallel. A subset of the tokens in the draft are accepted (and the rest rejected) based on a statistical method to guarantee that the final output follows the distribution of the large model. In this work, we provide a principled understanding of speculative decoding through the lens of optimal transport (OT) with membership cost . This framework can be viewed as an extension of the well-known maximal-coupling problem. This new formulation enables us to generalize the speculative decoding method to allow for a set of k candidates at the token-level, which leads to an improved optimal membership cost. We show that the optimal draft selection algorithm (transport plan) can be computed via linear programming, whose best-known runtime is exponential in k . We then propose a valid draft selection algorithm whose acceptance probability is (1 -1 /e ) -optimal multiplicatively. Moreover, it can be computed in time almost linear with size of domain of a single token. Using this new draft selection algorithm, we develop a new autoregressive sampling algorithm called SpecTr , which provides speedup in decoding while ensuring that there is no quality degradation in the decoded output. We experimentally demonstrate that for state-of-the-art large language models, the proposed approach achieves a wall clock speedup of 2.13X, a further 1.37X speedup over speculative decoding on standard benchmarks.

1 Introduction

Autoregressive language models have shown to achieve state-of-the-art results in several natural language tasks [2, 5, 25, 26]. During inference, given a context x t := x (1) , x (2) . . . , x ( t ) , an autoregressive model M b generates successive tokens x ( t +1) , x ( t +2) , . . . via temperature sampling [1, 10], where the next token x ( t +1) is drawn from the temperature-scaled distribution M b ( ·| x t ) . If the temperature is zero, i.e., greedy decoding, the next token is determined by the maximum likelihood method i.e., x ( t +1) = arg max x ∈ Ω M b ( x | x t ) , where Ω is the domain of a single token also referred to as the vocabulary. The sampling approach can be further combined with other sampling primitives such as nucleus sampling [16] and topk sampling [9, 22].

All these approaches are autoregressive decoding 2 methods, where tokens are generated serially one after another, which can be slow or even prohibitive in several applications [23]. Hence, several techniques have been proposed to improve the speed of decoding. Before we proceed further, we first present some notations and a simplified computational model.

Notations. We use x i : j to denote the sequence x ( i ) , x ( i +1) , . . . , x ( j ) and when i = 1 , we simply use x j = x 1: j . x ( i ) denotes the i -th entry of x . Subscripts are used to distinguish between different sequences, e.g., x t 1 and x t 2 denote two sequences of length t . We use [ n ] to denote the set { 1 , . . . , n } .

A simplified computational model.

Standard inference. Given a context x t , with O ( t 2 ) computation and O (1) time, an autoregressive model M b can compute M b ( y | x t ) , the (temperature-scaled) probability of all possible next tokens y ∈ Ω .

Parallelization along the time axis. Given a context x t , with O ( t 2 ) computation and O (1) time, an autoregressive model M b can compute M b ( y | x i ) , for all y ∈ Ω and i ∈ { 1 , 2 , . . . , t } .

Parallelization along time and batch axis. Let K be the maximum batch size that can be used during the inference of the autoregressive model. Given several contexts, x t 1 , x t 2 , . . . x t K , with O ( Kt 2 ) computation and O (1) time, an autoregressive model M b can compute M b ( y | x i j ) , for all y ∈ Ω , i ∈ [ t ] , and j ∈ [ K ] . 3

The above computation model shows that parallelizing along time and batch axes does not increase the computation time. It is a simplified characterization of the typical hardware, such as TPUs and GPUs, used in neural network inference. Previous approaches also assume similar computational model to devise faster decoding algorithms [19, 4]. In practice, there will be some overhead depending on hardware, implementation and resource utilization. In Appendix E, we experimentally verify that the theoretical gains are largely preserved for a large transformer model in practice. We also note that there are efficient transformer architectures, which reduces the computation cost from O ( t 2 ) to O ( t log t ) (see [24] for a detailed survey). Such approaches are orthogonal to the focus of this paper, and they can be easily combined with our approach.

Broadly speaking, multiple previous approaches proposed to guess a few possible future tokens using an efficient model. They then compute several conditional probability distributions from the large model based on the guesses. Computing the distributions takes O (1) time due to parallelization along the time axis. The guessed tokens are then accepted or rejected based on a statistical method such that the accepted tokens are effectively samples from the large model. This guarantees that there is provably no degradation in the quality of the decoded output compared to that of the large model. When the guesses are plausible under the large model, multiple tokens will be accepted, leading to a larger gain in latency improvement. We will further characterize the acceptance probability as a function of the closeness of the distributions of large model and the small model. While this approach incurs the same computation cost as vanilla decoding (under the simplified computational model assumed in this paper), it can significantly improve decoding latency due to parallelization.

The goal of this work is to provide a principled understanding of the above approaches and discuss optimality conditions and algorithmic improvements. We start by providing a more formal overview of speculative decoding and related works.

2 Previous works and speculative decoding

Previous approaches make use of parallelization along the time axis to provide speedups. They first predict multiple tokens and validate if these multiple tokens can be generated by the model with the corresponding sampling or decoding scheme. For greedy decoding, multiple tokens can be predicted by a separate model [23], aggressive decoding [11], or retrieval augmented text [28]. For sampling, recently [19, 4] proposed an algorithm called speculative decoding, and we provide an overview of this algorithm in the rest of the section. Suppose we have access to a computationally-inexpensive draft model M s , which predicts the next token given the context, and the predictions of M s are

This paper will

be liked for its

Context

Sample

one

draft

from the

small model

Compute conditional

probabilities by the

large model

Draft selection via

maximal coupling

Add to the context

This paper will

be

liked

by

Next iteration

(if not E.O.S.)

be

liked for by

close to that of M b for most contexts. Suppose we have obtained prefix x t . The next iteration of the speculative algorithm can be broken down into three steps (see Fig. 1 for an illustration).

Draft construction. The draft model is used to efficiently and 'speculatively' sample L tokens, ˜ x ( t +1) , . . . , ˜ x ( t + L ) . We keep the conditional probabilities on the next token M s ( y | x t , ˜ x t +1: t + i ) for each i < L and ∀ y ∈ Ω .

Conditional probability computation. After observing the samples, we compute the conditional distributions M b ( y | x t , ˜ x t +1: t + i ) for each i ≤ L and ∀ y ∈ Ω in parallel (along time axis) in O (1) time.

Draft selection. Validate and select first L ′ of the L tokens and set x ( t + i ) = ˜ x ( t + i ) for i ≤ L ′ given the draft sequence and the conditional probabilities from both models. Sample a token ˜ x ( t + L ′ +1) from a residual distribution as a correction to the rejected token. 4

After this step, we use x t + L ′ +1 1 as the next context and sample the next few tokens using speculative decoding iteratively. For a complete statement of the algorithm, we refer the readers to [19, 4]. The crux of the above steps is draft selection, which given a draft sequence and the conditional probabilities from both models, selects a valid sequence such that the output has the same distribution as that of the large model. In speculative decoding, this is achieved via recursively applying a token-level maximal coupling algorithm, which is provided in Algorithm 1. Note that for the draft selection, Algorithm 1 is applied where p is the conditional distribution of the draft model M s ( · | x t ) and q is the conditional distribution of the large model M b ( · | x t ) (which may be further conditioned on the newly decoded tokens).

Algorithm 1 Token-level maximal coupling

Input: Distributions p, q , Draft sample X ∼ p .

1: Compute the residual distribution p res where ∀ x ∈ Ω , p res ( x ) = q ( x ) -min { p ( x ) ,q ( x ) } 1 -∑ x ′ min { p ( x ′ ) ,q ( x ′ ) } .

2: Sample η ∼ U (0 , 1) .

3: if η ≤ min ( 1 , q ( X ) p ( X ) ) then

4: Return Y = X . { Accept the draft token. }

5: end if

6: Return Y ∼ p res . { Sample a corrected token from the residual distribution. }

Algorithm 1 returns a random variable Y which either is the accepted input X or a sample from the residual distribution p res , which is defined in Step 1 of Algorithm 1. The algorithm is recursively applied as long as the draft tokens are accepted to select the first L ′ ≤ L tokens from the draft model. For the first rejected token, the sample Y from the residual distribution is used as a correction . Previous works showed that if X ∼ p , then Y ∼ q [19, 4]. In the case of the draft selection, this means that the output of the algorithm is distributed according to M b ( · | x t ) , which is exactly the

This paper will

Context

Sample

multiple

drafts

from the

small model

Compute

all

conditional

probabilities by the

large model

Draft selection via an

OTM transport plan

Add to the context

This paper will

be

liked

by

three

Next iteration

(if not E.O.S.)

be

liked

by all three

be liked for its

be liked by all

not get good reviews

be read by four

desired outcome. Furthermore

where d TV is the total variation distance between p and q . The closer p and q are in d TV, the higher the chance of Pr( Y = X ) , and fewer the number of serial calls to the larger model. In the ideal case, if p = q , then Pr( Y = X ) = 1 , i.e., the draft token is always accepted, and when used for speculative decoding we have L ′ = L. Together with the extra sampled token 5 from M b , L + 1 tokens are obtained in one iteration. In such a case, based on our computational model (Section 1), assuming the decoding time of draft model is negligible, the speedup is ( L +1) times.

3 Our contributions

From a theoretical viewpoint, the speculative decoding algorithm raises multiple questions.

What is the relationship between speculative decoding and the broader literature of sampling in statistics?

Is speculative decoding optimal in an information-theoretic sense?

Speculative decoding uses parallelization along time to speed up decoding; would it be possible to use parallelization along batch (number of drafts) to further improve decoding speed?

We provide answers to all the above questions in this work. We first relate the problem of speculative decoding to the broader and well-studied discrete optimal transport theory through a token-level coupling problem (Section 4). With this connection, it becomes clear that the token-level draft selection is the optimal solution for optimal transport with indicator cost function and also related to the problem of maximal coupling [8]. Based on the connection to optimal transport, we show that one can further speed up the decoding by parallelizing along the batch axis by using multiple drafts from the draft model (Section 5).

More precisely, we formulate the token-level draft selection problem as a discrete optimal transport problem with membership cost, which is referred to as OTM. Discrete optimal transport can be solved with a linear program, but the number of variables is exponential in batch size, which can be prohibitive. To address this, we propose a valid transport plan that can be efficiently computed. Moreover, it achieves a (1 -1 /e ) -approximation of the optimal acceptance probability (Section 6).

With the theoretically motivated algorithms and guarantees, we circle back to speeding up decoding and propose a new algorithm called SpecTr and theoretically show that it can be used to derive valid sequences from the large model with better speedups (Section 7). See Fig. 2 for an illustration of SpecTr . Compared to speculative decoding (Fig. 1), the main difference lies in the number of sampled drafts sampled from the small model and the selection algorithm that selects a valid sequence from multiple draft sequences. We remark here that the latter requires completely new statistical

This paper will

be

liked

for

its

be

liked

by

all three

not

get

good

reviews

be

read

by

four

Figure 3: An example run of the sequence-level draft selection in SpecTr with L = 4 and 4 draft sequences. In the first step, there are 4 drafts tokens, and the token-level draft selection algorithm selects the word ' be ' which appeared thrice. Note that all tokens following ' be ' are valid draft tokens from the small model. In the second step, there are 3 drafts and the selection algorithm selects ' liked '. The next token-level selection algorithm will have two drafts (' by ' and ' for ') and it selects ' by '. Finally, there is only one draft following ' by ', and the selection algorithm doesn't select it and outputs ' three ' as a correction. The process ends and a total of 4 tokens are generated.

tools, and the connection between the token-level draft selection and OTM is critical for obtaining valid transport plans with good guarantees. We view this as one of the main contributions of the work. Similar to speculative decoding, there is provably no degradation in the quality of the decoded output compared that of the large model.

We then experimentally demonstrate the benefit of our approach on standard datasets (Section 8). More precisely, we show that for state-of-the-art large language models, SpecTr achieves a wall clock speedup of 2.13X, a further 1.37X speedup over speculative decoding on standard benchmarks.

4 Token-level draft selection and optimal transport

In this section, we focus on the draft selection step of SpecTr . We start by considering the case when L = 1 , which is a token-level draft selection problem. In particular, given context x t , let X 1 , . . . X k be a collection of draft tokens sampled from the small model, e.g., sampled i.i.d. from M s ( · | x t ) . Note that by our assumption of the computation model, we could compute the following conditional probabilities from the large model in parallel ( along time and batch axes):

The goal of the draft selection algorithm f : Ω k → Ω is to output Y = f ( X k ) , whose distribution follows M b ( · | x t ) , and hence is a valid sample from the large model. Moreover, when Y ∈ { X 1 , . . . , X k } , we could sample an extra token from M b ( · | x t , Y ) without calling M b since we have already computed the conditional probabilities M b ( · | x t , Y ) . Hence we would like to maximize the probability that we accept one token from the set of drafts.

When L > 1 , the drafts are sequences sampled from M s , a sequence of token-level draft selection algorithms could be used along the time axis to select a valid sequence from the M b . See an example in Fig. 3. The full details about the sequence-level selection algorithm is provided in Section 7.

The reminder of the section will be focused on the token-level draft selection problem. From the above discussion, there are the two main goals of the draft selection problem.

Validity. The output token is always a valid token from the large model i.e., its distribution follows the conditional probability of the large model. This guarantees that there is no quality degradation compared to the large model.

Maximizing acceptance. The higher the probability that we accept a draft token, the more serial computation we can save through parallelization, and hence better speedup.

Before proposing our framework to achieve the above goals, we would like to first discuss the technical challenge of draft selection with multiple draft tokens. One attempt is to sequentially apply the acceptance phase of Algorithm 1 (line 3 - 5) to each draft token X i with p = M s ( · | x t ) and q = M b ( · | x t ) . However, this approach would not guarantee that the final accepted token is from the desired distribution. To see this, consider the example of p = Ber (1) and q = Ber (1 / 2) . 6 Then we have ∀ i = 1 , . . . , k, X i = 1 and each of them will be accepted with probability 1 / 2 . After

applying Algorithm 1 to all X i 's, the probability of getting a 1 will be at least 1 -1 / 2 k and hence the output distribution would not be Ber (1 / 2) for k > 1 . Therefore the algorithm does not produce valid samples, which is a requirement of the draft selection problem.

In this work, we conduct a principled investigation of the draft selection problem, and show that these two main goals could be captured by the framework of optimal transport with a properly defined cost function. Next we define optimal transport formally and then connect it to draft selection with one draft. The generalization to multiple drafts is provided in Section 5.

Coupling and optimal transport. To simplify notations, we assume Ω is a discrete domain.

Definition 1 (Coupling) . For two probability distributions P over X and Q over Y , we say a joint distribution π supported over X × Y is a coupling between P and Q if ∀ x, y, π ( x, y ) ≥ 0 ,

We use Π( P, Q ) to denote the set of all possible couplings between P and Q .

When it is clear from context, we will overload notation and refer to the probabilistic mapping f π : X → Y introduced by the conditional probability π ( y | x ):= π ( x, y ) /P ( x ) as a coupling, which is also referred to as a transport plan from P to Q [27]. In this paper, we will set P to be the distribution of the draft tokens and Q to be the target distribution of the output token. In this case, the f π is a valid draft selection algorithm. Formally, this is stated in the claim below.

Claim 1. For all π ∈ Π( P, Q ) , let f π be the probabilistic mapping defined above . If X ∼ P , then f π ( X ) ∼ Q .

In this paper, we will design selection algorithms by finding valid couplings between the draft distribution and target distribution to guarantee validity of the output tokens.

Definition 2 (Optimal Transport (OT) [27]) . For a cost function c : X ×Y → R + , the transportation cost of a coupling is defined as:

The optimal transport plan is the coupling π ∈ Π( P, Q ) that minimizes the transportation cost.

Speculative decoding with one draft token. With these definitions in place, we can see that with X = Y = Ω , the domain of the tokens and P = p, Q = q , we recover the speculative decoding objective with one draft token using the cost function of indicator cost , which captures the resampling cost, defined below:

glyph[negationslash]

∀ x ∈ Ω , y ∈ Ω , c ( x, y ) = ✶ { y = x } . The transportation cost of the coupling will be C ( π ) = E X,Y ∼ π [ ✶ { Y = X } ] = P X,Y ∼ π ( Y = X ) . This optimal transport cost is known to be glyph[negationslash]

glyph[negationslash]

glyph[negationslash]

which is achieved by the maximal coupling between p and q stated in Algorithm 1 [8]. And hence speculative sampling achieves the optimal cost with one draft token.

5 Optimal transport with multiple draft tokens

In this section, we generalize token-level selection to allow for multiple drafts. More formally, let X = Ω k for some k ∈ N + , which is the space of k draft tokens from Ω and Y = Ω , which is the space of the final sampled token from the desired distribution. To characterize the resampling cost with multiple draft tokens, we use the cost function of membership cost , defined below:

✶ where S ( x ) = { o ∈ Ω | o appears in x } denotes the set of distinct elements in x . When k = 1 , it recovers the indicator cost mentioned before. The transportation cost of the coupling is

We will also refer to the above cost C ( π ) as the rejection probability due to its probabilistic interpretation. And similarly, α ( π ):=1 -C ( π ) = P X,Y ∼ π ( Y ∈ S ( X )) will be the acceptance probability.

From now on we will use membership cost as the default cost function and refer to the optimal transport solution as optimal transport with membership cost (OTM). We use π ∗ to denote the coupling that minimizes this cost π ∗ = arg min π ∈ Π( P,Q ) C ( π ); 7 and the cost C ( π ∗ ) is referred to as the optimal transport cost between P and Q . We use α ( P, Q ) = 1 -C ( π ∗ ) to denote the corresponding optimal acceptance probability.

Draft selection with i.i.d. draft tokens. In this paper, we will mainly focus on the case when the draft tokens are i.i.d. samples from a base distribution. 8 Let p, q be supported over Ω and the goal is to obtain one valid token from q given k i.i.d. samples from p . For SpecTr with context x t , we have p = M s ( · | x t ) and q = M b ( · | x t ) . We set P = p ⊗ k , a product distribution whose marginals are all p , and Q = q . The OT problem we want to solve is the following:

We overload notation and denote the optimal acceptance probability as α k ( p, q ):= α ( p ⊗ k , q ) = 1 -C ( π ∗ ) . To better understand the quantity, we state a few properties about α k .

Lemma 1. (Appendix A.2) The optimal acceptance probability statisfies the following properties.

Monotonicity. For any p, q and k ≥ 1 , α k ( p, q ) ≤ α k +1 ( p, q ) .

Consistency. If ∀ x ∈ Ω , q ( x ) /p ( x ) is bounded, we have lim k →∞ α k ( p, q ) = 1 . Else, lim k →∞ α k ( p, q ) = ∑ x ∈ Ω { p ( x ) > 0 } q ( x ) .

✶ The above properties demonstrate that for a large k , the value of α k can become large. Hence increasing k could increase the acceptance probability, leading to further speedups. We now focus on computing the optimal transport plan and the optimal acceptance probability.

OTMvia Linear programming. Optimal transport in discrete domain has been studied extensively [17, 21, 14], and it is shown that the optimal transport problem is equivalent to the following linear programming problem:

The linear program in (4) has | Ω | k +1 variables and | Ω | k + | Ω | equality constraints (see Definition 1). Linear programming can be solved in time polynomial in the number of variables and constraints [7, 21, 18], 9 implying the following lemma.

Lemma 2. Given p, q over Ω , the solution to Eq. (3) can be computed in time O ( | Ω | O ( k ) ) .

We refer to the optimal coupling obtained above as OTMk and denote it as π OTM -k . When k = 1 , there is a closed form expression for the optimal acceptance cost (see Eq. (1)), whereas for larger values of k , we are unaware of a general closed form expression. In Appendix A.1, we provide an information-theoretic upper (and lower) bound, which is tight up to a multiplicative constant of 1 -(1 -1 /k ) k ≥ 1 -1 /e .

While solving OTM in Eq. (4) gives the plan with optimal acceptance probability, to the best of our knowledge, the best-known runtime will be exponential in k , which can be prohibitive when either the vocabulary size | Ω | or the number of draft tokens k is large. 10 In the next section, we will present a selection algorithm that can be efficiently computed and show that it achieves an acceptance probability of at least (1 -(1 -1 /k ) k ) α k ≥ (1 -1 /e ) α k .

6 Draft selection via k -sequential selection

In this section, we present a sequential selection algorithm (K-SEQ), an approximate solution 11 to the optimal transport problem in Eq. (3), which can be efficiently computed in time almost linear in | Ω | and logarithmic in k . The algorithm is presented in Algorithm 2.

Algorithm 2 k -sequential selection algorithm (K-SEQ).

Input: Distributions p, q , samples X 1 , . . . , X k ∼ i.i.d. p . ρ ∈ [1 , k ] : division factor.

1: Let β p,q ( ρ ) = ∑ x ∈ Ω min( p ( x ) , q ( x ) /ρ ) and p acc = 1 -(1 -β p,q ( ρ )) k . Compute p res where

2: for i = 1 , 2 , . . . , k do

3: Sample η i ∼ U (0 , 1) .

4: if η i ≤ min ( 1 , q ( X i ) ρ · p ( X i ) ) then

5: Return Y = X i . { Return the i th draft token. }

6: end if

7: end for

8: Return Y ∼ p res . { Sample a corrected token from the residual distribution. }

At a high-level, the algorithm goes over all k draft samples generated from p sequentially, and decides on whether to accept each X i based on the ratio q ( X i ) /p ( X i ) . The algorithm output the first accepted sample or result from a residual distribution p res if none of the samples is accepted. To guarantee that the the final returned token is a valid sample from q , we choose an appropriate ρ ∈ [1 , k ] and accept X i with probability min(1 , q ( X i ) / ( ρ · p ( X i ))) instead of min(1 , q ( X i ) / ( p ( X i ))) as in Algorithm 1. In Theorem 1, we show that with appropriately chosen ρ 's, Algorithm 2 is indeed valid transportation plans from p ⊗ k to q . Moreover, to find the best transportation plan within the family, we only need to search over a single parameter ρ , which reduces the computation cost significantly. We also show that searching over this sub-family of couplings won't decrease the optimal acceptance probability by a multiplicative constant. The performance of Algorithm 2 is stated in Theorem 1.

When ρ ≥ ρ ∗ , the coupling π K-SEQ ρ in Algorithm 2 is a valid transport plan from p ⊗ k to q . When ρ = ρ ∗ , we have

Moreover, ρ ∗ can be computed up to accuracy δ in time O ( | Ω | log(( k -1) /δ )) .

We provide the proof in Appendix C.1. In Appendix B, using a few canonical examples of distributions, we plot the acceptance probability of K-SEQ and compare it with the optimal acceptance probability α k . It can be shown that K-SEQ could have a strictly worse acceptance probability compared to the OTM solution for certain cases while there also exist non-trivial cases where K-SEQ achieves the optimal acceptance probability.

Concurrent and recent work of [20, 29] has proposed another efficient algorithm for the draft selection phase. To the best of our knowledge, there is no optimality guarantee proved for their proposed algorithm. In Appendix B.3, we present its acceptance probability empirically for the canonical case of Bernoulli distributions, and show that both our proposed algorithms (OTM and K-SEQ) have a higher acceptance probability.

7 SpecTr: Application of OTM in autoregressive sampling

In this section, we describe how OTM can be used to speed up auto-regressive sampling, which we refer to as SpecTr sampling. Similar to speculative decoding, each iteration of SpecTr can be decomposed into three phases (Fig. 2):

Draft set construction. Given current context x t , use the draft model sample a set of K draft sequences with length L , denoted by S = { z L ∼ M s ( · | x t ) } . We keep the conditional probabilities M s ( y | x t , z i ) for all y ∈ Ω , i ≤ L and z L ∈ S .

Conditional probability computation. Compute the conditional probabilities on the next token for the large model M b ( y | x t , z i ) for all y ∈ Ω , i ≤ L and z L ∈ S in parallel.

Draft selection. Select first L ′ of the L tokens and set x ( t + i ) = z ( i ) for i ≤ L ′ and some z ∈ S given the set of draft sequences and the conditional probabilities from both models. Sample a token from a residual distribution as a correction to the rejected tokens.

The conditional probability computation step takes O (1) when | S | is not large based on our simplified computations model. We mainly focus on the draft set construction phase and draft selection phase.

Algorithm 3 Draft selection with multiple candidates (DraftSelection).

Draft set with i.i.d. draft sequences. Given context x t , a natural way to come up with a set of K drafts is to independently sample K draft sequences from M s ( · | x t ) , i.e.,

The draft set construction method in (7) can be generalized to a prefix-tree based algorithm. However, this generalized version did not perform better in our experiments. We include this construction in Appendix D for completeness.

Draft selection with multiple candidates. We present the sequence-level selection algorithm given a set of draft sequences in Algorithm 3. We assume the conditional probabilities on the next token are available given any prefix in the candidate set since they are computed in parallel in the second phase, and won't list them as inputs explicitly in Algorithm 3.

A sample run of the algorithm is presented in Fig. 3. The algorithm proceeds in a recursive fashion. Given prompt x t and a candidate set S sampled from M s ( · | x t ) , the algorithm first computes a token-level draft selection algorithm f π : Ω | S | → Ω which is a transport plan from M s ( · | x t ) ⊗| S | to M b ( · | x t ) . Then f π is applied to the set of first tokens of the draft sequences in S to obtained a valid token Y from M b ( · | x t ) . If Y is not the last token ( L ≥ 2 ), we filter out sequences in S whose first token is not Y and denote the remaining sequences as S next and feed it to the algorithm with context ( x t , Y ) and draft length L -1 . This goes on until we have L = 1 or S next = ∅ .

In this case when Y is the last token ( i.e., L = 1 ) and Y ∈ S , we have the choice to sample an additional token M b ( · | ( x t , Y )) since this conditional probability is already computed in the second phase. Due to the property of the token-level selection algorithms and the autoregressive structure of language models, it can be shown that Y is always a valid sample from M b ( · | x t ) . Let L ′ be the number of decoded tokens in one iteration. Note that this is a random variable in the range [1 , L +1] .

The formal quality guarantee is stated in Theorem 2. We present the proof in Appendix C.2.

Theorem 2. Assume all drafts in the set S are generated from the small model with input x t , or more precisely, ∀ z ∈ S,

Let ( x t , Y τ ) be the output of Algorithm 3 where τ is the length of the newly decoded tokens, then it satisfies that Y 1: τ is distributed according to M b ( · , · , . . . · ︸ ︷︷ ︸ | x t ) . More precisely, For any

τ dots τ 0 ∈ [1 , L +1] , and any τ 0 -length, sequence o τ 0 = ( o (1) , . . . , o ( τ 0 )) ∈ Ω τ 0 , we have

8 Experiments

We empirically evaluate SpecTr and compare it with two methods: (1) the baseline auto-regressive decoding; and (2) speculative decoding with K = 1 . Note that all three methods effectively generate samples from the same baseline large model, and hence the quality of the two speculative decoding methods is provably neutral to that of the large model. Thus, we will only focus on measuring the speedup in our experiments. In the simplified computation model, we made the following assumptions: (1) Decoding time from small models is negligible compared to decoding from the small model; (2) Parallelization along the batch and time axis doesn't increase the time for a serial call to the large model. With these, the theoretical speedup compared to baseline decoding will be the average number of decoded tokens per serial call, which is called block efficiency [19], defined below

However, in real deployment of the SpecTr algorithm, the actual end-to-end (wall clock) speedup is further impacted by the following aspects. (1) The decoding time for M s might not be negligible; (2) Parallelization along the batch and time axis might increase the time for a single call to M b ; (3) Overhead due to the implementation of additional functionalities in SpecTr such as the draft selection algorithm and switching between models. These factors will depend on how the algorithm is implemented and optimized. In our experiment, we consider both the block efficiency, and average wall clock speedup with our implementation of SpecTr .

We first present the performance of our algorithm and compare it to speculative decoding using state-of-the-art PALM-2 models with prompts from the one-billion language benchmark (LM1B) [3] . In Appendix E, we use a pair of smaller transformer models to break down different affecting factors mentioned above. In Table 1, we use PALM-2-Gecko and PALM-2-Bison as the small model and

large model, respectively [13, 12]. The wall clock speedup is normalized by the wall clock latency of baseline autoregressive decoding. The time we log include all above mentioned aspects. In the considered parameter configurations, the wall clock speedup increases as K and L increases. As seen from the table, the actual wall clock speedup is smaller than the theoretical speedup of block efficiency, which is consistent with what we expected. Importantly, the benefit from SpecTr outweighs these overheads. In particular, when L = 8 and K = 8 , our proposed SpecTr algorithm has a speedup of 2.13x, a further 1.37x increase compared to speculative decoding ( K = 1 ).

References

David H Ackley, Geoffrey E Hinton, and Terrence J Sejnowski. A learning algorithm for boltzmann machines. Cognitive science , 9(1):147-169, 1985.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robinson. One billion word benchmark for measuring progress in statistical language modeling. arXiv preprint arXiv:1312.3005 , 2013.

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.

Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311 , 2022.

Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In C.J. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K.Q. Weinberger, editors, Advances in Neural Information Processing Systems , volume 26. Curran Associates, Inc., 2013.

George B Dantzig. Linear programming. Operations research , 50(1):42-47, 2002.

Frank Den Hollander. Probability theory: The coupling method. Lecture notes available online (http://websites. math. leidenuniv. nl/probability/lecturenotes/CouplingLectures. pdf) , 2012.

Angela Fan, Mike Lewis, and Yann Dauphin. Hierarchical neural story generation. arXiv preprint arXiv:1805.04833 , 2018.

Jessica Ficler and Yoav Goldberg. Controlling linguistic style aspects in neural language generation. arXiv preprint arXiv:1707.02633 , 2017.

Tao Ge, Heming Xia, Xin Sun, Si-Qing Chen, and Furu Wei. Lossless acceleration for seq2seq generation with aggressive decoding. arXiv preprint arXiv:2205.10350 , 2022.

Google AI. Introducing PaLM 2, 2023. https://blog.google/technology/ai/ google-palm-2-ai-large-language-model/ .

Google PaLM-2 Team. PaLM 2 technical report, 2023.

Wenshuo Guo, Nhat Ho, and Michael Jordan. Fast algorithms for computational optimal transport and wasserstein barycenter. In Silvia Chiappa and Roberto Calandra, editors, Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics , volume 108 of Proceedings of Machine Learning Research , pages 2088-2097. PMLR, 26-28 Aug 2020.

Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. Flax: A neural network library and ecosystem for JAX, 2023.

Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. arXiv preprint arXiv:1904.09751 , 2019.

Leonid V Kantorovich. On the translocation of masses. In Dokl. Akad. Nauk. USSR (NS) , volume 37, pages 199-201, 1942.

Yin Tat Lee and Aaron Sidford. Path finding methods for linear programming: Solving linear programs in o (vrank) iterations and faster algorithms for maximum flow. In 2014 IEEE 55th Annual Symposium on Foundations of Computer Science , pages 424-433. IEEE, 2014.

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.

Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang, Rae Ying Yee Wong, Alan Zhu, Lijie Yang, Xiaoxiang Shi, Chunan Shi, Zhuoming Chen, Daiyaan Arfeen, Reyna Abhyankar, and Zhihao Jia. Specinfer: Accelerating generative large language model serving with speculative inference and token tree verification, 2023.

Ofir Pele and Michael Werman. Fast and robust earth mover's distances. In 2009 IEEE 12th international conference on computer vision , pages 460-467. IEEE, 2009.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

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

Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys , 55(6):1-28, 2022.

Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, HengTze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, et al. Lamda: Language models for dialog applications. arXiv preprint arXiv:2201.08239 , 2022.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. LLaMA: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023.

Cédric Villani et al. Optimal transport: old and new , volume 338. Springer, 2009.

Nan Yang, Tao Ge, Liang Wang, Binxing Jiao, Daxin Jiang, Linjun Yang, Rangan Majumder, and Furu Wei. Inference with reference: Lossless acceleration of large language models. arXiv preprint arXiv:2304.04487 , 2023.

Chao Zhang Yuhui Li and Hongyang Zhang. Eagle: Lossless acceleration of llm decoding by feature extrapolation, 2023. https://sites.google.com/corp/view/eagle-llm .

A Properties of optimal transport cost

A.1 Information-theoretic upper (and lower) bound of α k .

Below we provide an information-theoretic upper (and lower) bound in Lemma 3, which is tight up to a multiplicative constant of 1 -(1 -1 /k ) k ≥ 1 -1 /e . The proof is presented in Appendix A.3. For the case of k = 1 , the upper bound matches the optimal acceptance probability.

Lemma 3. For any two distributions p, q and ∀ k ≥ 1 , we have

where

In Appendix B, we plot α k as a function of k for a few simple pairs of ( p, q ) 's as illustrative examples. We note that the upper bound in Lemma 3 is tight for examples considered in Appendix B.

A.2 Proof of Lemma 1

We first prove monotonicity . By definition,

Moreover, for any π ∈ Π( p ⊗ k , q ) , we can construct π ′ ∈ Π( p ⊗ k +1 , q ) by setting

i.e., adding and independent sample from p to X k .

Hence we have

Next we prove consistency . We start with the case when ∀ x ∈ Ω , q ( x ) /p ( x ) < ∞ . To prove this, we will show that Algorithm 2 with ρ max = max x ∈ Ω q ( x ) /p ( x ) statisifies

Since α ( π K-SEQ ρ max ) ≤ α k ( p, q ) , the above equation implies lim k →∞ α k ( p, q ) = 1 . Notice that by Lemma 4 and Theorem 1, π K-SEQ ρ max is a valid coupling, and

where β p,q ( ρ ) = ∑ x ∈ Ω min( p ( x ) , q ( x ) ρ ) ≥ 1 /ρ max > 0 . Taking k →∞ concludes the proof.

For the case when q ( x ) /p ( x ) is unbounded, there exists x ∈ Ω such that q ( x ) > 0 and p ( x ) = 0 . Let

Let x 0 be such that p ( x 0 ) > 0 . We define q ′ such that glyph[negationslash]

Then we have d TV ( q, q ′ ) = p off , and hence by subadditivity of transport cost,

Moreover, we have ∀ x ∈ Ω , q ′ ( x ) /p ( x ) < ∞ . Hence

A.3 Proof of Lemma 3

For the upper bound, it would be enough to show that for any π ∈ Π( p ⊗ k , q ) , and any Ω 0 ⊂ Ω , we have

For the lower bound, we show that K-SEQ achieves an acceptance probability of at least (1 -(1 -1 /k ) k )¯ α k ( p, q ) , see Eq. (11), implying the lower bound guarantee.

B Comparison between α ( π K-SEQ ρ ∗ ) and α k for simple examples.

We illustrate the acceptance probabilities for our proposed token-level selection algorithms using a few simple examples and plot them in Figures 4 and 5. The analysis for these simple distributions is presented in Appendix B.1 and Appendix B.2.

glyph[negationslash]

Pairs of Bernoulli distributions. Let Ber ( b ) be a Bernoulli distribution with probability b of getting a head. In Figure 4, we plot the acceptance probability comparison between OTMk and K-SEQ for different Bernoulli distributions q = Ber ( b ) as a function of k when p = Ber (0 . 25) . Note that when p = q ( b = 0 . 25 ), the acceptance probability is always one for both methods. When p = q , the acceptance probabilities for both methods increase as k increases before they reach one. When b = 0 . 1 or 0 . 75 , K-SEQ has a worse acceptance probability compared to the OTMk algorithm. When b = 1 , the two algorithms have the same performance.

2.5

5.0

7.5

10.0

12.5

15.0

k

0.5

0.6

0.7

0.8

0.9

1.0

Acceptance probability

b

=0.1,

k

-seq

b

=0.1, OTM

b

=0.25,

k

-seq

b

=0.25, OTM

b

=0.5,

k

-seq

b

=0.5, OTM

b

=0.75,

k

-seq

b

=0.75, OTM

2

4

6

8

10

k

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

Acceptance probability

dq

=

d

/2

dq

=

d

/3

dq

=

d

/4

dq

=

d

/5

Pairs of uniform distributions. Let U ( d ) denote a uniform distribution over [ d ] . In Figure 5, we plot the optimal acceptance probability for different uniform functions q as a function of k . For these distributions, it can be shown that K-SEQ achieves the optimal acceptance probability α k . Hence only α k is plotted. Observe that all acceptance probabilities are monotonically increasing and tend to one when k →∞ , as stated in Lemma 1.

B.1 Calculations for α k .

In this section, we provide a sketch of optimal acceptance probability calculations for results in Figures 4 and 5.

Figure 4: Ber ( p ) and Ber ( q ) .

Setting Ω 0 = { 0 , 1 } in Lemma 3 yields the upper bound. For the lower bound observe that since Ω = { 0 , 1 } , { y / ∈ S ( x k ) } < 1 if and only if x k is 0 k or 1 k . Hence,

Consider the transport plan π given by π (1 k , 1) = min( p k , q ) , π (1 k , 0) = p k -min( p k , q ) , π (0 k , 0) = min((1 -p ) k , 1 -q ) , and π (0 k , 1) = (1 -p ) k -min((1 -p ) k , 1 -q ) . It can be checked that this is a valid transport plan. To see this matches the upper bound on the optimal cost from Lemma 3, notice that

If p k ≤ q and (1 -p ) k ≤ 1 -q , then the above equation simplifies to 1 and (10) also simplifies to 1 . If p k > q and (1 -p ) k ≤ 1 -q , then the above equation simplifies to 1 + q -p k and (10) also simplifies to the same quantity. Similarly, the proof applies for p k ≤ q and (1 -p ) k > 1 -q .

Figure 5: p = U ( d ) and q = U ( d/r ) . The optimal acceptance probability is

We first prove α k ( U ( d ) , U ( d/r )) ≥ 1 -(1 -1 /r ) k by a construction. Let S ( X k ) be the set of unique symbols in X k . Consider the following transport plan, where Y is drawn uniformly from S ( X k ) ∩ [ d/r ] and draws a new uniform sample from [ d/r ] if S ( X k ) ∩ [ d/r ] = ∅ . Observe that since U ( d ) is uniform over [ d ] , this is a valid transport plan and furthermore, glyph[negationslash]

The upper bound follows by setting Ω 0 = [ d ] \ [ d/r ] in Lemma 3.

glyph[negationslash]

B.2 Acceptance probability of K-SEQ for the example in Figure 5

In this section, we show that for the example in Figure 5, K-SEQ achieves the optimal acceptance accuracy. In this case, p = U ( d ) and q = U ( d/r ) . Recall that the optimal acceptance probability is

For U ( d ) and U ( d/r ) , we have

And hence solving 1 -(1 -β ( ρ )) k = ρβ ( ρ ) gives ρ ∗ = r (1 -(1 -1 /r ) k ) . And be Theorem 1, we have

And the equality holds since this is an upper bound for any coupling.

B.3 Comparison to multi-round rejection sampling in [20, 29]

In this section, we compare our proposed draft selection algorithms (OTM and K-SEQ) to the multiround rejection sampling algorithm (MULTI-ROUND) in concurrent and recent work of [20, 29] (see Algorithm 1 in [29]) using the example of Bernoulli distributions. As Figure 6 demonstrates, both our proposed algorithms outperform their algorithm. The advantage of OTM is demonstrated by the fact it is the optimal algorithm under the validity guarantee of the final accepted token. Our proposed efficient algorithm K-SEQ also outperforms MULTI-ROUND for the considered examples. We leave a systematic comparison of the algorithms as future work.

2.5

5.0

7.5

10.0

12.5

15.0

k

0.5

0.6

0.7

0.8

0.9

1.0

Acceptance probability

b

=0.1,

k

-seq

b

=0.1, OTM

b

=0.1, multi-round

b

=0.5,

k

-seq

b

=0.5, OTM

b

=0.5, multi-round

b

=0.75,

k

-seq

b

=0.75, OTM

b

=0.75, multi-round

C Analysis of SpecTr

C.1 Proof of Theorem 1

We start by proving the following lemma on ρ ∗ .

Then we have Let ρ ∗ be the solution to Eq. (6) . Then when d TV ( p, q ) ∈ (0 , 1) ,

f ( ρ ) is monotone in ρ in [1 , ∞ ) ;

Proof. It would enough to prove the followings: (1) f ( ρ ) is monotone in ρ in [1 , ∞ ) ; (2) f (1) ≥ 0 ; (3) f ( k ) ≤ 0 ; (4) f ( max x q ( x ) p ( x ) ) ≤ 0 .

To see (1), since β p,q ( ρ ) is decreasing in ρ , so is 1 -(1 -β p,q ( ρ )) k . Moreover, ρβ p,q ( ρ ) = ∑ x min { ρp ( x ) , q ( x ) } , which is non-decreasing in ρ . Hence we have 1 -(1 -β p,q ( ρ )) k -ρβ p,q ( ρ ) is decreasing.

To see (2), note that when ρ = 1 , β p,q ( ρ ) = 1 -d TV ( p, q ) . Hence we have

When ρ = k , (3) holds since for x ∈ [0 , 1] , we have 1 -(1 -x ) k ≤ kx . Moreover, when ρ = max x q ( x ) p ( x ) > 1 , we have β p,q ( ρ ) = 1 /ρ and (4) holds since

Next we prove Theorem 1, we will break the proof into four parts: (1) computation efficiency; (2) π K-SEQ ρ is a valid transport plan; (3) acceptance probability; (4) optimality guarantee of π K-SEQ ρ ∗ .

Computation efficiency. Note that Lemma 4 immediately implies that ρ ∗ can be computed up to arbitrary accuracy δ in time O ( | Ω | log(( k -1) /δ ) using binary search over [1 , k ] .

Valid transport plan. We next prove that π K-SEQ ρ is a valid transport plan when ρ ≥ ρ ∗ . By Lemma 4, when ρ ≥ ρ ∗ , we have 1 -(1 -β p,q ( ρ )) k ≥ ρβ p,q ( ρ ) . Recall that p acc = 1 -(1 -β p,q ( ρ )) k , and

∀ x ∈ Ω , we have

this implies p res ( x ) ≥ 0 for all x ∈ Ω . Moreover,

Hence p res is a valid distribution. It remains to show that the marginal of Y is q . We first compute the probability of the output Y = x . Note that probability that Y = X 1 is

Similarly, probability that

Hence,

Hence

Therefore, glyph[negationslash]

glyph[negationslash]

Summing over all symbols x yields

Hence we have

Acceptance probability. The acceptance probability holds since

Optimality guarantee of π K-SEQ ρ ∗ . It can be seen that β ( ρ ) is decreasing in ρ , and so is 1 -(1 -β p,q ( ρ )) k . Hence we have

where

The statement holds since f ( x ) = 1 -(1 -x ) k kx in monotonically decreasing when x ∈ (0 , 1 /k ] and f (1 /k ) = 1 -(1 -1 /k ) k , lim x → 0 + f ( x ) = 1 .

Moreover, ∀ x ≥ 0 , kx ≥ 1 -(1 -x ) k . Hence we have

where the last inequality is due to the upper bound in Lemma 3 with Ω 0 = Ω .

C.2 Proof of Theorem 2

Weprove the theorem via induction. When L = 1 , τ ∈ { 1 , 2 } . Let k = | S | . Since for the first step, f π in Algorithm 3 is a valid transport plan from M s ( · | x t ) ⊗ k to M b ( · | x t ) . We have Y 1 ∼ M b ( · | x t ) , which completes the proof when τ = 1 . When τ = 2 , we have Y 2 ∼ M b ( · | x t , Y 1 ) as stated in Step 5 of Algorithm 3. Hence the statement holds.

glyph[negationslash]

Suppose the theorem holds for L = glyph[lscript] ≥ 1 , we next prove that it holds for L = glyph[lscript] +1 . Let Y τ x t be the output sequence given context x t . When τ = 1 , since for the first step, f π in Algorithm 3 is a valid transport plan from M s ( · | x t ) ⊗ k to M b ( · | x t ) , we have Y 1 ∼ M b ( · | x t ) . When τ > 1 , S next = ∅ and by the assumption in Eq. (8), S next contains k ′ = | S next | drafts from M s ( · | ( x t , Y 1 )) with length glyph[lscript] . Let Y τ ′ x t ,Y 1 be the output sequence given context ( x t , Y 1 ) , by the induction assumption, we have for any τ 0 ∈ [1 , glyph[lscript] +1] , and any τ 0 -length, sequence o τ 0 = ( o (1) , . . . , o ( τ 0 )) ∈ Ω τ 0 , we have

Note that in this case τ = τ ′ + 1 , and for any ( τ 0 + 1) -length sequence o τ 0 +1 = ( o (1) , . . . , o ( τ 0 ) , o ( τ 0 +1)) ∈ Ω τ 0 +1 , we have

Combining the two cases, we complete the proof.

D Candidate set construction via a prefix-tree

As discussed in Section 1, the size of the draft set S is constrained by the number of parallel computations that can be supported in the hardware. Hence it is important to design the draft set carefully to allow for a longer sequence of accepted candidate sets. In addition to the i.i.d. draft set selection approach listed in Section 7, we present an algorithm that samples a draft set that forms the leaves of a prefix tree. Given a draft set size K , the algorithm can be specified by a sequence of parameter ( k 1 , k 2 , . . . , k L ) satisfying ∏ L i =1 k i = K .

The algorithm starts with a root node with sequence x 1: t and forms a prefix tree of depth L . At depth i ∈ [1 : L -1] , each node is expanded by a factor of k i +1 and each of its children will contain a sequence that satisfies: (1) Its prefix agrees with the sequence in the parent node; (2) The next token is sampled from the conditional probability given the prefix in small model. These child nodes will be at depth i +1 and the process goes until it hits depth L . We give a detailed description of the algorithm in Algorithm 4.

Algorithm 4 Draft set selection via prefix-tree.

E Additional experiments

In this section, we perform a detailed investigation of different factors that affect the speed of SpecTr with smaller transformer models. We train decoder-only transformer models on the LM1B

dataset based on the example provided in the FLAX library [15]. For the draft model, we use transformer models with 2 M , 6 M and 20 M parameters, and for the large model we use a 97 M parameter transformer model.

We first provide a verification of the computational model introduced in Section 1 by reporting the latencies of using the large model to compute the probabilistic distributions with parallelization over time and batch axes. As shown in Table 2, the latency stays roughly constant in these setting.

Similar to Table 2, we report relative latency when parallelizing across the time and batch axes using the small 6 M draft model in Table 3. In Table 3, the reported relative latencies are relative to the large 97 M model to get a sense of the relative cost of sampling multiple drafts with the small model compared to the large model.

To see how the size of size of the draft model will affect the block efficiency, we also include results for varying draft model sizes with the same 97 M large model for LM1B in Table 4. These draft models were produced by either halving ( 2 M ) or doubling ( 20 M ) the original 6 M draft model's number of layers, embedding dimension, MLP dimension, and number of attention heads. As expected, the larger draft models improve all speculative methods' block efficiency with SpecTr maintaining the best performance across all draft model sizes.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes SpecTr, a new speculative decoding framework for efficient Transformer and LLM decoding. \nSpecTr uses a combination of a small, fast model for generating draft samples and a larger, more accurate model for scoring and validating them in parallel. \nCompared to the prior works, the paper presents a formulation of the speculative decoding process through the lense of optimal transport theory and demonstrates the potential for further acceleration by generating multiple drafts in parallel along the batch axis. \nThe authors offer an approximate solution to efficiently solve the optimal transport problem at scale, resulting in significant latency improvements compared to baseline methods and previous speculative decoding approaches.\n\nSTRENGTHS:\nThe paper presents theoretical justifications for the methods that it proposes. Furthermore, it is the first attempt to incorporate parallelization along batch into the speculative sampling framework for better efficiency.\n\nWEAKNESSES:\n1. While the paper proposes theoretical justifications for the proposed method, it rather lacks proper evaluation. Since accelerating decoding processes is a very practical area, the author should provide a more thorough analysis of the end-to-end latency and text generation performance. In particular:\n   * (a) The paper lacks an evaluation of the impact of the proposed method on text generation performance. Tables 2 and 4 in the paper only focus on latency, and there is no comparison in terms of text generation quality (e.g. measured by metrics like BLEU scores, etc.) compared to the baseline or prior speculative methods. Latency numbers without these performance metrics make it difficult to assess the effectiveness of the proposed method.\n   * (b) The number of decoded tokens per serial call (in Table 2) is not directly indicative of the actual latency and does not necessarily represent the actual speedup. That is, it serves as a proxy measurement that may or may not yield the same degree of latency improvement. This is because there are factors such as framework overheads [1], hardware utilization, and others that can potentially impact the latency. For instance, it is difficult to state if processing 3 tokens (as in Table 2) in a single call is truly 3 times more latency efficient than processing them sequentially in 3 separate calls. Therefore, the claimed 3X speedup, and the additional 1.36X improvement over speculative decoding mentioned in the abstract, can be misleading.\n   * (c) It is unclear whether the proposed method introduces any additional run-time overhead compared to the prior speculative sampling method. If there is indeed extra overhead involved, it could potentially reduce the gap between the two methods, making the improvements presented in Table 2 less significant (thereby making the 1.36X improvement over the prior method misleading as well).\n   * (d) The latency overhead of running smaller models concurrently remains uncertain, since Table 2 only presents the efficiency of running the large model. For instance, a longer sequence length (L) can decrease runtime costs for the large model by increasing the number of tokens decoded per call as stated in the table. However, at the same time, it also increases the small model's cost as more tokens are generated but rejected/discarded. The paper lacks discussions around this point. While the authors presented the latency overhead of 10-15%, it is not clear under which particular setting it was measured and how the hyperparameters (K and L) affects this value.\n   * (e) Taken together, the authors should provide the end-to-end inference latency and the text-generation quality/performance measure to make the evaluation convincing.\n\n2. Including a concluding paragraph in the paper would enhance its professionalism. The authors should at least wrap up the paper with a one-paragraph summary.\n\n[1] The Framework Tax: Disparities Between Inference Efficiency in Research and Deployment, https://arxiv.org/pdf/2302.06117.pdf\n\nQUESTIONS:\nN/A\n\nLIMITATIONS:\nPlease see the weakness section", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies speculative decoding, a technique to increase inference efficiency in a large autoregressive model by sampling a set of candidate tokens (a *draft*) from a smaller (thus, faster) model, which is then scored according to the conditional distribution of the original larger model. The scores of the larger model are obtained in parallel, thus potentially resulting in a significant speedup depending on how well the smaller model's conditional distribution is close to the larger model's. The main contributions are twofold. First, the authors cast the speculative decoding problem as a maximal coupling problem in optimal transport. Then, they propose to use multiple drafts (i.e. multiple sets of tokens completing the context), which is solvable in exponential time in the number of drafts. Thus, a novel algorithm is proposed that runs linearly in the number of drafts and provably solves the problem up to a $(1-1/e)$ factor of the optimal acceptance probability.\n\nSTRENGTHS:\n1. To my knowledge, translating the speculative decoding problem in the context of optimal transport is not very surprising but novel. Furthermore, by connecting the problem to the field of optimal transport, the paper opens new possibilities for research, due to the maturity of the theory of optimal transport, which the authors actually use to derive an extension of the original speculative decoding algorithm to allow for multiple drafts and discuss its feasibility in terms of algorithmic complexity.\n\n2. The proposed algorithm to approximately compute the transport plan (Algorithm 2) nicely follows the problem casting performed earlier and is backed by a sound theory (Theorem 2) describing how the solution (in terms of acceptance probability) compares to the intractable case of exact solution in exponential time. I have quickly skimmed through the appendix and although not an expert in OT, the theory seems correct.\n\nWEAKNESSES:\n1. The idea of extending the speculative decoding algorithm in [15]  to multiple drafts is simple and elegant, but sometimes the authors add some unnecessary details. For instance, it seems trivial and very intuitive that by having multiple drafts, the acceptance probability increases with the number of drafts (hence Lemma 1 seems unnecessary). \n\n2. The authors test their proposed method solely on LM1B. I would have liked to see a larger set of experiments, covering at least a more significant subset of those performed in [1].  \n\n3. Algorithm 1 is wrong, as it always returns at line 7. I guess the authors want it to return at line 5 as well. Also, some quantities are not defined yet when Alg.1 is introduced (such as $\\mathcal{X}$), which complicates a bit the readability.\n\nMinor: \n\n3. In Figures 1,2,3 the size of ticks, labels, and lines should be increased. \n\n4. Typos: line 251: \"Gvien\". Line 217-218 the sentence \"To control the probability of accepting an $x \\in \\Omega$ with probability larger than $q(x)$.\" does not seem to be grammatically correct. Line 44 \"a several contexts\".  Line 120: \"an discrete\"\n\n[1] Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. arXiv preprint arXiv:2211.17192, 2022.\n\nQUESTIONS:\nNone.\n\nLIMITATIONS:\nThe authors address the limitations. I do not see any potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents speculative decoding, wherein a smaller language model is used to approximately sample from a large one, akin to\na type of accept/reject MCMC sampler.\n\nThe idea is that it is slow to sequentially sample from a large model of interest, but joint probabilities can be computed in parallel across the time dimension. At the same time, we could have access to many parallel copies of a smaller not-as-good model, and it's quicker to generate e.g. 100 tokens from the small model than to sample from the large one. We generate e.g. 100 tokens for K such smaller models, and accept one of the K sequences with some probability.\n\nThe paper is mostly a theory paper that gives bounds on the transport cost, between the large model distribution and the sampling distribution resulting from the accept/reject sampling composed with proposing from the smaller models.\n\nOn the plus side, the paper is a theory paper that I think has some nice implications for practitioners. I would love to see some more empirical work exploring this. The quality of the math / theory / bounds is high and the authors did a nice job breaking down the OT parts for the typical ML language modeling reader. \n\nOn the minus side, while a light set of experiments is totally fine for a theory paper, I believe most readers of the paper would benefit more if there were also some initial results of the actual sample quality from the presented algorithm. There seem to be results on latency and accept/reject probability (and there's some theory connecting accept/reject probability to closeness of sampled distribution to desired large model distribution), but there seem to be no direct results evaluating the samples resulting from the new algorithm. \n\nIn short, I tentatively accept but think the minus needs to be addressed and will change the score if there is no discussion or clarification on this. Note that light evaluation is totally fine since the theory is good. And it is even okay if the quality of the samples according to any of the metrics is possibly below competitive, since there is always room to improve these algorithms once they are established.\n\nSTRENGTHS:\n- novel algorithm for quicker sampling making use of smaller models and MCMC-like concepts to sample from a larger model\n- nice framing in terms of optimal transport that I think the typical ML + LM reader could take something positive away from\n- rigorous characterization theoretically of the algorithm in terms of transport cost bounds, accept/reject rates, etc.\n\nWEAKNESSES:\n- Table mismatch: e.g. referenced table 8 in main text but no table 8 in main text or appendix\n- Most noticeably, there seem to be no direct reports of quality of samples from the resulting approximate algorithm. There are tables for latency and for accept/reject rates (the latter is in appendix but could probably go in main text). And there is theory characterizing that the transport cost will not be too bad, etc. But, there are no direct metrics quantifying the original large model samples vs the proposed algorithm's samples (those produced by using the smaller models' sampling and larger models' scoring). Seems like any of the basic LM metrics would be warranted here\n\nGlad to be corrected if I have misunderstood.\n\nMinor:\n- Small type \"Gvien\" in source line 251 of the PDF.\n\nQUESTIONS:\nSee Strengths/Weaknesses\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-14"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a novel and efficient decoding algorithm for autoregressive large language model, SpecTr, which is an extension of speculative decoding. Given a large model M_b, it can only output one word at a time when it decodes, however, the cost of a serial call is quite expensive. The previous method speculative decoding can alleviate the slow decoding problem, which uses a much smaller model M_s to decode a single segment of length L and allows M_b to calculate its likelihood in parallel. The optimal transport plan between the distributions of the two models is used to determine whether to accept a certain part of the segment, which makes sure the samples has the same distribution as M_b. This approach has the advantage of reducing the number of serial calls of M_b, improving decoding efficiency. However, if the acceptance probability is too low, the reduction may not be significant.\n\nThe SpecTr algorithm allows to utilize K i.i.d. segments sampled by M_s to enhance acceptance probability. Concretely, the previous optimal transport problem is modified to transport the K-product distribution of M_s to the distribution of M_b. Since solving the new problem has an exponential cost with respect to K, the paper introduces the k-sequential selection algorithm with a 1-1/e approximation ratio with an appropriate cost.\n\nIn numerical experiments, they use an M_b model with 97M parameters and an M_s model with 6M parameters tested on LM1B, speculative decoding achieves an average of decoding 2.3 tokens per serial run of M_b, while the proposed method in this paper can decode 3 tokens, showing a significant improvement.\n\nSTRENGTHS:\nThe paper is well-written and self-contained, making it accessible even to someone like me who is not familiar with the related work on language model decoding. The proposed algorithm, SpecTr, combines the practical need for parallelization of large models in real-world scenarios. Moreover, it is mathematically concise, intuitive, and insightful. The paper also provides a basic theoretical analysis of the properties of the algorithm. The numerical experiments show a significance improving of their methods. \nI think this is a good submission.\n\nWEAKNESSES:\nThe numerical experiments on real datasets seem insufficient, as they only compare with speculative decoding when K=1, which may not be fair: using larger K would require more computational resources. It would be worth considering if there is a more equitable way to compare the methods, e.g. extend speculative decoding to use more resources in some crude way.\n\nQUESTIONS:\nThe first line of Algorithm 1:\nInput: ... \"X~iid p\" is confusing, since you have only one sample here, \"iid\" is unnecessary.\n\nLIMITATIONS:\nNo", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-26"}