SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ziteng Sun*
Google Research, New York zitengsun@google.com
Ananda Theertha Suresh*
Google Research, New York theertha@google.com
Jae Hun Ro
Google Research, New York jaero@google.com
Ahmad Beirami
Google Research, New York beirami@gmail.com
Himanshu Jain
Google Research, New York himj@google.com
Felix Yu
Google Research, New York felixyu@gmail.com
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.
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} \coloneqq x(1), x(2) \ldots, x(t)$ , an autoregressive model  $\mathcal{M}_b$  generates successive tokens  $x(t + 1), x(t + 2), \ldots$  via temperature sampling [1, 10], where the next token  $x(t + 1)$  is drawn from the temperature-scaled distribution  $\mathcal{M}_b(\cdot |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 \in \Omega} \mathcal{M}_b(x|x^t)$ , where  $\Omega$  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 top- $k$  sampling [9, 22].
All these approaches are autoregressive decoding<sup>2</sup> 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), \ldots, 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_1^t$  and  $x_2^t$  denote two sequences of length  $t$ . We use  $[n]$  to denote the set  $\{1, \ldots, n\}$ .
A simplified computational model.
- Standard inference. Given a context  $x^t$ , with  $O(t^2)$  computation and  $O(1)$  time, an autoregressive model  $\mathcal{M}_b$  can compute  $\mathcal{M}_b(y|x^t)$ , the (temperature-scaled) probability of all possible next tokens  $y \in \Omega$ .
- Parallelization along the time axis. Given a context  $x^t$ , with  $O(t^2)$  computation and  $O(1)$  time, an autoregressive model  $\mathcal{M}_b$  can compute  $\mathcal{M}_b(y|x^i)$ , for all  $y \in \Omega$  and  $i \in \{1, 2, \dots, 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_{1}^{t}, x_{2}^{t}, \ldots, x_{K}^{t}$ , with  $O(Kt^2)$  computation and  $O(1)$  time, an autoregressive model  $\mathcal{M}_b$  can compute  $\mathcal{M}_b(y|x_j^i)$ , for all  $y \in \Omega, i \in [t]$ , and  $j \in [K]$ .<sup>3</sup>
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  $\mathcal{M}_s$ , which predicts the next token given the context, and the predictions of  $\mathcal{M}_s$  are
In this work, we use the words sampling and decoding interchangeably to refer to the process of sequentially generating tokens from a language model.
3When the assumption holds, one could naively batch multiple decoding contexts, which improves decoding throughput, but not the latency of each context.
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  $\mathcal{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).
1. Draft construction. The draft model is used to efficiently and "speculatively" sample  $L$  tokens,  $\tilde{x}(t+1), \ldots, \tilde{x}(t+L)$ . We keep the conditional probabilities on the next token  $\mathcal{M}_s(y \mid x^t, \tilde{x}^{t+1:t+i})$  for each  $i < L$  and  $\forall y \in \Omega$ .
2. Conditional probability computation. After observing the samples, we compute the conditional distributions  $\mathcal{M}_b(y\mid x^t,\tilde{x}^{t + 1:t + i})$  for each  $i\leq L$  and  $\forall y\in \Omega$  in parallel (along time axis) in  $O(1)$  time.
3. Draft selection. Validate and select first  $L'$  of the  $L$  tokens and set  $x(t + i) = \tilde{x}(t + i)$  for  $i \leq L'$  given the draft sequence and the conditional probabilities from both models. Sample a token  $\tilde{x}(t + L' + 1)$  from a residual distribution as a correction to the rejected token.<sup>4</sup>
After this step, we use  $x_1^{t + L' + 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  $\mathcal{M}_s(\cdot \mid x^t)$  and  $q$  is the conditional distribution of the large model  $\mathcal{M}_b(\cdot \mid 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\sim p$  1: Compute the residual distribution  $p^{\mathrm{res}}$  where  $\forall x\in \Omega ,p^{\mathrm{res}}(x) = \frac{q(x) - \min\{p(x),q(x)\}}{1 - \sum_{x'}\min\{p(x'),q(x')\}}.$  2: Sample  $\eta \sim U(0,1)$  3: if  $\eta \leq \min \left(1,\frac{q(X)}{p(X)}\right)$  then 4: Return  $Y = X$  . {Accept the draft token.} 5: end if 6: Return  $Y\sim p^{\mathrm{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^{\mathrm{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' \leq 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 \sim p$ , then  $Y \sim q$  [19, 4]. In the case of the draft selection, this means that the output of the algorithm is distributed according to  $\mathcal{M}_b(\cdot \mid x^t)$ , which is exactly the
4See Algorithm 1 for definition of the residual distribution. When  $L^{\prime} = L$  , no token is rejected. The residual will just be the conditional probability  $\mathcal{M}_b(\cdot \mid x^{t + L})$  , which gives an extra decoded token.
Figure 2: One iteration of SpecTr. 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. See Fig. 3 for a more detailed run of the draft selection step.
desired outcome. Furthermore
$$
\operatorname * {P r} (Y = X) = \sum_ {x \in \Omega} \min (p (x), q (x)) = 1 - d _ {\mathrm {T V}} (p, q),
$$
where  $d_{\mathrm{TV}}$  is the total variation distance between  $p$  and  $q$ . The closer  $p$  and  $q$  are in  $d_{\mathrm{TV}}$ , the higher the chance of  $\operatorname*{Pr}(Y = X)$ , and fewer the number of serial calls to the larger model. In the ideal case, if  $p = q$ , then  $\operatorname*{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 from  $\mathcal{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.
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
<sup>5</sup>When  $L' = L$ ,  $x(t + L + 1)$  is sampled from  $\mathcal{M}_b(\cdot \mid x^{t + L})$ .
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.13\mathrm{X}$ , a further  $1.37\mathrm{X}$  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, \ldots, X_k$  be a collection of draft tokens sampled from the small model, e.g., sampled i.i.d. from  $\mathcal{M}_s(\cdot \mid 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):
$$
\mathcal {M} _ {b} (\cdot \mid x ^ {t}) \quad \text {a n d} \quad \forall i \in [ k ], \quad \mathcal {M} _ {b} (\cdot \mid x ^ {t}, X _ {i}).
$$
The goal of the draft selection algorithm  $f: \Omega^k \to \Omega$  is to output  $Y = f(X^k)$ , whose distribution follows  $\mathcal{M}_b(\cdot \mid x^t)$ , and hence is a valid sample from the large model. Moreover, when  $Y \in \{X_1, \ldots, X_k\}$ , we could sample an extra token from  $\mathcal{M}_b(\cdot \mid x^t, Y)$  without calling  $\mathcal{M}_b$  since we have already computed the conditional probabilities  $\mathcal{M}_b(\cdot \mid 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  $\mathcal{M}_s$ , a sequence of token-level draft selection algorithms could be used along the time axis to select a valid sequence from the  $\mathcal{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 = \mathcal{M}_s(\cdot \mid x^t)$  and  $q = \mathcal{M}_b(\cdot \mid 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 = \operatorname{Ber}(1)$  and  $q = \operatorname{Ber}(1/2)$ . Then we have  $\forall i = 1, \dots, k$ ,  $X_{i} = 1$  and each of them will be accepted with probability  $1/2$ . After
This paper will
$^6$ Ber(b) denotes a Bernoulli distribution with the probability of seeing a head  $b$ .
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  $\operatorname{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  $\Omega$  is a discrete domain.
Definition 1 (Coupling). For two probability distributions  $P$  over  $\mathcal{X}$  and  $Q$  over  $\mathcal{Y}$ , we say a joint distribution  $\pi$  supported over  $\mathcal{X} \times \mathcal{Y}$  is a coupling between  $P$  and  $Q$  if  $\forall x, y, \pi(x, y) \geq 0$ ,
$$
\forall y \in \mathcal {Y}, \sum_ {x \in \mathcal {X}} \pi (x, y) = Q (y), \qquad \text {a n d} \qquad \forall x \in \mathcal {X}, \sum_ {y \in \mathcal {Y}} \pi (x, y) = P (x).
$$
We use  $\Pi(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_{\pi} : \mathcal{X} \to \mathcal{Y}$  introduced by the conditional probability  $\pi(y \mid x) := \pi(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_{\pi}$  is a valid draft selection algorithm. Formally, this is stated in the claim below.
Claim 1. For all  $\pi \in \Pi(P, Q)$ , let  $f_{\pi}$  be the probabilistic mapping defined above. If  $X \sim P$ , then  $f_{\pi}(X) \sim 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: \mathcal{X} \times \mathcal{Y} \to \mathbb{R}_+$ , the transportation cost of a coupling is defined as:
$$
C (\pi) = \mathbb {E} _ {X, Y \sim \pi} [ c (X, Y) ].
$$
The optimal transport plan is the coupling  $\pi \in \Pi(P, Q)$  that minimizes the transportation cost.
Speculative decoding with one draft token. With these definitions in place, we can see that with  $\mathcal{X} = \mathcal{Y} = \Omega$ , 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:
$$
\forall x \in \Omega , y \in \Omega , \qquad c (x, y) = \mathbb {1} \left\{y \neq x \right\}.
$$
The transportation cost of the coupling will be  $C(\pi) = \mathbb{E}_{X,Y\sim \pi}[\mathbb{1}\{Y\neq X\}] = \mathbb{P}_{X,Y\sim \pi}(Y\neq X)$ . This optimal transport cost is known to be
$$
\min  _ {\pi \in \Pi (p, q)} \mathbb {P} _ {X, Y \sim \pi} (Y \neq X) = \sum_ {x \in \Omega} \min  (p (x), q (x)), \tag {1}
$$
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  $\mathcal{X} = \Omega^k$  for some  $k\in \mathbb{N}_{+}$ , which is the space of  $k$  draft tokens from  $\Omega$  and  $\mathcal{Y} = \Omega$ , 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:
$$
\forall x \in \Omega^ {k}, y \in \Omega , \quad c (x, y) = \mathbb {1} \left\{y \notin S (x) \right\},
$$
where  $S(x) = \{o \in \Omega \mid 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
$$
C (\pi) = \mathbb {E} _ {X, Y \sim \pi} [ \mathbb {1} \{Y \notin S (X) \} ] = \mathbb {P} _ {X, Y \sim \pi} (Y \notin S (X)). \tag {2}
$$