Multi-Draft Speculative Sampling: Canonical Decomposition and Theoretical Limits
Authors: Ashish J Khisti
MohammadReza Ebrahimi
Hassan Dbouk
Arash Behboodi
Roland Memisevic
Christos Louizos
Venue: ICLR 2025 Spotlight (2025)
Keywords: speculative decoding
multi draft speculative sampling
large language models
weighted importance sampling
optimal transport
Primary area: foundation or frontier models, including LLMs
TL;DR: We show the optimal multi-draft speculative sampling can be decomposed into a two-step solution: importance sampling followed by (single-draft) speculative sampling, and provide an explicit expression for the optimal acceptance probability.

Abstract: We consider multi-draft speculative sampling, where the proposal sequences are sampled independently from different draft models.  At each step, a  token-level draft selection scheme takes a list of valid tokens as input and produces an output token whose distribution matches that of the target model. Previous works have demonstrated that the optimal scheme (which maximizes the probability of accepting one of the input tokens) can be cast as a solution to a linear program. In this work we show that the optimal scheme can be decomposed into a two-step solution: in the first step an importance sampling (IS) type scheme is used to select one intermediate token; in the second step (single-draft) speculative sampling is applied to generate the output token.  For the case of two identical draft models we further 1) establish a necessary and sufficient condition on the distributions of the target and draft models for the acceptance probability to equal one and 2) provide an explicit expression for the optimal acceptance probability.  Our theoretical analysis also motives a new class of token-level selection schemes based on weighted importance sampling. Our experimental results demonstrate consistent improvements in the achievable block efficiency and token rates over baseline schemes in a number of scenarios.

Published as a conference paper at ICLR 2025
MULTI-DRAFT SPECULATIVE SAMPLING:
CANONICAL DECOMPOSITION AND THEORETICAL LIMITS
Asshish Khisti 12
M.Reza Ebrahimi
Hassan Dbouk!
Arash Behboodi'
Roland Memisevic
f Christos Louizos'
Qualccomm Al Research* 2 University of Toronto
ABSTRACT
We consider multi-draft speculative sampling, where the proposal sequences are sampled independently from different draft models. At each step, token-level draft selection scheme takes a list of valid tokens as input and produces an output token whose distribution matches that of the target model. Previous works have demonstrated that the optimal scheme (which maximizes the probability of accept- ing one of the input tokens) be cast solution to linear program. In this work we show that the optimal scheme can be decomposed into two-step solution: the first step an importance sampling (IS) type scheme is used to select one intermediate token; in the second step (single-draft) speculative sampling is applied to generate the output token, For the case of two identical draft models we further 1) establish necessary and sufficient condition on the distributions of the target and draft models for the acceptance probability to equal one and 2) provide an explicit expression for the optimal acceptance probability. Our theoretical analysis also motives a new class of token-level selection schemes based on weighted importance sampling. Our experimental results demonstrate consistent improvements in the achievable block efficiency and token rates over baseline schemes in a number scenarios.
INTRODUCTION
The transformer architecture (Vaswani et al., 2017) has revolutionized the field of natural language processing and deep learning. One of the key factors contributing to the success story of transformers, as opposed to prior recurrent-bace 2014), is their inherent basee architectures (Hochreiter and Schmidhuber, 1997; Chung et al., parallelization due to the attention mechanism. This allows for massive scaling and lead to the development of state-of-the-are Large Language Models (LLMs) (Touvron et al., 2023; Achiam et al., 2023; Brown et al., 2020; Chowdhery et al., 2023) which have demonstrated remarkable performance across wide range of tasks. Despite their parallelizable training, LLM inference is scquential, owing to their auto-regressive nature. This limits their text- generation to one token per forward pass, which is known to be memory-boue (Shazeer, 2019).
To alleviate the memory- nature of rugo-regressive decoding of LLMs, speculative decoding (Chen et al., 2023; Leviathan et al., 2023) leverages an arbitrary smaller language model (draft model) that generates multiple candidate tokens in auto-regressive manner. The L.M (target model) is then used to score all the tokens in the draft in parallel, and the draft tokens are verified through sequence of token-level rejection sampling which guarantees that the final sequence follows the same distribution as that of the target model. In order for speculative decoding to be beneficial, the combined cost of auto-regressively sampling from the draft model and parallel verification via the target model should be smaller than auto-regressively sampling from the target model. Intuitively, this requires that the draft model distribution resembles that of the target model, which can be measured via the acceptance rate of the speculative decoding process, i.e., the rate at which we accept/reject draft tokens.
Joint first authors Joint last authors Qualcomm AI Research is an initiative of Qualcomm Technologies, Inc. Correspondence a kh at (ece. uttronto..o and ebr ahiimi @ct quaa comm. com.

---

Published as a conference paper at ICLR 2025
A large number of works speculative decoding (Sun al., 2024b; Jeon et al., 2024; Miao et al., 2024; Sun et al., 2024a) have emerged recently in an effort to further improve decoding efficiency. The authors in (Sun al., 2024b) propose SpecTr. multi- extension where the draft model generates K candidate token sequences (which could be sampled in a batch) for each time-step (as opposed to one). The authors consider token-level selection scheme with the objective of maximizing the probability of accepting some token in the set of available tokens. They demonstrate that this problem can be cast into the framework of optimal transport and solved using linear program. However due to complexity reasons, the authors instead propose modified sequential rejection sampling scheme.
1.1 MAIN CONTRIBUTIONS
We revisit the optimal transport framework introduced in (Sun et al., 2024b) and introduce a canonical decomposition
demonstrate that the optimal acceptance probability can be achieved by a two-step scheme: the first step involves selecting a token from the input set using a type of importance sampling: the second step involves speculative sampling using the selected token and the target distribution.
For the case of K" = 2 identical draft models, we establish an analytical expression for the optimal acceptance probability. In this setting, we also establish a necessary and sufficient condition for the acceptance probability to equal one. We also discuss numerical evidence for the case of more than 2 drafts.
We propose a new token- token- selection scheme based on weighted importance sampling. To enable a faster implementation, we present heuristic approaches that reduce computation while also penalizing the acceptance probability.
We present experimental results involving the OPT model over a variety of tasks. We com- pare the performance our proposed schemes with baselines and demonstrate improvements in the block efficiency and token rates.
2 BACKGROUND AND RELATED WORK
Auto-regressive sampling from LLMs inherently sequential and memory-bound (Shazeer, 2019). Several approaches have been proposed in the literature to accelerate LLM inference (Shazeer, 2019: Jaszezur al., 2021; Frantar al., 2022: Frantar and Alistarh, 2023; Stern et al., 2018; Chen et al., 2023; Leviathan et al., 2023; Jeon et al., 2024; Sun et al., 2024b; Miao et al., 2024). Model compression techniques, such as quantization (Frantar al., 2022; Bondarenko et al., 2024) and sparsification (Jaszezur et al., 2021; Frantar and Alistarh, 2023) have been shown to reduce the overall complexity of LLMs at the expense of some degradation in decoding quality.
For lossless LLM inference acceleration, speculative decoding (Chen et al., 2023; Leviathan et al., 2023; Stern et al., 2018) has emerged as a promising and orthogonal alternative. Earlier works on greedy decoding draft and prediet multiple tokens by augmenting the base LLM (Stern et al., 2018) or aggressive decoding (Ge et al., 2022). However, LLM text-generation often requires sampling with non-zero temperature from the generated logits. To that end, speculative decoding (Chen et al., 2023; Leviathan et al., 2023) was proposed. In speculative decoding, auto-regressive sampling delegated to smaller language model (draft model) that generates multiple candidate tokens. The LLM (target model) is then used to score all the tokens in the draft in parallel, and the draft tokens are verified through a sequence of token-level rejection sampling. Speculative decoding guarantees that the final sequence follows the same distribution as that of the target model. The performance of speculative methods highly depends on the choice of the draft model. Zhou al. (2023) use knowledge distillation (Hinton et al., 2015) to better align the draft and target models which results in higher token acceptance rates.
More recently, the works of Sun et al. (2024b); Miao et al. (2024); Jeon et al. (2024) extend speculative decoding to the multi-draft setting where the draft model(s) generate multiple token sequences per time-step. Specifically, Sun et al. (2024b) formulate the token-level draft selection problem as a discrete optimal transport problem with membership cost and propose SpecTr:: a new decoding algorithm that allows for multiple candidates for each token in the draft. A related setting also studied in Miao et al. (2024); Jeon et al. (2024) where the authors consider a token tree based construction for improving the draft sequences well as token-level selection method
2

---

Published as a conference paper at ICLR 2025
different form Sun et al. (2024b). Instead of using a dedicated draft model, Cai et al. (2024) propose augmenting the target model with extra decoding heads that can concurrently draft multiple tokens. The extra heads are fine-tuned using auramtter-effiiien methods, and can be added to any pre-trained target model. Orthogonally, Sun et al. (2024a) study block-level verification in the single-draft setting as a block-level optimal transport problem. They propose computationally-fffiiee algorithm that optimally solves the block-level transport problem, and report speedups over prior token-level verification (Leviathan al., 2023).
TOKEN-LEVEL OPTIMAL DRAFT SELECTION: THEORETICAL ANALYSIS
We focus on token-level optimal draft selection framework introduced in Sun et al. (2024b). For sake of completeness we review the speculative sampling schemes in Leviathan et al. (2023); Chen et al. (2023) in Appendix A. We assume that = 2......... n denotes the vocabulary of tokens and at given step, say t, S {X1..... Xx J. denotes the K valid tokens under consideration. Each of these tokens is generated in an i.i.d. fashion from a distribution p(-) determined by the underlying draft model and the context sequence u' E 52' i.e, for each y E S2, have p(y) (yly)'),, where M, denotes the distribution generated by the small (draft) model. In a similar fashion we let q(.) be MM(yl'') where Mr denotes the
the distribution over 52 associated with the large model i.e. q(y) distribution generated by the large model. Note that we do not explicitly indicate the sequence u' when discussing p(-) and q(-).
is fixed and common to both models throughout our analysis.
Given an input S ~ 1I K p(X;) consisting of K candidate tokens (Xi,..., Xx), a token-level selection rule (TLSR) is a conditional distribution P(-(S) over Ω. A valid TLSR must satisfy the constraint that for each 2 E Ω, Es P(=)S)p(SS) = q(2). A natural metric to optimize for TLSR is the probability that one of the tokens is accepted i.e., if Z ~ P(-SS) denotes the output of the TLSR, then we wish to maximize Pr(Z E S).
Problem (Optimal Token Level Selection Rule) Given distributions p(-) and g(-) find a valid TLSR that maximizes the probability of acceptance: P(acc) = Pr(Z S) and let P (acc) be the optimal value.
Problem 1 was studied Sun et al. (2024b) and shown to be an instance of optimal transport, which
can be cast as a linear program. The authors used this framework to establish the optimality of speculative sampling (Chen al., 2023; Leviathan al., 2023) in the case of a single draft i.e.. the authors established an information theoretic upper bond on P* In this
K 1. For K >>> work, we revisit Problem and develop new insights into the structure of the optimal solution. In fact, establish that the optimal solution in the case of multiple drafts has natural connection to importance sampling (Tokdar and Kass, 2010). For the case of K = 2 drafts we exactly characterize P* (acc) and state necessary and sufficient conditions on p(-) and q(-) for P* (acc) to equal 1.
We begin by defining a family of schemes that we will refer to as importance weighted sampling.
Definition (Importance Weighted Sampling) An importance weighted sampling scheme takes as input the set of candidate tokens S = {XI...., }k} and outputs token Y₁ E S defined by the conditional distribution:
{ ( σk),
y E {s1, ... 5K
Pr(Yr = N| = =1ik)=
(0.
(I)
y E fa...... Fk
where Eyen By ..... σk) 1 for each FT:K E 9 and 0 I By IK) IK) I 1.
Note that instead of considering the probability over the value of the selected token in (1), one can instead consider the probability of selecting an index i between {l....,A K} i.c.. Pr(I = i ).1k). Such a distribution maps to (1) by simply summing over all indices where J; = y. We note that the form in (1) will be more convenient in the sequel. Also note that the classical importance sampling scheme (Tokdar and Kass, 2010) corresponds to the case where Pr(I X| xI.K) q(x:)) (xxx) However the family of schemes in Definition 1 is not
We also consider the case when the tokens generated from different distributions. See Remark 1 as well as the experimental results in Section 5.
3

---

Published as a conference paper at ICLR 2025
Draft X1 pl.)
Importance Weighted Sampling
2 X2 pl.)
Y1 pr())
Z ~ g())
Speculative Sampling
:
(21 (K)
D
Xx p())
Step
Step 1
Figure I: Optimal Approach for Multi-Draft Speculative Sampling
restricted to such a choice and we treat (z zx) as free parameters that can be optimized. While importance weighted sampling scheme may not lead be valid TLSR, explained next, it is a key building block in canonical decomposition for token selection.
Our first result is decomposition for the optimal token level selection rule that establishes connection to the importance weighted sampling in Definition 1. The proof is in Appendix B.
Theorem (Optimal Acceptance Probability and Canonical Decomposition) Let P (ace) be the acceptance probability for the optimal token level selection rule in Problem 1. Then we can express:
K Iplon)
Pa(cc) = max (F)())
Σ min
((v).
Σ
By (x(kk)
(2)
VER
F1umE En
t=1
where the maximum is over By (x1.K) for each {a 121,..., TK KKKM and
S such that (x1.K) S
VaL.K E S Σ (x(xk))= VER
furthrmoree
By(zik)=0, (ZI.....,
(4)
In addition, if ( (21.k) denotes the parameters that achieve the maximum in (2), then (acc)
can be attained by a two step canonical decomposition: in the first step, given the list of input tokens (x1,....
apply Importance Weighted Sampling in Definition I with parameters 3 (I1....,IK to output an intermediate token y E [I1....,22) IK} }; in the second step we apply single-drafi speculative sampling scheme (Chen el al., 2023; Leviathan el al., 2023) on the selected token y to generate the final output token.
Figure I illustrates the proposed system in Theorem 1, where the first step involves importance weighted sampling to output an intermediate token and the second step involves speculative sampling. This approach requires computing the optimal Bt (x1:k). In practice one can sub-optimal choices that are faster to compute, as will be discussed in the sequel.
Remark 1 Although our theoretical development focuses the case of identical drafts, our result in Theorem naturally extends when the the K tokens are instead sampled from a joint distribution i.e., p(z1,...Zk) This is discussed in Section B.l in the supplementary material. In particular Theorem naturally extends to the setting when the K tokens are sampled independentl from different distributions S ~ II (X.) as in Miao et al. (2024) as well as to the case when the tokens are sampled without replacement (Jeon el al., 2024).
We next build upon Theorem
establish new analytical results for the optimal acceptance proba- bility involving K = 2 drafts. Our first result is a characterization of the necessary and sufficient condition on the draft and target distributions p(-) and q(-) respectively that leads to P* (accept) =
Theorem With K = 2 drafts, necessary and sufficient condition for p* (acc) = in the Definition I the following:
Σ (((x) 2 BES
Σ p(a)) WE8
AA CQ Ω
(5)
4

---

Published as a conference paper at ICLR 2025
Note that the acceptance probability can equal 1 even when p(- and g(-) are not identical. Thus when the distribution of the draft model is close to the target model but not equal the acceptance probability equal 1. This is in contrast to the case of K = 1, where it is known that the acceptance probability can only equal when p()) and q(+) are identical distributions (Sun et al., 2024b). Furthermore to the best of our knowledge, previously proposed schemes for the multi-draft setting, such Spec'Tr (Sun et al., 2024b) and SpecInfer (Miao et al., 2024) based on modified rejection sampling also require p(.) = g(-) for the acceptance probability to be 1. Theorem is interesting in the context of two-step system in Fig. 1. In this case, the output of importance weighted sampling block Y matches the target distribution q(+) and the second step involving speculative sampling is not needed.
Example Consider S {1,2) and let the draft and target distributions be given by p (pt.Pa) P2) and q = (q), 92) respectively. We assume K = 2 drafts. In this case (5) reduces to 91 2 p and
Q2 2 P2 If P1 = P2 0.5 then it follows that 1 (acc) 1 if and only if 0.25 S qi S 0.75. In
contrast for the optimal scheme for K = draft we have P* (accc) = 1 only when q1 = Q2 0.5.
The proof of Theorem 2 in Appendix C involves analyzing the output distribution pr(-) of the Importance Weighted Sampling Scheme in Theorem and demonstrating that a feasible choice of (y(21, exists and sets (-) = q(-) when the condition (5) is satisfied. The proof is based on the Fourier-Motzkin (FM) elimination technique (Ziegler, 2012). However direct application of such a technique to satisfy the constraints q(i) = (i) for each i E Ω becomes intractable. Our key idea to demonstrate that instead considering a relaxation of the form q(i) 2 pr( leads to the same solution as the equality constraints and is amenable analysis using Nourier-Motzkin elimination. We explain this further with an example involving I = (1,2,3) in Appendix C.
The core technical challenge in the proof of Theorem 2 is in determining whether system of linear equations has a non-negative solution. Such problems have been studied previously in the literature, with Chernikova (1964); Dines (1926) providing algorithm. Such considerations lead to geometrie viewpoint involving polyhedral cones which we discuss in Appendix D. We explain how the double-description method (Fukuda and Prodon, 1995) for finding dual representations of polyhedral cones can be used to numerically verify the necessary and sufficient condition for the acceptance probability to equal 1. In fact this approach was used to verify analogous conditions to Theorem 2 for up to K- = 6 drafts and all alphabets of size [S2] S 14, although we only provide an analytical proof of the condition for K drafts in this paper. The reason for this is that our key step i.e., Lemma 2 in the Appendix that makes use of Fourier-Motzkin elimination, does not easily generalize to the case of K drafts.
Our final result is an explicit expression for the optimal acceptance probability for the case of K = 2 drafts.
Theorem For K drafts and for a draft distribution p(-) and target distribution q(-) and arbitrary token alphabel Ωt the acceptance probubility P* (a(cc)) for the optimal token level selection rule given by:
(I
Σ p(s) leES
Pa (*(ccc) = min
Σ g(s)+
Σ p(s) .
+ Σ p(s) LEES
(6)
BES
where S = Z \; S is the complement o S.
To the best of our knowledge the result in Theorem 3 was not known before. Upper bounds on p* (ace) are presented in Sun et al. (2024b), which are not necessarily tight. In contrast (6) provides an exact expression for the acceptance probability for the case of K = 2 drafts when Xi and X₂ are independently sampled from p(-). The proof of Theorem 3, presented in Appendix E, applies climination to the linear program presented in Theorem to characterize an analytical solution in the case of K = 2 drafts. The proof builds upon the proof of Theorem but requires elimination of additional variables.
the
5

---

Published as a conference paper at ICLR 2025
1,00
1.00
0.95
0.95
0.50
0.90
Pr(acosp))
Pr
0.85
0.85
0.80
0.75
0.75
Optinsal (Thm. 3)
Optinaal (Thm. 3)
0.70
SpecIr
0.70
SpecIY
Speclsder
SpecTnfer
0.65
0.65
1/6
2/6
3/6
4/6
1/6
2/6
3/6
4/6
5/6
92
Figure 2: Numerical evaluation of Pr(accept) for the optimal scheme (Theorem two baseline schemes - SpecTrr (Sun et al., 2024b) and SpecInfer (Miao et al., 2024). For sake of illustration we select alphabet Q {1, 2, 3) and p 1/3,1/3, 1/3). The left plot sets q = [1/3, 92,
well as
92. 2/3 q2l while the right plot sets [1/6, 5/6 22 where is varied on the x-axs.
Remark Note that (6) can be expressed as:
2
Σm(s) BES
(rcce))
min SCR
Eg(s)
g(()
+
FES
which leads us to conjecture that the optimal acceptance probability in the general case of K drafts is attained by replacing the exponent of 2 in the second term in (7) to K.
We provide numerical evaluation of the optimal acceptance probability in Fig. 2. For sake of illustration we assume that f2 is of size three, and assume p [1/3, 1/3, 1/3.1// 1/3]. We consider 92l in the right plot. The value of (2 is varied on the r-axis. We compare the optimal acceptance probability in Theorem 3 with two baseline schemes SpecTr (Sun al., 2024b) and SpecInfer (Miao et al., 2024). We observe that the optimal acceptance probability can equal 1 for a wide range of 92. This is consistent with Theorem 2. In contrast the baseline schemes seem to achieve an acceptance probability of 1 only in the special case when Q2 1/3 so that = [1/3, 1/3, 1/3].
[1/3, 42, 2/3 92. in the left plot and [1/6, 92, 5/6
4 FASTER IMPORTANCE WEIGHTED SPECULATIVE SAMPLING
We discuss fast approaches for computing (-) in canonical decomposition in Fig. 1. Ideally wish to select B,( (-) that maximize the acceptance probability in (2), but this involves a computationally expensive linear program. Fortunately achieve high acceptance probability. In practice the distribution of the target and draft model is often concentrated over small number of tokens. It has also been observed that sampling from a high probability set such as the top-p set (the set of high probability tokens with aggregate probability exceeding threshold) leads more coherent outputs (Meister al., 2023). After such sampling the effective alphabet size, i.c., the number of tokens with non-zero probability, is generally small, We report some measurements of the effective alphabet size for the OPT model Appendix H. This motivates us to develop some approaches for speeding up proposed solution by reducing the number of variables required optimization. Throughout this section we consider the case when the drafts have identical distribution and discuss the case of non-idenne non-identical draft distributions in Appendix I. We focus on the case of K = 2 drafts and then discuss how to tackle the general case.
develop procedures for faster computation that still
Linear Program for K 2 drafts: We first consider the case of K 2 drafts and revisit the linear program that needs to be solved to compute the optimal acceptance probability in (2). We will then explain how to reduce the variables in the linear program for faster computation. Let X1 and X2 denote the input tokens and Y denote the selected token. Note that when X, = X2 = i, we have that (i, 1. Furthermore when the draft models have identical distributions, from symmetry,

---

Published as a conference paper at ICLR 2025
we have (y(i.)) (). It is more convenient to introduce a new set of variables way which are defined as:
Wi = Pr(Y = [X1, [x) = {i.j}),
=
(8)
i.e., WAy denotes the probability that the output token is i given that we see the pair {i,}} at the input in any order. We next discuss how to formulate a linear program involving the variables was to maximize the acceptance probability in (2). For convenience we let p (PI.....PM Px) as the probability vector for the draft model and q = (q.....,9h qn) as the probability vector for the target model. The acceptance probability i.e., the objective in (2), is given by Σ min (pr(t). (:). where
(b()) = PK
Σ i=1,igh
denotes probability that the selected token is Y k. Furthermore for our definition of wi, in (8), the associated constraints are: Wiy + Wy = and 0 S Wi,j S 1. This optimization can be cast as a linear programming problem with ((n²) variables which may be slow in practice. We next discuss techniques to reduce the number of variables in optimization without significant loss in the acceptance probability.
Faster Approximations to Linear Program: In order to propose a faster approximation, we note that when the draft and target distribution are concentrated over a few tokens, the linear programming solution will not be sensitive to most choices of wt,j- As a result one can heuristically set most of the variables, We refer to this method as the truncated LP scheme and present the pseudocode in Algorithm 2 in Appendix G. Assume that the vocabulary Q (1.2.....; n has the tokens sorted in decreasing order i.e., Q1 př 2 42 p2... .... 2 In Pn- We partition Ω into two sets Ω {1, 2, .......)) .) and f2g = {s 1..., n... where is design parameter select. We fix a subset of weights as follows:
Way = [1, i E R1,j E E
(10)
11, i E Ω2 E Ω2 <
while we leave the weights Wa,, for i < j and i,j E U2 as free parameters. The intuition behind the choice of weights in (10) is that in these cases we prefer token i over token j to increase pr(1) further, which is in turn can decrease the difference between di and (i). Note that (9) reduces to:
p P
= 1,ik 2pspauk. + Eis+1 2p:pk, k: E S2 Σ 1i=k1 2p:Pk.
(k)
(11)
k S2
Our objective to maximize reduces to Ei= min(pr(k) (k), over the variables Wh,-- Thus the number of variables is reduced to O(s²). We further show in Appendix F that if P* (acc) is the optimal acceptance probability associated by applying the linear program over all O(z?) weight
variables and P(acc) is the acceptance probability for the truncated program then:
Placc) 2 P*acc)=
Σ ((x)))
(12)
.Ef
Thus if 22 is selected so that the penalty term is small then the decrease in the acceptance probability can be kept small.
In the experiments we observed that for well-trained target models the drop in accuracy is negligible even for small values of s. Thus by appropriately truncating the number of variables to optimize the linear program we expect to have faster implementation. We discuss the computational complexity of the proposed method in Appendix G.
Our second proposal is to truncate the alphabet when performing importance sampling. In particular let 24 S Ω be a high probability subset of Z under the target distribution. We re-normalize the target distribution to q so that it is supported entirely over 2o and use this distribution in proposed scheme to generate an output Y q(-). Since we want the output to follow the distribution 9-)). we perform the following post-processing. We accept Y with probability Pa Σ Ezen, g(z) and reject with probability p. selected with probability proportional to g(z). This ensures that the output token follows the original distribution q(-). We refer to this scheme as the truncated alphabet scheme.
Pa. If rejected, output token from Ω Go such that any Ω Pl is
7

---

Published as a conference paper at ICLR 2025
To tackle the case Kf > 2 drafts we propose group the input tokens into groups of size and then apply the two-draft importance sampling scheme in multi-ster stage manner. For example if 5 X2, X3) and K 3 we first apply the fast importance weighted sampling to the group {X1, X2) to oulput an intermediate token Y1 with distribution say pi (.).). Then we apply importance weighted sampling the input (Y1, X3), where the tokens now have non-identical distributions, and produce an output token Y to which speculative sampling is applied.
Remark 3
5 EXPERIMENTAL RESULTS
Setup. We conduct experiments using an instance A100 GPU with 80GB memory. We use the OPT models (Zhang et al., 2022), where the draft model has 125 million parameters and the target model has 13B parameters. For evaluation purposes we consider the datasets associated with the XSum (Narayan al., 2018), Databricks-Dolly-155 (Conover et al., 2023) and the WMT18 (Bojar et al., 2018) tasks. All experiments were done over arbitrarily chosen seeds and the performance was averaged over these.
5.1 EXPERIMENTS WITH TOP-P SAMPLING
In this section we report experiments LLM tasks when top-p sampling is used with p 0.95. In our first set experiments Fig. 3, we consider case of K 2 draft models, which share same weights but use different temperatures for token generation. We set the temperature of the target model to 1.0, and one that of the draft models to 1.2 while we vary the temperature of the other draft model between the range of 1.0 to 2.4. In all our experiments generate 5 tokens per call of the draft model.
Our baseline scheme is the single-draft speculative sampling scheme (Leviathan al., 2023; Chen et al., 2023) when we only use a single draft with sampling temperature of 1.2. The other two schemes are the SpecInfer (Miao et al., 2024) and our proposed importance sampling (IS) scheme. In the IS scheme we employ both truncated L.P (with = as the truncation parameter) and truncated alphabet a size of 40 tokens) as discussed in section 4. Fig. 3 we report the performance achieved by the different schemes across the three tasks. The top plots report the block efficiency, which is the average number tokens that are accepted per use the draft model (Leviathan et al., 2023). The block efficiency achieved by the single-draft scheme is shown by the horizontal dotted red line in each plot, while block efficiency of the scheme and SpecInfer are shown using the dark blue and light blue bars. We observe that the IS scheme consistently outperforms SpecInfer
IS
Speciader
Single Draft SD
Dolly
XSum
WMT
Block Efficiencyy
1.5
1.0
0.0
Over Single lnggle-drall draft SD
Token Rate Improvenee
09%
.008
-10%
Draft 2 Temperature
Draft
Temperature
Draft 2 Temperature
Figure 3: Performance comparison of different multi-draft schemes. The temperature of the first draft models set to 1.2, while we vary the temperature of the other draft.
8

---

Published as a conference paper at ICLR 2025
across all three tasks. In fact when the temperature of the second draft is increased, the improvement in the block efficiency for SpecInfer is rather negligible compared to the single draft baseline. On the other hand proposed IS scheme is able to achieve consistent improvements. The bottom plots show the percentage improvement in the token rate with respect to the baseline single draft scheme. We observe that when the temperature of the second draft increases, the gains negative for the SpecInfer scheme as the computational time for the second draft does not translate into a sufficient improvement in the block efficiency. On the other hand our proposed importance sampling scheme shows a consistent improvement the token-rate over the baseline schemes due to improved block efficiency, despite the additional compute from the second draft. In the experiments involving the XSUM and WMT tasks we also measure the ROUGE-L (Lin, 2004) and BLEU (Papineni et al., 2002) scores respectively, which are reported the Appendix J.3. In Appendix J. we provide a similar experiment when both drafts use identical temperatures. We are able to also include comparisons with the SpeeTr scheme (Sun et al., 2024a) in that experiment. Likewise in Appendix J.2 we present similar experiment involving K drafts.
In Table I we study the effect of choosing different parameters for vocabulary truncation and LP truncation discussed in Section 4. We use the same parameters as in the previous experiment but fix the temperature of the second draft to 2.0. As we increase the size of the vocabulary S4g from 10 to 40 see improvements in the block efficiency as well the token rates. Beyond 40, appears that the gains in the block efficiency saturate and the token rate decreases. We also find that the block efficiency not sensitive the choice of the parameter the LP program and choosing s yields the best token rate.
Table 1: Effect of LP Truncation and Alphabet Truncation
Token Rate improvement to SD
Block Efficiency
10
1.98
0.03
-0.57 3.38%
20
2.00
= 0.04
1.00 + 3.08%
Vocab. Truncation ((9o))
40
2.05 0.04
6.63 + 3.18%
50
2.03 0.05
3.22 + 3.39%
5
2.05
0.04
6.63 + 3.18%
A--Truncation Threshold (s)
10
2.04 + 0.05
1.52 + 3.47%
15
2.04 0.04
1.74 + 2.36%
In Figure 4, we fix the temperature of the two draft models to 1.0, and vary the temperature the target model in the range of 0.2 to 1.0. The rest of the parameters remain the same. We again report the block efficiency and the improvement in the token rate the single draft scheme with the same draft temperature over the Dolly task. As the two drafts have identical temperature we are also able to include comparisons with the SpecTr scheme. We again observe that our proposed method generally attains the best performance.
IS
Spec/r
SpecInfer
Single-Draft SD
3.0
100
Token Rate Improvement Over Single-drafe SD
Block Elficiency
2.5
2.0
1.5
39%
1.0
0%
0.5
-2%
0.0
0.2
0.4
0.8
1.0
Target Temperature
Target Temperature
Figure 4: Performance comparison of different schemes on the Dolly task, while we vary the temperature of the target model and keeping the temperature the two drafts 1.0.

---

Published as a conference paper at ICLR 2025
5.2 EXPERIMENTS WITH TOP- SAMPLING
We next consider experiments when top-k sampling is applied to LLM tasks. We again use the same OPT models as in the previous subsection for both draft and target models. As before, also sample a total of 5 tokens during each call to the draft model. In the Dolly task we use a sampling temperature of 1.0 for both the target and draft models. In the XSUM task we use a temperature of 0.9 for the target and a temperature of 0.3 for the draft model.
In the first experiment, reported in Table 2, we compare the toke-level acceptance probability across different methods. We first generate a sequence of tokens via auto-regressive decoding using the target model. At each step we compute the logits generated by the draft and the target model, which we then use to compute the acceptance probability for different schemes. We average the acceptance probability over 100 instances. In the importance sampling (IS) scheme, we use truncated linear program with cut-off threshold
5. We also consider the theoretically optimal acceptance probability for the case of = 2 drafts and also compute the natural extension of this expression for K = 4, discussed in Remark 2 (reported using the gray font). We use top-k sampling with k 5 for both models in this experiment. We report the acceptance probabilities for the case K = 2, 4, 8 drafts. In applying the IS scheme to handle K > 2 drafts,
the successive selection scheme discussed in Remark 3. Note that the acceptance probability increases as we increase the number of drafts in each case.
Table 2: Comparison of average acceptance probability across different tasks.
Scheme
XSUM
Dolly
K
= 4
K = 8
K
K
K 8
Optimal
0.5009
0.5226
0,5419
0.6384
0.6731
0,6962
IS SpecTTr
0.4933
0.5145
0.5333
0.6348
0.6691
0.6919
0.4889
0.5083
0.5263
0.6246
0.6560
0.6800
SpecInfer 0.4875
0.5058
0.5227
0.6202
0.6489
0.6722
In Table 3 we compare the block efficiencies for different methods using K = 2 and K = drafts. We apply top-k sampling with k = 10 and k = 5 and use a temperature 1.0 for both models. In our IS scheme, truncate the variables in the I.P to a higher block efficiency over prior works, the gains are relatively small. This can be explained by noting that this setting the acceptance probability of the baseline schemes is already close the theoretical limit, as noted in Table 2.
7. While our proposed scheme still achieves
Table 3: Block Efficiency achieved in the Dolly Task with top-k sampling
Sampling
Scheme
K = 2 drafts
K = 3 drafts
Block Efficiency
Loss
Block Efficiency Loss
IS
2.48 + 0.01
2.59 + 0.02
top-k (k = 10)
SpecTr SpecInfer
2.43 + 0.01
98%
2.55 + 0.01
98%
2.38 + 0.02
96%
2.49 + 0.02
96%
IS
2.52 0.02
2.63 + 0.03
top-k (k: 5)
SpecTr
2.48 + 0.02
98%
2.56 + 0.03
97%
SpecInfer
2.47 0.01
98%
2.55 + 0.04
97%
CONCLUSION
We revisit the problem of maximizing the acceptance probability for token-level multi-draft selec- tion (Sun et al., 2024b). We first present a decomposition result that demonstrates connection to importance sampling. For the case of K = 2 (identical) drafts we establish closed for solution for the acceptance probability and also present a necessary and sufficient condition for the acceptance probability to equal one. We propose a practical token selection scheme that mimics the theoretical decomposition principle and provides a way to tradeoff the computational speed with the acceptance probability. We present experimental results using OPT model and three tasks: Dolly, XSUM and WMT and demonstrate improvements in block efficiency and token rates over baselines schemes.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We thank Mingu Lee, Wonseok Jeon, and Babak Ehteshami Bejnordi for their help with this work.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
Ondiej Bojar, Christian Federmann, and Mark et al. Fishel. Findings of the 2018 conference on machine translation (WMT18). In Proceedings of the Third Conference Machine Translation: Shared Task Papers, pages 272-303, 2--3,, Belgium, Brussels, October 2018. Association for Com- putational Linguistics. doi: 10. 1/W///////11-6// W1 81-8- 1188--4011 01.
URL https: aaalanhooannttyyooooaaaao .rg/
Yelysei Bondarenko, Markus Nagel, and Tijmen Blankevoort. Quantizable transformers: Removing outliers by helping attention heads do nothing. Advances Neural Information Processing Systems, 36, 2024.
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 leamers. Advances in neural information processing systems, 33:1877-1901, 2020. 2020.
Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen, and Tri Dao. Medusa: Simple I1m inference acceleration framework with multiple decoding heads. arXiv preprint arXiv: 2401. 2024.
Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, -ran-Baptiste Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.01378, 2023.
NV Chemikova. Algorithm for finding general formula for the non-negative solutions of system of linear equations. USSR Computational Mathematics and Mathematical Physies, 4(4):151-158, 1964.
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. Journal of Machine Learning Research, 24(240):1-113, 2023.
Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks sequence modeling. arXiv preprint arXiv:1412.38 2014.
Mike Conover, Matt Hayes, Ankit Mathur, Jianwei Xie, Jun Wan, Sam Shah, Ali Ghodsi, Patrick Wendell, Matei Zaharia, and Reynold Xin. Free dolly: Introducing the world's first truly open instruction-tuned 1lm, 2023. URL https: http://wwwwaaaaabaaa..a.a..bb..a dat wekaabbriiikk..e. com/ waaa///log//0// ////// 12/dolly first open conmerciailaioooo www.erciall--vialeeeeeeeeee. hwwteaabeeeinssttue ins- won.ucttiiotteeee tund--- 11m.
George B Dantzig and B Curtis Eaves. Fourier-motzkin elimination and its dual. Journal of Combinatorial Theory, Series A, 14(3):288-297, 1973.
Lloyd L Dines. On positive solutions of a system of linear equations. Annals of Mathematics, pages 386-392, 1926.
Elias Frantar and Dan Alistarh. Sparsegpt: Massive language models can be accurately pruned in one-shot. In International Conference on Machine Learning, pages 10323-10337. PMLR, 2023,
Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint arXiv:2210.1332. 2022.
Komei Fukuda and Alain Prodon. Double description method revisited. In Franco-Japaness and Cranco-Chnesss conference combinatorics and computer science, pages 91-111. Springer, 1995.
11

---

Published as a conference paper at ICLR 2025
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.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv preprint arXiv:1503.0255: 2015.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):
1735-1780,, 1997.
Sebastian Jaszezur, Aakanksha Chowdhery, Afroz Mohiuddin, I Lukasz Kaiser, Wojciech Gajewski, Henryk Michalewski, and Jonni Kanerva. Sparse is enough in scaling transformers. Advances in Neural Information Processing Systems, 34:9895-9907,, 2021.
Wonseok Jeon, Mukul Gagrami, Raghavy Goel, Junyoung Park, Mingu Lee, and Christopher Lott. Recursive speculative decoding: Accelerating Ilm inference via sampling without replacement. arXiv preprint arXiv:2402. /4/60, 2024.
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.
Chin- Yew Lin. ROUGE: A package for automatic evaluation of summaries. In Text Summariza- tion Branches Out, pages 74 81, Barcelona, Spain, July 2004. Association for Computational Linguistics. URL https: / /www. www.aclwwwdt web.web.co. httgg//annhoooog////000000
Clara Meister, Tiago Pimentel, Gian Wiher, and Ryan Cotterell. Locally typical sampling. Transac- tions of the Association for Computational Linguistics, 1 11:102-121, 2023.
Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang, Zhengxin Zhang, Rae Ying Yee Wong, Alan Zhu, Lijie Yang, Xiaoxiang Shi, et al. Specinfer: Accelerating large language model serving with tree-based speculative inference and verification. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3. pages 932-999 949, 2024.
Shashi Narayan, Shay B Cohen, and Mirella Lapata. Don't give me the details, just the sum- mary! topic-aware convolutional neural networks for extreme summarization. arXiv preprint arXiv:1808.0 08745, 2018.
Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: method for automatic evaluation of machine translation. In Proceedings of the 40th annual meeting of the Association for Computational Linguistics, pages 311-318, 2002.
Noam Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.221500 2019.
Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit. Blockwise parallel decoding for deep autore- gressive models. Advances Neural Information Processing Systems, 31, 2018.
Ziteng Sun, Jae Hun Ro, Ahmad Beirami, and Ananda Theertha Suresh. Optimal block-level draft verification for accelerating speculative decoding. arXiv preprint arXiv:2403./0444 2024a.
Ziteng Sun, Ananda Theertha Suresh, Jae Hun Ro, Ahmad Beirami, Himanshu Jain, and Felix Yu. Spectr: Fast speculative decoding via optimal transport. Advances in Neural Information Processing Systems, 36, 2024b.
Surya T Tokdar and Robert E Kass. Importance sampling: a review. Wiley Interdisciplinary Reviews: Computational Statistics, 2(1):54-60, 2010.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones. Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017
12

---

Published as a conference paper at ICLR 2025
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre- trainaddd transformer language models. arXiv preprint wrXiv:2205.01066 2022.
Yongchao Zhou, Kaifeng Lyu, Ankit Singh Rawat, Aditya Krishna Menon, Afshin Rostamizadeh, Sanjiv Kumar, Jean-François Kagy, and Rishabh Agarwal. Distillspec: Improving speculative decoding via knowledge distillation. arXiv preprint arXiv:2310.08461, 2023.
Günter Ziegler. Lectures on polytopes, volume 152. Springer Science & Business Media, 2012.
Zny. Skeleton algorithm, 2018. Available at http:://www.iic uic rnn.ru/ zny/skcleton/.
13

---

Published as a conference paper at ICLR 2025
A SPECULATIVE DECODING BACKGROUND
We briefly review the proposal made in Leviathan et al. (2023); Chen et al. (2023). Let denote the set of all permissible tokens associated with a language model. Given context x² (a=11)...... .(t)) Ωt our target model produces conditional probabilities M,(yk=") for each E Ω. Following prior works and explicitly stated in Sun al. (2024b) we assume the following computational model:
1. Standard Inference: Given context æ², an auto-regressive model can output Mo(ylæ") for each y E Ω in (1) time.
2. Parallelization along time-axis Given context x², auto-regressive model can output Mo(ylz") for all i E [1,...,4] and each E Ω again in ((1) time.
Since the standard inference is a sequential process, the idea in Leviathan al. (2023); Chen et al. (2023) is to exploit parallelization along time-axis to speed up inference. This is accomplished using a draft model M, that is capable of generating multiple tokens with a much lower computation than Ma(.). The main steps are summarized below. We assume that context token x! is the input.
Draft Construction:
The draft model can efficiently sample L tokens in an auto auto-regressive manner, extending the context x' to (e(1)... ....,.......... sr(),, a(t + 1). 1). _itt + T.). In addition, keep the conditional probabilities Mylulæf, x(t1) for each 0 S S L and each Ω.
2. Conditional Probability Computation: Given the samples = generated by M, we
compute the conditional probabilities Mo(yl=1, atti) for each 0 S < L and each y E P using parallelization along time-axis.
3. Draft Selection We select first L' < L tokens and set x(t a(t i) for S L'. The selection is based on speculative decoding algorithm below.
Algorithm Speculative Sampling
1: Input: Distributions p(-) and q(-). Sample X p(-)
2: Compute residual distribution (xx) I- min(p (ar) (())'))))))
q(æ) m((l
3: Accept = False
4: With probability (x) = min (1,
g(X) p(x)
set Accept = True.
5: if Accept= True then
6:
y= X
7: else
8:
Y ~ <=""*(-)
9: end if
10: Return Y.
For convenience we denote p(v) = Ms(ulz') and d(v) = Mo(yl=". The selection procedure is as follows: Given context xf and candidate draft sequence +LL the algorithm in the first step sets p(ã++1) = ( and q(t++1) = Mo( respectively. It accepts ão41 with probability (x241) = min p(z-)))) If the token accepted then xxt + 1) = Ø and we proceed with 2 If the token is rejected then we sample ~ p= (.) and set æ(t + 1) = Y. At this step we initiate call to sample fresh tokens from the draft model with s++1 as input. This process continues until all L tokens have been accepted or a token is rejected. The correctness of the proposed algorithm follows through direct computation of Pr(Y = y). We also compute the probability of accept:
Pr(accept)
Σ
Pr(ace| X =x)x(())
(13)
DEN
Σ (p)(())((2))
(14)
sEX
Σ
min(p(z). drv(p.
(15)
TE
14

---

Published as a conference paper at ICLR 2025
where (-) is the total variational distance. Note that Pr( Pr(accept) is key metric that determines the probability that a token from the draft model is accepted. The higher the value of Pr( Pr(accept), the more efficient is the algorithm. In this paper also study Pr Pr(accept) in the multi- setting.
B PROOF OF THEOREM
We will consider the case when there are K drafts i.e., XI,..., "i....,KK'. are sampled i.i.d. from a draft model with distribution p(-), while the target model has distribution of q(-). We assume the alphabet ? = MM for some arbitrary M.
Analysis of Importance Weighted Sampling Scheme: We first consider the family of importance sampling schemes followed by speculative sampling and derive the acceptance probability. Assuming denotes selected sample importance sampling, let:
(xf))
E 11,,... 3K
Pr(Y = x| = x;{) = {0.
(16)
Φ f TKK
where Ey σK) for each zf 22 and s (21,....... TK) 1 It follows that
K (x(x))
Pr(Y = = Σ
p(,))
(17)
Flyont EQ
Σ
By(xf) Iy ( σk)
II p(x))
(18)
where L xK) denotes the indicator function that equals if y E { RK} and equals 0 otherwise. Note that (18) follows since ", 0 if y $ (kk)
By applying speculative sampling to the selected sample X, the probability of acceptance is given by:
PM-18 (accept 1) = Σ
Pr((, = =y))
(19)
yU
Σ E
min
g(v),
Σ
3y(x1) Ig .,, rK)
II p(z:)
(20)
Thus within the proposed class of importance sampling schemes, we can formulate our objective as:
Øx) II p(o.)) }} (21)
K
max (x Thy
min
q(v),
Σ
By(at()
BEn
such that 0 5 (x(K < 1 for each y, II,... tK E S2 and
Σ
( 11, Vy E
(22)
EAA
and furthermore
,,(x1) = y (rI,,.., rz)
(23)
We the notation X short hand for XI-K. Similarly aif shrot hand for Tisk
15

---

Published as a conference paper at ICLR 2025
Analysis of Optimal Solution: We now consider the problem of optimizing the acceptance proba- bility for any given p(-) and q()) in the general setting. Following the framework in Sun al. (2024b), seek to find Py|X Xx (2)2),.......................... xx) for each y, ....... JK Zt such that maximize
Pr(accept 1) Pr(Y E (X1...... Xx1)
(24)
subject the marginal constraints Py(():
((v) - Pr(Y = =y) = D>Prrre Pr(Y = y, X = =xf)
Exrxx (
II p(x)).
(25)
x
Next we consider:
((v)= Σ
PrIXF (ula²² K) II p(x.)
#tEn*
Σ atEn
PrIX (ylaf (1
xK)
p(x:)
+ Σ
(ulaf
xx)
II p(xs)
(26)
2 EN
1=1
Σ
PrIx (yla (o .......... rK)
II p(m))
(27)
#ten*
where (xx) = 1 - 1,(,,,..,,æ denotes the complement of I. Now note that:
Pr(Y E (XI,,..,,K,
Σ Pr(Y (XI...., .x). |x x (x( =x
(28)
ot ENK
K
Σ
Σ
Pr Pr xx (yl= (21.......... II(z))
(29)
of EAK yER
Σ Σ VER of ENK
Py xx
................................ FK)
(30)
[I(tss)
min
((v),,
Σ
Prixp(ylat (ol.......
................................ xK)
(31)
yE
x Enx
where we use (27) which implies that for any feasible Py xx (y)xf();
K xx) II p(x.) S q(v)
Σ 2f ED
Py xx (512)) )
(32)
is satisfied.
Upper Bound the optimal acceptance probability: We now establish an upper bound on (31) and show that it coincides with the acceptance probability optimized in the importance weighted sampling scheme (21).
For each σ E K let us define
D(a = Σ
PrIx (gl (21......
(33)
VEQ
16

---

Published as a conference paper at ICLR 2025
and furthermore with ((1,,....... denoting the number of unique elements in zf"
Prix (wl=f)
y E { [1
rk}.
D(af) > D(af)) = 0,
D(2f)
Prixg (ylof)
y E {o IK].
(34)
[a [.)
y $ fa ( vK).
Note by construction that for each of Q
Σ
prixf (ul=f)
(35)
wen
and
Pr|x (vl=)) =0, y f21,,.., oK}
(36)
and furthermore:
Pr|i (y/=f) [ (21,..... rx) 2 PYIXE (y ) [,(11,...... (KK)..., $1,..., TK (37)
Substituting (37) into (31) have that for any feasible PrIxF( (-) there exists 0yl (-) satisfy-
ing (35) and (36) such that:
K
Pr(Y& E (x1)
& Σ min BEN
((y). 2f EAK
Σ
PrIX (ulz) (21 xK)
II p(x.)
(38)
It thus follows that that optimal acceptance probability in the general case is upper bounded by optimizing the (38) over PYIx (y)= satisfying (35) and (36). But this problem precisely coincides with the optimization the proposed class of IS schemes stated in (21)-(23), thus establishing the optimality of the latter.
B.1 EXTENSION BEYOND IID SETTING
The proof in Theorem I assumed that ø1... FK are sampled form the same underlying distribution p(-). Here provide natural extension when the tokens are sampled are not sampled from product distribution but instead from a joint distribution: p(I),...Ix).
Theorem Let P* (acc) be the acceptance probability for the optimal token level selection rule when S ~ p(x1,....... R.). Then we have
P*(acc)
max (9,(#1:))
Σ Emin
min <lv), Σ
By(xik): p(oxk))
(39)
EE
FI EQ
where the maximum is over By (x1.K) such that S (T(KK S and
Σ (x(21k)= =1 Vy E S
(40)
EDK
furthermree
By(ank)= =. O. "II,...,KKK)
(41)
Furthermore if (3((k)))) denotes the paramelers that achieve the maximum in (39), then P* (acc) can be attained by two step approach as follows: in the first step, given the list of input tokens (21....,IK] we apply Importance Weighted Speculative Sampling in Definition with parameters 3y (7 to output an intermediate token y E {I1...., .......,2. X J; in the second step we apply single-draft speculative sampling scheme (Chen et al., 2023; Leviathan et al. 2023) on the selected token y to generate the final output token.
[x)
The proof of Theorem 4 identical to the proof of Theorem I. We note that replacing the distribution of S from I p(X;) to the joint distribution p(= Ix). does not affect any of the steps as we did the property that the joint distribution factorizes in the product.
17

---

Published as a conference paper at ICLR 2025
C PROOF OF THEOREM 2
We first consider the special case of f = =1,,,33 to illustrate the key ideas. We then proceed with the proof.
Example 2 Consider the case when Q = 21, and let p (p1,P2,P3 P3) and q (q) 92, qa) denote the draft and target model distribution for the current token of interest. We again assume K 2 drafts. Let X1 = and X₂ j denote the pair of input tokens and Y denote the output of the importance weighted sampling scheme in step 1 in Fig. 1. Since X1 p(.)) and X2 p(.) it is clear the the optimal TLSR does not depend on the order of X1 and X2 but only on the unordered set (X1,X2) and let (i.j} denote the realization. Let 04, Pr(Y-= = i, ( (X1, X2) } (i,j3) denote the probability of the event that the (unordered) input tokens are {i.j} and the output token is Y = i. Similarly let aj, = Pr(Y = j, (X1, X2] = {i Note that Qi,i P¹ must hold, when X1 = X2 = i, clearly Y = i in the Importance Weighted Sampling scheme. Note that P (acc) = 1
requires that Pr(Y 1) 9. for each E Ω. This results in the following system of linear equations:
a1 = pi + 01, + 113
Q2 = 22 + 02.1 + 0223,
q3 = P3 + 03,1 + 03.2
(42)
subject to 05.. + jz. 2p.9; and 0 S aij S 2p:Pj. We prove that (5) provides a necessary and sufficient condition that the above system of linear equations has a feasible solution.
Our initial attempt was to directly apply Fourer-Motzkin (FM) elimination technique Ziegler (2012); Dantzig and Curtis Eaves (1973) to (42). However a direct application of FM elimination does not appear be tractable for arbitrary sized alphabets, as the elimination of each variable introduces large number of inequalities. Our key observation is that (42) is equivalent to the following relaxed set of inequalities:
Q1 2 p + 01. + a1.3.
92
2 22 02,1 + 0223,,
93 2 13 + 0,,, + 03,2
(43)
with the same conditions on 0i.j as before. A solution to (42) exists if and only if a solution to the relaxation (43) exists. Indeed contradiction, suppose that a solution (43) exists with strict inequality in one of conditions. Then summing over all the inequalities and using 0i.j + a,,i = 2p.P; gives q1 + 92 + > (p1 +p)) However since p and q are probability vectors both sides should sum to 1, leading to a contradiction. Our second key idea is to augment the system of inequalities in (43) with the following additional inequalities:
Φ1 92 2 (p1 + p2)) + 01.3 + 022,3+
@ + 03 2 (p) + pa)² + 01,2 + 03,2;
02 + 4 (p2 + pa)2 + 012,1 + C13,1
(44)
Note that the inequalities in (44) are redundant and follow by simply adding each pair of inequalities in (43) and using 04 + O,j = 2p:Pg. However applying FM eliminations simultaneously over the expanded system of inequalities involving (43) and (44) is surprisingly tractable. In fact we show that applying FM elimination for eliminating each 04. (and by extension aj,z) simply involves dropping that variable in the system of inequalities (43) and (44), For example eliminating 01,2 (and simultaneously (22,1) 1) in the first step is equivalent to:
41 2 pi 01,3; 92 2 p2 022,3, 93 2 p3 0:3,1 + 033,
(45)
a1 + Q2 2 (p) p2)² + 01,3 + 02,3, a1 a 2 (pl + p2)² 03,22 Q2 + q3 2 (p2 2a + 013,1
(46)
Eliminating all Oi, in this fashion establishes that feasible solution exists if and only if qi 2 pt and
9, + Qk 2 (p; + px)² for j,ijj E Ω and j + k. This is precisely the condition in for an alphabet
of size [[2] 3.
We now proceed with the proof of the result.
Setting of Linear System of Equations and its Relaxation: Following the simplified notation in the main text for the case of K = 2 drafts, we let q = (9I.....9.) qn) be the target model distribution and p (P1,.... .., Pr.) be the draft model distribution. Also recall that define 0g,j Pr(Y
i, (X,,XX) X2) = {(,j}) as discussed in the main text. In order to match the output distribution Pr(Y =) to the target distribution, need satisfy the following system of linear equations:
18

---

Published as a conference paper at ICLR 2025
91 pi 01,2 + ... 01,
(47)
42 p 02 + .... + 022
(48)
(49)
On Pn = 0n,1 +... + an
(50)
where 02. 2 0 and 04, + 0j,
2p;p; 2ps.j for each i + ml.
We instead consider a relaxed system of inequalities:
91 p 2 011, + 011m 2
(51)
92 P2 2 02,1 + 02,2
(52)
(53)
On Pn Z On,1 + ..... Onnnn-
(54)
where Ciij 2 0 and 0æj + دzب = 2p,P; = 2p:.j for each + j f1....,n). We note that the nn. system of inequalities (47)-(50) has a solution and only if the system of inequalities (51)-554) has solution. Indeed, for contradiction assume that one of the inequalities in (51)-(54) is a striet inequality. Then summing over the left and right hand sides and using at + aaja 2ptpy we get that
2
Σ 4s
Σ
Pb
(55)
which is contradiction both sides sum to 1. Thus it suffices to consider the system of linear inequalities.
Augmented System of Inequalities: Instead of the original system of inequalities (51)-(54), consider an augmented system of inequalities defined as follows.
Lemma Our original system (51)-(54) has a solution if an only if the following system has solution:
2
Σge 9: -
Σ
Ps
ΣΣ Csst
S [l....,n}
(56)
SES
IEE
#ES LESt
for ag 2 and asst + 01,8 for s,t E (1,2....,n)) nn with +
To establish this, (51)-(54) and sum s S:
Σ (q. pš)
as.j
(57)
#ES j=1.fs8
Σ Obgtt +
Σ
Oggt
(58)
8ES
LESt
tES1[s)
algtt +
(59)
SES
SES (ESE(S(()
ΣΣ
Ogst +
Σ
(a.t + as,))
(60)
SESIESS
Ogt +
Σ
2p
(61)
8ES tESe
(8.))ESx51155
follows that:
Σ
11 2
Σ
2ps. 2 ΣΣ
Osst
(62)
#ES
#ES
(s,)
#ES LES
19

---

Published as a conference paper at ICLR 2025
Σ
9
Σ Ps
ΣE
1gg
(63)
#ES
sES tES
as required. The other inclusion follows by simply setting S {i} for each Induction Argument
We will prove the following by induction.
Lemma 2 Lel
V == ((i, (r): (r,,,,)
(64)
denote the indices (with in ju for all k ,,...,T) FM elimination. Then the remaining constraints are given by:
T) of the variables eliminated after rounds of
Σ %
Σ Ps
ΣΣ O.d &((s,t)) AS S (I.....n) SESIES
(65)
Remark 4 When all the variables have been eliminated the right hand side in (65) will equal foror any choice S S 21,2....,7 n. and we will recover the result Theorem
Note that the base case with V, = {.} immediately follows from (56). We will assume that the variables di and dg eliminated for q {I....,r 1} and the associated Fourier-Motzkin (FM) conditions are given by:
Σg. g0 #ES
Σ Ps
EE aee I(((stt & VA-1), AA S (1...., n). *ESTES"
(66)
At step climinate the variable a,drr and o,,, and will show that (65) is satisfied. In applying the FM elimination, we only need to consider those inequalities in (66) where either ai a appears the right hand side. The remaining equations will not be affected in this step of FM elimination and replacing V,-1 with Vr will not have any effect there. Any such inequality will be associated with choice of S where either both i, and j. belong to S neither and j. belong to S. Thus we have:
Σg. #ES
%
Σ Ps
2 ΣΣ a.s I((s,t) & Vr),
BES
PESIES-
VS
S (I....,m i, S & j. ir &S& S.
(67)
The FM elimination will only consider those inequalities in (66) where either ai..j or a appears in the right hand side. The inequalities where ai.., appears the right hand side is associated with those subsets S of (1....,7) n where i, E S) and j. & Si. Likewise the inequalities in (66)
where J,,,, is appears jr E S2 and i, & Sz. Thus the FM elimination applied to variables air and a will consider the following system of equations:
the right hand side are associated those subsets S2 E 21,2,..., n) where
Σ
Σ
2
Σ
Σ
Og I((8, V,-1),
SES:
BES
BES; LES;
VSA C (1...., m), in E S1.j. S1.
(68)
Σ Us
Σ Ps
Σ
ass I((s,t) & V,-1),
SESa
SES
SE tESS
VS2 C (1...., n).jr Jr E S2, ir & S2.
(69)
Ci, + ajiht 2pi,...
(70)
20

---

Published as a conference paper at ICLR 2025
Jr 2 O, jr 2 O.
(71)
Accounting for (71) and using the fact that V, V ((ir:))) (rr,,))) we immediately have that:
Σ 9s
Σ
Ps
2
ΣΣΣ sES; LES
Ustt [(((,1) $ Vr),
#ES1
+ES;
VS1 _
[1,...,71
n,, in E Sr, $ St,
(72)
Σ 9s
Σ Ps
a,t [(s, & V.),
#ES2
BES2
SES, tESS
VS S {1,
n). Jr E S fr $ S1.
(73)
In addition the FM elimination procedure is required to combine every possible inequality in (68)
with every possible inequality in (69) and eliminate d, and ji,.., by applying (70). For a specific choice of S₁ and Sn the inequality we consider is of the form:
2
2
Σ 9s + sES;
Σ9 2s-
Σ .ES
Σ
Ps
Ps
sES,
Sc
Σ
Σ
Oxst
I((s,t) & V,-1)
+
Σ SES² NESS
Σ
ast I(((stt $ Vy-1).
(74)
#ES1
tES;
We will show that this inequality is redundant as it is dominated by the set of inequalities in (67). Let R- S n S2 and T = = S2. Note that i, & R and j. & R. Now consider the left hand side of (74).
Σ aES1
9. +
Σ
Σ
Ps
Σ Ps NoESZ
8ES2
SES
2
= Σ
9s +
as
+ 2
Σ
90
Σ
Ps +
Σ
Σ #ES
Pe +
Σ
Pa
#ES
|R
SES, R
sER
$ES.
RR
SER
/R
SER
(75)
2
2
Σ 4 +
Σ
Σ Ps
Σ MER
Σ
Σne 9s
9
Ps
-
Ps
R
BES, JR
BEST
IR
WE
Σ
Ps
Σ
Ps
2
Σ Ps
Σ
Ps
(76)
RR
ER
NES
RR
NE
2
2
2
Σ
1s
Σ
Σ
Ps
Σ Ps
Σ
Ps
-
Σ
Ps
Σ 2s LE
BE
SESIYR
BER
$ESO
SESI
RR
-2
Σ P&
Σ
Σ P.
Σ
Ps
2
Σ
Ps
BES,YR
RE
BES,
BESTiR
RE
+2
Σ Ps
Σ Ps YR
(77)
(R
IR
2
90
Σ
+
9.
Σ Ps
+2
Σ
Σ
Ps
(78)
SET
SET
SER
LER
#ES2iR
BESI IR
21

---

Published as a conference paper at ICLR 2025
We now consider the right hand side of (74). We recall that with T= S1 U22 and R. S1 n the following relations that can be easily established using Venn diagram of sets SJ and Sy:
Si T U (S2 \ R),
(S2 R) = (
(79)
S = T U (Si \ R),
(S \ R) = {-}
(80)
R T" USi \ US2 \ R), (Si (S) \ R) = -1
(81)
T = R (S:) \ R) (S2) \ R)
(82)
Now consider the following:
Σ Σ
att I((,,t) $ V+-1)
SES, 1EST
Σ
Σ
Ort [((8,)) & V.-1))+ + Σ
Σ
Oxt T((8,t) V.-11
(83)
SES, IET"
SES, (R
Σ
Σ
I((, @0 (V,-1)+ + Σ
Σ
Oss I((8,)) $ V.-1)
DESARIETT
\R IET
SER IET
Σ
Σ a,t II((s,t) V + Σ
Σ ast I(((st) $ v,-1)
(84)
SEST 1R
SER IES:)R
where we use (79) in (83).
In similar fashion we can express,
ΣΣ
(((( & V--1)
#E8, IES;
Σ
Σ Cee I[(8,)) $ VV-1)++ + >>>
ass I((8,1) & V.-1)
IET
SER IET
a,t I(((st) $ VV-1)) + Σ Σ a,t I((s,t) 1.-1) (85)
RRIS
RR
Combing (84) and (85) and re-arranging terms, we get that:
Σ
Σ
T(((s & + Σ Σ ast (((s,1)&VV--
asa
sES, LEST
SES; LEST
Σ
Σ Ougt
I((sst) V,-1)
asst I(((st) (V+-1)+
Σ Σ
asit I((s, 1) $ V--1)
sES, E
SER 1ET
#E5 IET
Σ
Σ
Σ
Σ ass 1((stt) t) $ Vソ-1) + Σ Σ agtt I((sst) $ V--1)
ag,t +
SER (R
1ER
SER IET
Σ
Σ
agp I((s,t) & Vy-1) + Σ Σ as, I((s,t) $ v.-1)
(86)
SES, RESAR
SESARIESI
= -Σ
agt I((stt) $ 12-1)+ + Σ Σ
age I(((st) $ Vy-1)
SET TET
SER IER
Σ
Σ 0gt I((s,t) $ Vn-1)++ + Σ
Σ
Ot. T((stt) $ Vr-1)
(87)
sES, AR LESAR
SE RR
Σ E Ogt I((,,) vv) Σ Σ asd I(((s,) ((s, V.)
SET LET
sER IER
Σ
(aL, as.)) I((s, f) $ 1--1)
(88)
SES RESS
where we use (81) and (82) in (87) as well as the fact that [[(s t) & V,-1)= = I[(t,s) $ v,-1) as the pair (s, () and (1, is eliminated simultaneously. In (88) we the fact that T contains both i, and
22

---

Published as a conference paper at ICLR 2025
Jr while R contains neither Łr. and jr-1 and hence ai.,,. Jr,., do not appear in the first two terms in (88) so that V,-1 can be replaced by V,. Combining (78) and (88) it follows that the FM elimination for choice of S1 and S2 leads to:
Σn.
Σ ps
Σ
40
ΣPo BER
+2
Σ
ps
Σ Ps (ESI
SET
NET
FER
SESiiR
2 Σ
Σ
Cns I(((s,)) & V.)+
Σ E
ass
I[((st) $ V+)
NET T
SERIER
+ Σ
Σ
(ar + as,)) I(((s,t) $ V,-1)
(89)
sESi
(R
RESAR
Note that this condition equivalent to:
9
Σ
Ps
E
aes I[((stt) & Vr)
SET
GET
SET TET
Σ 4s
Σ Ps
ΣΣ BERIERS
aas I((s,t) (1))
BEE
E (ar, + agg)) I(((st) $ vV-1)) )
+
2
Σ
Ps
Σ $ESii
Ps
2 0. (90)
AASTVIESSY
We now show that this condition is redundant it is implied by other conditions. Since T and R satisfy the conditions in (67) we already have that:
Σ 4
Σ
Ps
ΣΣ NET LET
ass 1((s,t) $ V.)
(91)
SET
SET
Σ 9x
Σ Ps
Σ
Σ
Ggst I((sst) $ V+)
(92)
LER
IE
Further since the sets (S, R) and Sy
R are disjoint follows that:
2
Σ
Pt
Σ Ps
Σ
Σ
(ar, + a..e)) [[((s))) $ VV-1)
(93)
(ES
#ES /R
VES LES,/R
Σ
2psPt
Σ
Σ
(at, + as I[((stt) $ VA-1)
(94)
1ES, R
SES,
R
sES \R R
Σ
Σ
(2p-Pt
(au,s + a.,)) I((s,t) $ V V.-1) 2 20
(95)
/R
where we use the fact that by construction aa,, + Ot, 2p-P:. It thus follows that the condition (90) implied by other conditions already presented in the FM elimination and thus redundant. Since our choice S and S2 is arbitrary it follows that every combination of the form (74) is redundant and the only equations that remain upon elimination of atr.jr and a.j,,, are given by (67), (72) and (73). This completes the induction step in Lemma 2 and the proof.
D CONNECTION BETWEEN THEOREM AND POLYHEDRAL CONE REPRESENTATION
We consider the case of = (1, 3, for sake of concreteness. We discuss how the characterization P2, Pa) denote of Pt (acc) = 1 is related to dual representation of polyhedral cone. Let p (p,,P2,P
23

---

Published as a conference paper at ICLR 2025
the draft probability and q (91, 92. 9s) denote the target probability vector. As before define az,j = Pr(Y = i, Xx1 = {i,j}). ). We need to solve the following system of equations:
91 pp 0122 0,,3
(96)
Q2 P 0221 0223
(97)
93 P3 = 03,1 + 03,2
(98)
subject to the conditions that oij + 0j = 2p:P; and 0 < aij S 2p:P3- Using the fact that
Q1 42 + Q3 1 and P1 + P2 + Ps 1, suffices the consider the following system equations:
01,2 + 01,33 = q1 - p
(99)
02,1 + 02,3 = 92 - P2
(100)
01,2 02,1 2p1,2
(101)
01,3 + 03,1 = 2p1
(102)
22,3 + 03,,2 2p2.3
(103)
with the additional requirement that aij 2 We will represent this system of equations in matrix
form. Our variables of interest are x [a112, 01,,3 022, 022,3, 03,1, 03,2] 2 0. Our equality
constraints can be expressed in the following form:
A A b,
X X20
(104)
where
1, 1,1,0,0,0,0
01,2
p
0,0, 1, 1,0,0
01,3
- p2
A
1,0,1,0,0,0
0g21
b
2p1,2
(105)
0, 0,1,0,0,1,0
0223
2p1.3
0,0,0, 0,0,,
03,11
2p2,3
(33,2
Upon application of Farakas' Lemma follows that the system (104) has solution if and only if every y that satisfies y" A 2 0 also satisfies y 2 0, where b depends on and q as in (105). Let us define
I. 0. 1, 0, 0
1, 1,0,0,1,0
0,1,1,0,0 B B
(106)
0,1,
0,0,0,1,0 0,0,0,0,7
and note that the set
B = (y By 2 0}
(107)
denotes a polyhedral cone in R We need to show that for each y E we must have that > 2 The representation (107) is the so-called hyperplane representation of the code each of B
defines a hyperplane. We would like to find an equivalent generator representation the form:
R
{z = RA, X 0)
(108)
The Minikowski- Weyl Theorem (Fukuda and Prodon, 1995) guarantees that for every B in (107) there exists in (108) of finite dimensions such that = R. Furthermore the double-description method is algorithmic way of computing R given B and vice versa. Using the package skeleton for double description (Zny, 2018) we could show that for the B matrix in (106) the associated R matrix is given by:
Is
1
1
0
R
1
0
(109)
0
-1
-1
-1
1
24

---

Published as a conference paper at ICLR 2025
where Is is 5 5 identity matrix. The generator representation in (108) is convenient as in order to show that (104) has a feasible solution, it suffices to show that R R²b 2 Indeed substitution of (109) and (105) yields
b
b
@ + 02 (pi + 12)²
+ 02 - (pi + p2)²
R²b-
41 + pf + 2p1.2 2p1,3
=
92 + 93 - (22 + ps )
(110)
42 + p2 + 2p1.2 2p2,3
91 + 93 + (11 + ps)²
Q1 92 + pl
2p1.2 1 2p1,3 3 + 2
P3
2
In the last step we use the fact that Eq: =
It thus follows that R > if and only if 2 p and qi + q; 2 (p: + P2) holds as stated in Theorem 2. Thus this approach provides an alternative proof for Theorem for the case of [92]
We did not however find a simple approach to analytically compute the generator representation R from the hyperplane representation for arbitrary dimensions. On the other hand we used the numerical implementation of the double
description method to compute B and R. for the case of up-to K 6 drafts and [S2] < 14 and demonstrate that the natural counterpart of our result Theorem 2 appears to valid all these cases.
D.1 INTUITION FOR THE CASE WHEN K AND [ 3
We consider the case of Ω 21, and K = 3 drafts. We let q= (11: 42 (2) and p (P1, P2 P2.Pa) denote the probabilities of the target and draft models. We build upon the ideas in Example 2 in Section C and provide some intuition that the necessary and sufficient condition for the acceptance probability to equal 1 is:
9i > P 9 2 +
Vi.j E (1,2,3),
We let (1) Pr(Y- = 1, X2 = (,), i.e., the probability that the selected token equals 1 and the input tokens are {1,1,2] in any order. Likewise we let 01,1,2(2) Pr(Y
2, (X1, X2, X3) = 1, 2}). Note that (1)
3pfp2 is the probability of observing the input tokens to equal 1,, 2} in some order. We define ox,s,g (1) and (5, for other values of
i in similar fashion. Finally we let (1) Pr(Y = i, {X1, X2, {1, 3)) for each
i E 3 be the probability that token i selected and the input tokens are {1, 2, 3] in some
order. We observe that Σ, (1) 6p1P2p3. We note that the probability that the selected token
Y = 1 is given by:
(1(1) = p¹ + (1) + (1, + (111((1)
+ 0133,(11 + a122,(11
(111)
Here the right hand side denotes all the events that a token appears among the put tokens and that token 1 is selected. Following the same relaxation in Example 2 show that the acceptance probability equals 1 if
a 2 p1( = př + 001,1,211 + +1.22( + (1,,,3(1) + 0133,(1) + (1.
(112)
2/ (2) = pI + a1,11 + 01.2222) + (2) + 42,3,3 + (1)
(113)
gs (3) = p + (1)1s(33 + (1) + (3) + 023,3(3) + (1.2.33)
(114)
Following Example 2 we also consider the following augmented system of inequalities:
+ 2 (pi + p2)x +.,, a1,, + (1)
+ (2) + 22) + 21,,,,3
(115)
41 + 8 Iv (p1 p2) + +11,1,(11) + 0,,,,, 01,2,(1) (3) + 22 (3) (3)
(116)
42 + 93 > (p2 p3)$ 3 (1) + 21,,22
01,1,333) a1,,3(2) +
01,,33 21,,233) (117)
We note that (115) is direct consequence of adding (112) and (113) and using the condition that (f) + (5) = 3p The inequalities (116) and (117) follow in a similar manner. Although these inequalities are redundant their presence aids in simplifying the Fourier-Motzkin elimination as in the case of K = 2 drafts. We do not considering summing all the inequalities in (112)-(114) as
this implies ạ 42 93 (p1 + p2 + p,)3, which holds trivially both sides are equal to 1. In
25

---

Published as a conference paper at ICLR 2025
order to find conditions when the probability equals 1, need apply Motiee- elimination to the system of inequalities in (112)-(117). where the variables satisfy the conditions that
(1i + aas 3p
(118)
and
0 +123(1)+01,23(2)) - 0122, = 6p122P3
(119)
and all the variables non-negative. By following elimination procedure we can show that, similar to Example 2 and Lemma 2, the elimination of each variable in this augmented system simply amounts to dropping from the right hand side in the system of inequalities (112)-(117). Thus the
intuition behind our conjecture is that when one considers an augmented system of inequalities, only
terms of the form (E Pa) K appear and lower order terms do not appear. After application of the Fourier-Motzkin elimination the required condition follows.
E PROOF OF THEOREM 3
Pn] be the distribution of the draft model and q = [gl,-..: 2n] be the distribution of the target model. Our optimization problem can be expressed as follows:
As in the proof of Theorem 2, we let =
pi.....PPp
maximize
to
(120)
t S min
9t, P. + Σ
Easl
(121)
i
os + 0j,i 2p.py, 0 S at 1.
(122)
In order to solve this linear program analytically we introduce an additional variable = satisfying single inequality S t + t..... We provide the range of feasible feasible values of 2 and pick the maximum. Following the techniques used in the proof of Theorem 2 we have the following Lemma:
Lemma 3 Upon applying (122), we have the following system of inequalities with Ω nI....,m}:
elimination technique to eliminate variables 04, in (120)-
t S 9ú, E Ω
(123)
2
Σ 44
Σ
Pi
+2
Σ Pi YES
Σ Pi) liES
AA 22
(124)
iES
VES
2
ti-
(125)
We will defer the proof of this lemma after the main proof. We will use (123)-(125) to establish the following step by induction.
Lemma 4 Suppose that we apply Fourier Motzkin elimination to eliminate variables t .i....,ty 1--1 in (123)-(125). Let Q = {l.....j 1} and 922 = {j....,, ) be partition of Q. Then we have
2
2 5 Σ 9:
Σn +++
Σ
Pt
+2
Σ Pi
Σ P1)
iES
iev
(EESU
ViESLUVe
AS
C
V S2
22 \, yr
S2z IV
(126)
Σ t <
Σ Pi
+2
Σp,
p:
ΣP:)
Σ Pi
VS C l2
Se 02 \
(127)
IES
LIES
NEE
LiES
1 Ji
Vi E S2
(128)
26

---

Published as a conference paper at ICLR 2025
Note that this results implies the main result by setting Q2 S2 and f22 {-}
have:
2 -
Σ ai +
Σ
Pi
Σ iES
Pi
Σ
Pi
(129)
iESt
FESe
We first consider the base case: 1. In this case Ω
{ is the empty set and f2 Ω. Thus
S = - = }.} and V _ f2 and V = f2 V. In this case (126) reduces to:
Σ Pc NEVe
+2
Σ Pi NEV
Σ Pi
(130)
KVV
and (127) and (128) have Ω2 = Q2 Note that (127) and (128) are equivalent to (123) and (124). It
thus suffices to show the equivalence between (130) and (125). To show that the condition (130)
implies (125) it suffices to set V = S and V = {-) To show that (125) implies (130), for any
C Q, we can express:
+
Σ 4 iEVe
(131)
1EV
2
Σ t
+
Σ
Pt
+
Σ FeV
Ps
Σ
Pt
(132)
iEV
LEVe
LiVee
where we use (124) the second term. This completes the proof of the base case.
climination on eliminate ..,.., - leads to (126)-(128) with Q = (1..... . - 1} and f22 = [j..... n). We want to show that upon applying Fourier-Motzkin elimination to eliminate łj. we reduce the system of inequalities again (126)-(128) with D2 (I......) j) and 522 {j + .n].
For induction we assume that for some j > the application of Fourier-Motzee
Let consider those V C P2 [j....n) n where & V in (126). Each such V S Ω22 = 1j+ ] as j E V. Since the variable l does not appear in the right hand side (126), the Fourier-Motzkin elimination will not modify the inequality. We can reinterpret (126) as:
Σ 1E8'
Σ 4+
Σ Pi +2
Σ Pi
Σ Ps
'EV'
VES'UV'
AS'
R & S', V' C 22 Stc S2 S v
= S22 L
(133)
Next consider the case (126) where j E V. In order to apply Fourier- elimination, we express V- = [j) U where V C S2 {j + n.. We explicitly consider the variable I in (126) below.
2
x S
Σ %i
+
Σ "EV'
44 Hatty
Σ NESEJVE
+2
Σ VESUVV
pi
Σ Pi )²)
(134)
iES
iES"UPe
first combine (134) with the inequality
l S qg and introduce S2
On U [/). S2 = 92 \ (ih.
S' = S (j) and S
22 S V
VV
(5) C S2 and D 02
1" to have:
Σ 9 iES'
Σn+ h
Σ iES
P
+2
Σ Pi
Σ iiES'SUe
Pi
LEV"
VEE''U"'
VS'
C V
= S D >'== 2 '
(135)
Note that (133) and (135) recover all the upper bounds on in the induction step for (126). We further need to show that the Fourier-Motzkin elimination does not introduce any further inequalities during the elimination of t.. In particular with V consider combining:
2 S
Σ di iE8
Σ
+
Σ ViESEUVe
Pi
+2
Σ Pi ViSSU
Σ LiES
Pa):
(136)
iEV
27

---

Published as a conference paper at ICLR 2025
with the inequality:
Σ iEW
4 S
Σ
Pi
+ 2
Σ Di
Σ Pi We
(137)
NEW
NEW
where W _ f22 = Ui.....". n. and j E W. Defining W1 = W
V and = wn V we have:
EG=
Σ
Ps
+
Σ
Pi
+
E
Pi
Σ
Pi
Σ Pi iENi
(138)
iEW
NEW!
VEll
NEWI
EMM
2
Σ
Σ
Σ
Σ
Pr
+
Pi
+
Pi ViEU
NEW!
iell
iEW!
+
Σ
Pi
Σ Pi
+2
Σ Ps
Σ Pi iEDIW
(139)
NiE
iEDIW
2
Σ
Ps
+
Σ
Pt
+
Σ Pi
Σ
Ps +
Σ Pi
NEW!
VEW
iEW!
iEM
iED(
+
Σ
Ps
Σ
P
(140)
iEll
W
2
Σ
Pi
+
Σ
Dt
+
Σ
Σ
+2
Σ pi
Σ
NEW!
VEU
iEW!
im
iERii
NEE
(141)
Next we consider (136):
2
2 S Σ
9: +
Σ
4
+
Σ
t1 +
Σ
ps
iES
iEVi
+2
Σ Ps
Σ Pi
+
2
Σ
Pi
Σ
Pt
(142)
JESUM
ViES-UVe
ViEW
NES-UVE
where we use the fact that V = V1 UU. Adding (142) to (141) and eliminating t, where i E U, we get:
+
Σ
t
Σ
+
Σ t, +
Σ LE
Ps
iEW1
EE
iEVI
+ 2
Σ Pi
Σ
Pi
+2
Σ Pi
Σ Ps
VESUV
VEES-UV
iiEl
2
+
Σ Pi
Σn
+2
Σ
Pi
Σ
+2
Σ
Pi
Σ
(143)
NEWi
NEW
NEWN
iEnii
iENR
2
Σ
9:
Σ
t +
Σ
P:
Σ
Pi
iES
iE
LEES-UEE
NEU
+
Σ Ps
+2
Σ
Σ
NESUU
iES.UN
iEDi
28

---

Published as a conference paper at ICLR 2025
2
+
Σ NEW
Pi
+2
Σ pi
Σ Pi iE
VEWi
(144)
Next note that SUV, 5Ω W. This follows since f2 = Ωi U and S S Q2, V. W W and Vi UW = by definition as V1 = V W. Thus the application Fourier-Motzkin elimination with
V y Uul gives:
+
Σ
Σ
qi
Σ
f
Σ
Pi
+2
Σ Pi
Σ iESELVY
Pi
iEW1
iES
iEVi
iESSUVS
VESSUN
+
Σ
Pi
+2
Σ
Pi
Σ
Pi
(145)
!EW!
NEW!
iENA
However the above inequality is consequence of the following:
2 <
Σ
4
Σ
Pt
+ 2
Σ
Σ Ps iESTUVE
iES
IEVI
iES.UVI
ViESUNi
Σ 4 <
Σ JEWi
Pi
+2
Σ ieWi
Pi
Σ Pi
(146)
!EW!
where Vi S. Ω2 = Ω2 Vi. C Q and S C 24 \ S and Wi C Q2- which are already implied in the induction step. Thus conclude that each combination of the form (136) and (137) redundant and need not be included the next step of the Fourier-Motzkin elimination. This concludes the analysis of the upper bound on z in (126).
It remains to establish the induction for (127) and (128) i.e., upon elimination of &j results in
2
E
Σ
Pi
+
Σ iES
P:
Σ ViES=
Pi
VS S S2 S == 22 \
(147)
IES
LiES
t4 S qi
Vi
Ω
(148)
where 02 1......... n). Naturally every inequality (147) and (148) is already contained in (127) and (128) where j $ S. So we only need to show that the application of Fourier-Motzkin elimination remove any other inequality does not result in any additional inequality. Note that the elimination of ty simply involves combining each inequality with > 0. Thus any inequality in (127) where C l2 with j reduces to:
Σ
t
Σ iES
+ 2
Σ LiES
Ps
Σ VE.S
(149)
SESY()
We show that (149) is weaker than
Σ 4
Σ
Di
+2
Σ pi
Σ Pi
(150)
NES
(ESTUU
NESHU)
(ES-UU)
which is already contained in (127) and hence redundant. In particular consider the right hand side of (149):
Σ (ESHU
Pi
P
+2
Σ pt
P
Σ 1
(151)
iESIU)
29

---

Published as a conference paper at ICLR 2025
2 PJ +
2ps
Σ
Pi
+
Σ
Pi
+2
Σ
pi
Σ P
(152)
iES(
iESYis
iESifs}
ViES=
2
IV
Σ
Pi
+
Σ Pi (ES(i)
Σ iESTUi>
(153)
(ESMI)
which implies that (149) is indeed weaker.
Thus we have completed the induction step. Continuing the induction to eliminate all variables t,,...,t results in
2
Σ
qi
Σ VEES
Pi
+2
Σ iES
Pi
Σ VE
Pi
AA
E
(154)
claimed. It now only remains establish the proof of Lemma 3 which will do. As considering the elimination of 04,j it suffices to consider the following inequalities:
t S p. +
i 1,2,...,
(155)
j=1
0. + + = PpbPj
0 S 03 1
(156)
We will show the following by induction. Suppose that at step > = 1 let
V. (()1,)1), (T,, ii), (f--1, (fr- 1-1
(157)
denotes the indices (with is I jiz) of variables that are eliminated using the Fourier-Motzkin elimination. Then the resulting system of inequalities is given by:
2
Σ < #ES
Σ Ps
+
ΣΣ ast I(((s V.) + ΣΣ
2p.Pt I((s,() E V,) (158)
.ES
BES LES*
#ES LES*
for all S C Ω = (1, 21, n). n.). For the base case, consider the case whenr = 1 i.e., V, = (.). The condition (158) reduces to:
2
Σ SS
Σ
Ps
ΣΣ agsi
AS S
(159)
L
#ESIESE
We show that (159) equivalent to (155). Indeed setting S = {i} in (158) and using ag,, = 2p.Pt
recovers (155) for each i =1,2.... We will show that the conditions (155) and (156) also imply (158). Note that for any C s:
Σ
1/4
S
Σ
+
Osst
(160)
SES
SES
SES i=
Σ
Σ
Olaii +
Σ
asd
(161)
IES0
(ES)(s)
Oggi + Σ
Cz,5
(162)
SES
&ES iES
#ES iESH(a)
P
Σ
O'sii +
Σ
(ass, + (s,.)
(163)
#ES
8CS LES
(s.))SxSsi1s
Σ
P4
2
Σ
Ossi +
Σ
2ps.i
(164)
#ES
#ES $ES
(s,))S8xsibb
2
Σ
Ps
ΣΣ
Og
(165)
SES
#ES FES
30

---

Published as a conference paper at ICLR 2025
We thus recover (159) from (155). This establishes the base case.
For the induction step, let us assume that we have eliminated all os., where the indices (i.j) are the set V, and that (155) is satisfied. We consider climination of indices (ir, Jr) and (jrs &r) associated with oi and jj.,rr'
E S AES
Σ LEES
ΣΣ asss i(( v.)++ SESiES"
ΣΣ 2psp; I(((s,) V.) (166)
Ps
+
SES
We need to show that upon elimination of an and az,t, using Fourier-Motzkin climination the resulting system of inequalities given by:
2
E t = #ES
Σ LES
+
ΣΣΣ
Σ
Ces I(((s,) & V++1) Σ Σ NES VESE
2p.p: I((s,t) E Ve+1), (167)
RES VES
with
V++1 = (11,,11). (1, (Ir, (Jr, 1,)
(168)
We note that in the Fourier-Motzkin elimination step have only consider those inequalities where either oi,.j. or Og.,,, appears on the right hand side of (166). This is equivalent to having ir S and jr S or J. E S and ir E S For those S that do not satisfy either condition, we immediately have (167). If the selected S follow either of these cases, combining (166) with atrúr S 2p,,Pj, and aj.., S Pp., we reduce (167). At this point all the equations in (167) have been recovered. Nevertheless Fourier-Motzkin elimination requires us to also consider all pairwise equations where S2 S. Ω i E S1 and jr & S1 and i, & S2 and jr E S2:
Σ to xES1
Σ SESI
Ps
ag Ve,, v SESI SES
Σ
Σ
2p.p: I(((, E V,) (169)
sES; LES?
Σ to SES
Σ ag,,)((s,)))) i) $ V+)+ Σ Σ 2p.p I((s, V,) (170) iESS
Ps
S
SES₁
1ESS
The Courier-Motzker elimination step requires us combine (169) and (170) and use dt Pirr Oirije Ojrir 2 0 to eliminate 0irrj- and Ojrrir in the induction step.
ta
Σ to #ES
Σ
+
Σ to
Σ Ps
WES
8ES2
HE
ag ((s, & V,)
2p.Ps I(((s)) V.)
SESi XESS
SESi iES;
+
Σ
Og De,TT((s ((s, $ V+) + Σ Σ
2p.P: I((s,t) E V,)
(171)
SES: YESS
#E.S XESS
We will show that each such inequality is redundant and already implied by the set of equations already established in (167).
Let R S1 S& and T = S S2. Note that i, & R and jr & R. First following the same steps
leading to (78) we can show that:
2
2
Σ to
Σ
Ps
Σ
to
Σ
Ps
SES
SES
sES,
S
Σ
Ps
+
Σ Ps
+2
Σ Pe
Σ Po
SET
BET
sER
SE
MES,
(172)
31

---

Published as a conference paper at ICLR 2025
Next, following the steps leading to (84) we have that:
Σ
Σ
asi
I((s,)) $ V.) + 2p-p: I[((s, i) E V.)
BES, iESt
Σ asi
I((s,)) $ V,) + 2psps - 1((ss)) E V2)
SES TET
Σ
E
03i ((s, $ V.)+ + 2ppp I(((s,)) i) V.)
(173)
SEST iES
E
Σ
O.. -..,))((s))))))))) $ V,) - 2p.p. ((ss,) E V,)
SE8, RRIETT
ΣΣ asd 1((s,)) $ V.) + 2p.p. [(((s)) i) V2)
SERiET
Σ
E
0ss I((, $ V,) + 2p.p, [I((s,)) i) E V,)
sESi R /R
+ Σ
Σ
O,1 --.(s((()))))))))))(( & v,) + 2p.p. [((ss i) E V,)
(174)
SRR RR
and, likewise,
E
Σ
ass I((s, & V.) + 2p.p: - I[((s,) E V,)
*ES, JES;
Σ aas (s) V.) + 2p.p: [l((s)) i) E v,)
#ES, (R T
Σ
I((ss)) i) $ v.) + 2p.pr I((s.)) E V,)
R IET
Σ
Σ
Cgs I((ss i) & V.)) + 2p.p: I((s, i) E V,)
SEAAR iE VR
Σ
Σ
((s & V.) 2p.p: I((s, E V+)
(175)
RER iR
Next we combine the terms to get:
ΣΣ a.s I(((,)) $ V2) + 2p.ps I[((s))) E v.)
SES, FEST
Σ sE ieS
Σ
Ces 1((s,)) & V,) + 2p.p. I((ss v,)
Σ
Σ
as,i {((s,))(V,)))) + 2p.p. I((s, i) e V,)
SES, \R 1ET
ΣΣ
Σ
a.si
I(ss V.) + 2p.p. I((ss E V.)
SER
Σ
Σ
O..j I((, $ V.) + 2p.p: I((s, ([(ss)){()))) E }.1
REET
Σ
Σ
agK I((s,)) $ V,) 2psp: I((s,)) E V,)
SER YR
Σ
Σ
a.p [(((s $ V.) + 2psp. I[(ss i) E V+)
iET
Σ
Σ
O.. T[(s, $ V.) + 2p.P: T((s, 1) E v.)
(176)
\R
32

---

Published as a conference paper at ICLR 2025
E
Σ asi [((, $ V.) 2pspi [[((s))) v.)
#ES R
Σ
Σ
Qsii 1((8, $ v.)) + 2p.p. II((s, e V,)
(177)
#ES RESI R
asd 1((s, V,) + 2p.p. [(((,)) V,)
SET
E sER
OBs I(((s)) $ V.) 2p.p: I[((s i) E V,)
Σ
(as,3 a1,)) I((s,t) i) $ V,) + 4psp, I(((st) V.)
(178)
EE R RE
asi I((s,t) $ V,) + 2psp, I(((s,) E V,)
NET T
Oui I[((s & V.) 2p.0: [(((s,)) i) E V,)
#RR E
Σ
2p-p; + 2p-ps I(((s))) E V,)
(179)
VE
iES
(180)
Thus the resulting inequality from Fourier-Motzkin elimination is given by:
Σ Ps
Σ t,
Σ
Ps
+
Σ
Ps
Σ Po (ESiiR
RET
WET
SER
HER
BES21R
asii i((s, $ V.) + 2pspi 1((8,1) V2)
SET iET
Ous (((, $ V.) + 2p,p, II((s; i) E V+)
sER I
+
Σ
Σ
2p.p. + 2p.p: 1((ss i) E V,)
(181)
R
Next note that since (rr E T and (ir.j)) & R, the inequalities:
2
Σ
1.
<
Σ BET
Ps
a.s 1((s,)) $ V.) + 2p.p, I[((s,)) E V,)
SET
sER iE
2
Σ t
Ps .ER.
Σ
Σ
Σ
a I[((s i) $ V.) + 2p.p: I((s, i) E V.)
(182)
*ER
SER iER=
are already constructed in the induction step. Also clearly (181) is implied by these since:
(
2
Σp, P+
Σ P.
Σ
Σ
2psPs
(183)
LESSIR
PESLiR
is an identity since S₁ \ cR and S₂ Motzkin climination is redundant and we have completed the induction step and in turn established Lemma 3.
R are disjoint. Thus each such inequality form the Fourier-
F PROOF OF EQUATION (12)
First we consider the non-truncated program and let wiy be the variables for i, E with < maximize the objective:
j that
Σ min pı (1))
(184)
33

---

Published as a conference paper at ICLR 2025
where
pr(1) = pi +
Σ
Wis
(185)
j=1
Note that we have
PM(ace) =
Σ min (ainggi pr (1) v
min +
Σ
(186)
1=++1
For the truncated linear program we have for each E {1, 83
pi(1) = p +
Σ
Wp:
Σ
2p:Pi
(187)
s=++1
and for i >> 8:
pr(1) pi +
2p-Pj
(188)
j=
We consider potentially sub-optimal choice of weights Wa wai i,j E {1.....s) 8} for the truncated linear program. Note that
pr(1) 2 (1),
Viss
(189)
and
(1( 2 p
viss. s.
(190)
As a result, using (186) we have:
P(acc)
min(q:ppr(
(191)
min
ge (r()) +
min(k: pp
(192)
n
2 P' (acc) Σ
9: +
min (q.
(193)
i=0+1
ims+1
1
Paaaa))
(g: p
(194)
G COMPUTATIONAL COMPLEXITY OF TRUNCATED LP
We study the computational complexity of the truncated linear program (see Algorithm 2) and propose a variation that could improve it further. Let Q = 21,2....,,7 denote the vocabulary size under consideration. Note that the truncated linear program has two steps:
Sort the candidate tokens based on 4 D.a This requires Onn logn) computations. We select that largest tokens and identify this set as Ωn
Apply linear program for the tokens in Ω which is the size of s. There are Oss²) variables and standard implementation of the linear program involves complexity of (s)))
Thus the overall complexity of our proposed scheme is O(s + n log n) where is the alphabet size. In practice we keep the size of 22 relatively small e.g., we set s = 5 in many of our experiments. This avoids slowdown arising from the 5 term in the linear program. Secondly, in practice, Ω. is a subset tokens from the original vocabulary after top-p sampling. So even when the original vocabulary is large, the value of in practice is relatively modest. In Appendix H we present the histogram of the alphabet size after top-p sampling to illustrate this point. Thus the Onn n) is an acceptable cost for implementation.
34

---

Published as a conference paper at ICLR 2025
Algorithm 2 Truncated LP
Input: Threshold 8, Input tokens X1, X2 sampled independently from p(.)
2: Output: Selected token Y,, output distribution pr(-).
3: Order vocabulary = {1, 2, ......... r1,, sorted in decreasing order with (g: - -p.).
4: Set f2 {1, {1....,5 5} and f2g {s ,.....,?. n).
5: For i, E f22 or E Q and j E f2 set wi, in (10)
6: For i.j E $2,, compute wij
solution to linear program:
Maximize: Σi=1 min(g: Pr (i)), where P1 (i) is defined in (11)
Constraints: Wi.j 2 0, wi, + Wy, =1
7: Compute pr(-) according to (11).
8:
if X = X2 then
9:
Set Y1 X1.
10: else
Let {X1, X2) [i.j} and <
12: Set Y1 = with probability Wi,j. and Y, = j otherwise.
end if
G.1 VARIANT OF TRUNCATED LP
We now present variation of the LP that avoids the O(n l) cost from sorting Ω. The proposed method is as follows:
Select s tokens with largest values of 9: - p This requires O(n computations. We identify this set as Q and the put the remaining tokens in Ω2.
We propose the following heuristic choice of weights:
[1, E R1.j E 92
Wi.j
(195)
11/2,
i S2.j 92 }
Note that this results in the following distribution of the selected token:
1agk 2p.PaWp. + + 2p.Pa, I E Q2
(k)
(196)
P
Lith PiPk.
f2
We leave the weights way for i < and i,j E Ω as free parameters.
Apply linear program for the tokens in Ω which is the size of s. There (ss²) variables and standard implementation the linear program involves complexity (s)))
can be seen that the overall complexity of the proposed method is O(n s + s By following the same steps as in Section F can also be verified that the resulting probabilities in (196) are also satisfy (12) yielding the same theoretical guarantee.
H HISTOGRAM OF ALPHABET SIZE FOR OPT MODEL
We present additional evidence that after top-p sampling the effective alphabet size can be significantly reduced. We consider the OPT model and report the histogram of the alphabet size under the draft and target models for the XSum task in Fig. 5, the Dolly task in Fig. 6, and the WMT task Fig. The histogram is truncated to 40 tokens make the figures clearer.
35

---

Published as a conference paper at ICLR 2025
00
C0
00
001
000
Figure 5: Truncated Histogram for OPT Draft and Target Models for the effective alphabet size after top-p sampling with p=0.95 on XSum dataset
dutegrme Dely Daleor
00
40
000
Figure 6: Truncated Histogram for OPT Draft and Target Models for the effective alphabet size after top-p sampling with p=0.95 on Dolly dataset
36

---

Published as a conference paper at ICLR 2025
Hatggrec
Hanggarc
01
Figure 7: Truncated Histogram for OPT Draft and Target Models for the effective alphabet size after top-p sampling with p=0.95 on WMT dataset
37

---

Published as a conference paper at ICLR 2025
I LP AND FAST LP VERSION FOR NON-IDENTICAL DRAFT DISTRIBUTIONS, K=2
For the of K 2 drafts explain how the importance weighted sampling scheme and its faster variants can be extended when the two tokens are sampled independently but from different distribution i.c. X1 ~ pi(-) and X2 p2))). We P1 (pi,1..... Pi.nn) and P2 (P2,1 P2.a) denote the distributions of the draft models to sample X₁ and X2. We let q = (f),.... 9n) denote the target distribution.
The order of the tokens matters and accordingly for i < J, we define:
Wijj = Pr(Y = X1 i, i, x2= Waij
wi,j = Pr(Y = =|| X, = Wj.i Pr(Y i| = j. = بنز = - /)et Pr(Y j|X1 X, = =)
(197)
(198)
If Y denotes the selected token, then considering all cases where token appears as one of the input tokens, we have:
(i) = P2, P1. Way + E
W1,
Σ j=l
PL.Pe.
E jmit1
Pi,
jmit1
(199)
We need to find way and W that maximizes Σx=1 min (g: pi)).. This is a linear program in variables Wiy satisfying S Way S The truncated version L.P obtained by sorting the tokens in Ω based on qi P1, again considering sets Ω = (1,2, ....,..... 8) and f2 = {ss + 1,...,n, We treat Wa,j as variables that need to be optimized i.j E S2 If E f2 and j E S22 we set wa = 1. If both i E f2 we set Wii i < and 0 i > j. The resulting distribution is given follows. For E (8.
n).
pr(1) P1,122 + Σ
PI.iP2 Di +
Σ
u uss +
Σ ==1
Pi,j
j=l
j=
P1.P2,WW3s +
Σ
(P1.72, + p1,))2))
(200)
jms41
and for i s I,... n. have:
pr(() P1,i22, +
(Pi.P224 + P1.P2)))
(201)
j=i+1
Upon following the sequence of steps leading to (194) we can show that
Place) 2 (*ccc)- - Σ (gs -P..P2.)
(202)
The truncated alphabet scheme can be applied in similar fashion by considering high probability subset On S Ω and only keeping those input tokens that belong to Ωg. We generate truncated distributions Pi (-) and p2(-) and apply the linear program on these followed by speculative sampling using the target distribution q-))
38

---

Published as a conference paper at ICLR 2025
J ADDITIONAL EXPERIMENTAL RESULTS
J.1 TWO DRAFT MODELS WITH IDENTICAL TEMPERATURES
Here, we consider the case where identical draft models are used generate candidate tokens sequences, We compare the performance of our method with SpecTr (Sun et al., 2024b) and SpecInfer (Miao et al., 2024), well as single-draft speculative sampling (Leviathan al., 2023; Chen et al., 2023), where we use the same temperature for the draft model.
set the temperature of the target model to 1.0, while we vary the temperature of the K = draft models between the range of 1.2 to 2.4, and we report the performance achieved by the different schemes across the three tasks discussed earlier. The top plots report the block efficiency, which is the average number of tokens that are accepted per use of the draft model (Leviathan et al., 2023). The bottom plots show the percentage improvement in the token rate with respect to the baseline single draft scheme.
In Figure
IS
SpeeTr
SpecIafer
Sigule-drafi SD
Dolly
XSum
WMT
2.0
Block Effrciency
1.5
1.0
0.0
Token Hate Improvemen Over Single le draft SD
4055
309g
2055
106%
1.8
2.4
1.2
1.4
2.0
2.4
12
1.8
2.0
2.4
Drafts Temperature
Drafts Temperature
Drafts Temperature
Figure 8: Performance comparison of different multi-drafe schemes, while we vary the temperature of the two draft models.
We observe that the IS scheme consistently outperforms the both SpecTTr and SpecInfer across all three tasks. In fact when the temperature of the second draft is increased, the improvement in the block efficiency of both SpecTrr and SpecInfer is rather negligible compared to the single draft baseline, while the token rate are in fact lower than the single draft baseline. On the other hand our proposed IS scheme is able to achieve consistent improvements in all the three tasks. In the experiments involving the XSum and WMT tasks we also measure the ROUGE-L (Lin, 2004) and BLEU (Papineni et al., 2002) scores respectively, which are reported in the Appendix J.3.
Table 4 compares the block efficiencies for different multi-draft speculative sampling methods using K to K drafts when all the drafts are identical and use sampling temperature of 1.2. We again see a consistent improvement by the proposed importance sampling scheme.
Table 4: Block efficiency achieved the Dolly task for different number of draft models.
Scheme
K
K 3
K 4 K
K
IS
2.13 + 0.05
2.22 0.05
2.26
0.05
2.27 + 0.05
2.28 + 0.06
SpecInfer
1.76 + 0.04
1.86 + 0.05
1.95 + 0.05
2.00 0.04
2.04
+ 0.05
SpecTr
1.77 + 0.04
1.89 0.05
1.96 0.05
2.03 + 0.06
2.08 + 0.04
39

---

Published as a conference paper at ICLR 2025
J.2 THREE DRAFT MODELS
In this section, we consider the case of K = 3 draft models that use different temperatures for token generation. As stated before, in case of non-identical draft models, the only plausible multi-draft sampling scheme for comparison is the SpecInfer scheme (Miao et al., 2024). We set the temperature of the target model to 1.0, and the temperature of two of draft models to 1.2 and 1.4, while we vary the temperature of the first draft model between the range of 1.0 to 1.6. As observed in Figure 9, the proposed IS scheme outperforms SpecInfer and Single-draft speculative sampling method in terms of block efficiency and achieved token rate.
IS
SpecInfer
Single Draft SD
2.5
15%
Token Rate Improtement Over Single-draft SD
Block Efficiency
10%
1.5
0.0
-5%6
1.6
1.6
Temperature of Draft 1
Temperature of Draft 1
Figure 9: Performance comparison of different schemes on the Dolly dataset with K = 3 drafts. We vary the temperature the first draft model and keeping the temperature the other two drafts to 1.2 and 1.4, with the target model having temperature of 1.0. The single-draft baseline uses draft temperature of 1.2.
J.3 ROUGE-L AND BLEU SCORES
In the experiments involving the XSum and WMT tasks measure the ROUGE-L (Lin, 2004) and BLEU (Papineni et al., 2002) scores respectively. Table 5 and Table 6 show the ROUGE-L and BLEU scores for the case of identical draft models, corresponding the experiment in Figure 8 in Section J.1.
Table 5: ROUGE-L scores the XSum task across various decoders and sampling temperatures.
Draft Temp.
1.2
1.4
1.6
2.0
2.4
Decoder
IS
0.186 + 0.004
0.188
0.002
0.191 - 0.003
0.186 - 0.004
0.187 + 0.003
Signle-draft SD
0.190 = 0.006
0.185 = 0.005
0.190 = 0.004
0.186 = 0.003
0.186 + 0.004
SpecInfer
0.184 + 0.004
0.190 0.002
0.187 + 0.001
0.186
0.003
0.186 + 0.004
SpecTr
0.188 + 0.002
0.182 + 0.006
0.188 + 0.001
0.185 + 0.006
0.188 + 0.001
Table 6: BLEU scores on the WMT dataset across various decoders and sampling temperatures.
Draft Temp.
1.2
1.4
1.6
2.0
2.4
Decoder
IS
0.037 + 0.002
0.038 + 0.004
0.034 + 0.002
0.039 + 0.003
0.039 + 0.002
Signle-draft SD
0.036 + 0.000
0.037 + 0.003
0.038
0.004
0.037
0.003
0.038 0.002
SpecInfer
0.035 + 0.003
0.039
+ 0.004
0.035
= 0.003
0.034 + 0.009
0,036 + 0.003
SpecTr
0.039 0.001
0.037 0.001
0.039 0.001
0.036 0.002
0.035 0.001
Similarly, Table 7 and Table 8 show the ROUGE-L and BLEU scores for the case of non-identical draft models, corresponding the experiment in Figure 3 in Section 5.1.
40

---

Published as a conference paper at ICLR 2025
Table 7: ROUGE-L scores the XSum task across various decoders and sampling temperatures.
Temperature
Draft
1.2
Draft
1.2
1.6
2.0
2.4
N/A
Decoder
IS
0.187 + 0.004
0.189 + 0.007
0.189 + 0.001
0.191 + 0.002
SpecInfer
0.184 + 0.004
0.190 + 0.003
0.185 = 0.006
0.189 + 0.006
Single-draft SD
0.190 + 0.006
Table 8: BLEU scores on the WMT dataset across various decoders and sampling temperatures.
Temperature
Draft
1.2
Draft
1.2
1.6
2.0
2.4
N/A
Decoder
0.036 + 0.003
0.035
0.002
0.036
0.002
0.035
0.002
SpecInfer
0.035 + 0.003
0.038 + 0.005
0.041 = 0.002
0.040 = 0.002
Single-draft SD
0.036 + 0.000
K NOTATIONS
Table 9: Table notations summarizing symbols and their descriptions.
Symbol
Description
Fia or x
Sequence tokens T1, T2, "1 3
K
Number of draft models
Vocabulary of tokens
Size of the vocabulary, [
The set valid draft tokens under consideration at every step
p
Distribution of the draft model any step. with pt denoting the probability of token
Distribution of the target model at any step, with denoting the probability of token Draft model
M.
M
Target model
nt
Context sequence as a sequence of t tokens ?
Pd (acc)
Optimal acceptance probability
By ( (z FK)
Conditional probability that token y is selected given the input tokens Ti, JK Optimal Conditional probability that token is selected given the input tokens 21, ....... UK
8g [
Y
The selected token from the importance weighted sampling step The final output token after the combined importance weighted sampling and speculative sampling step
7
pr(-)
Distribution of the output token in the importance weighted sampling step Index of the selected token n the importance weighted sampling step Conditional probability that token Joint probability that token i is selected and the input tokens are and j Joint probability that token i selected and the input tokens are i, j. and
I
Ws.
selected given the input tokens are i and j
aij
aij.k
LP-Tnnncation Threshold parameter
Qo
High probability subset tokens of the vocabulary Ω
41


Reviews:
[{"review_id": "73dSt334YZ", "reviewer": "Reviewer_aVvF", "summary": "Previous papers showed the optimality for a token selection rule when one token was chosen from the draft model. The authors do the same for the case when two tokens are selected. The authors decompose the problem of token selection into two steps --- first select one of K tokens using an importance weighted sampling technique, and then using the optimal single token selection proposed in previous papers. Finding the optimal token selection rule then reduces to finding the best importance weighted sampling technique.\n\nThe authors then show a necessary and sufficient condition for the existence of an optimal token selection rule where two tokens are selected from the draft. They then present an algorithm to solve the resulting linear program faster than the O(n^2) standard solution.\n\nThe authors compare the Importance sampling algorithm with SpecInfer on various tasks, and show that IS comes closer to optimal than SpecInfer.", "questions": "I would like to see a comparison to the results in https://openreview.net/forum?id=9KxnxWOBA5", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The results in the paper are quite original and significant, to my knowledge, which is a bit limited. The experiments show that the proposed algorithm works well in practice.", "weaknesses": "I would have liked the presentation in Section 4 to be a little clearer - request you to write out the full linear program, and then explain the truncation a bit more - what is the truncated program computing? \n\nIn the experiment section, please include results for Temperature < 1, which is much more common in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "18CFrPE7DJ", "reviewer": "Reviewer_6EsD", "summary": "The paper studies the problem of multi-draft speculative sampling and makes new observations and provides a new algorithm. Of the new observations, Theorems 2 and 3 standout which characterize the optimal acceptance rate for two drafts and further show when the acceptance rate is one. The new algorithm exploits certain structure in the optimal transport problem and provides a clever approximate solution.", "questions": "- In Figure 3, experiments are conducted for draft temperatures > 1.0. However typically, target model temperature is often set to 1.0 or less than 1.0. Given this, I was expecting draft temperatures also to be <= 1.0. I was wondering how these draft temperatures were chosen?\n\n- Would it be possible to say how the gap between the proposed approach and previous approaches vary as the target model temperature varies from 0.0 to 1.0?\n\n- Theorems 2  and 3 provide a nice characterization of acceptance probability. However, given two distributions p, q, evaluating the condition seems to require 2^{vocabulary size} number of computations. Is there a faster way to evaluate these quantities?\n\n- There are some approximate ways to solve the optimal transport by Sinkhorn methods e.g., https://amsword.medium.com/a-simple-introduction-on-sinkhorn-distances-d01a4ef4f085. I was wondering if these ideas are applicable here?\n\n- Architectures in ML, typically refer to neural architectures. Given several works which propose different drafting methods for speculative decoding, I was thinking \"canonical architectures\" refers to some new drafters. It might be good to clarify this distinction early on in the paper.\n\n- Can you please expand on what token rate means? Is it the number of tokens / second?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The paper is well written and easy to follow. Theorems 2 and 3 are novel and the experimental results are convincing.", "weaknesses": "See questions section below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MEOML05pZ8", "reviewer": "Reviewer_V7r6", "summary": "This paper studies the optimal token-level draft selection for a speculative sampling. In particular, speculative sampling is a method for accelerating autoregressive sampling by using smaller \"draft\" models to generate candidate tokens and a larger model to determine which tokens to accept. The key problem studied in this paper is how to propose tokens to optimize the probability they will be accepted. Importantly, increasing the acceptance probability leads to less calls to the larger model which in theory could increase the tokens per seconds. The key theorem of this paper is to characterize the necessary and sufficient conditions when two tokens are generated by the draft model. Interestingly, unlike the one token case, this paper shows that one can reach probability 1 acceptance even if the draft and large models have different token distributions. Finally, the paper ends with (i) heuristics to practically leverage the proposed sampling scheme, e.g., truncating the token distributions (ii) a heuristic scheme for leveraging the proposed importance sampling to larger sequences and (iii) empirical evidence of the effectiveness of the approach for K=2 and larger.", "questions": "- Could you give some intuition about the conjecture in remark 2. It says that the exponent over the sum of the p-weights should match K (the number of tokens to generate per round). It's not at all clear to me why these should be so tightly coupled and why other lower order terms wouldn't appear.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "This paper provides a strong theoretical contribution to a very important topic. Speculative decoding serves as a refreshing and novel direction for accelerating LLM inference. This work puts the multi-draft model selection problem on better theoretical footing.\n\nThe results seem to be sound from my reading and no obvious errors were found.", "weaknesses": "By and large the biggest complaint I have with this draft is the background exposition. Prior to reading this I was not familiar with speculative decoding. While I understand space is tight, many key concepts such as the role and form of accepting probability in the speculative decoder were not clearly explained. This is a *key* aspect of the work, and the draft would strongly benefit for giving it more treatment.\n\nTo be honest, I couldn't understand anything the first time I read the paper and had to go back to the original speculative sampling paper [1] before I could piece things together.\n\nNevertheless, after reading, I now understand that this paper is targeting the acceptance criteria as it determines the number of candidate tokens used via a Bernoulli like-process. \n\nThis plus the strong theoretical analysis are enough for me to lean towards acceptance.\n\n[1]Accelerating Large Language Model Decoding with Speculative Sampling,  https://arxiv.org/abs/2302.01318", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "riye1MHlJk", "reviewer": "Reviewer_yfZp", "summary": "This paper studies the problem of multi-draft speculative sampling, where the optimal scheme can be cast as a solution to a linear program. The key idea behind this work is that this optimal scheme can be decomposed into a two-stage solution. Based on this idea, the authors provide some theoretical findings towards the two-stage solution. Motivated by the analysis, the authors propose a new class of token-level selection scheme based on weighted importance sampling, which shows consistent improvement in various scenarios.", "questions": "1. Can proposed truncated LP combined with SOTA speculative decoding (e.g. Medusa [4] and Eagle [5]) methods for better performance?\n2. Why the authors do not discuss and compare a multi-draft selection method RRS [6]?\n3. Could you please discuss the limitations of this paper?\n\n[4] Cai, Tianle, et al. \"Medusa: Simple llm inference acceleration framework with multiple decoding heads.\" *arXiv preprint arXiv:2401.10774* (2024).\n\n[5] Li, Yuhui, et al. \"Eagle: Speculative sampling requires rethinking feature uncertainty.\" *arXiv preprint arXiv:2401.15077* (2024).\n\n[6] Jeon, Wonseok, et al. \"Recursive speculative decoding: Accelerating llm inference via sampling without replacement.\" arXiv preprint arXiv:2402.14160 (2024).", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The introduction of the two-step scheme are simple and intuitive, making it easy to understand.\n2. The authors provide a solid theoretical analysis for token-level optimal draft selection.\n3. The truncated LP algorithm is novel to the field of speculative decoding and theoretically better than SpecTr.\n4. Extensive experiments, including different datasets, different decoding temperature and different draft models demonstrate the effective of proposed method.", "weaknesses": "1. I think the contributions of this work is a little weak. Given that speculative decoding shows poor performance with batch size larger than 1, the multi-draft sampling is usually not an appropriate option for speculative decoding. I think the authors should emphasize the contributions to real-world speculative decoding applications. \n2. While I appreciate that the authors add some examples for better explanation, I think the authors should add a clear notation section with distinct description.\n3. For the truncated LP draft selection algorithm, there exists an operation of sorting, which takes a complexity of $O(n\\log n)$ and may influence the final speedup. The authors should report its real-world time consuming and compare to vanilla multi-draft selection, especially the empirical results on models with large vocab size, e.g., Llama 3.1 with vocab size 128k. Besides, for multi-draft scenarios with $K > 2$, the authors propose a multi-stage manner, which may affect the final speedup as well.\n4.  I think the authors should conduct experiments on more benchmarks (e.g. HumanEval [1], GSM8K [2] and MT-bench [3]) for a fair comparison. Besides, the authors should report the real-world speedup to demonstrate the effectiveness of proposed method.\n\nI know that conducting a wide range of evaluation experiments is costly and time consuming, and I will increase my score if the authors clearly address my concerns.\n\n[1] Chen, Mark, et al. \"Evaluating large language models trained on code.\" *arXiv preprint arXiv:2107.03374* (2021).\n\n[2] Cobbe, Karl, et al. \"Training verifiers to solve math word problems, 2021.\" *URL https://arxiv. org/abs/2110.14168* (2021).\n\n[3] Zheng, Lianmin, et al. \"Judging llm-as-a-judge with mt-bench and chatbot arena.\" *Advances in Neural Information Processing Systems* 36 (2023): 46595-46623.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)