SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We will also refer to the above cost  $C(\pi)$  as the rejection probability due to its probabilistic interpretation. And similarly,  $\alpha (\pi)\coloneqq 1 - C(\pi) = \mathbb{P}_{X,Y\sim \pi}(Y\in S(X))$  will be the acceptance probability.
From now on we will use membership cost as the default cost function and refer to the optimal transport solution as optimal transport with membership cost (OTM). We use  $\pi^{*}$  to denote the coupling that minimizes this cost  $\pi^{*} = \arg \min_{\pi \in \Pi(P, Q)} C(\pi)$ ;<sup>7</sup> and the cost  $C(\pi^{*})$  is referred to as the optimal transport cost between  $P$  and  $Q$ . We use  $\alpha(P, Q) = 1 - C(\pi^{*})$  to denote the corresponding optimal acceptance probability.
Draft selection with i.i.d. draft tokens. In this paper, we will mainly focus on the case when the draft tokens are i.i.d. samples from a base distribution. Let  $p, q$  be supported over  $\Omega$  and the goal is to obtain one valid token from  $q$  given  $k$  i.i.d. samples from  $p$ . For SpecTr with context  $x^t$ , we have  $p = \mathcal{M}_s(\cdot \mid x^t)$  and  $q = \mathcal{M}_b(\cdot \mid x^t)$ . We set  $P = p^{\otimes k}$ , a product distribution whose marginals are all  $p$ , and  $Q = q$ . The OT problem we want to solve is the following:
$$
\min  C (\pi) \quad s. t. \quad \pi \in \Pi \left(p ^ {\otimes k}, q\right). \tag {3}
$$
We overload notation and denote the optimal acceptance probability as  $\alpha_{k}(p,q)\coloneqq \alpha (p^{\otimes k},q) = 1 - C(\pi^{*})$ . To better understand the quantity, we state a few properties about  $\alpha_{k}$ .
Lemma 1. (Appendix A.2) The optimal acceptance probability satisfies the following properties.
- Monotonicity. For any  $p, q$  and  $k \geq 1$ ,  $\alpha_k(p, q) \leq \alpha_{k+1}(p, q)$ .
- Consistency. If  $\forall x \in \Omega$ ,  $q(x) / p(x)$  is bounded, we have  $\lim_{k \to \infty} \alpha_k(p, q) = 1$ . Else,  $\lim_{k \to \infty} \alpha_k(p, q) = \sum_{x \in \Omega} \mathbb{1}\{p(x) > 0\} q(x)$ .
The above properties demonstrate that for a large  $k$ , the value of  $\alpha_{k}$  can become large. Hence increasing  $k$  could increase the acceptance probability, leading to further speedups. We now focus on computing the optimal transport plan and the optimal acceptance probability.
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  \sum_ {x \in \Omega^ {k}} \sum_ {y \in \Omega} \pi (x, y) \mathbb {1} \{y \notin S (x) \} \quad s. t. \pi \in \Pi (P, Q). \tag {4}
$$
The linear program in (4) has  $|\Omega|^{k+1}$  variables and  $|\Omega|^k + |\Omega|$  equality constraints (see Definition 1). Linear programming can be solved in time polynomial in the number of variables and constraints [7, 21, 18], implying the following lemma.
Lemma 2. Given  $p, q$  over  $\Omega$ , the solution to Eq. (3) can be computed in time  $O(|\Omega|^{O(k)})$ .
We refer to the optimal coupling obtained above as OTM- $k$  and denote it as  $\pi^{\mathrm{OTM} - k}$ . When  $k = 1$ , there is a closed form expression for the optimal acceptance cost (see Eq. (1)), whereas for larger values of  $k$ , we are unaware of a general closed form expression. In Appendix A.1, we provide an information-theoretic upper (and lower) bound, which is tight up to a multiplicative constant of  $1 - (1 - 1 / k)^{k} \geq 1 - 1 / e$ .
While solving OTM in Eq. (4) gives the plan with optimal acceptance probability, to the best of our knowledge, the best-known runtime will be exponential in  $k$ , which can be prohibitive when either the vocabulary size  $|\Omega|$  or the number of draft tokens  $k$  is large.[10] In the next section, we will present a selection algorithm that can be efficiently computed and show that it achieves an acceptance probability of at least  $(1 - (1 - 1 / k)^k)\alpha_k \geq (1 - 1 / e)\alpha_k$ .
The existence of optimal coupling in discrete domain is well-known, e.g., see [27]. When the optimal coupling is not unique, we use  $\pi^{*}$  to denote one of the optimal couplings.
The above generic formulation immediately allows generalization to more complex draft selection strategies, such as sampling  $k$  tokens without replacement, or using a different drafting distribution for each draft.
To our best knowledge, the best practical computation bound (through interior-point method) is  $O(|\Omega|^{3k})$  [21] and the best theoretical computation bound is  $O(|\Omega|^{2.5k})$  [18].
For discrete OT, Sinkhorn algorithm could be used to solve an entropy-regularized version of OT, which has a better computation complexity [6]. However, the computation cost of the algorithm will still have a linear dependence on  $|\Omega|^k$ , which can be prohibitive.
6 Draft selection via  $k$ -sequential selection
In this section, we present a sequential selection algorithm (K-SEQ), an approximate solution<sup>11</sup> to the optimal transport problem in Eq. (3), which can be efficiently computed in time almost linear in  $|\Omega|$  and logarithmic in  $k$ . The algorithm is presented in Algorithm 2.
Algorithm 2  $k$ -sequential selection algorithm (K-SEQ).
Input: Distributions  $p,q$  samples  $X_{1},\ldots ,X_{k}\sim i.i.d. p.\rho \in [1,k]$  : division factor. 1: Let  $\beta_{p,q}(\rho) = \sum_{x\in \Omega}\min (p(x),q(x) / \rho)$  and  $p_{\mathrm{acc}} = 1 - (1 - \beta_{p,q}(\rho))^k$  . Compute  $p^{\mathrm{res}}$  where  $\forall x\in \Omega ,p^{\mathrm{res}}(x) = \frac{q(x) - \min\left\{p(x),\frac{q(x)}{\rho}\right\} \frac{p_{\mathrm{acc}}}{\beta_{p,q}(\rho)}}{1 - p_{\mathrm{acc}}}$  (5)   
2: for  $i = 1,2,\dots ,k$  do   
3: Sample  $\eta_i\sim U(0,1)$    
4: if  $\eta_{i}\leq \min \left(1,\frac{q(X_{i})}{\rho\cdot p(X_{i})}\right)$  then   
5: Return  $Y = X_{i}$  . {Return the ith draft token.}   
6: end if   
7: end for   
8: Return  $Y\sim p^{\mathrm{res}}$  . {Sample a corrected token from the residual distribution.}
At a high-level, the algorithm goes over all  $k$  draft samples generated from  $p$  sequentially, and decides on whether to accept each  $X_{i}$  based on the ratio  $q(X_{i}) / p(X_{i})$ . The algorithm output the first accepted sample or result from a residual distribution  $p^{\mathrm{res}}$  if none of the samples is accepted. To guarantee that the final returned token is a valid sample from  $q$ , we choose an appropriate  $\rho \in [1,k]$  and accept  $X_{i}$  with probability  $\min (1,q(X_i) / (\rho \cdot p(X_i)))$  instead of  $\min (1,q(X_i) / (p(X_i)))$  as in Algorithm 1. In Theorem 1, we show that with appropriately chosen  $\rho$ 's, Algorithm 2 is indeed valid transportation plans from  $p^{\otimes k}$  to  $q$ . Moreover, to find the best transportation plan within the family, we only need to search over a single parameter  $\rho$ , which reduces the computation cost significantly. We also show that searching over this sub-family of couplings won't decrease the optimal acceptance probability by a multiplicative constant. The performance of Algorithm 2 is stated in Theorem 1.
Theorem 1. Let  $\beta_{p,q}(\rho) = \sum_{x\in \Omega}\min \left(p(x),\frac{q(x)}{\rho}\right)$  and  $\rho^{*}$  be the solution to the identity below.
$$
1 - \left(1 - \beta_ {p, q} (\rho)\right) ^ {k} = \rho \beta_ {p, q} (\rho). \tag {6}
$$
When  $\rho \geq \rho^{*}$ , the coupling  $\pi_{\rho}^{\mathrm{K - SEQ}}$  in Algorithm 2 is a valid transport plan from  $p^{\otimes k}$  to  $q$ . When  $\rho = \rho^{*}$ , we have
$$
\alpha (\pi_ {\rho^ {*}} ^ {\mathrm {K - S E Q}}) \geq (1 - e ^ {- 1}) \alpha_ {k} (p, q).
$$
Moreover,  $\rho^{*}$  can be computed up to accuracy  $\delta$  in time  $O(|\Omega| \log((k - 1)/\delta))$ .
We provide the proof in Appendix C.1. In Appendix B, using a few canonical examples of distributions, we plot the acceptance probability of K-SEQ and compare it with the optimal acceptance probability  $\alpha_{k}$ . It can be shown that K-SEQ could have a strictly worse acceptance probability compared to the OTM solution for certain cases while there also exist non-trivial cases where K-SEQ achieves the optimal acceptance probability.
Concurrent and recent work of [20, 29] has proposed another efficient algorithm for the draft selection phase. To the best of our knowledge, there is no optimality guarantee proved for their proposed algorithm. In Appendix B.3, we present its acceptance probability empirically for the canonical case of Bernoulli distributions, and show that both our proposed algorithms (OTM and K-SEQ) have a higher acceptance probability.
7 SpecTr: Application of OTM in autoregressive sampling
In this section, we describe how OTM can be used to speed up auto-regressive sampling, which we refer to as SpecTr sampling. Similar to speculative decoding, each iteration of SpecTr can be decomposed into three phases (Fig. 2):
<sup>11</sup>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  $x^t$ , use the draft model sample a set of  $K$  draft sequences with length  $L$ , denoted by  $S = \{z^L \sim \mathcal{M}_s(\cdot \mid x^t)\}$ . We keep the conditional probabilities  $\mathcal{M}_s(y \mid x^t, z^i)$  for all  $y \in \Omega, i \leq L$  and  $z^L \in S$ .
2. Conditional probability computation. Compute the conditional probabilities on the next token for the large model  $\mathcal{M}_b(y\mid x^t,z^i)$  for all  $y\in \Omega ,i\leq L$  and  $z^{L}\in S$  in parallel.
3. Draft selection. Select first  $L'$  of the  $L$  tokens and set  $x(t + i) = z(i)$  for  $i \leq L'$  and some  $z \in S$  given the set of draft sequences and the conditional probabilities from both models. Sample a token from a residual distribution as a correction to the rejected tokens.
The conditional probability computation step takes  $O(1)$  when  $|S|$  is not large based on our simplified computations model. We mainly focus on the draft set construction phase and draft selection phase.
Algorithm 3 Draft selection with multiple candidates (DraftSelection).
Input: Input sequence  $x^t$ ; draft sequence length:  $L$ ; draft sequences  $S = \{z_i^L \mid i \leq |S|\}$ .  
1: Compute a transport plan (using linear programming in Lemma 2 for an optimal solution or Algorithm 2 for a suboptimal solution) from  $\mathcal{M}_s(\cdot \mid x^t)^{\otimes |S|}$  to  $\mathcal{M}_b(\cdot \mid x^t)$ , denoted by  $\pi_t$ .  
2: Get the multi-set of next token-level drafts:  $S_z = \{z_i(1)\}_{i \in [|S|]}$  and compute  $Y = f_{\pi_t}(S_z)$ .  
3: if  $L = 1$  then  
4: if  $Y \in S_z$  then  
5: Sample  $Y' \sim \mathcal{M}_b(\cdot \mid (x^t, Y))$ .  
6: Return  $(x^t, Y, Y')$ . {Sample an extra token if the last token is accepted.}  
7: else  
8: Return  $(x^t, Y)$ . {Stop and return the corrected token and previous accepted tokens.}  
9: end if  
10: end if  
11: Let  $S_{\mathrm{next}} = \{z^{2:L} \mid z \in S \text{ and } z(1) = Y\}$  be the set that consists of sub-sequences of the candidates that agree with the selected next token.  
12: if  $S_{\mathrm{next}} = \emptyset$  then  
13: Return  $(x^t, Y)$ . {Stop and return the corrected token and previous accepted tokens.}  
14: else  
15: Return DraftSelection(( $x^t, Y$ ,  $L - 1$ ,  $S_{\mathrm{next}}$ ). {Keep the draft token and proceed to the next time step.}  
16: end if
Draft set with i.i.d. draft sequences. Given context  $x^t$ , a natural way to come up with a set of  $K$  drafts is to independently sample  $K$  draft sequences from  $\mathcal{M}_s(\cdot \mid x^t)$ , i.e.,
$$
z _ {1} ^ {L}, z _ {2} ^ {L}, \dots , z _ {K} ^ {L} \sim_ {i. i. d.} \mathcal {M} _ {s} (\underbrace {\cdot , \cdot , \dots} _ {L \text {d o t s}} \mid x ^ {t}). \tag {7}
$$
The draft set construction method in (7) can be generalized to a prefix-tree based algorithm. However, this generalized version did not perform better in our experiments. We include this construction in Appendix D for completeness.
Draft selection with multiple candidates. We present the sequence-level selection algorithm given a set of draft sequences in Algorithm 3. We assume the conditional probabilities on the next token are available given any prefix in the candidate set since they are computed in parallel in the second phase, and won't list them as inputs explicitly in Algorithm 3.
A sample run of the algorithm is presented in Fig. 3. The algorithm proceeds in a recursive fashion. Given prompt  $x^t$  and a candidate set  $S$  sampled from  $\mathcal{M}_s(\cdot \mid x^t)$ , the algorithm first computes a token-level draft selection algorithm  $f_{\pi} : \Omega^{|S|} \to \Omega$  which is a transport plan from  $\mathcal{M}_s(\cdot \mid x^t)^{\otimes |S|}$  to  $\mathcal{M}_b(\cdot \mid x^t)$ . Then  $f_{\pi}$  is applied to the set of first tokens of the draft sequences in  $S$  to obtain a valid token  $Y$  from  $\mathcal{M}_b(\cdot \mid x^t)$ . If  $Y$  is not the last token ( $L \geq 2$ ), we filter out sequences in  $S$  whose first token is not  $Y$  and denote the remaining sequences as  $S_{\mathrm{next}}$  and feed it to the algorithm with context  $(x^t, Y)$  and draft length  $L - 1$ . This goes on until we have  $L = 1$  or  $S_{\mathrm{next}} = \emptyset$ .
In this case when  $Y$  is the last token (i.e.,  $L = 1$ ) and  $Y \in S$ , we have the choice to sample an additional token  $\mathcal{M}_b(\cdot \mid (x^t, Y))$  since this conditional probability is already computed in the second phase. Due to the property of the token-level selection algorithms and the autoregressive structure of language models, it can be shown that  $Y$  is always a valid sample from  $\mathcal{M}_b(\cdot \mid x^t)$ . Let  $L'$  be the number of decoded tokens in one iteration. Note that this is a random variable in the range  $[1, L + 1]$ .
The formal quality guarantee is stated in Theorem 2. We present the proof in Appendix C.2.
Theorem 2. Assume all drafts in the set  $S$  are generated from the small model with input  $x^t$ , or more precisely,  $\forall z \in S$ ,
$$
\forall i \in [ 1, L ], \quad z (i) \sim \mathcal {M} _ {s} (\cdot \mid x ^ {t}, z ^ {i - 1}). \tag {8}
$$
Let  $(x^t, Y^\tau)$  be the output of Algorithm 3 where  $\tau$  is the length of the newly decoded tokens, then it satisfies that  $Y^{1:\tau}$  is distributed according to  $\mathcal{M}_b(\underbrace{\cdot,\cdot,\ldots}_{\tau \text{dots}} \mid x^t)$ . More precisely, for any  $\tau_0 \in [1,L+1]$ , and any  $\tau_0$ -length, sequence  $o^{\tau_0} = (o(1),\ldots,o(\tau_0)) \in \Omega^{\tau_0}$ , we have
$$
\Pr \left(Y ^ {\tau_ {0}} = o ^ {\tau_ {0}} | \tau = \tau_ {0}\right) = \Pi_ {i = 1} ^ {\tau_ {0}} \mathcal {M} _ {b} (o (i) \mid x ^ {t}, o ^ {i - 1}).
$$
We empirically evaluate SpecTr and compare it with two methods: (1) the baseline auto-regressive decoding; and (2) speculative decoding with  $K = 1$ . Note that all three methods effectively generate samples from the same baseline large model, and hence the quality of the two speculative decoding methods is provably neutral to that of the large model. Thus, we will only focus on measuring the speedup in our experiments. In the simplified computation model, we made the following assumptions: (1) Decoding time from small models is negligible compared to decoding from the small model; (2) Parallelization along the batch and time axis doesn't increase the time for a serial call to the large model. With these, the theoretical speedup compared to baseline decoding will be the average number of decoded tokens per serial call, which is called block efficiency [19], defined below
$$
\text {B l o c k} := \frac {\text {T o t a l n u m b e r o f d e c o d e d t o k e n s}}{\text {N u m b e r o f s e r i a l c a l l s t o} \mathcal {M} _ {b}}.
$$
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  $\mathcal{M}_s$  might not be negligible; (2) Parallelization along the batch and time axis might increase the time for a single call to  $\mathcal{M}_b$ ; (3) Overhead due to the implementation of additional functionalities in SpecTr such as the draft selection algorithm and switching between models. These factors will depend on how the algorithm is implemented and optimized. In our experiment, we consider both the block efficiency, and average wall clock speedup with our implementation of SpecTr.
We first present the performance of our algorithm and compare it to speculative decoding using state-of-the-art PALM-2 models with prompts from the one-billion language benchmark (LM1B) [3]. In Appendix E, we use a pair of smaller transformer models to break down different affecting factors mentioned above. In Table 1, we use PALM-2-Gecko and PALM-2-Bison as the small model and
Table 1: Experimental results on the LM1B dataset with PALM-2-Gecko as the small model and PALM-2-Bison 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	1	4	2.4	1.67
SpecTr	8	4	3.1	2.08
Speculative	1	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. In particular, when  $L = 8$  and  $K = 8$ , our proposed SpecTr algorithm has a speedup of  $2.13\mathrm{x}$ , a further  $1.37\mathrm{x}$  increase compared to speculative decoding ( $K = 1$ ).
Authors thank Asaf Aharoni, Kwangjun Ahn, Badih Ghazi, Sanjiv Kumar, Teodor Marinov, Michael Riley, and NeurIPS reviewers for helpful comments and discussions.
[1] David H Ackley, Geoffrey E Hinton, and Terrence J Sejnowski. A learning algorithm for boltzmann machines. Cognitive science, 9(1):147-169, 1985.
[2] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[3] Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Philipp Koehn, and Tony Robinson. One billion word benchmark for measuring progress in statistical language modeling. arXiv preprint arXiv:1312.3005, 2013.
[4] Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean-Baptiste Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.01318, 2023.
[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.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 (http://websites.math.leidenuniv.nl/probability/lecturenotes/CouplingLectures.pdf), 2012.
[9] Angela Fan, Mike Lewis, and Yann Dauphin. Hierarchical neural story generation. arXiv preprint arXiv:1805.04833, 2018.
[10] Jessica Ficler and Yoav Goldberg. Controlling linguistic style aspects in neural language generation. arXiv preprint arXiv:1707.02633, 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.10350, 2022.
[12] Google AI. Introducing PaLM 2, 2023. https://blog.google/technology/ai/google-palm-2-ai-large-language-model/.
[13] Google PaLM-2 Team. PaLM 2 technical report, 2023.
[14] Wenshuo Guo, Nhat Ho, and Michael Jordan. Fast algorithms for computational optimal transport and Wasserstein barycenter. In Silvia Chiappa and Roberto Calandra, editors, Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pages 2088-2097. PMLR, 26-28 Aug 2020.
[15] Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. Flax: A neural network library and ecosystem for JAX, 2023.
[16] Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. arXiv preprint arXiv:1904.09751, 2019.
[17] Leonid V Kantorovich. On the translocation of masses. In 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 o (vrank) iterations and faster algorithms for maximum flow. In 2014 IEEE 55th Annual Symposium on Foundations of Computer Science, pages 424-433. IEEE, 2014.
[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 Jia. 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 IEEE 12th international conference on computer vision, pages 460-467. IEEE, 2009.
[22] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019.
[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, HengTze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, et al. Lamda: Language models for dialog applications. arXiv preprint arXiv:2201.08239, 2022.
[26] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Roziere, Naman Goyal, Eric Hambro, Faisal Azhar, et al. LLaMA: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 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.04487, 2023.
[29] Chao Zhang Yuhui Li and Hongyang Zhang. Eagle: Lossless acceleration of llm decoding by feature extrapolation, 2023. https://sites.google.com/corp/view/eagle-llm.
A Properties of optimal transport cost
A.1 Information-theoretic upper (and lower) bound of  $\alpha_{k}$ .
Below we provide an information-theoretic upper (and lower) bound in Lemma 3, which is tight up to a multiplicative constant of  $1 - (1 - 1 / k)^k \geq 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  $\forall k \geq 1$ , we have
$$
(1 - (1 - 1 / k) ^ {k}) \cdot \bar {\alpha} _ {k} (p, q) \leq \alpha_ {k} (p, q) \leq \bar {\alpha} _ {k} (p, q),
$$