SpecTr: Fast Speculative Decoding via Optimal Transport
Authors: Ziteng Sun
Ananda Theertha Suresh
Jae Hun Ro
Ahmad Beirami
Himanshu Jain
Felix Yu
Venue: NeurIPS 2023 poster (2023)
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
Ziteng Sun" Google Research, New York zitengsun@google.comm
Ananda Theertha Suresh" Google Research, New York theertha@google.c
Jae Hun Ro Google Research, New York jaeroßgoogle. com
Ahmad Beirami Google Research, New York beirami@google. com
Himanshu Jain Google Research, New York himj&google,
Felix Yu Google Research, New York felixyuggoogle. com
Abstract
Autoregressive sampling from large language models has led to state-of-the-ar results in several natural language tasks. However, autoregressive sampling gener- ates tokens one 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 tokens), and then score all tokens 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 guarantee that the final output follows the distribution of the large model. In this work, we provide prin- cipled 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-come maxinul-coupling problem. This new formulation enables us to generalize the speculative decoding method to allow for a set of & 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 We then propose a valid draft selection algorithm whose acceptance probability is (1 - -optimal multiplicatively. Moreover, it can be computed in time almost linear with size develop
domain of single token. Using this new draft selection algorithm, new autoregressive sampling algorithm called SpecTr, which provides speedup decoding while ensuring that there is quality degradation in the decoded output. We experimentally demonstrate that for state-of-the-ar large language models, the proposed approach achieves wall clock speedup of 2. 13X,, further 1.37X speedup over speculative decoding on standard benchmarks.
1 Introduction
Autoregressive language models have shown achieve state-of-the-art results in several natural language tasks [2, 5, 25, 26]. During inference, given a context =':=x(1), ()))).,, autore- gressive model Mr generates successive tokens ([+1), ... via temperature sampling [1. 10], where the next token xx(1 + 1) is drawn from the temperature-scales distribution Ms((-l²)). If the temperature is zero, i.c., greedy decoding, the next token determined by the maximum likelihood method i.c., rtt+1(+++)))) arg max Mp(z(x"), where Ω is the domain of single token also referred to as the vocabulary. The sampling approach can further combined with other sampling primitives such nucleus sampling [16] and top-k sampling [9, 22].
'Equal contribution.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

All these approaches are autoregressive decoding" methods, where tokens are generated serially one after another, which can be slow or even prohibitive in several applicationn [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 2d to denote the sequence zri),, æ(i + T(j) and when i == 1, we simply use x² = وا z(ii denotes the i-th entry of æ. Subscripts are used to distinguish between different sequences, e.g, zi and zy denote two sequences of length t. We use [n] to denote the set {l....,7) n.).
A simplified computational model.
Standard inference.
Given a context x², with O(t²) computation and ((1) time, an autoregressive model Ma can compute MM(yl>="), the ((emperature-scaled) probability of all possible next tokens y E Ω
Parallelization along the time axis. Given context x', with ((1)) computation and (11) time, an autoregressive model M, can compute Mo(ylx"), , for all y E 52 and e (1,2...., tJ.
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¹, ri. rkk. with (Ktt²) computation and O(1) time, an autoregressive model Mb can
compute Ma(y) for all E S2, i E [4]. and j E [K].
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 large transformer model in practice. We also note that there are efficient transformer architectures, which reduces the computation cost from ((²²) to Ott 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 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 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- penxeivv draft model Ma. which predicts the next token given the context, and the predictions of M, are
PIn this work, generating tokens from a language model. 'When the assumption holds, one could naively batch multiple decoding contexts, which improves decoding throughput, but not the latency of each context.
the words sampling and decoding interchangably to refer the process of sequentially
2

---

This paper will
Compute conditional probabiliies by the large model
Sample one draft from the small model
Context
Nost decation (f not E. D.S.)
Add to context
Draft selection via maximal coupling
This paper will be liked by
be liked for by
Figure 1: One iteration of speculative decoding [19, 4]. Tokens in blue are decoded tokens from previous iterations, which are used as context for the current iteration. Tokens in red are drafts from the small model based on the context. The underlined tokens are the newly decoded tokens in the current iteration, where underlined red tokens represent tokens selected from the draft and underlined green token is selected from the residual distribution.
close to that of M2 for most contexts. Suppose have obtained prefix x'. The next iteration of the speculative algorithm can be broken down into three steps (see Fig. for an illustration).
1. Draft construction. The draft model is used to efficiently and "speculatively" sample L tokens, z(t + 1), ......., ,(l + L). We keep the conditional probabilities on the next token M2 xf, (4,4+1++)) for each i < L and Vy E Ω.
2. Conditional probability computation. After observing the samples, we compute the conditional distributions MMsyy x² aef,5'1:5t 1:t+)) for each i s L and Vy Ω in parallel (along time axis) in O(1) time.
3. Draft selection. Validate and select first L' of the L tokens and set xrt + i) z(t for & L' given the draft sequence and the conditional probabilities from both models. Sample token ztt L'' + 1) from residual distribution correction to the rejected token."
After this step, we use xi +++++1 the next context and sample the next few tokens using speculative decoding iteratively. For complete statement the algorithm, we refer the readers to [19, 4]. The crux of the above steps is draft selection, which given 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 token-level maximal coupling algorithm, which is provided in Algorithm 1. Note that for the draft selection, Algorithm I is applied where is the conditional distribution of the draft model M,( x*) and is the conditional distribution of the large model Ma x²) (which may be further conditioned on the newly decoded tokens).
Algorithm Token-level maximal coupling
Input: Distributions Draft sample X ~
1: Compute the residual distribution yes where Vz 9. (x)
q()) (p((m()((()))()) 1-E,, min(p(a' ).9((()))
2: Sample 1) ~ UO,, 1).
3: in' S. min (1, g(X)
then
p(X)
4: Return Y X.
(Accept the draft token.)
5: end if
6: Return Y y (Sample corrected token from the residual distribution.)
Algorithm returns random variable Y which either is the accepted input or a sample from the residual distribution pts which is defined in Step 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 correction. Previous works showed that X p. then Y ~ g [19, 4J. In the case of the draft selection, this means that the output of the algorithm is distributed according to My x'), which is exactly the
'See Algorithm for definition of the residual distribution. When L' = L. no token is rejected. The residual will just be the conditional probability Ms(-] x f.. which gives extra decoded token.

---

be ilked be (ked all not get good reviews be read four
This paper will
Compute conditional probabilitis by the large model
Semple multiple drafts from the small model
Context
Nest deraaion (f not O.S.)
Draft selection via OTM transport plan
context
This paper be liked by three
be (ikd by ell Uhree
Figure 2: One iteration of SpecTr. Tokens blue are decoded tokens from previous iterations, which used context for the current iteration. Tokens in red are drafts from the small model based on the context. The underlined tokens are the newly decoded tokens in the current iteration, where underlined red tokens represent tokens selected from the draft and underlined green token is selected from the residual distribution. See Fig. 3 for a more detailed run of the draft selection step.
desired outcome. Furthermore
Pr(Y = X)
Σ min(p(=). = dTv(p.)). #EQ
where drv is the total variation distance between and q. The closer p and g in drv, the higher the chance of Pr(Y = X), , and fewer the number of serial calls to the larger model. In the ideal i.e., the draft token is always accepted, and when used for case, if p q, then Pr(Y- = X) = L. Together with the extra sampled token' from Mi. L + 1 speculative decoding we have tokens are obtained in one iteration. In such a case, based assuming the decoding time of draft model is negligible, the speedup is (L + times.
computational model (Section 1),
Our contributions
From 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-theen theoretic sense?
Speculative decoding uses parallelization along time to speed up decoding; would possible to use parallelization along batch (number of drafts) further improve decoding speed?
We provide answers to all the above questions in this work. first relate the problem of speculative decoding to the broader and well-studied discrete optimal transport theory through 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 maximal coupling [8]. Based on the connection to optimal transport, we show that 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 discrete optimal transport problem with membership cost, which is referred to as OTM. Discrete optimal transport can be solved with linear program, but the number of variables is exponential in batch size, which can be prohibitive. To address this, we propose valid transport plan that can be efficiently computed. Moreover, it achieves (1 - apprrximaaiinn of the optimal acceptance probability (Section 6).
With the theoretically motivated algorithms and guarantees, we circle back speeding up decoding and propose 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 valid sequence from multiple draft sequences. We remark here that the latter requires completely new statistical
When L' = L. xtt L + 1) is sampled from Ma (-1 LL....

---

be
liked
for
its
liked
by
three
This peper will
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 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' 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 It select it and outputs "three" correction. The process ends and a total of 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-an large language models, SpecTr achieves a wall clock speedup of 2. 1.33, 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 token-level draft selection problem. In particular, given context x let X XI,... X be a collection of draft tokens sampled from the small model, e.g.,ampled i.i.d. from M.(- x²). Note that by our assumption of the computation model, could compute the following conditional probabilities from the large model in parallel along time and batch axes):
Ma x²)) and
E E [k]] Mal. | xx)))..
The goal of the draft selection algorithm f : 2 + Ω is to output Y f(x*), whose distribution follows My(- æ²), and hence valid sample from the large model. Moreover, when Y {X1,..., X+). we could sample an extra token from Ma(. x²,Y) without calling M, since we have already computed the conditional probabilities Ma(. v,,Y). Hence we would like to maximize the probability that we accept one token from the set of drafts.
When L>> > the drafts sequences sampled from M,,, sequence of token-level draft selection algorithms could be used along the time axis to select valid sequence from the MM. See an example in Fig. 3. The full details about the sequence-level selection algorithm provided in Section
The reminder the section will be focused on the toven-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 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 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 3-5, 5) to each draft token X; with p = MM x²) and q = Ma( x²). However, this approach would not guarantee that the final accepted token is from the desired distribution. To see this, consider the example of Ber(1) and we have Vi = I,.... k, X₁ and each of them will be accepted with probability 1/2. After
Ber(1/2). Then
BBer(b) denotes Bemoulli distribution with the probability of seeing head
5

---

applying Algorithm 1 to all Xi's, the probability of getting a 1 will be at least 1 - 1/2 and hence the output distribution would not be Ber(1/2) for k > Therefore the algorithm does not produce valid samples, which is 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 draft selection with one draft. The generalization to multiple drafts is provided in Section 5.
Coupling and optimal transport. To simplify notations, assume Ω discrete domain.
Definition (Coupling), For two probability distributions P over X' and Q over y we say a joint distribution π supported over X' is coupling between P and Q if Vz, y, (x, 2 0,
Vy E y,
Σ TEX
イ(. v) = Qu),,
and Vo E X,
Σ VEY
(1. v) = P(x).
We use II( 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
fx X + y introduced by the conditional probability π (y .);=,(<,,)) y) (()) as a coupling, which is also referred distribution of the draft tokens and Q to be the target distribution of the output token. In this case, the fx valid draft selection algorithm. Formally, this is stated the claim below.
transport plan from P to Q [27]. In this paper, we will set P be the
Claim 1. For all ι E (P,, Q), let fa be the probabilistic mapping defined above. If X P, then (X) ~ Q
In this paper, we will design selection algorithms by finding valid couplings between the draft distribution and target distribution guarantee validity of the output tokens.
Definition 2 (Optimal Transport (OT) [27]). For cost function e X ) - R4, the transportation cost of a coupling defined as:
C(T) = Ex,y ()]
The optimal transport plan is the coupling II().Q) that minimizes the transportation cost.
Speculative decoding with one draft token. With these definitions in place, we can see that with X' y S2, the domain of the tokens and P = p, Q q. recover the speculative decoding objective with one draft token using the cost function of indicator cost, which captures the resampling cost, detined below:
<(x, v) = (v + x). The transportation cost of the coupling will be C(T) Ex, [1 {Y + x3] Px,Y (Y++)) X). This optimal transport cost is known to be
Vx E Ω E Ω
min Px FX,Y x Σ minn(p(z) rell(p.g)
(1)
ZEQ
which is achieved by the maximal coupling between and 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' 2 for some k E N., which is the space of & 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, use the cost function of membership cost, defined below:
Yz E 2², y E Ω
c(x. v) = 1 {y & S(x)
where S(x) to Ω o appears in æ) denotes the set of distinct elements in When k = recovers the indicator cost mentioned before. The transportation cost of the coupling is
((T)) = Ex [L (Y & S(X)) Px, (Y & S(X)).
(2)

---

We will also refer to the above cost ( as the rejection probability due to its probabilistic interpre- tation. And similarly, a(r))== ((²) = =x,yy= E S()) 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 optimal transport with membership cost (OTM). We use n² to denote the coupling that minimizes this cost π = arg miD.. C(T);; and the cost C(T*) is referred to as the optimal transport cost between P and Q. We use oe Q) = (()") to denote the corresponding optimal acceptance probability.
Draft selection with i..dd draft tokens. In this paper, we will mainly focus on the case when the draft tokens are i..dd samples from base distribution." Let p. be supported over S and the goal is to obtain one valid token from g given i...d. samples from p. For SpecTr with context x', we have p (- x²) and q Ma(- x²). We set P p0* product distribution whose marginals are all p. and Q = q. The OT problem we want to solve is the following:
min C(T) s.l. π E (T(p),,))
(3)
We overload notation and denote the optimal acceptance probability as az (p..))=(((pp)) 1 C(T*). To better understand the quantity, we state a few properties about ag. Lemma 1. (Appendix A.2) The optimal acceptance probability statisfies the following properties.
Monotonicity. For any p. and 2 1, (p. S $atilp..)) q))
Consistency. Vr Q (xx/p(x)) is bounded, we have lime (p, = 1. Else, lims- 0n (p,9) = Eures {p(x) > 0} (xx).
The above properties demonstrate that for large k, the value of a can become large. Hence increasing & could increase the acceptance probability, leading to further speedups. We now focus on computing the optimal transport plan and the optimal acceptance probability.
OTM via 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: min
Σ (Tt tu $ S(x)} s.t. n E
(4)
EE0 VER
The linear program in (4) has [02|k++ variables and (S2 + |o] equality constraints (see Definition 1). Linear programming can be solved in time polynomial the number of variables and constraints [7, 21, 18)," implying the following lemma.
Lemma 2. Given g over Ω the solution to Eg. (3) can be computed in time ((12))(()))
We refer to the optimal coupling obtained above as OTM-k and denote it as π OTTM- When A = 1, there is closed form expression for the optimal acceptance cost (see Eq. (I)), whereas for larger values of k. we are unaware of a general closed form expression. In Appendix A.1, we provide an information- -theorrtttoot upper (and lower) bound, which is tight up to multiplicative constant of 1 (1-111/18) 2 1-11 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 prohibitive when either the vocabulary size [S]] or the number of draft tokens l is large. 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)* Jak 2 (1 1/e)au.
"The existence of optimal coupling in discrete domain is well-known, e.g., see [27]. When the optimal coupling is not unique, we use x² denote one of the optimal couplings.
"The ahove generic formulation immediately allows generalization to more complex draft selection strategics, such as sampling & tokens without replacement. or using different drafting distribution for each draft.
"To our hest knowledge, the best practical computation bound (through interior-point method) is O sk [21] and the best theoretical computation bound is O 2.5k) [18].
FoFo discrete OT, Sinkhorn algorithm could be used to solve an entropy-regularin version of OT, which has better computation complexity [6]. However, the computation cost of the algorithm will still have linear dependence [ which can be prohibitive.
7

---

Draft selection via k-sequential selection
In this section, present sequential selection algorithm (K-SEQ), an approximate solution' to the optimal transport problem in Eq. (3), which can be efficiently computed in time almost linear in [S2] and logarithmic in k. The algorithm is presented in Algorithm 2.
Algorithm 2 k-sequential selection algorithm (K-SEQ).
Input: Distributions samples X Xx mi.idd. P. E [1, division factor. I: Let (p) = Exeo
p(x)/)) and Pace = 1 (1 -p))( Compute p where
q(x) min {D(x), 2(a)
p Prialo)
Vo E (s)
(5)
1 - Pace
2: for .2....,, do
3: Sample 7/i (O, 1).
4: if 7J I min (1, o(X.)
then
px(().)
5: Return Y X4. {Return the ith draft token.) }
6: end if
7: end for
8: Return Y ~ 1/ (Sample 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; based on the ratio q ((x:)/p(X; The algorithm output the first accepted sample result from a residual distribution pds if none of the samples is accepted. To guarantee that the the final returned token is valid sample from 9. we choose an appropriate E [1, k] and accept X2 with probability min(1, (X.)/(p p(X))) instead min(1. ((X.)/ as Algorithm 1. In Theorem 1, we show that with appropriately chosen p's, Algorithm 2 is indeed valid transportation plans from oa to q. Moreover, find the best transportation plan within the family, we only need to search over a single parameter p. which reduces the computation cost significantly. We also show that searching over this sub-family of couplings won't 't decrease the optimal acceptance probability by multiplicative constant. The performance of Algorithm 2 is stated in Theorem 1.
Theorem 1. Let (P) = Exeo min (p(x), 2(x) and p" be the solution to the identity below.
1- (1 - Bp. (p)) = pB).a()).
(6)
When p 2 p², the coupling u SEQ in Algorithu is a valid transport plan from z* to q. When p", we have
a(=x-St (1 (1- e (p,
Moreover, p² can be computed up to accuracy in time O[[ log((k: 1)/6)).
We provide the proof in Appendix C.I. In Appendix B, using a few canonical examples of distri- butions, plot the acceptance probability of K-SEQ and compare it with the optimal acceptance probability a'g. It can be shown that K- 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- 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 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 SpecTrr: Application of OTM in autoregressive sampling
In this section, we describe how OTM can be used 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):
Note here that the solution still satisfies the constrains in Eq. (3), and hence is a valid transport plan. The term approximate here means that the solution is not the exact minimizer of the cost in Eq. (3).

---

1. Draft set construction. Given current context z', use the draft model sample a set of K draft sequences with length L, denoted by S {2" M,
x²)}. We keep the conditional probabilities My(y (y | =²) for all y U2 S I. and こ4 E
2. Conditional probability computation. Compute the conditional probabilities on the next token for the large model Ms(y x², for all y 22 S L and 24 E S in parallel.
3. Draft selection. Select first L' of the tokens and set art + i) = z(t) for S L' and some
z E 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 ((1) when |S] is not large based on our simplified computations model. We mainly focus on the draft construction phase and draft selection phase.
Algorithm 3 Draft selection with multiple candidates (DraftSelection).
Input: Input sequence w²; draft sequence length: L; draft sequences S = (=f | |S|]-
1: Compute a transport plan (using linear programming in Lemma for an optimal solution or
Algorithm 2 for a suboptimal solution) from MM,(- x1)015| to Ma x²), denoted by Tt.
2: Get the multi-set of next token-level drafts: S_ = (1, (e(sS) and compute Y fx, (Sa).
3: if L = 1 then
4: if Y S. then
5: Sample Y" ~ MMs( (x', ()),
6: Return (a. Y,,YY,Y Y"). (Sample an extra token if the last loken is accepted.)
7: else
8:
Return (x", YY. (Stop and return the corrected token and previous accepted tokens.
9: end if
end if
11: Let Suext = [222 E S and x(1) Y} be the set that consists of sub-sequences of the candidates that agree with the selected next token.
12: if Smext = then
13: Return Y). (Stop and return the corrected token and previous accepted tokens.)
14: else
Return DraftSelection((s) Y), I 1, Snext)) (Keep the draft token and proceed to the next
time step.)
16: end if
Draft set with i..d draft sequences. Given context x' a natural way to come up with set of K drafts is to independently sample K draft sequences from Ma(- x²), i.e.
* z .... zk miidd. Ma x²).
dots The draft set construction method in (7) can be generalized to 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 available given any prefix in the candidate since they are computed in parallel in the second phase, and won't 't list them as inputs explicitly in Algorithm 3.
A sample run of the algorithm is presented in Fig, 3. The algorithm proceeds in recursive fashion. Given prompt x and a candidate set S sampled from Ma x²), the algorithm first computes a token-level draft selection algorithm fa ! + which a transport plan from M.(- | (11) to MM(- zz)) Then fa is applied to the set of first tokens of the draft sequences in S to obtained valid token Y from Mp(- x²). If is not the last token (L 2 2), we filter out sequences first token is not Y and denote the remaining sequences as Spext and feed it to the algorithm with context (xf, f Y) and draft length
whose
1. This goes on until we have L = or Suext = V.
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 Ma(- (ar', Y)) since this conditional probability already computed in the second phase. Due to the property of the token-level selection algorithms and the autoregressive structure of language models, can be shown that Y is always valid sample from Ms(- x²). Let T/' be the number of decoded tokens in one iteration. Note that this is 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 drasts in the set S are generated from the small model with input x', or more precisely, Vz E S,
ViE E [1, L],
2(1) ~ MMl- 1 |x'22"" (01,22=1)
(8)
Let (x', YT) be the output of Algorithm 3 where T is the length of the newly decoded tokens, then it satisfies that Y is distributed according to Moli..... x²). More precisely, For any
dotr
TO E [1, 11, and any To-length, sequence ṭºh = =(o(1),... (.(1),..... o(To)) 9 we have
Pr = TT = TD) = [ Mo(o(i) t)
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 't increase the time for a serial call to the large model. With these, the theoretical speedup compared baseline decoding will be the average number of decoded tokens per serial call, which is called block efficiency [19], defined below
Total number decoded tokens
Block efficiency
Number of serial calls to MM
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, might not be negligible; (2) Parallelization along the batch and time axis might increase the time for a single call to M₁;
(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 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-arr PALM-2 models with prompts from the one-billion language benchmark (LMIB) [3]. In Appendix E, we use a pair of smaller transformer models break down different affecting factors mentioned above. In Table 1, we use HALM-2-Geckee and SLLM-2-Bison as the small model and
Table 1: Experimental results on the LMIB dataset with PALM-2-Geece as the small model and PALM--BBison as the large model. Results are averaged over 1000 test prompts and 3 random seeds.
Algorithm K
L
Block efficiency
Relative wall clock speedup
(normalized by baseline)
Baseline
1.0
1.0
Speculative
4
2.4
1.67
SpecTr
4
3.1
2.08
Speculative
8
2.9
1.56
SpecTr
8
8
4.0
2.13
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. partieular, when T= 8 and K = 8, our proposed SpecTr algorithm has speedup of 2. 13x, further 1.37x increase compared to speculative decoding (K 1).
10

---

Acknowledgements
Authors thank Asaf Aharoni, Kwangjun Ahn, Badih Ghazi, Sanjiv Kumar, Teodor Marinov, Michael Riley, and NeurlPS reviewers for helpful comments and discussions.
References
[1] David H Ackley, Geoffrey E Hinton, and Terrence J Sejnowski. A learning algorithm for boltzmann machines. Cognitive science, 9():1477-669 1985.
[2] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et Language models are few-shot learners. Advances in neural information processing systems, 33:1877-19011 2020.
[3] 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 131.:3003000 2013.
[4] Charlie Chen, Sebastian Borgeaud, Geoffrey Irving. Jean- -aa-ittt Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.013366 2023.
[5] 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.023/1 02311, 2022.
[6] 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.
[7] George B Dantzig. Linear programming. Operations research, 50(1):42-47, 2002.
[8] Frank Den Hollander. Probability theory: The coupling method. Lecture notes available online pdf), 2012.
(http://websiites. math. leidenuniv.
[9] Angela Fan, Mike Lewis, and Yann Dauphin. Hierarchical neural story generation. arXiv preprint arXiv: 1805.0883.3333 2018.
[10] Jessica Ficler and Yoav Goldberg. Controlling linguistic style aspects in neural language generation. arXiv preprint arXiv:1707.0266: 2017.
[11] Tao Ge, Heming Xia, Xin Sun, Si-Qing Chen, and Furu Wei. Lossless acceleration for seq2seq generation with aggressive decoding. arXiv preprint arXiv:2205.103500 2022.
[12] Google AI. Introducing PaLM 2, 2023. google-palm- 2-aa--aareeeeaaaaeea -awge-language--aeeee-ar- model/.
https://blog http::///booggggggee//eeeooooo//
[13] Google PaLM-2 Team. PaLM 2 technical report, 2023.
[14] Wenshuo Guo, Nhat Ho, and Michael Jordan. Fast algorithms for computational optimal trans- port 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.
[15] Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Mare Zee. Flax: A neural network library and ecosystem for JAX, 2023.
[16] Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. arXiv preprint vrXiv:1904.09755: 2019.
[17] Leonid V Kantorovich. On the translocation of masses. Dokl. Akad. Nauk. USSR (NS), volume 37, pages 199-201, 1942.

---

[18] Yin Tat Lee and Aaron Sidford. Path finding methods for linear programming: Solving linear programs in (vrank) iterations and faster algorithms for maximum flow. In 2014 IEEE 55th Annual Symposium on Foundations of Computer Science, pages 424 433.. IEFE, 2014.
[19] 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.
[20] 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 Jía. Specinfer: Accelerating generative large language model serving with speculative inference and token tree verification, 2023.
[21] Ofir Pele and Michael Werman. Fast and robust earth mover's distances. In 2009 JEEE 12th international conference on computer vision, pages 1600467.- IEEE, 2009.
[22] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
[23] Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit. Blockwise parallel decoding for deep autoregressive models. Advances in Neural Information Processing Systems, 31, 2018
[24] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys, 55(6):1-28, 2022.
[25] Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, Heng- Tze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, et al. Lamda: Language models for dialog applications. arXiv preprint arXiv::201. 08239,, 2022.
[26] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. LLaMA: Open and efficient foundation language models. arXiv preprint arXiv:2302.13977 2023.
[27] Cédric Villani et al. Optimal transport: old and new, volume 338. Springer, 2009.
[28] 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.044887 2023.
[29] Chao Zhang Yuhui Li and Hongyang Zhang. Eagle: Lossless acceleration of Ilm decoding by feature extrapolation, 2023. https: //st:///ssessss google wwt.cccc/c/c///c http///or/vvee/eeeeeeeeeee 11m.
12

---

A Properties of optimal transport cost
A.1 Information-theorrtie upper (and lower) bound of o'+
Below we provide an information-theoneiicc upper (and lower) bound in Lemma 3, which tight up multiplicative constant of - (1 - 1/k)* 2 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 Vk 2 1, we have
(1 (1 (1-(11/8k)) (x 4) S cx(p. q) S ax(p.))
where
(p(p,))
min faC
Σ min {9(0), p(u))}} yellg
Σ min p(..) x"
E
d(n)
nE
(9)
In Appendix B, plot ag function of ls for few simple pairs of (p. q) illustrative examples. We note that the upper bound in Lemma 3 is tight for examples considered in Appendix B.
A.2 Proof of Lemma
We first prove monotonicity. By definition,
04 (p.
min Prxx,ywa (Y 5(x*)) TET 9)
Prxxy (Y S(x+)
max #Ell(pkk.g) .4)
Moreover, for any II .), we can construct &' 11(p@t+1 .4) by setting k+ E g y E Ω x x'( v) ((* .v)((<(k + 1)),
i.e., adding and independent sample from p to X* Hence we have
(x(pp,)) = max RE .9)
P (Y S((*))
Prxetly (Y E S(X*))
max # 9))
Prx (V E s(x
max r
Prxsily (Y E s(x
max
r'Ell(p#t11
(p,41(p,0)
Next we prove consistency. We start with the case when Vo E S, (<))/p(x) < oc. To prove this, we will show that Algorithm 2 with Pmax maxzen (z)) /p(z)) statisifies
lim (-(TK- = 10000
KEO < an (p. q). the above equation implies lim (p, 1. Notice that by Lemma 4 and Theorem 1, T T--SFO is valid coupling, and
Since (TT (<T...EE)
a n K-SBQ) (T = 1 - (1 3p.0 (pmax))
where (p) - min(p(x), ((x)) 1/ 0. Taking it concludes the proof.
For the case when ((x)/((x) is unbounded, there exists Ω such that g(x) and p(x) 0. Let
Poff Σ (p(x)) 0) g(x).
13

---

Let Zo be such that p(ro) > 0. We define g' such that
0,
if p()) 0, if p(r) and o Do,
q(x);
((x) + Pott
i = 20-
Then we have drv(g, Poft, and hence by subadditivity of transport cost,
On (p(p11 4) 2 cu(p, - Poff-
Moreover, we have Ve (x))(())p(a) < oc. Hence
lim (p 2 lim (p d') Poft
Poft
Σ
8-000
1+000
sEN
A.3
Proof of Lemma
For the upper bound, it would be enough to show that for any π E IIp .). and any Sn C S2 we have
Pr {(( < Σ min 1 -
-()))*)
+
Σ
min p(aa), Σ
(v))
VERo
a En
vES(*
Pr((
E Σ Pr(x* = = = y) 1 S(x²))
= Σ
Σ
Y -) {u E S(x*)
yEla shEft
Σ
Σ
Pr(x*= v) 1 1 E S(x*)}
yEn ee
Σ VE
Pr Pr S(x"), + Σ
Σ
Pr = xx,Y y)
x En yES(u*) ynnj
Σ min (y E S(x*)) ) ((0)) + Σ
min (x* x²))
Σ
{(1)}
y
a *
NES(a*)nog
Σ yElo
min(1 - (1- (1-p)))) {())) + min
p(r:),
q(v)}.
z zkgnt
(
For the lower bound, we show that K-SEQ achieves an acceptance probability of at least (1 (1 (/k)))āk(e q). see Eq. (11), implying the lower bound guarantee.
B Comparison between a17 K-SEO) and Og for simple examples.
We illustrate the acceptance probabilities for our proposed token-level selection algorithms using few simple examples and plot them in Figures 4 and 5. The analysis for these simple distributions is presented in Appendix B. and Appendix B.2.
Pairs of Bernoulli distributions. Let Ber(b) be Bernoulli distribution with probability b of getting head. In Figure 4, we plot the acceptance probability comparison between OTM-k and K-SEQ for different Bernoulli distributions
Ber(b) as a function of k when Ber(0.25). Note that when p = g (b = 0.25), the acceptance probability is always one for both methods. When p * the acceptance probabilities for both methods increase A increases before they reach one. When = 0.1 or 0.75, KSSF has worse acceptance probability compared to the OTM-k algorithm. When 1, the two algorithms have the same performance.
14

---

10-
1.-
0.9
0.9
0.8
0.7
0.1.1. OTM
0.6
0.7-
b.25. OTM
6.5.5, k-sgq
0.4
d, 453
0.6
b 0.75, k-50q
0.3
d5
2,5
5.0
7.5
10.0
12.5
15.0
2
6
10
Figure 4: Acceptance probability compari- son between OTM-k and K- K-SEQ when Ber(0.25) and q Ber(b).
Figure 5: Optimal acceptance probability (ax)) as a function of lo when ((d) for d 120 and q = (dg).
Pairs of uniform distributions. Let (d) denote uniform distribution over [d]. In Figure 5, we plot the optimal acceptance probability for different uniform functions y as a function of For these distributions, it can be shown that K-SEQ achieves the optimal acceptance probability only ore plotted. Observe that all acceptance probabilities are monotonically increasing and tend to one when + stated in Lemma 1.
Hence
B.1 Calculations for O'k+
In this section, provide sketch of optimal acceptance probability calculations for results in Figures and 5.
Figure 4: Ber(p) and Ber(y). The optimal acceptance probability is
(Ber(y), Ber(z)) = min(g. (1 - p)*) + min(1 - 2, x²)) Setting So = {0, 1} in Lemma 3 yields upper bound. For the lower bound observe that since [0, 1 {y S(**)} 1 and only if * is 0 1 Hence,
(10)
ax (Ber(pp)) Ber(4))
(x** # {0tt }² }} + max{=(Y = 0, x =
+ =
= 1, 1
dd (1 p)* max
x = =Y==11,7 =1,= }}
Consider the transport plan  given by * =1* 1) min(p* k g). क
= - min(p² .4)
T0* 0) min((1 - p)* g). and π (0*, ( 1) = (1-p) p)* min((1 - 1 4). can be checked that this is valid transport plan. To see this matches the upper bound on the optimal cost from Lemma 3, notice that
1-1
- 1)*h max(=(V = 0, x = 0*) + ((Y = =
=1,x=1117 1h)}
r² -11 (1- -)* min(p" k,,
- (
If pt & and (1 - p)* < 1 - q. then the above equation simplifies to 1 and (10) also simplifies to If pd > q and (1 - f < 1 - q, then the above equation simplifies to 1 + q pd and (10) also simplifies to the same quantity. Similarly, the proof applies for pd S q and p)* 1-9.
Figure p = Udd and q= Uddrr). The optimal acceptance probability
first prove a ((()), (L)) = - - 11/))*
S(X*) be the set of unique symbols in X Consider the following transport plan, where is drawn uniformly from S(X*) n [d/r] and draws a new uniform sample from [d/r] if S(x*) ) n [d/r] = 0. Observe that since (d) uniform over [d], this a valid transport plan and furthermore,
1 -1-1/r** by a construction.
as (U),, (d/r)) 2 Pr(S(X)* n [d/r] + 0) The upper bound follows by setting Ro [d] \ [d/r] in Lemma 3. cu(U(d)UUd/r)) )
(1 1/r)
[d/r]
(1 -1- 1
15

---

B.2 Acceptance probability of K-SEQ for the example in Figure 5
In this section, show that for the example in Figure 5, K-SEQ achieves the optimal acceptance accuracy. In this case, = (d) and = (d/r). Recall that the optimal acceptance probability is
an(U(d), = (1 (1-1/r)*.
For (d) and (d/r), we have
1
3(p) Σ
min hin(p(z).g(<//9) max p}
=EId
And hence solving - 3(p))* p3(() gives r (1-1/r)) And be Theorem 1, have
ž 1 -1- =
And the equality holds since this is 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 multi- round rejection sampling algorithm (MULTI-ROUND) in concurrent and recent work of [20, 29] (see Algorithm 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 the optimal algorithm under the validity guarantee of the final accepted token. Our proposed for the considered examples. We leave systematic comparison of the algorithms as future work.
efficient algorithm K- K-SEQ also outperforms
1.0
Acceptance probabiliyy
0.9
0.1, k-seq
b= 0.1, OTM
0.8
b= 0.1, multi-round
b 0.5, k-seq
0.7
b= 0.5, OTM
b 0.5, multi-round
b= 0.75, k-seq b 0.75, OTM
b 0.75, multi-round
2.5
5.0
7.5
10.0
12.5 15.0
k
Figure 6: Acceptance probability comparison among OTM-k, K-SEQ and MULTI-ROUND when = (0r(0.)5) and Ber(b). When b = 0.25, all three algorithms achieve an acceptance probability of 1, and hence omitted in the plot.
C Analysis of SpecTr
C.1 Proof of Theorem
We start by proving the following lemma on p².
Lemma 4, Let
f(p) = 1- (1 (1-BBpa(pp)) *²) p3p.o (p).
Then we have Let n² be the solution Eq. (6). Then when dr 4) E (0, 1),
-(() is monotone in p in oo);
16

---

p² E minkin max, / l].
Proof. It would enough to prove the followings: (1) f(p) is monotone in p in [1, 00); ; (2) f(1) 2 0:
(3) f(k) & 0; (4) f( mmax, [()
piz) 0.
To see (1), since (p)g(p) is decreasing in P. so is 1 (1 (p.) Moreover, ppp.
Σ, min(pp(x), (()),, which is non-decreasing in p. Hence we have 1 - (1 pp pp
is decreasing.
To see (2), note that when 1,8p,(()) - drv (rv(p,)) 9). Hence we have
1 - [1 - Bpg = (pv(p.))) 2 drv(p.
When p k (3) holds since for E [0, have (1 2) S kr. Moreover, when
a(x) p max, p(z) 1, we have Bp-(p) = 1/p and (4) holds since
1 (p- -01 1/p)* <1 = p. 1/p.
L
Next we prove Theorem 1, we will break the proof into four parts: (1) computation efficiency; (2) X-SBQ is valid transport plan; (3) acceptance probability; (4) optimality guarantee of m K-SBQ T
Computation efficiency. Note that Lemma 4 immediately implies that p" can be computed up to arbitrary accuracy in time log(() 1) using binary search over [1, ks..
Valid transport plan. We next prove that xp SEQ is valid transport plan when p 2 p². By Lemma 4, when p 2 p², we have (1-Bp.e(p)) 2 pSp.a (p). Recall that Pace 1 (1-8p.(pp))
and
((x)) - min pzz) g(x))
Pace Spi(p))
Vz E Q2 (x)
Paoe
Vr E Ωt we have
min p(z), g(x)
Pace
1 - (1- Bpg(p))*
pp) g(z) S q(x),
(p)
this implies p=*())
(x) 2 for all x E Ω. Moreover,
((x) min p(x)) 2(wr)
Pacc
Σ #En
10 (z)
Sp.lle)
1.
Pace
Hence ysa valid distribution. It remains show that the marginal of Y is q. We first compute the probability of the output Y = r. Note that probability that Y = X1 is
p(Xi) min (1,2(xx)) min (p(X)), g(X1)
Hence
min p(x), g(x))
Pr(Y = X, =x) = min
P
Therefore,
p(x), a(x))
Pr(Y X1) Σ min
=(())
p
Similarly, probability that
Pr(Y X2 = z) Pr(Y 4 Xa/X + Xi) = (1 (p) min (p(s), g(x))
Hence,
17

---

Pr (Y I, one of X is accepted)
k-1
Pr Pr (X1,.... X, are rejected, X;+1 is accepted, and Xi41 x))
Σ (1 Ap.()))) p(x)) min (1, g(x)
R-1
p(z)p)
p(x)), g(x) Σ
min
8 (p))
1=0
p(z). g(x) 1 - (1 3pn(p))
min
(p)
Summing over all symbols yields
Pr(one of x is accepted) = Σ min {
{ p(x),
g(x)
(1- (p)( = -1 (p)(
P (p)
Hence have
Pr x) = Pr (Y I, one one of X is accepted) + (1 Pr (one of X is accepted)) ) p
= min p(x). g(x)) (1 (p)
(p)(p)
(1--11- (1-(1-Bpg())1)) q(x)) -min p(x), g(z) 1-11- 1-11-Ap((P))
(p))(p)
(1 (p)
= =(x).
Acceptance probability.
The acceptance probability holds since
a( x-S00
Pr(one x is accepted) (1 (1 (p))()))))
Optimality guarantee n
K-SBO can be seen that 3(p) is decreasing a. and so (1
(p)) Hence we have
a(xx-Se0) > 1- (1
)) 21 (pg(k))* (v Σ Emin()((=),((=)}
where
(p( 1 1 (1- (-Bpp(())
E -
kBp.g( (k)
The statement holds since f(x) 1-(1-a)² in monotonically decreasing when x 1/K) and f(1/k) = 1 - (1 1/ lim f(r) 1.
bx
Moreover, Væ: 2 0, kr 2 1 - (1 - xt Hence we have
> (1 -1-1/k)) Σ min min(kp(z).q((z)
(1 (1- - 1/b)* Σ
minf1 - - x( {(x)}
(11)
2
2 (1 (1 1/k)* )x (p. 4), nonumber
(12)
where the last inequality is due to the upper bound in Lemma 3 with Pg Ω.
18

---

C.2 Proof of Theorem 2
We prove the theorem via induction. When -= r E (1,2]. Let 1S]. Since for the first step, fo in Algorithm is valid transport plan from ( =²) to Mal. 1²)) We have Y1 Mo(. x²), which completes the proof when = 1. When r = 2, we have 22 ~ Mol. x², Vi) as stated in Step 5 of Algorithm 3. Hence the statement holds.
Suppose the theorem holds for = f 2 1, we next prove that it holds for L = t 1. Let Y be the output sequence given context x'. When r 1, since for the first step. fo in Algorithm 3 is valid transport plan from M, x² joo to MM x²), we have Y1 ~ ~M,(- | z¹). When 7 > 1, S Smext * and by the assumption Eq. (8), Snext contains k' = |Snext] drafts from M,(- (x*, Yi)) with length É. Let Y at,Y, be the output sequence given context (x*, Yi), by the induction assumption, we have for any E [1,
o(To)) E 2 we have Pr (Y = = -To) -IT Mo(o(() x², NY,00
& I], and any To-length, sequence 0" =
Note that in this case (o(1), ...). o(To), o(To + 1)) E f2T0+1, we have
r' + 1, and for any (To + 1)-length sequence o
Pr Pr = 0 Tot || = +1 = PX = o1) Pr (v = =000 = To)
-Mo(((1) x²) Mo + +1) x²)
=I xt o')
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 the hardware. Hence it is important to design the draft set carefully to allow for longer sequence of accepted candidate sets. In addition 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 prefix tree. Given a draft set size K, the algorithm can be specified by sequence of parameter (k1, kz,.................. ......... k2) satisfying I K K.
The algorithm starts with root node with sequence xht and forms a prefix tree of depth L. At depth i E [1 : - 1]. each node is expanded by a factor of k,+1 and each of its children will contain a sequence that satisfies: (I) 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 + 1 and the process goes until it hits depth L. We give detailed description of the algorithm in Algorithm 4.
Algorithm 4 Draft set selection via prefix-tree
Input: Input sequence z²; expansion factors at each level: (k1, ke,...,
.......,k) kL).
1: Su =xr).
2: for i = 0 0,2, 2.
do
3:
S =
4: for all seq E Sz do
5: Sample ki+1 i.i.d. tokens X1, X2,..., .,, XAan from Mo(. seq).
6: S S.+1 U { (seq, X,), i = 1,2,.. ....))))..
7: end for
8:
end for
9: Return St.
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 the LMIB
19

---

Table 2: Average latency with parallelization along the time axis and batch axis. We report average latency with standard deviation over 1,000 runs using a 97M transformer relative length 4 and batch = on GPU.
Relative latency
batch 1
batch 2
batch 4
batch 8
length =
1.00 + 0.16
1.01 = 0.15
1.06 + 0.10
1.10 = 0.16
length
1.01 = 0.18
1.09 = 0.25
1.10 0.09
1.42 0.4
Table 3: Average latency with parallelization along the time axis and batch axis. We report
age latency with standard deviation over 1,000 runs using 6M transformer relative to the 97M transformer with length 4 and batch
GPU.
Relative latency
batch 1
batch 2
batch 4
batch 8
length =
0.18 + 0.02
0.19 = 0.04
0.18 + 0.09
0.20 = 0.13
length = 0.17 0.04
0.19 + 0.05
0.16 0.02
0.18 0.04
dataset based on the example provided in the FLAX library [15]. For the draft model, we use transformer models with 2M, 6.M and 20M parameters, and for the large model we use a 97M parameter transformer model.
We first provide 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 Table 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 6M draft model in Table 3. In Table 3, the reported relative latencies relative to the large 97 model to get sense of the relative cost 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 97M large model for LMIB in Table 4. These draft models were produced by either halving (2M) or doubling (20M) the original 6M 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.
20

---

Table 4: Experimental results on the LMIB dataset with varying draft model sizes and the 97M transformer as the large model. All results are over 1000 test prompts averaged over three different random seeds and sampling temperature of 1.0 for both the draft and large models.
Draft model
Algorithm
K
Block efficiency
2MM Transformer
Baseline
1.00
Speculative
1
1.86 0.00
SpecTr
2
2.07 + 0.01
SpecTr
4
4
2.30
Spectr
8
2.56 0.01
Speculative
1
8
+.1
SpecTr
2
2.15
SpecTr
4
8
0.00
SpecTr
8
2.68 0.01
6MM Transtormer
Baseline
1.00
Speculative
1
2.21 0.01
SpecTr
2
4
SpecTr
4
2.74
SpecTr
8
2.99
Speculative
1
8
2.33
SpecTr
2
2.61
SpecTr
4
8
2.96 0.03
SpecTr
8
3.27 0.02
20MV Transformer
Baseline
1.00
Speculative
1
2.71#0.01
SpecTr
2
4
2.96 + 0.00
0.000
SpecTr
4
3.28 + 0.02
SpecTr
8
4
3.49
Speculative
1
3.02
SpecTr
2
8
3.48
SpecTr
4
3.85
SpecTr
8
8
4.15
21


Reviews:
[{"review_id": "s9CD68QUuZ", "reviewer": "Reviewer_TBLD", "summary": "This 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", "questions": "The first line of Algorithm 1:\nInput: ... \"X~iid p\" is confusing, since you have only one sample here, \"iid\" is unnecessary.", "limitations": "No", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The 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", "weaknesses": "The 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive comments on the novelty of the algorithm and theory, and acknowledging that SpectTr is mathematically concise, intuitive, and insightful. \n\n#### ***[Fair comparison with speculative decoding]***\nWe would like to point out that one of the main contributions of our work is to relate speculative decoding to the theory of optimal transport, which allows for the use of $k$ draws from the small model. We are unaware of other intuitive baselines on how to use more resources to speed-up speculative decoding due to the subtly involved in ensuring that the final sequence is still a *valid statistical draw from the large model*, which is our main contribution in this paper. We would be happy to conduct comparisons if we are missing some baseline extensions that you may have in mind.\n\n#### ***[Typo]***\nThanks for catching; we will fix it."}, {"review_id": "3a2pF7jfuE", "reviewer": "Reviewer_hEU2", "summary": "This 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", "questions": "See Strengths/Weaknesses", "limitations": "Yes", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- 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.", "weaknesses": "- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive comments on the novelty of the algorithm and soundness of the theory. Please see replies to specific questions below. \n\n#### ***[Quality of the samples]***\n(Identical to the same comment in the global response, provided here for completeness.) We would like to clarify that one of the biggest advantages of our proposed acceleration methods is that there is **provably** no drop in performance because our algorithm guarantees that the final outcome is a statistical draw from the large model. This holds for both the optimal solution from exactly solving the linear program, and the approximate solution $k$-seq. The “approximation” part in $k$*-seq* comes from the non-optimal acceptance probability, which leads to less decoded tokens per serial call and a hit in latency compared to the optimal solution.\n\nMore precisely, SpecTr guarantees that the final output sequence follows the exact same distribution as the output sequence from the large model as long as the token-level algorithm is a valid coupling between $p^{\\otimes k}$ and $q$ where $p,q$ are the conditional distributions on the next token from the small and large model respectively. Hence all metrics such as BLEU, avg. sentence likelihood, will be *neutral*. The formal statement is stated in Theorem 3 of the submission. The guarantee is the same as what is claimed in previous speculative decoding methods.\n\nWe will add more discussions in the revised version and improve the theorem statement to make this fact clearer.\n\n#### ***[Scale of the empirical evaluation]***\nWe agree that performing a more complete set of experimental evaluations would provide more evidence on the effectiveness of the method. For the rebuttal, we also evaluate our proposed method on the state-of-the-art PALM-2 models and report end-to-end latency improvements with respect to baseline autoregressive decoding and speculative decoding. Please see details in the global response and Table 1 in the attached PDF. We will add additional experimental results in the final version.\n\n#### ***[Table 8]***\nThanks for catching this! Table 8 was a latex error on our part. The experimental section refers to Tables 1, 3 (wall clock for large and small models) and Table 2 (Results on LM1B dataset).\n"}, {"review_id": "oxntqGvb1S", "reviewer": "Reviewer_jweu", "summary": "The 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. ", "questions": "None. ", "limitations": "The authors address the limitations. I do not see any potential negative societal impact.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "\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.  ", "weaknesses": "\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive comments, pointing out the typos, and suggesting ways to improve the paper. We will incorporate them in the final version. Please see inline replies below.\n\n#### ***[Lemma 1 seems unnecessary]***\nWe agree that the monotonicity part of  Lemma 1 follows from definitions, but the consistency part needs proof. We will clarify in the subsequent version. \n\n#### ***[Larger set of experiments]***\nFor the rebuttal, we have evaluated our method on state-of-the-art PALM-2 models and report end-to-end latency improvements. Please see details in the global response. We will add additional experimental results in the final version.\n\n#### ***[Algorithm 1 always returns at Line 7]***\nThank you for pointing out the typo. As you mention, Line 5 of the algorithm should say Return Y = X and accept = True. Will add other definitions in the Algorithm.\n\n#### ***[Typos & suggestions on font size]***\nThanks for these suggestions. We will incorporate them.\n"}, {"review_id": "VCWXj5rbmh", "reviewer": "Reviewer_DdMT", "summary": "The 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.", "questions": "N/A", "limitations": "Please see the weakness section", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The 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.", "weaknesses": "1. 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for agreeing on the novelty of the algorithm and the theoretical justification. Below we address the concerns on the evaluation of our method. \n\n#### ***[Quality of the final outputs.] (response to 1(a))***\n(Identical to the same comment in the global response, provided here for completeness.)  We would like to clarify that one of the biggest advantages of our proposed acceleration methods is that there is **provably** no drop in performance because our algorithm guarantees that the final outcome is a statistical draw from the large model. This holds for both the optimal solution from exactly solving the linear program, and the approximate solution $k$-seq. The “approximation” part in $k$*-seq* comes from the non-optimal acceptance probability, which leads to less decoded tokens per serial call and a hit in latency compared to the optimal solution.\n\nMore precisely, SpecTr guarantees that the final output sequence follows the exact same distribution as the output sequence from the large model as long as the token-level algorithm is a valid coupling between $p^{\\otimes k}$ and $q$ where $p,q$ are the conditional distributions on the next token from the small and large model respectively. Hence all metrics such as BLEU, avg. sentence likelihood, will be *neutral*. The formal statement is stated in Theorem 3 of the submission. The guarantee is the same as what is claimed in previous speculative decoding methods.\n\nWe will add more discussions in the revised version and improve the theorem statement to make this fact clearer.\n\n#### ***[End-to-end latency improvement.] (response to 1(b-e))***\n(Mostly the same with comment in the global response, provided here for completeness.) We agree that all factors mentioned by the reviewer would affect the effectiveness of the proposed approach in practical systems, and it is important to implement the method and report end-to-end latency comparisons including the delays caused by system overheads. To further demonstrate the effectiveness of our proposed approach, we conduct experiments on the state-of-the-art PALM-2 models [1] with PALM-2-Gecko and PALM-2-Bison (where Bison is a larger model) as the small model and large model, respectively. We report end-to-end (wall clock) latency comparisons between regular decoding, speculative decoding, and SpecTr. This includes the time to (parallelly) draw drafts from the small model, the time to verify the drafts with the large model, latency introduced by running the sequential rejection algorithm, and other system overheads. See the attached PDF file in the global response for detailed numbers. While we do see a smaller wall clock speed-up compared to the number of decoded tokens per serial call, our proposed method still achieves significant improvement in wall clock latency with respect to baseline decoding and speculative decoding. When $K =8$ and $L=8$, our relative wall clock speed-up over baseline is 2.13x (in contrast to 1.56x for speculative decoding over baseline), a further 1.37x improvement over speculative decoding.\n\nWe will add additional experimental results in the final version.\n\n[1] Google   AI. Introducing   PaLM   2, 2023. https://blog.google/technology/ai/google-palm-2-ai-large-language-model/.\n"}]

Author rebuttal:
We thank all reviewers for their detailed reading and encouraging comments about our submission. We are delighted to read the reviewers’ acknowledgement of the novelty of the algorithm (reviewer TBLD, hEU2, jweu, DdMT) and the soundness of the theory (reviewer hEU2, jweu). We will incorporate their suggestions to improve the presentation in future revisions of the paper.  Below we address a few common concerns raised by the reviewers. Each reviewer's individual questions will be answered in separate responses.

#### ***[End-to-end latency improvement.]***
We agree that it is important to implement the method and report end-to-end latency comparisons including the delays caused by system overheads. To further demonstrate the effectiveness of our proposed approach, we conduct experiments on the state-of-the-art PALM-2 models [1] with PALM-2-Gecko and PALM-2-Bison (where Bison is a larger model) as the small model and large model, respectively. We report end-to-end (wall clock) latency comparisons between regular decoding, speculative decoding, and SpecTr. This includes the time to (parallelly) draw drafts from the small model, the time to verify the drafts with the large model, latency introduced by running the sequential rejection algorithm, and other system overheads as suggested by reviewer DdMT. See the attached PDF file for detailed numbers. While we do see a smaller wall clock speed-up compared to the number of decoded tokens per serial call, our proposed method still achieves significant improvement in wall clock latency with respect to baseline decoding and speculative decoding. When $K =8$ and $L=8$, our relative wall clock speed-up over baseline is 2.13x (in contrast to 1.56x for speculative decoding over baseline), a further 1.37x improvement over speculative decoding.

We will add additional experimental results in the final version.

#### ***[Quality of the final outputs.]***
We would like to clarify that one of the biggest advantages of our proposed acceleration methods is that there is **provably** no drop in performance because our algorithm guarantees that the final outcome is a statistical draw from the large model. This holds for both the optimal solution from exactly solving the linear program, and the approximate solution $k$-seq. The “approximation” part in $k$-seq comes from the non-optimal acceptance probability, which leads to less decoded tokens per serial call and a hit in latency compared to the optimal solution.

More precisely, SpecTr guarantees that the final output sequence follows the exact same distribution as the output sequence from the large model as long as the token-level algorithm is a valid coupling between $p^{\otimes k}$ and $q$ where $p,q$ are the conditional distributions on the next token from the small and large model respectively. Hence all metrics such as BLEU, avg. sentence likelihood, will be *neutral*. The formal statement is stated in Theorem 3 of the submission. The guarantee is the same as what is claimed in previous speculative decoding methods.

We will add more discussions in the revised version and improve the theorem statement to make this fact clearer.

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


Meta-review: This paper presents speculative decoding, wherein a smaller language model is used to approximately sample from a large one, akin to a type of accept/reject MCMC sampler. This is a borderline paper. The AC had a thorough discussion with the SAC and the SAC agreed on the decision.
Decision: Accept (poster)