BanditSpec: Adaptive Speculative Decoding via Bandit Algorithms
Authors: Yunlong Hou
Fengzhuo Zhang
Cunxiao Du
Xuan Zhang
Jiachun Pan
Tianyu Pang
Chao Du
Vincent Y. F. Tan
Zhuoran Yang
Venue: ICML 2025 poster (2025)
Keywords: LLM Inference Acceleration
Speculative Decoding
Training-free
Hyperparameter Selection
Bandit Algorithms
Primary area: deep_learning->large_language_models
TL;DR: This paper proposes a training-free online learning framework to adaptively choose the configuration of the hyperparameters for speculative decoding as text is being generated.

Abstract: Speculative decoding has emerged as a popular method to accelerate the inference of Large Language Models (LLMs) while retaining their superior text generation performance. Previous methods either adopt a fixed speculative decoding configuration regardless of the prefix tokens, or train draft models in an offline or online manner to align them with the context. This paper proposes a training-free online learning framework to adaptively choose the configuration of the hyperparameters for speculative decoding as text is being generated. We first formulate this hyperparameter selection problem as a Multi-Armed Bandit problem and provide a general speculative decoding framework BanditSpec. Furthermore, two bandit-based hyperparameter selection algorithms, UCBSpec and EXP3Spec, are designed and analyzed in terms of a novel quantity, the stopping time regret. We upper bound this regret under both  stochastic and adversarial reward settings. By deriving an information-theoretic impossibility result, it is shown that the regret performance of UCBSpec is optimal up to universal constants. Finally, extensive empirical experiments with LLaMA3 and Qwen2 demonstrate that our algorithms are effective compared to existing methods, and  the throughput is close to the oracle best hyperparameter in simulated real-life LLM serving scenarios with diverse input prompts.

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
Yunlong Hou
Fengzhuo Zhang 1i
Cunxiao Du 2 Xuan Zhang Jiachun Pan
Tianyu Pang?
Chao 2 Vincent Y. Tan Zhuoran Yang"
Abstract
Step Profe:
Speculative decoding has emerged as a popular method to accelerate the inference of Large Lan- guage Models (LLMs) while retaining their supe- rior text generation performance. Previous meth- ods either adopt fixed speculative decoding con- figuration regardless of the prefix tokens, train draft models in an offline or online manner to align them with the context. This paper proposes a training-free online learning framework to adap- tively choose the configuration of the hyperpa- rameters for speculative decoding as text is being generated. We first formulate this hyperparame- selection problem
Adaptive speculative
Adaptive speculative decoding via a
Prefie:
Draft
secoding
Eagle-2
Verify
idecoding
vi
Eagle-2
BBQ
Suffix
Accept
rate
2/3
Sufti Tree
Which model should choose?
KINM
Proftc:
Adaprive speculative decoding via
PLD
Draft sufto: tree way
BBB
Eagle-2
Verify
tree method
Rast
Suffix Tree
Accept rate
1/3
Multi-Armed Bandit problem and provide a general speculative de- coding framework BANDITSPEC. Furthermore. two bandit-based hyperparameter selection algo- rithms, UCBSPEC and EXP3SPEC, are designed and analyzed in terms of a novel quantity, the stopping time regret. We upper bound this re- gret under both stochastic and adversarial reward settings. By deriving information -hherretccc impossibility result, it is shown that the regret per- formance of UCBSPEC is optimal up to univer- sal constants. Finally, extensive empirical experi- ments with LLaMA3 and Qwen2 demonstrate that our algorithms are effective compared existing methods, and the throughput is close the oracle best hyperparameter in simulated real-life LLM serving scenarios with diverse input prompts.
Figure 1. Given the prefix tokens and the candidate hyperparam- eter configurations (c.g., models), which configuration should be selected to decode the next tokens? We formulate this problem bandit problem and propose general framework BANDITSPEC.
tokens (Brown et al., 2020; Touvron et al., 2023). This autoregressive decoding approach involves multiple forward passes, with each pass generating one token sequentially. Consequently, this process can lead significant latency during inference.
Speculative decoding was introduced by Leviathan et al. (2023); Chen et al. (2023) to accelerate the inference of LLMs. The standard speculative decoding framework has been extended with improved performance since then. thorough overview presented at Appendix A. While the existing speculative decoding methods diverse, most previous works adopt fixed one across tasks, severely lim- iting their potential. For instance, when dealing with code debugging or srummar-checking tasks, the generated tokens are expected to resemble most of the input tokens. There- fore, retrieval-based speculative decoding techniques are preferred (Hu et al., 2024). contrast, for story generation tasks, we expect the generated tokens be more creative. Thus, draft model with high-temmerature parameter is preferred over retrieval- methods. The potential of these speculative decoding methods can only be exploited when the configuration of hyperparameters is well-aligned to the given task. There are existing works that attempt to achieve this goal, e.g., Zhou et al. (2024) distills the draft
1. Introduction
Large Language Model (LLM) is trained to predict the probability of the next token conditioned on all previous
"Equal contribution Work done as an associate member at Sea AI Lab PPoject Lead National University of Singapore 'Sea AI Lab "Singapore Management University 'Yale University. Cor- respondence to: Cunxiao Du <ensdunm gmmm, Zhuoran Yang aichuoran.yangee .tauuorrn.ang@@al...dd.... edl'.
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
model during inference. Furthermore, even when the choice of the draft model is optimized, the associated hyperparam- eters can still be refined. For instance, Liu et al. (2024)
Algorithm I CANONICAL DECODING
Inputs: initial prompt ptg = pt E X", target model P. Procedures:
and Huang et al. (2024) aim to optimize the speculation length in a training and training-free manner. Based on these observations, we ask the questions (see Figure 1): : given prefix prompts and candidate configurations of hy- perparameters, there theoretically sound framework to model and solve the hyperparameters selection problem? Is there any training-free method that can adaptively choose the hyperparameters such that the latency of speculative decoding can be minimized?
1: Set t =0 0.
2: while L + 0 and x2 + EOS do
3:
t t= 1.
4: 2 ~ P(- pt,-1)-
ptg - concat(pt,
6: end while
7: return pt,
In this paper, we answer these questions affirmatively. We adopt bandit framework to leverage its adaptivity in un- known environments to achieve this goal. Our contributions can be summarized follows.
respectively the t-th token and first t tokens. In the inference stage, LLMs use an additional temperature pa- rameter Y > 0 to predict the next token's probability as so softmax(7 log ((11),))) where softmax is the softmax operator. The results in our work hold for any 7 > 0, and we just denote softmax(7 log P) as P for ease of notation. When y > O, we sample the next token from the predicted distribution, which is called sampling (sampling decoding). When y 0, the next token will be the token that corresponds to the highest logit value; this is called greedy decoding. We note that the greedy decoding is deterministic, i.e., the token is sampled from a degenerate distribution. These two families of decoding methods can be unified as Algorithm I, where LLMs autoregressively generate tokens until the EOS token.
We formulate the hyperparameter selection problem in speculative decoding as bandit problem and propose gen- eral speculative decoding framework BANDITSPEC(ALE (ALG) (see Algorithm 3), where the hyperparameter selection al- gorithm ALG selects the hyperparameters to be deployed in each round of speculative decoding. The objective min- imize the stopping time regret, which measures the latency of ALG compared to that of the best hyperparameter.
Under mild stochastic and adversarial reward assumptions, we devise two hyperparameter selection algorithms, UCB- SPEC and EXP3SPEC, respectively. By deriving upper bounds on the stopping time regret, we prove that the infer- ence latency between the proposed algorithms and the best hyperparameter under given initial prompt vanishes asymp- totically. In addition, we show, via deriving an information- theoretic impossibility result, that the regret performance of UCBSPEC is optimal up to constants.
Speculative Decoding As shown in Algorithm 1, the au- toregressive decoding feature requires multiple forward in- ferences of LLMs P sequentially. To reduce the number of forward inferences, Leviathan et al. (2023); Chen et al. (2023) proposed the vanilla speculative decoding algorithm which implements a draft model Q to generate draft tokens and let the target model P verify them in parallel. For completeness, we present and describe the vanilla specu- lative decoding algorithm in Appendix C.I. This vanilla specualtive decoding is then extended by some existing works, e.g., Miao et al. (2023); Cai et al. (2024) organizes the draft tokens as a tree, which improves the number of accepted tokens. The speculative decoding algorithm con- tains many hyperparameters, e.g., the draft model Q, and the tree structure in Miao et al. (2023); Cai et al. (2024), Most existing works keep these hyperparameters fixed for all the tasks. Some other works optimize the draft model in an online or offline manner (Liu et al., 2023) and the size of the tree (Chen et al., 2024), which are designed for specific considerations. In contrast, our work aims to derive unified online hyperparameter selection algorithm that can be applied for any type of hyperparameters.
Extensive empirical experiments with LLaMA3 and Qwen2 are conducted to demonstrate the efficacy of the proposed framework. When the batch size is 1, the adap- tive selection of models via UCBSPEC and EXP3SPEC can greatly improve the latency, exhibiting competitive per- formance against the existing methods. Under simulated real-life scenarios where LLMs are implemented for diverse prompts simultaneously, the adaptive selection of specula- tion length via UCBSPEC achieves comparable throughput with the oracle best.
2. Preliminaries
LLM Decoding We denote an LLM as P: X - Ax. where X and X are the space of tokens and the space of all token sequences, respectively. Most LLMs pre- diet this conditional probability via predicting the log- its of the next token. Concretely, the LLMS predict log P(xel | 11),1), where Ie E X and F E X are
Multi-Armed Bandits The Multi-Armed Bandit (MAB) is a fundamental online decision- decision-making problem (see Algo- rithm for its dynamics). In its classical stochastic form, an agent chooses from K arms, each of which delivers a re-
2

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
Algorithm 2 SPECUATIVE DECODING SUBROUTINE (SpEcDEcSUB)
Algorithm 3 SPECULATIVE DECODING WITH BANDITS (BANDITSPEC)
Inputs: pt E X", target model P. the hyperparameters S, maximum speculation length L.
Inputs: arm selection algorithm ALG, initial prompt pla pt E X,, bandit configuration [P.S [Sthe(k] L). Procedures:
Procedures:
1: Call a standard speculative decoding algorithm with (pt, P.S, L).
1: L t= 0, Ho = 0, To 11,41.00 = 0.
while EOS & æI,, do
2: return the accepted and bonus tokens Fier, where 7 2
3: t=t 1.
4: Select a hyperparameter index L, = ALG(H4-1))
5: I = SPFCDDESSUR(PL P, SI,
ward sampled i.i.d. from an unknown but fixed distribution when pulled (Lattimore & Szepesvári, 2020). The goal is to select arms over I' rounds to maximize cumulative rewards. Two primary classes of algorithms- Cotorithms-UB-yeee UCB-type (Auer et al., 2002b) and sampling-basee methods (Russo et al., 2017)- have been developed and proven optimal in this setting. In the adversarial formulation, there are no assumptions on the reward distributions; rewards can evolve arbitrarily over time and may be correlated across arms (Auer al., 2002a). Several algorithms, such as EXP3 and EXP4 (Auer et al., 2002a), are known to achieve optimal performance under these conditions. In this work, we frame the hyperparameter selection problem as an MAB problem and develop algo- rithms tailored to both stochastic and adversarial settings.
6:
pt, = concat 21,2).
7: Hi = con (Ie, xI,,)).
8: end while
9: return ST pt. v) = t, PTST(ALL,p))))) = pt,-
length L, and the output is the accepted token sequence Ther E X We provide two examples of the hyperparame- ter sets here. (1) If we adopt the vanilla speculative decoding (Algorithm 6) as Line I, S can be different draft models Q : X + Ax and S is the set of all the provided draft models. We would like to choose a draft model according to its training context, c.g. math, creative writing, to decode the current prefix. Then the problem we consider is how to adaptively select a proper draft model for speculative de- coding via bandit algorithms. (2) If we adopt Medusa (Cai et al., 2024) as Line 1, can be different tree structures, and S is the set of plausible tree structures. In this problem, we would like to adaptively adjust the speculation tree structure according to the context.
:= := {1, J.. For finite set X, we denote the set of distributions supported on it as Ax {P: X' + [0, 1] ΣeEx P(z) 1, P(x) 2 0 for all x E X). The space of all finite length sequences whose components belong to X' is denoted as X', and we use II.1 E x² C X to denote a length-L L sequence. The Kullback-Leible (KL) divergence between two distributions P and Q is denoted as KL(P,Q) Q).
Notations:
Let [N]
With the help of SPECDECSUB, the speculative decoding with a bandit framework. BANDITSPEC, can be specified in Algorithm 3 and as illustrated in Figure 1. The handit config- uration v = (P, S = {S:).e L) consists of three com- ponents: the target model P. the set of K candidate hyper- parameter specifications S, and the maximum speculation length L. Each hyperparameter specification S. S corre- sponds to an arm in the bandit problem. Given a prompt pt. and an arm selection algorithm ALG, a hyperparameter spec- ification is chosen according to the history H4-1 in Line 4. Then SPEC is invoked with selected hyperparam- eters SJ, as input. The output of SPECDECSUR, F1,,1 is then adopted update the prompt (Line 6) and the history information (Line 7). The whole process stops when the EOS token appears in the prompt. We denote the number of calls to SPecDecSUB (the stopping time) and the gen- erated token sequence as ST pt, v) and plSTTALG,,t. pt.")) respectively. To minimize the decoding latency, we aim to design ALG to minimize ST(ALG, pt., v).. Since the position of the EOS token itself is a random variable, we would like
3. Bandits for Adaptive Speculative Decoding
In the section, we formally formulate the hyperparameter selection problem in speculative decoding using the parlance of multi-armed bandits. The goal of this online decision- making process is to decode as soon as possible, i.e., mini- mizing the latency of the LLM decoding. Different from the classical multi-armed bandit problem, this problem involves two stochastic processes that march at various paces. In fact, as described in Appendix C. each (vanilla) specula- tive decoding subroutine produces several tokens, where the number of accepted tokens itself is also a random variable. Thus, the selection of hyperparameters of each speculative decoding subroutine and the token generation processes are evolving different paces.
To put the problem in mathematically sound way, we first specify a general speculative decoding subroutine (SpecDecSUB) in Algorithm 2. The input of this sub- routine a prompt pt E X², a target model P, a specifica- tion of hyperparameters S, and the maximum speculation
We abbreviate the notation Fi, 1:1. as F1., represent the accepted tokens generated by St, at time step t.
3

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
to minimize the expectation of ST(A pt, v). The perfor- mance of ALG is measured via the stopping time regret
Stop Prefix:
Adaptive specuiative decoding via bandi: algorithm is an idea that that aims to strise
History Information:
(eLGA pt, := E [[SG p) - E [ST(ALG; (nt,))) pt, v) pt. v],,
pp., v]
(1)
select traft model according to UCB or EXP3
Eagle-2
Suffix Tree
uou
3 times
2 times
Accept 7 tokkns
Accept 2 tolions
where ALG; is the arm selection algorithm which adopts S; in all rounds, i.c., i = (H,, for all He and t, and d (pt, v) = argmin,g [x] E (ALG, pt, v) pt, v] denotes the index of the best hyperparameter for prompt pt under bandit configuration v.
Decis sior Making Statistics:
Confidencre
Cumulative
interval
UCB
EXP3
risk
For ease of notation, when and pt are clear from the con- text, (ALG pt, v), Reg (ALG, pt, v), and i" (pt, 2/ will be abbreviated as (T(G))), (ALG) and i", respectively. We will use BANDITSPEC(ALG) to specify the choice of ALG in Algorithm 3. For simplicity, we regard the bonus token as the last accepted token. Thus, the length of the ac- cepted tokens It,,t at each round, yr,,,, is between [1, L++1]
Figure 2. Illustration of our bandit model for choosing configura- tions to decode the next token, where UCB and EXP3 refer to UCBSPEC and EXP3SPEC, repectively.
Reg(ALG, pt, we want to devise an ALG to (approxi- mately) match the performance of ALG,,. In particular, if (egGAAG, pl, v) = o (([len)p)sTTaaao) }]), it implies that BANDITSPEC requires the same number of specula- tive decoding rounds as (ALLD) asymptoti- cally even though the information about S is not revealed at the beginning. In order words, BANDITSPEC(A (ALG) learns the identity of S,. quickly and the price for this learning process can be amortized over time. Additionally, when BANDITSPEC(ALC) is deployed over diverse prompt inputs, we expect significant acceleration of token generation compared to any fixed single speculative decoding method.
Before going the algorithm design and the theoretical analysis, we would like specify some important properties that are shared for any arm selection algorithm ALG and clar- ify the intuitions about our theoretical analysis. We denote the stopping time of the canonical decoding (Algorithm 1) and the generated sequence as Te and pt,,. respectively.
Proposition 3.1. For any arm selection algorithm ALG that selects an arm according to the history, the generated prompt (EST(ALG) is equal to pt, in distribution, i.e.,
Why do we consider stochastic and adversarial settings? To derive efficient algorithms and meaningful theoretical analysis, it is necessary make certain plausible assump- tions of the problem. For the BANDITSPEC problem, we need to model the stochasticity of the number of accepted tokens for each hyperparameter specification. We highlight that in real-world applications, they are far from identically and independently distributed. The stochastic case (Sec- tion 4) models it as random variables and only assumes that each hyperparameter will have a stationary mean accep- tance length (Assumption 4.1) without the independence assumption. The adversarial case (Section 5) removes this stationarity assumption and does not make any distributional assumption of the number of accepted tokens for each hyper- parameter. We highlight that there is no explicit adversary in the speculative decoding, but we model the stochasticity of the number of accepted tokens as the randomness from (imaginary) adversary.
(LST(AAG)) d pt,, and (en( A len(pt,.). (2) len(ptST(ALG)
The stopping time ST(AGG) can be bounded as
< <ST(A < len(PEsT(ALA)); a.s. (3)
L+1
The proof of Proposition 3.1 is provided in Appendix D.I. This proposition states that the distribution of the generated prompt the same as that of the prompt generated by Algo- rithm 1. The stopping time STLG) is equal to the length of the generated prompt up to a constant. To facilitate our theoretical understanding, we pose the following question.
Question: Whether it is possible to devise an arm selec- tion algorithm ALG to achieve sublinear regret in terms of the length of the generated token sequence, iee, is Reg(ALG,pt,, = ((B[len(pttsrrua))
Interpretation of the Desired Result. Given prompt pt and bandit configuration v, BANDITSPEC adpatively se- lects the hyperparameter via ALG and learns the context. The stopping time regret (1) measures how the stopping time of BANDITSPEC(ALO) compared to that of the (ag- nostic) best one B.NDITSPEC(ALG=:) By minimizing
4. Modeling Tokens Stochastically
In this section, we model the length of the accepted tokens as random variables.
Assumption 4.1 (Stationary Mean Values). There exist K values (Hi]e(K] [(1, L+1], such that conditioned on the
4

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
Algorithm 4 UCBSPEC
prompt pl E X², the expected length of the output sequence of the canonical decoding algorithm (Algorithm 1) is finite, i.e. B[len(pt,,)) < oc.
Inputs: number of hyperparameter specifications K, history
H1 (L,::1,,e)) s= confidence parameter 8.
Procedures:
This assumption states that the expected length of the gener- ated prompt is finite. In real-world applications, the length of the generated prompt is always finite due to the limits of computation and storage.
1: if K then return /++1 = (
2: Compute the lengths Yı,,. = len(X s,,s) .) for all s [t].
3: Set the statistics ((()) f [k], where
E. Y 1 = =
To state our main result, we denote the suboptimality gap between the best arm i* := argmax;ę[K] i4 and arm as A, := 14-- Ms. Define the hardness parameter H(pt, 1) := Eig; 1/(u-- Ai), which captures the diffi- culty of acceleration given the initial prompt pt and bandit configuration v.
74,1
1[1 =1], Pr,e
niat
L
1+ ni,t
Kta (1 + ng.e)
Cri
2
? n
1 + log
8
UCB = -b,,t + Cri,t-
Theorem 4.3 (Upper Bound), Under Assumptions 4.1 and 4.2, given any prompt pt X and bandit configura- tion (P,S = (Si L), the expected stopping time regret of Algorithm 3 with ALG Algorittmm 4 (UCBSPEC) is upper bounded as
4: return index 1++1 = argmax UCB
history H4-1 and the chosen arm 1, at time t, the expected number of the accepted tokens , -1]=
Reg((AG, (AG vt = I.2 log Bllen(pt.,.)) )1)
This assumption assumes that the conditional expectation of the number of accepted tokens for each hyperparame- ter is equal to a fixed number conditioned on the previous tokens. We emphasize that this assumption does not re- quire independence between the number of accepted tokens across implementations of SPECDECSUB, which would be unrealistic in real-world applications. More discussions are provided in Appendix B.
Theorem 4.3 answers the proposed question in Section 3 in the affirmative under Assumptions 4.1 and 4.2. To inter- pret the results of the theorem, for each hyperparameter S., it requires = O log E[(en(ptsT(ALL)) )1 /)) pulls to identify that S; is suboptimal under the current prompt pt and bandit configuration v/ resulting in n, token loss com- pared to the case where S₁- had been adopted. Additionally, this loss could be compensated by n,A:/14 pulls of S.,. which constitutes the final stopping time regret bound. The proof is postponed to Appendix D.2 with more discussions in Appendix B.2.
4.1. Upper Bounds for the Stochastic Case
Algorithm Design We design a UCB-type arm selec- tion algorithm UCBSPEC, as shown in Algorithm 4. To avoid additional terms, we call the aggregated algorithm, CANDITSPEC(UCBSSEEES as UCBSPEC. The full version of UCBSPEC is detailed in Algorithm 8.
4.2. Lower Bound for the Stochastic Case
We further provide an information- the regret under the greedy decoding strategy to indicate how the upper bound is in Theorem 4.3. More details and the proof of Theorem 4.4 are deferred to Appendix D.4.
lower bound of
This aggregated algorithm, UCBSPEC, is adapted from the classical UCB-1 algorithm in Auer al. (2002b). The main differences are the confidence radius design cr;, and the stopping rule. We highlight that the form of cr2,, is designed to fit the weak assumption of the number of accepted to- kens. In fact, the proof of the regret of UCB-1 assumes that the values of each arm are generated before the pull of arms (Auer et al., 2002b; Lattimore & Szepesvári, 2020), which bifurcates from practical L.M inference scenarios. In contrast, we remove this strong restriction. The stop- ping rule of UCBSPEC makes the analysis of our algorithm rather different from that of UCB-1. The stochasticity of the total number of arm pulls requires a novel regret decompo- sition analysis that is not presented in previous works.
Theorem 4.4 (Lower Bound), Given any sequence of initial prompts (pt )m= C Xint with len(pt - oC, m - Do and a bandit configuration v = S {S.) L), under Assumption D.4, the greedy decoding strategy and the dynamics represented in Algorithm 3. for any non- anticipatory and consistent arm selection algorithm ALG. the expected regret satisfies
lim inf Reg(ALG, pt, v) m00 log(len(pt..))
Mi--
kl."
Theoretical Analysis We first state an assumption. Assumption 4.2 (Finite Generation Length). Given any
where kl, == inf (KT, Ps): ExxPs(X] > >.
5

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
To provide more concrete example of the lower bound, consider the truncated geometric distribution (TGD) on [1, + 1I with parameter p E (0, 1), i.e.
Algorithm 5 EXP3SPEC
Inputs: number of hyperparameter sprciaccctttteespppcccccccttttot K, history
He = (1,,,
Procedures:
p (1 p). x 1,2 L.
Ps(x)
(4)
I: Compute the lengths Vi,,= = len(X1,,5) for all 8 [t].
pt,
x =+1+1.
2: Set the statistics for all E [K]
This TGD was considered in the seminal works on specula- tive decoding (Leviathan et al., 2023; Chen et al., 2023).
L+1- Yi.
Z,, = lfi = 11-
(6)
L Pt,
Proposition 4.5 (Tightness Result). Let STOD = {S>
Ps satisfies (4)). Let { C STGD and S, satisfies (4) with Pe (Line 5 Algorithm 2), then
3: Set learning rate 77. = Vlog K/(t (t)..
Set probability vector pt E A with for all i e [K]
liminf Reg(ALG, pt"",, v)
Pe- (1 pf-) 2 H( (1 p:-)
(-
1/1 D4-1 2i.0)
70000 log(len(pt.:.))
exp
Pt,
Σ}= exp (17tt E Zi.o)
Therefore, the upper and lower bound match up absolute constants and 1.((pp.)) p.- (1-p)) factor. In particular, if pe- E
5: return hyperparameter index Le+1 Pt.
(2-1/4,1), they match up to absolute constants and L.
The proof is deferred to Appendix E.3. Proposition 4.5 in- dicates UCBSPEC is optimal up to constants and L when considering the TGD. In other words, the additional spec- ulative decoding rounds of UCBSPEC not only achieves Olog ([len(pt,,)] )I compared to ALG,,, but is also among the best possible for any arm selection algorithm (up to constants).
time step be yi.t = (X..) We assume (yr,e) kEjK,, is fixed by the environment before the algorithm starts.
The bandits problem with this assumption is often referred to the oblivious adversarial bandits in the online learning works (Auer et al., 2002a; Lattimore & Szepesvári, 2020). It admits more general and practical setups compared to the stochastic MAB. We find the greedy decoding strategy aligns more closely to this setup the sense that the gen- erated tokens by the models are (potentially) fixed given the initial prompt. Hence, we present our result under the greedy decoding strategy in this section.? Given prompt pt E X and bandits configuration v, the stopping time regret (1) of an arm selection algorithm ALG becomes
For the tightness of our algorithm, according to Note 15.3 in Lattimore & Szepesvári (2020), kl, O( when A; is small. This indicates the dominating terms in the upper bound in Theorem 4.3 match the lower bound in Theorem 4.4 up to (possibly instance-dependent) constants. Futherfmore, because the truncated geometric distribution is more close to a sub-exponential family distribution, es- pecially when T. large, bandit algorithms built upon UCB1 (Auer et al., 2002b) are generally loose in some fac- tors. In order to close the gap between the upper and lower bounds completely, KI-UCB (Garivier& & Cappé, 201 1) can possibly be adapted to this problem out of theoretical in- terest. However, on the practical side, KL-UCB demands solving an optimization problem at each round, which can be cime-consuming during implementations. Thus, does not perfectly align with our ultimate goal of LLM inference acceleration.
Reg B[ST(ALG)] min ST(L,,,) (E(K]
(5)
It is worth pointing out that under the greedy decoding strategy, the stopping time of any proposed algorithm can still be random due to the internal randomness embedded in the algorithm. For instance, the choice of hyperparameter St, in Line 5 in Algorithm 5.
Algorithm Design We present our arm selection algorithm in Algorithm 5, which is an abridged version of the full ver- sion BANDITSPECI (EXTTSPECEEXP333SS (EXP3SPEC) delineated in Algorithm 9. This algorithm modifies the anytime EXP3 algorithm (Latti- more & Szepesvári, 2020) to suit the speculative decoding application. In terms of the algorithm design, the main dif- ference lies in the change of the stopping rule. We highlight that while the stopping time of the algorithm is random,
5. Modeling Tokens Adversarially
In this section, we weaken Assumption 4.1 in Section 4 and consider a more general case. Specificallyy, we make the following assumption on the number of accepted tokens.
Assumption 5.1 (Adversarial Mean Values). Let the num- ber of accepted tokens generated by hyperparameter S; at
OOrr analysis can also be extended to cover the sampling de- coding strategy (see Remark D.3).
6

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
the anytime feature of Algorithm 5 does not require any information about the time horizon. This is achieved by the vanishing sequence of learning rates (7h HeN which can be elegantly adapted to the unknown stopping time. With regard to the analysis, previous works (Auer et al., 2002a; Bubeck et al., 2012; Lattimore & Szepesvári, 2020) only consider the gap between the cumulated rewards over the same fixed horizon T, i... | ( Mt., .l..el- In contrast, we need to upper bound the stopping time regret in (5) where the baseline ALG, and any proposed algorithm ALG have different termination times in general. Thus, the analysis is much more involved.
BANDITSPEC, along with UCBSPEC and EXP3SPEC. In the first experiment, the candidate hyperparameters are dif- ferent draft models. In the second experiment, the candidate hyperparameters are different speculation lengths, where real-life I.LM serving scenarios are simulated with diverse input prompts. Additional experimental results on memory utilization and additional experiments on larger models and different hardwares are provided in Appendix G. The code is accessible via https: / /github. com/ sa sg/ Bandit
6.1. Experiment with Draft Models
Theoretical Analysis To ease the analysis, we make an assumption on the stopping time of Algorithm 5.
Experimental Setups We adopt the open-sourced LLaMA3-8B-Instrue (Dubey et al., 2024) and Qwen2- 7B IB-Instruct (Yang et al., 2024) as the target models. The commonly-used existing speculative decoding meth- ods PLD (Saxena, 2023), Rest (He et al., 2024), Suffix Tree (Oliaro et al., 2024; Hu et al., 2024) and Eagle-2 (Li et al., 2024b) are adopted as the baselines. Among these baselines, PLD, Rest, and Suffix Tree represent the non- parametric (or model-free) speculative decoding methods, whereas Eagle-2 represents the speculative decoding meth- ods that utilize smaller draft models. Each of these methods corresponds to an arm in our problem.
Assumption 5.2 (Stopping Time assumption). Given :=
prompt pt E X and configuration v, let i* [rgmingg(k]] ST(A),;) We assume that ST(GG ST(ALG,-) almost surcly.
In speculative decoding, when the initial prompt is given, there generally exists a hyperparameter that has the highest acceptance rate in most rounds compared to the rest of the hyperparameters. As bandit algorithms will explore those suboptimal hyperparameters, the termination time falls behind that of the optimal hyperparameter. Therefore, Assumption 5.2 is satisfied in practical applications.
The experiments are carried out on Spec Bench (Xia et al., 2024), Alpaca (Taori et al., 2023), Code Editor (Guo et al., 2024) and Debug Bench (Tian et al., 2024). Among these benchmarks, Spee Bench and Alpaca encompass multiple topies, while Code Editor and Debug Bench focus on coding tasks, a representative scenario for specialized models.
Theorem 5.3. Under Assumptions 4.2, 5.1 and 5.2, given (P,S- = (ett) L). the expected stopping time regret of Algo- rithm 3 with ALG = Algorithm 5 (EXP3SPEC),
any prompt pt X* and bandit configuration
Reg(ALG, pt, 2L min { { len(pt,.) K log K,
We record the number of accepted tokens for each specu- lative decoding step, as well as the wall-time for generat- ing each complete response. The Mean Accepted Tokens (MAT) and the throughput (Tokens/s) are computed. These two metrics are widely adopted in the speculative decoding community and are positively correlated (Xia et al., 2024). In particular, Tokens/s measures the actual latency during decoding. The experiments are conducted on a single A100 and set batch size as 1.
2LK log K
min ST(ALG;) log K kE[K]
Theorem 5.3 also provides an affirmative answer to the question posed in Section 3. The first term in the minimum provides a worst-case guarantee. Even if all hyperparame- ters in S are not good or K is large, EXP3SPEC will stop at no more than (len((pt, )) time steps after S;. termi- nates. The second term is an instance-dependent bound in terms of hyperparameters S. Specifically, when the best hyperparameter S,. has small stopping time, EXP3SPEC will scale as ST(AG,;-) + ((VST(ALG..))) This upper bound suggests that the number of speculative decoding rounds of EXP3SPEC is almost the same as that of the best hyperparameter configuration ALG,-.
Experimental Results We report the results of our experi- ments in Table 1. The proposed adaptive speculative decod- ing framework BANDITSPEC exhibits superior performance compared to existing methods in the datasets we consider. In particular, the best performance measured by Token/s is always achieved by the proposed framework. We note that although the non- non-parametric methods are worse than Eagle- in average, they are effective on a portion of prompts. Our proposed methods, UCBSPEC and EXP3SPEC, automati- cally adapt to different prompts, i.e., suffering from a small stopping time regret on each prompt. Thus, they achieve better performance than all the methods that only use a fixed model. On Debug Bench, UCBSPEC can even achieve
6. Experiments
In this section, we conduct two sets of experiments to demonstrate the efficacy of the proposed bandit framework
7

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
Table Empirical Comparison between the proposed algorithms the existing works, measured by Mean Accepted Tokens (MAT) (
and Tokens/s (1). The best result is highlighted in bold, while the second best result is underlined. The proposed algorithms demonstrate unequivocal superior performance compared with existing methods.
Spec Bench
Alpaca
Code Editor
Debug Bench
Methods
MAT(1)
Tokens/s(1)
MAT(T)
Tokens/s(t)
MAT(f)
Tokens/s(f)
MAT(1)
Tokens/s(1)
LLaMA3-8B-Instruc
Vanilla
1.00
35.73
1.00
35.92
1.00
36.32
1.00
36.89
PLD
1.46
43.96
1.53
53.06
2.13
82.61
1.67
82.76
Rest
1.29
40.67
1.48
52.40
1.33
51.32
1.29
48.49
Suffix Tree
1.83
55.10
1.71
64.02
2.30
90.21
2.13
77.56
Fagle-2
3.94
98.15
4.04
110.00
4.79
128.76
4.78
119.12
EXP3SPEC
3.65
102.10
4.23
120.38
4.36
137.29
4.50
132.25
UCBSPEC
3.98
105.72
4.35
125.78
4.83
138.27
4.60
135.34
Qwen2-7B-Instruct
Vanilla
1.00
38.71
1.00
39.32
1.00
39.30
1.00
39.57
PLD
1.55
52.44
1.42
58.41
1.89
64.56
2.15
70.49
Rest
1.31
46.42
1.47
59.01
1.31
53.79
1.22
50.51
Suffix Tree
1.96
68.42
1,46
62.60
2.18
85.75
2.49
101.47
Eagle-2
3.64
97.82
3.61
104.43
4.88
138.58
4.79
126.01
EXP3SPEC
3.76
107.36
3.83
113.90
4.90
160.41
4.86
151.73
UCBSPEC 4.13
112.33
3.93
114.20
4.92
161.35
5.10
151.37
improvements of 13% for LLaMA3 and 19% for Qwen2. Moreover, as UCBSPEC demonstrates better performance under almost all benchmarks with both two target models, this suggests that speculative decoding in real-life environ- ments tends to be closer to the stochastic (Assumption 4.1) compared to the adversarial reward case (Assumption 5.1).
while a lower acceptance rate can lead wasted compu- tation resources of GPU. Additionally, the execution time of the draft model contributes to overall decrease in throughput. Given these confoundingly interrelated fac- tors, along with latent variables such as the acceptance rate (which is unknown before verification and depends on the input prompts), we adopt a sundit-basee approach to model the current throughput as the reward. Specifically, we em- ploy UCBSPEC to dynamically adjust the hyperparameter 7/, the speculation length, to maximize the throughput, i.c., the number of generated tokens per second. We set the maximum speculation length L as 4, and 7 takes values in {0. where 7 = corresponds to the canonical decoding (Algorithm 1). As the first experiment suggests UCBSPEC is more in line with the real-life speculative decoding environment than EXP3SPEC, only evaluate UCBSPEC in this experiment. The experiments are con- ducted on a single A100.
Remark 6.1. The adversarial setting can be regarded as a means of comparison to the stationary setting. Prior to this work, was a priori unclear how use MAB to improve speculative decoding. Should one employ stochastic, ad- versarial or even more generalized model? We consider range of such MAB models and do a comparison among them to provide the community with guide on which MAB model is best suited to the speculative decoding problem. As the empirical performance of UCBSPEC better than EXP3SPEC (Table I), implies that real-life scenario tends be benign and may be more aligned with the stationary mean assumption.
Specifically. we use LLaMA3-8B Instruct and Qwen2-7B Instruct as the target models and adopt Fagle-1 (Li et al., 2024a), the current state-of-the-art model, as the draft model. We do use Eagle-2 (Li al., 2024b) because does not support batch inference. For evaluation, we adopt Al- paca (Taori al., 2023) as the test set, as it covers various topics, thereby simulating realistic setting with diverse acceptance rates. To approximate real-world conditions, we randomly sample prompts from the test set to form a batch for inference, with batch sizes ranging from to 50. As
6.2. Experiment with Speculation Lengths
Experimental Setups In addition to improving the latency when batch size is our proposed algorithms also improve the throughput in real LLM serving scenarios with different batch sizes. In practical serving environments, speculative decoding does not always yield performance gains due to variations in batch size and acceptance rate. As the batch size increases, the system rapidly becomes compute-bound. bound,
8

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
choose appropriate hyperparameters to accelerate LLM in- ference under realistic assumptions. Both theoretical guaran- tees and extensive experiments are provided to demonstrate that adaptive speculative decoding via bandit algorithms can boost the performance of existing methods in a training-free manner. For future work, we would like to point some direc- tions, improving the performance of the current algorithms.
1.8
oracle best
UCBSpec
1.8
1.4
Throughput
1.2
40
80
120
160
200
240
280
320
360
400
440
480
Semple Index (sorted)
Therefore, another direction to design hyperparameter selection algorithms that can achieve the (near) optimal balance between these two goals based on practical needs.
(a) Target model: LLaMA3
Thimughyan Improvment
1.0
orable best
UCBSoes
Structured Bandits Our current framework is based on the standard K-armed bandit model. However, broader classes of bandit problems with additional structures- such as linear bandits (Abbasi-yadkori et al., 2011) and Lipschitz bandits (Magureanu et al., 2014)-can also be considered. This aligns more closely with practical scenarios, where the number of hyperparameters can be large, and the value of K may be very high when modeling the problem as a K Karredd MAB. By leveraging such structures in MABs, we can expect to identify better hyperparameterss more efficiently, thereby further accelerating the optimization process.
1.5
1.4
1.2
mymmm
0.8
80
200
820
440
480
Sample India (sortad)
(b) Target model: Qwen2
Figure 3. We compare throughtput improvements with different speculative decoding lengths ? [4] and the canonical decoding (y 0). The performance of UCBSPEC approaches that of the best hyperparameter across all samples for both target models LLaMA3 and Qwen2. The sample indices sorted according to the best arm improvement for a clear demonstration.
Robust bandits and Non-stationary bandits As indicated by the experimental result, the real-life speculative decoding environment is closer to the stochastic reward case (Assump- tion 4. than the adversarial reward case (Assumption 5.1). Therefore, one direction for future work is to consider the settings "in between", e.g., robust bandits in the presence of adversarial corruptions (Ding et al., 2022; Zhong et al., 2021), non-stationary bandits (Cao et al., 2019; Besbes et al., 2014; Hou et al., 2024) where the mean number of accepted tokens can vary across time. These settings are more benign than the adversarial reward assumption and can be exploited to accelerate the inference.
our evaluation metric, we measure the throughput improve- ment relative to the canonical decoding (non- specullaaltt baseline. Our result is averaged over 16 independent runs to smoothen the hardware-dependent factors.
Experimental Results The results are presented in Figure 3, where we reorder the 500 sample indices in ascending order of the performance of the best hyperparameter (blue line) for easy comparison. Otherwise, the lines in this figure will not be largely monotonic. Here the "worst" and "best" lines are calculated among results of y E {1...., 4] in hindsight. Thus, we call the "best" line as the oracle best. Firstly, since the optimal hyperparameter y varies with different input prompts for either target model, fixing single hy- perparameter is suboptimal, e.g., in Figure 3 (b), the best hyperparameter changes from = 1 (light green) to 2 (green) at sample index around 80; and the original Eagle- 1 (Li et al., 2024a) (? = 4 in purple) is even inferior to the canonical decoding (7 in grey) for sample indices less than 80. This necessitates the use of adaptive hyperparam- eter selection. Next, UCBSPEC demonstrates competitive throughput performance, outperforming the second- best hy- perparameter in most cases and closely approaching the (varying) oracle best across experiments. These benefits are obtained thanks to the adaptivity of BANDITSPEC.
Contextual Bandits Another direction is to explore con- textual bandits, where the environment reveals additional information that can be leveraged to reduce the learning burden (Luo et al., 2018; Kato & Ariu, 2021).
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There many potential societal consequences of our work, none which we feel must be specifically highlighted here.
This work is supported by funding from the Singapore Ministry of Education Academic Re- search Fund (AcRF) Tier I grants under grant numbers A-8002934-00--8 and A-80009880000000. This research is also supported by the National Research Foundation, Singa- pore under its Al Singapore Programme (AISG Award No: and is part of the programme DesCartes which is supported by the National Research
Acknowledgementss:
7. Conclusions and Discussions
In this work, we propose a MAB framework together with two hyperparameter selection algorithms that adaptively
SGGG2-PDD-022-08--40
9

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
Foundation, Prime Minister's Office, Singapore under its Campus for Research Excellence and Technological Enter- prise (CREATE) programme.
Bubeck, S.,
N., et al. Regret analysis of stochastic and nonstochastic multi-armed bandit prob- lems. Foundations and Trends@ in Machine Learning, 5 (1):1-122, 2012.
References
Cai, T., Li, Y., Geng, Z., Peng, H., Lee, J. D., Chen, D., and Dao, T. Medusa: Simple Ilm inference acceleration framework with multiple decoding heads. arXiv preprint WrXiv:241.100772 2024.
Abbasi-yadkori, Y. Pál. D., and Szepesvári, C. Improved algorithms for linear stochastic bandits. In Advances in Neural Information Processing Systems, volume 24, pp- 2312-2320, 2011.
Cao, Y., Wen, Z., Kveton, B., and Xie, Y. Nearly optimal adaptive procedure with change detection for piecewise- stationary bandit. In Proceedings of the 22nd Interna- tional Conference on Artificial Intelligence and Statistics, pp. 418-427. PMLR, 2019.
Agrawal, S. and Goyal, N. Analysis of thompson sampling for the multi-armed bandit problem. Conference on learning theory, pp. 39-1. JMLR Workshop and Confer- ence Proceedings, 2012.
Chen, C., Borgeaud, S., Irving, G., Lespiau, J.-B., Sifre, L., and Jumper, J. Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.01318, 2023.
Agrawal, S. and Goyal, N. Near-optimal regret bounds for thompson sampling. Journal of the ACM (JACM), 64(5): 1-24, 2017.
Antos, A., Grover, V., and Szepesvári, C. Active learning in heteroscedastic noise. Theoretical Computer Science, 411129-30)2712-22723 2010.
Chen, Z.,, May, A., Svirschevski, R., Huang. Y., Ryabinin, M., Jia, Z., and Chen, B. Sequoia: Scalable, robust, and hardware-aware speculative decoding. ArXiv, abs/2402.12374, 2024. URL https: semant wew.cciscchaa.cocco org /Corg//rrrsss. / //api.semantic 267751265.
Arora, R., Dekel, O., and Tewari, A. Online bandit learning against an adaptive adversary: from regret to policy regret. In Proceedings of the 29th International Coference on International Conference Machine Learning, pp. 1747- 1754, 2012.
Ding, Q., Hsieh, C.-J., and Sharpnack, J. Robust stochastic linear contextual bandits under adversarial attacks. In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, Proceedings of Ma- chine Learning Research, pp. 7111-7123. PMLR, 28-30 Mar 2022.
Auer, P., Cesa-Bianchi N., Freund, Y., and Schapire, R. E. The nonstochastic multiarmed bandit problem. SIAM journal on computing, 32(1):48-77, 2002a.
Du, C., Jiang, J., Yuanchen, X., Wu, J., Yu, S., Li, Y., Li, S., Xu, K., Nie, L.,, Tu, Z., et al. Glide with a cape: A low- hassle method accelerate speculative decoding. arXiv preprint urXiv:2442.00088 2024.
Auer, P., Fischer, P., and Cesa-Bianchi, N. Finite-time analysis of the multi-armed bandit problem. Machine Learning, 47:235 2566, 2002b.
Besbes, O., Gur, Y., and Zeevi, A. Stochastic multi-armed- bandit problem with non-stationary rewards. In Ghahra- mani, Z.., Welling, M., Cortes, C., Lawrence, N., and Weinberger, K. (eds.), Proceedings of the 27th Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The llama 3 herd of models. arXiv preprint 24XY::2007.218.. 2024.
Garivier, A. and Cappé. O. The kl-ucb algorithm for bounded stochastic bandits and beyond. In Proceedings of the 24th Annual Conference on Learning Theory, vol- ume 19 of Proceedings of Machine Learning Research. pp. 359-376. PMLR, 09-11 Jun 2011.
Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Ka- plan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert- Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D. M., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., teusz Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. Language mod- els are few-shot learners. ArXiv, 0ix/2005.14163 2020. URL https:/ org//CorpuSDD:218977.
Gordon, G. Regret bounds for prediction problems. In Proceedings of the 12th Conference on Learning Theory, pp. 29 40, 1999.
Guo, J., Li, Z., Liu, X., Ma, K., Zheng, T., Yu, Z. Pan, D., LI, Y., Liu, R., Wang, Y., Guo, S., Qu, X., Yue, X., Zhang, G., Chen, W., and Fu, J. Codeeditorbench: Evaluating code editing capability of large language models, 2024.
10

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
He, Z., Zhong, Z., Cai, T., Lee, J., and He, D. REST: Retrieval-basee speculative decoding. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume I: Long Papers), pp- 1582-1595. Association for Computational Linguistics, June 2024.
Li, Y., Wei, F., Zhang, C., and Zhang, H. EAGLE-2: Faster inference of language models with dynamic draft trees. In Empirical Methods in Natural Language Processing, 2024b.
Liu, X., Hu, I.,, Bailis, P. D., Stoica, I., Deng, Z., Cheung, A., and Zhang, H. Online speculative de- coding. ArXiv, abs/2310.07177, 2023, URL https: 2 263835233.
Hou, Y., Tan, V., and Zhong. Z. Almost minimax optimal best arm identification in piecewise stationary linear ban- dits. Advances in Neural Information Processing Systems, 17:128967-122044 2024.
Liu, X., Daniel, C., Hu, L., Kwon, W. Li, Z., Mo, X., Che- ung, A., Deng. Z., Stoica, I., and Zhang, H. Optimizing speculative decoding for serving large language models using goodput. arXiv preprint arXiv:2406. 14066, 2024.
Hu, Y., Wang, K.. Zhang, J., Li, C., and Chen, H. Sam decoding: Speculative decoding via suffix automaton. arXiv preprint arXiv:2411. 2024.
Luo, H., Wei, C.-Y., Agarwal, A., and Langford, J. Ef- ficient contextual bandits in non-stationary worlds. In Proceedings of the 3 Conference On Learning Theory, volume 75 of Proceedings of Machine Learning Research, pp. 1739-1776 PMLR, 06-09 Jul 2018.
Hu, Z. and Huang, H. Accelerated speculative sampling based on tree monte carlo. In Forty-first International Conference on Machine Learning. 2024.
Huang, K., Guo, X., and Wang, M. Specdec++: Boost- ing speculative decoding via adaptive candidate lengths. arXiv preprint arXiv:2405.19715 2024.
Magureanu, S., Combes, R., and Proutiere, A. Lipschitz bandits: Regret lower bound and optimal algorithms. In Conference on Learning Theory, pp. 975-999, PMLR, 2014.
Kato, M. and Ariu, K. The role of contextual information in best arm identification, 2021.
Miao, X., Oliaro, G., Zhang, Z., Cheng, X., Wang, Z., Wong, R. Y., Chen, Z. Arfeen, D., Ab- hyankar, R., and Jia, Z. Specinfer: Accelerat- ing generative Ilm serving with speculative inference and token tree verification. ArXiv, 099781, 2023. URL https: /api. http:/appp....aaaa.e.ee.a bemanticcchoolee htg//ooppaaD:228740000
Khisti, A., Ebrahimi, M. R., Dbouk, H., Behboodi, A., Memisevic, R., and Louizos, C. Multi-draft speculative sampling: Canonical architectures and theoretical limits. In The Thirtentth International Conference on Learning Representations, 2025.
Kveton, B., Szepesvari, C., Vaswani, S., Wen, Z., Latti- more, T., and Ghavamzadeh, M. Garbage in, reward out: Bootstrapping exploration in multi-armed bandits. In Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp- 3601 3601-3610. PMLR, 09-15 Jun 2019.
Oliaro, Gi, Z., Campos, D., and Qiao, A. Suffixecoding: A model-free approach to speeding up large language model inference. arXiv preprint arXiv:241 0./24/1..4497. /.04975, 2024.
Robbins, H. Some aspects of the sequential design of exper- iments. Bulletin of the American Mathematical Sociery, 58:527-535, 1952.
Lai, T. L. and Robbins, H. Asymptotically efficient adaptive allocation rules. Advances in applied mathematics, 6(1): 4 4-22, 1985.
Russo, D., Roy, V., Kazerouni, A., and Osband, I. tu- torial thompson sampling. ArXiv, abs/1707.02038, 2017. URL https: / //api. semanticscholar. org/ ntg/CorppusDD:322991
Lattimore, T. and Szepesvári, C. Bandit algorithms. Cam- bridge University Press, 2020.
Leviathan, Y., Kalman, M., and Matias, Y. Fast inference from transformers via speculative decoding. In Inter- national Conference on Machine Learning, pp. 19274- 19286. PMLR, 2023.
Saxena, A. Prompt lookup decoding, November 2023,
URL https://github. com/ apoorvumang/ prompt-lookup-decodi
Sun, H., Chen, Z., Yang, X., Tian, Y., and Chen, B. Tri- force: Lossless acceleration of long sequence generation with hicrarchical speculative decoding. arXiv preprint arXiv:2404. /9/22 2024a.
Li, Y. Wei, F. Zhang, C., and Zhang, H. Eagle: Speculative sampling requires rethinking feature uncertainty. arXiv preprint XrXiv:2401. 15077, 2024a.

---

Yang, A., Yang. B., Hui, B., Zheng, B., Yu, B., Zhou, C., Li, C., Li, C., Liu, D., Huang, F., Dong, G., Wei, H., Lin, H., Tang, J., Wang, J., Yang, J., Tu, J., Zhang, J., Ma, J., Yang, J., Xu, J., Zhou, J., Bai, J., He, J., Lin, J., Dang, K., Lu, K., Chen, K., Yang. K., Li, M., Xue, BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
Sun, Z.,, Mendlovie, U., Leviathan, Y., Aharoni, A., Beirami, A., Ro, J. H., and Suresh, A. T. Block verification celerates speculative decoding. In Workshop on Efficient Systems for Foundation Models I/@ ICML2024, 2024b.
findings-acl.456. URL https: / /tc//accccac/ lnthhlloyh org/2024 (ttg/2224.finiimae fin.i4.fiiinnng-oee
456.
Sun, Z., Suresh, A. T., Ro, J. H., Beirami, A., Jain, H., Fast speculative decoding via optimal transport. Advances in Neural Information Processing Systems, 36, 2024c.
and Yu, F. Spectr:
M. Ni, N., Zhang, P., Wang, P., Peng, R., Men, R., Gao, R., Lin, R., Wang, S., Bai, S., Tan, S., Zhu, T., Li, T., Liu, T., Ge, W., Deng, X., Zhou, X., Ren, X., Zhang,
Taori, R., Gulrajani, I. Zhang. T., Dubois, Y., Li, X., Guestrin, C., Liang, P., and Hashimoto, T. B. Stanford alpaca: An instruction-folloiingg Ilama model. hitps:/ hempp:///itinn //gpttgttt//////// com/tatsu-Lab/ stanford_alpaca 2023.
X., Wei, X., Ren, X., Liu, X., Fan, Y., Yao, Y., Zhang, Y. Wan, Y., Chu, Y., Liu, Y., Cui, Z., Zhang, Z., Guo,
Z.,, and Fan, Z. Qwen2 technical report, 2024. URL org. abs// /2t../44...102 10671.
betps://arxio
Tian, R., Ye, Y., Qin, Y., Cong, X., Lin, Y., Liu, Z., and Sun, M. Debugbench: Evaluating debugging capability of large language models, 2024.
Yin, M., Chen, M., Huang, K., and Wang. M. A theoretical perspective for speculative decoding algorithm. arXiv preprint arXiv:241 1.00841, 2024.
Touvron, H., Martin, L., Stone, K. R., Albert, P., Almahairi,
Zhong, Z., Cheung, W. C., and Tan, V. Probabilistic sequen- tial shrinking: best arm identification algorithm for stochastic bandits with corruptions. In Proceedings of the 38th International Conference on Machine Learning, vol- ume 139 of Proceedings of Machine Learning Research, pp. 12772- PMLR, 18 24 Jul 2021.
A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D. M., Blecher, L.., Ferrer, C. C., Chen, M., Cucurull, G., Esiobu, D., Fernandes, J., Fu, J., Fu, W., Fuller, B., Gao, C., Goswami, V., Goyal, N., Hartshorn, A. S., Hosseini. S., Hou, R., Inan, H., Kardas, M., Kerkez, V., Khabsa, M., Kloumann, I. M., Korenev, A. V., Koura, S.S Lachaux, M.-A., Lavril, T., Lee, J., Liskovich, D., Lu, Y., Mao, Y., Martinet, X., Mihaylov, T., Mishra, P., Molybog, I., Nic, Y., Poulton, A., Reizenstein, J., Rungta, R., Saladi, K., Schelten, A., Silva, R., Smith, E. M.. Sub- ramanian, R., Tan, X., Tang, B., Taylor, R., Williams, A., Kuan, J. X., Xu, P., Yan, Z., Zarov, I., Zhang, Y., Fan, A., Kambadur, M. H. M.. Narang, S., Rodriguez, A., Stojnic, R., Edunov, S., and Scialom, T. Llama 2: Open founda- tion and fine-tuned chat models. ArXiv, 0hs/2307.092883 2023. URL https: wttppppp..iiissc...aeceeamoemaee Brg//oorpsII::2555 :5:995009:
Zhou, Y., Lyu, K., Rawat, A. S., Menon, A. K., Ros- tamizadeh, A., Kumar, S., Kagy, J.-F,, and Agarwal, R. Distillspec: Improving speculative decoding via knowl- edge distillation. In The Twelfth International Conference on Learning Representations, 2024.
Wan, R., Wei, H., Kveton, B., and Song, R. Multiplier bootstrap-based exploration. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 35444 35490. PMLR, 23 29 Jul 2023.
Wang, J.. Wang, J., Athiwaratkun, B., Zhang, C., and Zou, J. Mixture-of- agents enhances large language model capabilities. arXiv preprint S4Kv:24000046992 2024.
Xia, H., Yang, Z., Dong, Q., Wang. P., Li, Y., Ge. T., Liu, T., Li, W., and Sui, Z. Unlocking efficiency in large lan- guage model inference: A comprehensive survey of spec- ulative decoding. In Ku, L.-W., Martins, A., and Sriku- mar, V. (eds.), Findings of the Association for Computa- tional Linguistics ACL 2024, pp. 7655-766 7671,, Bangkok, Thailand and virtual meeting, August 2024. Association for Computational Linguistics. doi: 10. 18253/
12

---

BANDITSPEC:
Adaptive Speculative Decoding via Bandit Algorithms
A. Related Works
Speculative Decoding Speculative decoding is proposed in Leviathan et al. (2023); Chen et al. (2023), where the draft model only generates a single chain of draft tokens. Then a line of works extends the chain structure to the tree structure (Miao et al., 2023; Cai et al.. 2024; Du et al., 2024; Li al., 2024a; Hu & Huang, 2024). In these works, the draft tokens are organized as connected tree. To further improve the number of accepted tokens, previous works propose to generate tokens in batch manner, i.e., the draft tokens are organized as multiple disconnected parts. SpecTr (Sun et al., 2024c) views this problem from the optimal transport perspective and derives the algorithm that is optimal up to a multiplicative constant. Khisti et al. (2025) derives the canonical form of this problem and design the relaxed optimization algorithms. All these algorithms verify the draft tokens in a token by token manner. Sun et al. (2024b) proposes to verify all the draft tokens as a whole block, which further boosts the acceleration ratio. Sun et al. (2024a) proposes to fit the speculative deciding into hierarchical structure, which multiple draft models with various sizes are generating and verifying tokens. The smaller model will generate more tokens. This fine- ine-grained behavior improve the overall performance of the system. Liu et al. (2023) design algorithms to update the draft model parameters in online manner, which makes the draft model adaptive to the current context. Liu et al. (2024) and Huang et al. (2024) aim to optimize the speculative length in a training and training-free manner (more discussions on SpecDec++ (Huang et al., 2024) are provided in Appendix B.4). Chen et al. (2024) optimizes the hyperparameters related to the hardware by dynamic programming in an offline manner. We also note that there are series of non--orameen can-aaraetric speculative decoding algorithms (Hu et al., 2024; Oliaro et al., 2024), i.e., the draft model itself does not require any training procedures. Yin et al. (2024) derives the theoretical analysis of the speculative decoding.
Multi-Armed Bandit The multi-armed bandit problem is fundamental topic in decision theory and reinforcement learning, with various algorithms developed to address the exploration-exxloitaaioo trade-off. The standard stochastic K K-armed bandit problem is firstly introduced by Robbins (1952) and then studied by Lai & Robbins (1985). There has been a major theoretical advancement with the introduction of Upper Confidence Bound (UCB) algorithms (Auer et al., 2002b). Various algorithms have been proposed to achieve improved theoretical guarantees and practical performance since then (Garivier & Cappé, 2011: Bubeck et al., 2012). Beyond UCB-type algorithms, sampling- Sampling (Agrawal & Goyal, 2012; 2017; Russo et al., 2017) and sampling via bootstrap (K et al., 2019; Wan et al., 2023), have also exhibited strong empirical performance with provable regret bounds. Furthermore, the problem has been extended to the adversarial settings where the rewards are no longer stochastic (Auer et al., 2002a; Bubeck et al., 2012). We refer to Lattimore& & Szepesvári (2020) for a comprehensive introduction of the Multi-Armer Bandit problem.
algorithms, such as Thompson
B. Additional Discussions and Remarks
B.I. Discussion on the Assumptions
On the theoretical side, the stationary mean assumption (Assumption 4.1) is strictly weaker than the i.i.d. assumption. In particular, the number of accepted tokens can depend on the generated tokens. Therefore, this assumption is aligned with real-world scenarios in which different decoding steps are correlated. Furthermore, the basic Multi-Armed Bandits (MAB) model can be generalized to contextual bandits and non-stationary bandits. The proposed BANDITSPEC framework provides a basic template to apply these more general MAB setups to speculative decoding. Our formulations under the stationary/adversariall mean assumptions are just basic setups and we leave the more general/elaborate setups as future research.
On the experimental side, our experimental results (Table 1) indicate, the performance of UCBSPEC significantly outperforms one of the best speculative decoding methods, Eagle-2 (Li et al., 2024). This corroborates the stationary mean assumption in our formulation.
B.2. Discussion UCBSPEC
We comment that UCBSPEC is among the simplest UCB-type algorithms in the sense that only the empirical means and UCB's need to be maintained, and the hyperparameter to be selected can be directly determined via the UCB's.
In contrast, Thompson Sampling (Agrawal & Goyal, 2012) and KL-UCB (Garivier & Cappé, 2011) generally achieve better empirical regret bounds than UCB1 (Auer et al., 2002b). However, Thompson Sampling requires maintaining the posterior distribution and sampling from it to select the arm to pull; whereas KL-UCB involves solving an optimization problem for arm selection. These steps add additional complexity to the algorithms.

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
Given our goal of accelerating LLM inference, the simplicity of UCBI I is more in line with this objective. Therefore, propose UCBSPEC, which redesigns the confidence radius and the stopping rule of UCBI adapt specifically to the speculative decoding application.
B.3. Discussion Adaptive Adversary
We consider the oblivious adversary in this paper where the numbers of accepted tokens generated by all hyperparameters all time steps, i... (,1,))ejK))ee)) are fixed before the decoding process starts. One may be interested in considering the adapative adversary, where the environment (adversary) can choose the number of accepted tokens generated by S₁, based on (part of) the history information H4-1 and / (Arora et al., 2012). This adversary is more malicious than the oblivious one and the regret is expected to be even larger than the current one in Theorem 5.3. As our empirical experiments suggest that the practical scenario aligns more closely with the stochastic reward assumption (Assumption 4. 4.1) and deviates from the oblivious adversarial reward assumption (Assumption 5.1), we believe it is unnecessary to consider the adaptive adversarial reward.
B.4. Discussion on SpecDec++ (Huang al., 2024)
We compare the proposed methods with an existing adaptive speculative decoding method, SpecDec++ (Huang et al., 2024), in this section.
SpecDec++ (Huang et al., 2024) is adaptive in choosing the speculation length, achieving good performance compared to the vanilla speculative decoding method (Leviathan et al., 2023; Chen et al., 2023). It trains an acceptance probability prediction head and stops drafting new tokens when the predicted rejection probability reaches certain threshold.
We compare it with the proposed methods as follows:
Firstly, we highlight that our proposed method training-free which can be deployed easily along with existing off-the-shelf methods. In contrast, SpecDec++ focuses on training of an acceptance prediction head. Currently, SpecDec++ is only available when using LLMA-22Chht- as the draft model and LLaMA-2- CAAMAA--Chht-7000 ChBt-70B the target model (bfloat 16).
Secondly, the proposed BANDITSPEC framework considers the more general hyperparameter selection problem that goes beyond merely the speculation length. Therefore, "orthogonal" to SpecDec++ in the sense that any methods with (or without) SpeeDec++ can also be candidates for the hyperparameter in our framework, e.g. {Eagle-2, LLLMMA-2-Chat-7B] with SpecDec++ can also be regarded arms (if they available).
C. Additional Details
In this section, we provide more details that complement the main paper.
C.1. Vanilla Speculative Decoding
For completeness, we present and describe the vanilla speculative decoding algorithm (Leviathan et al., 2023; Chen et al., 2023) in Algorithm 6 in this section.
We introduce some notations first. For any nonnegative function f X' - R4, with Ezex f(x) > 0, we define the distribution induced by it as Norm(f(-)) Fr)Erexf(e f(x²). The positive part of function is denoted as [f())]+ max(0, f()))
Speculative decoding implements draft model Q to generate draft tokens and let the target model P verify them in parallel. In practice, the draft model is much smaller than the target model. Thus, the draft token generation (Line 1) can be achieved in a short time. Then we let the target model only forward inference once with these draft tokens as inputs (Line 2). The verification procedures (Lines to 9) are designed to guarantee that the output tokens 21. is distributed as the target model P. Here, the additional (r + 1)-st accepted token is also called the bonus token.
2

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
Algorithm 6 Vanilla Speculative Decoding
Inputs: base model P, draft model Q, prefix pt, maximum speculation length L
Procedures:
I: Generate L draft tokens æ1.z via i; ~ Q(- pt, 215-1) for E [Z].
2: Set r = 0, and calculate the values of P(ū; pl, for E [L] in parallel.
3: for == l....... do
4: Sample r; (n 1]).
5: rf min PT;; pt, 1 pt, (1=11) then
6: Set and 3; 5z.
7: else
8:
24 1/ 21--))))
9: break.
10: end if
11: end for
if L then sample #L+1 P(: (pt,
13:
return F1.++1-
C.2. Dynamics of MAB
We provide description of the dynamics of MAB in Algorithm 7.
Algorithm/ 7 Dynamics of MAB
Inputs: K arms, time horizon T.
1: Hu = 00.
2: for = 1,1,2 1, 2,, T do
3: Agent adopts an algorithm select 4e based on H4-1-
4: Environment reveals the reward X ft, the agent.
5: Hi concat(H4--) (I, x).
6: end for
The goal minimize the cumulative regret
T
max (E[K]
X1,
X1,1
where the expectation taken w.r.t the randomness in the rewards (for the stochastic setup) and the possible internal randomness in the arm selection algorithm.
C.3. Full Description of BANDITSPEC with UCBSPEC
We provide the full description BANDITSPEC(ALG) with ALG UCBSPEC in Algorithm
3

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
Algorithm 8 BANDITSPEC(UCBSPEC) (Full version of UCBSPEC)
Inputs: initial prompt pig pt X", bandit configuration (P,S (8.) L).
Procedures:
1: Hy 0, Lu 1, 1o.00 0.
2: while F.OS ot,,t do
3:
f
4: if S K then
5:
I = -Round-Robin)
6: else
7: Select index 1 = argiax UCB.,-1
8:
end if
9:
Observe Xt, = F1,1 Sp L) and Yi. yla len(X),,)).
10:
pt, concat(pt,-1 x1,,), H4 concat(He-1 (Ie, ()).)).
11:
Update the statistics (ūr.))e(k)) {crG,)eE]K]. where
Σ. Y Hcal(1a
n4,t =
Σ
i1s =1, Prg
ni,t
L
1 + Riit
K1 + Mi,e)
CI₁ =
2
n₂.
1 + 2oo
6
UCB, = Pitt + CTi,
12: end while return 4, pt,
C.4. Full Description of BANDITSPEC with EXP3SPEC
We provide the full description BANDITSPEC(ALG) with ALG EXP3SPEC Algorithm 9.
Algorithm 9 BANDITSPE((EPP3SPE)) (Full version EXP3SPEC)
Inputs: initial prompt pla pt E X", bandit configuration (P,S = (S:heiK L), learning rates f/:
T-K -T*
N.
Procedures:
1: = 0, Hn = 0, In = 1, 1 = 0.
2: while EOS It,tt do
3: i -=6 --1
4: Set probability vector ps A [K] with
- s1-1 2i.0
exp
i4 [K].
Pts
Σx= exp
(17t 5 21.0)
5: Select hyperparameter index I pt-
6: Observe X F1,t SPEEDDECSUB((LL L) and yr.t len(X
7: pt, (pt Xx,,), H4 = concat (H4-1, (Jr, X),,)))
8: Set the statistics
L 1 Aort
Zu 1 = 1.}
Vi E [K].
L PL.i
9: end while
10: return pl,

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
D. Proofs of Main Results
D.I. Proof Proposition 3.1
To prove Proposition 3.1, we note that we only need to prove
(LL d 7. and len(PEST(ALA)
(ALG)) I (en(PfiST(Aa)): a.s.
L+1
The other results are implied by these two results. For equality, we note that this is already proved by Theorem in Yin et al. (2024), where the equality holds for any specification of the hyperparameters. For the inequality, we note that each implementation of SPecDecSUB generates least one token and most tokens. Thus, the inequality holds almost surely.
D.2. Proof of Theorem 4.3
Theorem 4.3 (Upper Bound). Under Assumptions 4.1 and 4.2, given any prompt pt X and bandit configuration V (P, ((ihetk), L), the expected stopping time regret of Algorithm 3 with ALG Algorithm (UCBSPEC) is upper bounded as
Reg(ALG, pt, v) = O (H((p,, v) L log E [[len(pt,,]]]
Proof of Theorem 4.3. Our proof of Theorem 4.3 consists of three steps.
Reward and Stopping time decomposition.
Construction of the high probability event.
Concluding the proof.
As we mentioned in Section the main difference lies at the two aspects:
Firstly, the stopping time is now random, which depends the generated tokens. This cause trouble when decompose the reward/regret, as both the rewards and time horizon depend on the history. We tackle this problem in Step by making use of the martingale structure of the rewards sequence.
Secondly, we consider the problem under Assumption 4.1, where the number of accepted tokens can be dependent. This practical as LLM generates tokens an autoregressive manner. In contrast, under the commonly seen assumption for the K-armed MAB, the rewards are i.i.d. and can be regarded as they have been sampled before the algorithm starts (see Chapter 4 in Lattimore Szepesvári (2020)). Thus, cannot be used under Assumption 4.1. We solve this problem in Step 2, by adopting the so-called self nerf-nordd confidence bounds (Abbasi-yadkorf
bound (Lemma F:2) can be directly applied, which
al., 2011).
Step 1: Reward and Stopping time decomposition.
By the property of speculative decoding in (2) and (11), for any algorithm ALG,
ST(ALG)
E [len(pt.)] = E [leerptesraae))] =
We wish to decompose the expected total token sequence in terms of each hyperparameter i e [K] in the first step, i.c.,
ST(LG))
K = Σ ps E [IIITTAAL))
E
Y1.t
1=1
*We clarify that only the tokens up the EOS token be appended to prefix token sequence in practice. Therefore, the actual
number of accepted tokens is between [1, Y Yigguee ST(LL))] While this mismatch can be solved, using Y VisrrauyySST(uuo as the number of accepted tokens the last round will introduce an error most to the token sequence length len which is an error most the stopping time (ALL))). This error is negligible compared to the other values in the stopping time regret. Therefore, assume EOS token only appears at the accepted tokens for the sake of simplicity.
5

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
The standard regret analysis adopts Wald's equation to decompose the expected cumulative regret, or equivalently the stopping time. However, both the stopping time (ALL) and YI,, depends on the history under our problem setup, Wald's equation fails. We propose new and general approach to decompose the reward.
Step 1.1: We first prove that M, := Ei=1 Y1, MI, n = 0, 1, 2,... is martingale with respect to {Fr where
My == Fn == d(Hn).
By the definition of martingale, we only need to show (1) B|[MM.]] < oc, and (2) B|Ma+1 Mn.
(1) M,]] oct As the number of the accepted tokens at each round is bounded as Yi. [1, + 1] almost surely and μ1, E [1, + 1], we have [Y,,, -vs, S L. Then the triangular inequality |M I Ei= (i,e pu.l S L < o0 indicates that
M,ll < co.
(7)
(2) M Mn: The conditional expectation of M,+1 can be calculated via tower property
B F.. M,,, + E [Va41,,41 Mlata Hun [
Hn]
Ma
(8)
where the last equality results from Assumption 4.1.
Based on (7) and (8), Mn, nl 0, 2,2, ... is martingale with respect to (
Step 1.2: We then prove E [MST(AL = 0 via Doob's optional stopping lemma (Lemma F.1).
We have already showed that Σi= Yi,, i ,4,, 1,2.... is martingale with respect to (H, ) In order to apply Lemma F.1, we firstly verify the prerequisites listed in Lemma F.1 condition (b): (1) E[ST(ALG)) < co, and (2) there exists c E R. such that M, M4-1/| F(-1] I c almost surely for t $ ST(LGG).
(1) (ALG) oo: According to the property of speculative decoding (3) and Assumption 4.2, we have that [[ST[ALG)] Ellen(pt,,]] )] 20.
(2) E M M1-1l Fi-1l S As M, M Yı, Mr, and |Yr, Mr, S L almost surely, it holds that E|[,M; M-1 F:-1l & L. Taking e = L finishes the verification.
Therefore, condition (b) in Lemma F.1 is satisfied and we obtain
ST(ALG)
Yi.t per. O.
Step 1.3:
We show that E
===1
ST(AG))
Vieit]
B[nL,ST(ALO]-
We firstly note that
(TLL)
ST(ALG)
E
(L 1)E[ST(ALG)] oc and E
Σ
11.
S ()ESSAAA))) < co,
so the expectations of E Yi,t and Et=l μz, exist and are finite (integrable).
ST(ALG)
ST(LL))
Furthermore, we have
(ALG)
ST Σ
E
Σ
P1.
E
11 = H4 = Σ
μs E [N.SST(ALa]
(10)
1=1
t=1
d=1
where n4.ST(ALG) (10) imply
2tm1
ST(LG
1[I,
i) by definition. Because
ST(LLG) YI,, and _t=1 Pr. integrable, (9) and Etm1
ST(ALG)
STYALG) E
(ALG)
ST(AGG)
=
E
Σ
YI
H4,
+ E
Σ F1 = Σ Ms E [[ALST(AL)]
(11)
[=1
tm1
6

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
This equality decomposes the cumulative reward (and stopping time) in terms of the arms.
Step 2: Construction of the high probability event.
We then derive the concentration property for the number of accepted tokens. Define the good events:
& := (ii, E luu Cri.,, Mi + cra, Vi E [K] at round t).
Since random variables supported on (a, b] is (b- //2/24/2/// Sub--aaauaaaa and 1 pt, S Vscit pt, + Mr, under problem setup, According to Lemma F.3, we obtain
8
P(E)) S n2 8
P (EE) 2
and
(12)
t2
where 8 is a confidence parameter that will be specified later. We remark that Lemma F.3 from Abbasi-yadkori et al. (2011) adopts self-normalized concentration bound for the martingale sequence, which generalizes the standard i.i.d. reward assumption in the K-armed MAB problem.
As result, we can now bound the number of times arm i pulled any round 2 K. Conditional on the good event &,, we have Pi,r Jau Cri,t: i Py + craa and arm i will not be pulled if cri, A./2. By adopting Lemma F.4, when arm i selected at time t 1, must hold that
2L ni. 4+
LK t²
A 1 + 2 log
(13)
Step 3: Concluding the proof. According to Step 1,
(T(L)
E len(pt,.) ) Elee(p(ssTraa)) E
YIos Σ pi E (.sTTaa))
Therefore,
(ALG) =
Σ (E(K)
Ms E [TE.ST(AM]] + Σ
A; E [MGST(ALa]] 145- E [ST(GA))
<E[K]
El E[N.ST(ALL
(14)
P4-
Under the UCBSPEC algorithm, we have
ST(AGG)
1f 1, = Ri,e-1 54+
2L2
LKtz
}
E
1 + 2 log
A?
A
ST(ALG))
A
2L
LK(t 1)2
}}]
E
Σ
1 I i, Mi.11 > 4 +
1 2lg
+ K
A?
4,d
(=K+1
ST(ALC)
2L
LK (ST(ALG))
)
E
4 +
2 2log
Σ 1(8))
K
A7
A
1=
212
LK (E[ST(ALG)))),
4
(1 + 2 log
Σ
P())) + K,
Pl
AZ
A
[=K+1
212
LK (E[een(pt,,)))
))
22
4
(1 2log
+
+ K,
A
A
6

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
where the first inequality results from (13), the second inequality utilizes Jensen's 's inequality, and the last inequality adopts the property of speculative decoding (3) and the upper bound the error probability of good event (12) in Step 2. Taking 6 = 1/2 in the above bound concludes the proof of this theorem.
D.3. Proof of Theorem 5.3
Fix any [ the baseline algorithm is set to be ALG,, i.e., Algorithm 3 implements Line 4 of Algorithm 3 with hyperparameter Ss only. Let ALG EXP3SPEC. we assume the BANDITSPEC repeats the while loop in Algorithm 3 until max(ST ST(L)))) To avoid any confusion, we restate the algorithm for the purpose of analysis in Algo- rithm 10. Algorithm 10 takes ALG and ALG, as an input and stops until plsT(aLS) and plsT(aLes) are generated. The two token sequences up to EOS token are output at the end of the algorithm.
Algorithm 10 B.NDITSPEC(EXPPSSEC)) (For analysis purpose)
Inputs: initial prompt ptp = pt E X", speculative decoding configuration v/ (P,S = [Ei L), stopping time oc, baseline hyperparameter S., initial prompt Ptg pt, stopping time 7
Procedures:
0, Ho = 0, To = 1,400.0 = 0, =ho0
2: while r = 0o or 7 = oc do
3: / = -+1
4: // Procedures of the original EXP3SPEC
5: if oo and EOS E -1,-,t11 then
6: т t and pt, = pt.
7: end if
8: Set probability vector Ps A [ with
-1] E1-1 2n.0)
cxp
log K 1. K
Ptj
with learning rate T) =
Vi E [K].
(15)
Ej=1 exp (n1p E1=1 Z,,0)
9: Select hyperparameter index 1 ~ Pt.
10: Observe X1,t = F1,, SPEEDEESSUB(Ppt... P,Sr, L) and fr.t = len(X
pt, 1 H4 concat (H (Le, (),)))
12: Set the statistics
L + 1 Vitt
Z,, = 1li 1
Vi [K].
(16)
L Pt.i
13:
// Procedures of the baseline ALG, if n and EOS & 2t = then
14:
15:
b
and p = pt
end if
17: Observe X x SPECIDECSUR(PL2... P, I.) and wot Iee( (X,.).
18: ptş = cconcat(pt] Xh,,)
19:
end while
20:
return ST(LLG) T, piST(aL)) = pt, and ST(ALG.) = Th. piST(ALL")) = pt....
Theorem 5.3. Under Assumptions 4.2, 5.1 and 5.2, given any prompt pt X and bandit configuration (P,S
(St)i L), the expected stopping time regret of Algorithm 3 with ALG Algorithm 5 (EXP3SPEC).
Reg (AGG pt, v) 2L min
/Ien(pt...) log K,
2LK log K +
min ST(LG,,) K log K KE(K]
8

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
Algorithm 11 Dynamics of the OLO Problem with Full Information Feedback
1: Ho =0.
2: for 1 T........,T do
3: Selects Pt E [ based on H4-1-
4: Observes the loss vector & and suffers loss P. 4.
5: Hr = concal Concal(He-1) (pt,ft): Et)).
6:
end for
Proof of Theorem 5.3. ease of presentation, we present Algorithm 10, where the while loop in BANDITSPEC is repeated until mAL(ST(ALG), ( (TT,),))).
Our analysis of Algorithm 10 is novel compared the standard analysis of EXP3 algorithm (Auer et al., 2002a; Lattimore & Szepesvári, 2020). It requires more technical manipulations due to the fact that the termination times of the baseline algorithm ALG, and ALG are different and random, and that our goal is to minimize the stopping time regret (5).
For theoretical analysis, we regard Algorithm 10 as an instantiation of the Tollow-the-RggullriieeeL..add (FTRL) algo- rithm (Gordon, 1999 Lattimore & Szepesvári, 2020).
The proof is decomposed into steps:
Connection between FTRL and Algorithm 10: we firstly introduce FTRL and the problem where it is applicapable. The shared features and differences are highlighted.
Transformation of the stopping time regret: the stopping time regret is related to the regret under FTRL framework. In this case, the minimized regret by FTRL. can be translated to the stopping time regret.
Regret decomposition: the FTRL regret is decomposed for easier processing.
Upper bound each term in the decomposed regret: we upper bound each term in the decomposed regret. The main difficulty is to deal with the difference in the time scales (TLG) and (TLG) and the randomness in (ALG))). Specifically, because both the loss vectors (Zihee and the stopping time ST(ALG) are random, taking expectation of the cumulative loss within (TAALL,)),S(AALA)) is non-trivial. We devise Lemma D.1 and Lemma D.2 to deal with this problem.
Conclusion the stopping time regret: we aggregate the results in the previous steps and derive the final bound for the stopping time regret.
Step 1: Connection between FTRL and Algorithm 5.
We denote A [ as the K-dimensional probability simplex.
FTRL is often used in the Online Learning Optimization Problem (OLO). We firstly provide brief introduction to OLO that operates on A [K] Let E ..... E R be sequence of unknown loss vectors. The dynamics of OLO problem is stated in Algorithm 11. Given time horizon T N, the agent (or algorithm) aims to minimize the bass-based) regret
MAX PEAIK] [K] 1=1
Σ Pp. (ec-p)Ter T E..
(17)
where Pt is the action taken by the agent at time step t, is some fixed baseline action in A [K and the maximum operator indicates the agent is competing with the best fixed baseline. The FTRL algorithm minimizes (17) by taking the action Pt = Argmin (p) at time step i, where &, A + is defined as
I-1
2x(x) = F.(T) +
Σ Tl
lo
and A[K] R, Vt N, are some convex functions.

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
In the following, we illustrate the connection between FTRL and Algorithm 5 (or Algorithm 10). We let & 2 (Z,,..., 2k.)) FF)x) = f nh with F(x) A X]] + R and F(x) := Σi= (z; log Iz) +log K + Furthermore, some calculation indicates the action taken by FTRL is exactly Pt as in (15), i.e.,
P: argmin (p)))
|
Therefore, Algorithm is indeed an instantiation of FTRL in terms of the algorithm design.
Under our problem setup, the difference lies at the target of the algorithm. Instead of minimizing the corresponding regret
T Pp- - p) T 2
Σ
max PEA
we aim at minimizing the (loss-based) regret defined on two different time scales
ST(AG) Σ
STAALG,) Σ t=1
Reg (ALG) := E
P.
T 2
min E SE(K]
e 2
(18)
t=1
where 6 E R is an one-hot vector with the jth coordinate being 1, and the expectation taken w.r.t the internal randomness within ALG and Z. We highlight again that ST(AG random whereas ST(ALG;) fixed under the greedy decoding strategy.
Step 2: Transformation of the stopping time regret.
The stopping time regret (5) and Reg(ALG) in (18) may look different at first sight. We demonstrate that these two notions of regret be transformed into another up some constant factors.
We firstly simplify Reg(ALG). Let Zi,t
(3i,
L)/L. According to the definition of pt in (15) and Z, in (16),
(TGA
Σ
T 2
ST(AG)) Σ
L+1
Vi
L+1
STT(AGG)
len(ptsT(MLa)
L
L
Furthermore, note that 2, is an unbiased estimator for 2ı and (T(L)G;) is a fixed real number,
ST(ALG,)
Σ
eT2,
e
ST(ALG.)
Σ
21, =
L+1
ST(ALG;)) len(ptsTTALa.)
(19)
L
1=1
1=1
Hence, Reg(ALG) is simplified as
L 1
len(ptisT(ALa)
L+1 (
len len(ptsrrata.))
(eG(A IE
ST(AGG)
min iE(K]
STAALG)))
T,
L
T.
+
(E[ST(ALC)] - min ST(ALG;))
L
>E[K]
where we adopt the property of speculative decoding (2) in the last equality. This indicates that
+ 1
Reg(ALG)
Reg ALG).
(20)
Step Regret decomposition.
Based on the previous two steps, we now upper bound Beg(A This notion regret distinguishes from the standard regret analysis due the difference the time scales ST(L)G) and ST(ALG,).
10

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
For simplicity, we define the Bregman divergence induced by convex function / - [ with () = f(a) f(b) - (a b) T (bb. Fix i E [K], we now decompose this empirical regret w.r.t.
R+ as Dy(: [KK A [K] +
STYALG)
AAG,)
Pe 2
Σ
Zz.
1=1
ST(ALG)
(AT(A
(ALO))
Σ
PN+1) T2,)
Σ
Pit1 2.
Σ
E 2x
+
STOLG)
ST(AGG)
Σ
(pr
1D-+P41) Z1,)
Σ
(Pt1 Feei((et)) (x(((11) F.(Pet1)
ST(ALG)
(ALG)))
ST(AG..)
Σ
e Z
Σ
0,
+
Σ
c, 21)
1=1
E=1
1=1
ST(LLG)
(T
(pr -O--P4+1 Po++) T2,
+
Σ
(P41((E41) - (P+4()))2)
ST(ALG)
(F.(pe))) P(p1(pe11)
-FsT Fi
ST(ALG)
ST(ALG.)
( (PST(ALA)+1) - (()( +
Σ
e, T,
21
c c
E=
ST(LE)
- (pt-P(+2) Dr(peti,Pe)
+ FST(AAL)++ - F111
(0)
(0)
ST(G
(FrlDem)) Fet1()e + (ST (PST(ALA)+ (ST((LA)+1 (e:)
(t)
ST(A
ST(A.C;)
Σ
6 2.
Σ cT
20
(21)
(1)
where the last inequality adopts the fact that Do, (a, (a, = DD. (a,b) = Dr(a, b)/1 and the inequality
(p.) (p(+1)) = -Da, (p++1,P)) (pt+1 - pp:) (p:) S -Do, -Do,, (p)+1,P)).
Here (p++1 - pt) (p:) S 0 results from the choice of pt = argmin
p prgmin px(p) and the first-order optimization
condition.
Step 4: Upper bound each term in the decomposed regret.
In this step, upper bound each term in (21). We comment that (D) and (9) require to attend to the randomness in (T(ALG) and the different time indeces. This issue will not be encountered in the conventional scenario where ST(ALG) (A(L).
Upper bound (O): will show that B((D) < E E ST(ALG) n//2] almost surely.
Recall the definition of 2, in (6), 2, only has a positive value at the I, coordinate. We divide the problem into two cases: (1) Pt. PLt1, < and (2) Pt.r, PL41,1, 2

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
Case (1): Pt. Ptt1. < 0. Because the Bregman divergence is always non-negative, hence,
(0) S (Pt.), PP+111)) Zr,, 00 0.05
2pr.
Case (2): Pt,l, 1,+ 2 0. Note that F(x) is Legendre function on A (K)- By invoking Lemma F.8,
(pe Pc+1)) T Zz Dr(p(+1.Pt)
Th Ze
1/t 2 H
where H, = (2)((a) and 4: Pr + (1 - (): Pt+1 for some E [0, 1]. Furthermore, V²F(2) is K x K diagonal matrix with (V2F(g)),,, = 1/qr,,, vi E [K]. Therefore,
1
Ieat Pi.
T1 1
4t
S
Pt.l,
2
Pi.1
2pr,
To conclude the two cases, holds almost surely that
ST(AGG)
(0)
S
Σ
Me
2pr.
Lastly, by adopting Lemma D.I which is proved in Appendix E.I, we obtain
(T(C)
K
E
[[)]]
S
E
Σ
718
Lemma D.1. Under Assumption 4.2, consider the learning rates T) defined in Algorithm 5 (or Algorithm 10), it holds that
ST(LG))
ST(AGG)
E
K
E
Σ
Pt.ee
Upper bound (o): Because F(z) is non non-negative on A [K] and F(ei)) = log K, hence,
F(es)
log K
log K TAST(A
([ST(LG+11)
VIST(AG)+11
where we manually set +)ST( = 7TT((A)))
Upper bound (t): Recall that F(x) Ex (zx (n logæ; x;) + log K + 1, F(x) is non-negative for any æ E A[K]
Additionally, f)- = / /og K/(K t),t E [ALG))))) is a decreasing sequence. Therefore,
(T(L Σ
F(Pet1)
F(pe+1)) S 0, a.s.
(1)
T/
Mj++1
Upper bound (3): Since +1 is the minimizer of ()T())))1IP)) we have (t) S 0.
Upper bound (1) Under Assumption 5.2, ST(ALG) 2 ST(()); almost surely. We further prove that E{())
<
[ALGG) - STAL,,,) for = i", which we summarize in the following lemma with proof postponed to Appendix E.
Lemma D.2. Under Assumption 4.2 and 5.2, consider Z, defined in Algorithm 5, it holds that
ST Σ
E
T 7
[[ST(ALG)) - ST
SAST
+1
12

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
Step 5: Conclusion of the stopping time regret. Aggregating the upper bounds for each terms in (21) and taking expectation,
(ALG))
K
log K (SST(AL))
Reg (LG))
E
Ne
+
+ +[[ST(ALG)] ST(ALG;-).
2
By substituting the learning rate values into the equation,
Reg(ALG) < 2 V /ST(ALG). log K
[[[T(AAL)) - ST(LL,,-)
< 2 v [E[ST(ALG]] K log K E[ST(ALG)) S
2 /len(pl..) K log K + E[ST(ALG)) ST(ALG,)))
where the second inequality adopts Jensen's 's incquality and the last equality holds due to ST(ALG) & len(pt,)) almost surely as in (3), According to the regret transformation (20),
Reg(ALG) S 2L
(TLGG) log K S 2L / Vlen(pt..) K log K.
(22)
Furthermore, by solving the quadratic function in terms of (ALG), i. i.e.,
Reg(ALG) < 2L V (Reg(AAGG) ST(ALG;-)) K log
we obtain
Reg((AG) 4L2 K log K 2L
min ST(G)))) K log K.
(23)
kE[K]
Aggregating (22) and (23) concludes the proof of this theorem.
Remark D.3 (Sampling Decoding under Adversarial Mean Values). Since the tokens can be regarded as fixed given the initial prompt and the hyperparameter configurations under the greedy decoding strategy, we consider the greedy decoding strategy under the adversarial mean values assumption (Assumption 5.1). If one wishes consider the sampling decoding strategy, the proof of Theorem 5.3 can be adapted to it. Specifically, this switch of decoding strategy mainly influences (19), the proofs of Lemma D.I and Lemma D.2. We can depend on Doob's 's Optional Stopping Theorem (Lemma F.1) solve this problem, just like what we have done to prove (28) and replacing the condition (1) therein by [ALG( < E [plen(pt,,)) < co. The rest of the proof can go through in a similar manner. In the end, we can arrive at a similar result, iee.
Reg(ALG, v) S 2L min 2LK log K
min ((ST(AL,;)) log K, /Bleen(pt,,)] log K (E,K)
D.4. Proof of Theorem 4.4
Under the greedy decoding strategy, the problem is alleviated in two aspects. Firstly, given the target model P and the initial prompt pt. the total length len(pt,,) is (potentially) determined. While the total length is determined, it is worth noting that the number of accepted tokens at each round (Line 5 in Algorithm 3) is still random. Additionally, under the dynamics presented in Algorithm and given the history H4, there is a one-to-one mapping between the accepted tokens X1.t and its length Yi,t-
Since the lower bound is established in terms of class of algorithms over set of initial prompts, we adopt X to denote the set of initial prompts and adopt San to denote the set of all hyperparameter specifications that can be selected to constitute S. To further ease the problem, augment Assumption 4.
Assumption D.4. We assume that Given any bandit configuration = (P,S = (Stherk L) and initial prompt pt E Vi conditional on the history H4-1 and the selected arm I, at round t, the distribution of the length of the accepted tokens PP pt, H4-1, = i) =
(s.(.), Vi E [K].
For any two hyperparameter specifications S, E Sall, we have KI.(Ps, Psr) < co.
13

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
We consider the class of arm selection algorithms which non-anticipatory and consistent.
Definition D.5 (Non-anticipartory Algorithm). An arm selection algorithm ALG is non-anticiputory if ALG(- | (Hi), Vt E N.
Hi) e
Definition D.6 (Consistent Algoirthm). An arm selection algorithm ALG consistent over a class of bandit configurations
and prompt set initiol if for all v A and any sequence of initial prompts (pt X Amitat with len(pt + 00, m oo, and for all (0, 1].
RALG pt"", v)
lim
O.
m2 700
len(pi (1p )
)==1 C Xint with len() - oc, + xo and a bandit configuration v (P.S = {Si (E(K) L), under Assumption D.4, the greedy decoding strategy and the dynamics represented in Algorithm 3, for any non-anticipatory and consistent arm selection algorithm ALG, the expected regret satisfies
Theorem 4.4 (Lower Bound). Given any sequence of initial prompts (pt
Reg(ALG, pt,
lim inf
m0000 log( (pt())
kl.
where kl, := ( Ex~ps[A] > 3.
Proof of Theorem 4.4. The proof consists of three steps:
Divergence decomposition: similar to the reward decomposition in the upper bound proof, the divergence decomposition cannot be done as the time horizon (ALG)) is a random stopping time. tackle this problem in the first step.
Lower bound on Ey, (+ adapt the standard trick lower bound the expected number of times arm has been chosen.
Conclusion of the proof.
Step Divergence decomposition. The divergence decomposition suffers from the same issue as the reward decomposition, ie.. the stopping time (ALG)) depends on the history. We adopt the same trick as in the reward decomposition step to overcome this issue and the result is summarized in Lemma D.7 whose proof is postponed to App. E.
Lemma D.7. Under Assumption D.4, given two bandit configurations u = (P,S = (S L) and o' = (P,S' = (5)) L) which only differ the hyperparameter specifications, for any pt E
X initial and algorithm ALG,
(L PALLL,,,t,,,,,,,p..... PALG. Ea
Ps;)
where PL.c,pt.v (resp. Ptta,pt.)) is the probability measure induced by (ALG, vt, (resp. (ALG, vt, defined on the a-algebra (HH))
Step 2: Establishment for the lower bound of E,, (TLA))) Given algorithm ALG, bandit configuration E A, prompt
pt (pr and any 0, construct K alternative bandit configurations v; = (P.S: {St,) K L) for i # it with
S1, S, S. 11j # + S 1(j = 1).
where S² in satisfies that its mean ps; > with (L Ps:) - kl, + In other words, under bandit configuration v, only S, changes into S and arm becomes the best arm. As the bandit configurations only differ in the hyperparameter selection, we adopt the shorthand notation P,, E, and Reg(v) )) for Pt.c, EaLo,pt.r and Reg(ALG, pt, v) respectively when there is no risk of confusion.
14

---

E, (T,LL)) BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
According to Lemma D.7, for any pt E (p("))=11 m==1
[L,Fe,Frl = E (ASTATT(ALa)) KL(Ps,, Ps/) S EE (MA,ST(AL)) (kl, + te).
By Lemma F.6, with A, (n,,ST(ALO) > (ALL)))
P.[A,] + P., [A ス 2 exp ( Ey (mn,S(((( [L(Ps.. Pst))
(24)
exp
Ry (M,,T((LL) (kl, + E))
Thus, by adopting (14), the expected reward under y can be lower bounded as
A
2 A
Reg(v) 2
By (n
Ey (TI.SSST(ALL)
|A. PP[A
(25)
P4-
A 1 (ST((A)) A. Po|A]] 2
A; len(pt
M4 2(L + Fr[A.))
Mr 2
Similarly, under the alternative bandit configuration Vi.
Us. Ma MSt
Ms p4-
len(pt,.)
Reg(v:)= 2
Eeg (M,,ST(AUU))
2
Pers [A.]
(26)
1s,
2(L +
Combining (24), (25) and (26).
Reg(v) + Reg(u:)
Ai US, - Mi-
len(pt,e)
min
(Pr[As] + (.,(4)))
Ui.
PS;
2(L + 1)
A:
Ms; pi*
len(pt,))
2 min
exp ( ( EE.lnAsT((LO)) (kl, + E))
ps- As,
4(L + 1)
which holds for any pt E (p Rearranging the terms, have
Eg (
lim inf
(pt(
log (min A,,
Meg- }) log (4(I.+ - log(Reg(u) + Reg(u:)
1
14. Ast (kl; + 8) (og(len(pt..))
lim inf
kl, +
m0+00
1
kl, +
where the last equality follows from the definition of consistent algorithm. Because >> 0 is arbitrarily chosen, by sending -> 0, we obtain the lower bound for E,, (TAA))(
lim inf ()
2 kl,
(27)
Step Conclusion of the proof. Aggregating (14) and (27).
Reg(ALG, pt v)
lim inf
m00
log(len(pt..))
kl
This concludes the proof.
15

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
E. Supporting Propositions
E.1. Supporting Lemmas for Theorem 5.3
Lemma D.1. Under Assumption 4.2, consider the learning rates тι defined Algorithm 5 (or Algorithm 10), it holds that
ST(A
STL)
E
K
E
Σ
Pt,
Proof of Lemma D.I. Similar to the proof of Lemma D.2, we adopt Doob's Optional Stopping Theorem (Lemma F.1) to show
(TTCL Σ
Tle
7/
0.
(28)
Pt.
According to condition (b) in Lemma F.1, we only need to show that (1) E[ST(ALG)) < oc, and (2) there exists e E RR such that for all < ST(ALG), E[]n/Pr. - mel 1[ <.
Condition (1): Given Assumption 4.2, it holds that len(pt,.) oo under the greedy decoding strategy. Therefore, we have [[ST[ALG) len 0o.
Condition (2): Note that
K
E Pt.
K
mil
H1-1
< E Ne
H4-1 + K ne S
E
1(4
i} Ho-1 K m
Pt.le
Pt.
1=1
$ 2K S 2 K log K.
Therefore, Condition (2) is satisfied and (28) is established. In addition, by using Assumption 4.2,
STT(LG)) Σ
K
S (TT / K Klog K x0.
So E STTALO)
nh] exists. Lastly,
(TIG
ST(LG))
STT(LGG)
E
K
17t
+ E
Pt..
Pt.e
which indicates E> t=1
(TLA)
T/. .c..pr..] exists.
In conclusion, by adding K D
ST(ALG)
ne] on both sides of (28), the desired result is obtained.
Lemma D.2. Under Assumption 4.2 and 5.2, consider Z, defined in Algorithm 5, it holds that
(TAL Σ
E
e 2
[[[ST(ALG) ST(ALG.-).
SALG
)+1
Proof of Lemma D)2. the two expectations below exist and
ST
(T(G)
Σ
e
2
E
Σ
Zz,t
(29)
16

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
then it holds that
(TTL
(ALG))
(ALC)
ST(ALG.))
E
21
EE
Σ
2
E
Σ
zit
Σ
(ALG)
E
Zi,t
-[ST(AAG))]
(TTLLG)))))SAAA)))))
ST
The desired result can be obtained.
Therefore, we aim to prove (29). Since both ST(ALG) and 2, are random, the obstacle is that we cannot directly take expectation of the summand. We will firstly prove
(TLL) Σ
E
4 2
Zi,t
=
(30)
by Doob's Optional Stopping Theorem (Lemma F.I). According to condition (b) in Lemma F.I, we only need to show that (1) B[ST(ALG)) oc, and (2) there exists I such that for all f ST B[le, (Z, 21] H4-1] <
Condition (1) holds as shown the proof Lemma D.I. We only need to show Condition (2). Note that B|2] = 21.
E[le] (2: 10---,] =:] H4-1] - E B 2, 94-1. T4]
He-1 + e. 22g,t S 2.
Therefore,
ST(AGG) =t=1 e, 7, Zi,t is well-defined and (30) is established.
We then prove the two expectations exist. Note that all involved variables are positive, we only need to show
(T(LGG)
E
Zizt
00.
This can be obtained by noticing that 2 E [0, 1] for t E N.
ST(L)G)
24,1
[[S(AAL))]
<0.
Combining the above with Condition 2, have
(T(A)))
ST(ALG)
ST(ALG)
E
Σ
eT
=E
Σ
T
TTa zt
+ E
Σ 1=1
cT
t=1
t=1
which means E [> In conclusion, by adding E [[>ST(M
(TLG)))
E; 2.] exists.
1t=
(TLG))
zi,e] on both sides of (30),
(T(A
(A(L)))
E
Σ
E
2
E
Σ
Zi,1
This finishes the proof of (29).
E.2. Proof Lemma D.7
Lemma D.7. Under Assumption D.4, given two bandit configurations (P,S {S.) L) and v' (P,S'
(S L) which only differ the hyperparameter specifications, for any pt E Nimitiall and algorithm ALG,
K
(L PALG,pe,rr) =
Σ
EaLa,tt Ensapt.wli..TT/AA KLL Ps:
17

---

Psn. (Yi,,) BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
where P.as,pt. (resp. (LL..pt.)) is the probability measure induced by (ALG, (resp. (ALG, pl, v)) defined the a-algebra {a(H)))irr
Proof. As the bandit configurations only differ in the hyperparameter specifications, we adopt the abbreviated notation Py and Par for the induced probability Pttc,pt,rr and pt..,pt, respectively. We use PAta()) to denote the output distribution of the arm selection algorithm ALG in Line 4 in Algorithm 3.
With the bandit configuration // = (P, L). the probability of HST(ALO) is
(ALG)
STTA
(A
II
PALO( Pt)PP. It)
II ==1
Puc 14-1. (Y,,))
t=1
Similarly, under the bandit configuration v' = (P.S',
(T(L)
Par (HST(ALA))
= II t=1
PaLalLe Ht) (Y1,)).
Therefore, holds that
Pe (HSSTAAA)))
ST(AGG)
Psr, (Yrea)
log
log Pst (Yra.)
Per (HST(ALO))
Because KL(Ps,, Ps;) < 00, Vi E [K] under Assumption D.4, the divergence between IP, and PPr can be rewritten as
STYALG))
Pr (HST(AAA)))
KL(P,,, PP)) E, log P.. (HST(A))))
E,
log
< 20.
Pst (Y
E=
We then prove that
ST(ALG)
Ps.. (Yr,,)
E,
log
Σ
E (KTT Pss)
(31)
Psy. (Y1,))
The proof is composed by two arguments:
Argument 1:
ST(A
Psn. (Y1.2) llog Ps; (Yr,z)
E
KL(Pss, Pss 0.
(32)
Argument 2:
(T(L)))
E
KL(Pst. Pst
(KlTsTTuLa)KKL(PPS Ps:).
(33)
If the arguments true, by summing up (32) and (33), obtain the desired result (31).
We prove the two above Arguments.
Argument 1. Let
Psi, (X1.))
M, :
log
, 1.2....
psi (X1ea)
18

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
and My := O.
We firstly prove that (M is martingale w.r.t (Hx)m-ui (1) [I),I] < cc, and (2) E. [Mx+1 Hn] Mn.
(1) I, < oc. According to Assumption D.4, for any i E [K]. KL (ss. Pst) < co, this indicates there exists E R such that
Ps. (z) log (x) (x)
<c<00, [K]
(34)
This indicates
psi (X + KL(Ps Psi >1 <oc.
psi (X1rt)
E, M, -
Ey
log
(35)
[=
(2) Ey [M++1 Hn] = Mn By adopting the tower property.
Psints (Y (11++1m+1)
Ey Mxx11
[m]] = Mn + Ey
Ep
log
(T(PSt St,41 Psyn12 ) Hos Tnt1|
Hn
Psi (V (1141m41)
= M
(36)
From (35) and (36), (Mr) is martingale w.r.t. (Hn)nEN- Additionally, we will adopt Doob's Optional Stopping Theorem (Lemma F.1) on MST(ALo)- The prerequisites are verified as follows:
(1) ([LG() < oct By Assumption 4.2, (TLG))) is stopping time w.r.t. (Hn)neN with E|ST(ALG)) S len(pt < co.
(2) there exists ē e IR, such that E|[ Mn. H.l < for any n ST(ALG): According (34),
Ps., (Xi,t)
EIJM,+11 - Mnl Hn] S Ev log Psy (X1,,)
Hn + Ev KL(Ps,, (si.s..Pssi)))
2e < oo.
Taking ε 2c finishes the verification.
Therefore, the prerequisites in Lemma F.1 (b) are satisfied. By invoking Lemma F.1, (32) is established. Argument 2. Firstly, by (34),
ST
||
E, Σ
KL(Ps., Ps,
s S. [ST(AL))] :c oc.
ST(ALG) So Et=1
(L.(Ps,, Ps₁,) has finite expectation. Furthermore, it can be observed that
(T(A
(T(L)))
E>
KL(Ps PsS,)
E,
114 = ()sL(Ps..ss)
E,
Therefore, (33) is proved.
This concludes the proof of this divergence decomposition lemma.
19

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
E.3. Proof of Proposition 4.5
Proposition 4.5 (Tightness Result), Let STGD {S: Ps satisfies (4)). Let {8: STGD and St satisfies (4) with p (Line 5 in Algorithm 2), then
Reg pt"",
2 H(pt, v). Pi- Pi- 11-11-p/) p/-)
lim inf m0000
log(len(pt::)))
(1 - (1-Pi-)
Therefore, the upper and lower bound match up absolute constants and a p1-(1-pp; factor. In particular, if E (2-1/1, they match up to absolute constants and L.
L2(1-ppo)
1/1 ,1),
Proof of Proposition 4.5. Given any S E S with parameter p.
L+1
pt 1- pb+1
1g = x Ps(x) =
x p (1 p) (1+
1 1-p
Note that if ps 2 pú, we have p > Pt. Therefore,
1 - pl+
p. L+1
(p - pi) + (0. 1+1 - pp--1 + -:pp [+1
pp: 1111)
(37)
1-9 p 1 Pi
(1 p) p.)
(p pı) + (p)
1+1
(p - -i)p -p²)
- (1-p)(1 - -pi)
(1 p) - pi)
In addition, for any i E [K].
KT(Ps,, Ps) =
L+1
Ps. (1) log (x)
Pi - P I+1
log
Pi
+ (1 - p,))og log Pi
Ps(x)
p
1 -
By utilizing log x S x for 00
(1- pp- (pi- - pp)
KL(Ps,, 5 P: Pi-
(1 p1
p - Ps
(38)
1 Pi P
1-
r(1 pi)(1 P)
Combining (38) and (37),
(L(Ps,, Ps) < (us - p1) - p1:)(1 p)(1 -p-
(MS - 14:3
p(1 - p2--2
p(1 //1 p)
According to the definition of kl; in Theorem 4.4,
(His - 14:)
A?
kl. S
Pi. Pi- (1.)
Pe (1 - p/: / (L-pp)
Pr
Thus, the regret is lower bounded by
Reg(ALG, (ALG, ptm
Pi- - pf-
inm m=
2
Σ
log(len (pt()e(pt177)))
Mir A;
(1 Pr-
Furthermore, pf E (2-1/4,1) 1),
lim inf Reg(ALG, "t',.ptt..)
L/2 87 Pi- A
(pt())))pt?))
20

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
F. Supporting Lemmas
Lemma E.1 (Doob's optional stopping, Theorem 3.8 in Lattimore & Szepesvári (2020)). Let =F= (Fi) be a filtration and (Xt)sen be an a-dapte martingale and T F-stopping time such that at least one of the following holds:
(a) There exists an E N such that Pf(r > n] = 0.
(b) R(r] < oo, and there exits a constant IR such that for all t. E, B||x++1 - X. Fi) S almost surely on the event that T > t.
(c) There exists constant such that [Xtrrl almost surely for all t N.
Then X- almost surely well-defined, and [[X+] = E[Xu]. Furthermore, when (Xt) is super/sub-marrtingee rather than martingale, then equality is replaced with lexx/greater-than, respectively.
bound, Fact 1 in Auer et al. (2002b)). Let XI,..., Xn n be random variables with common range [0, 1] and X XI....,Xp-- Xn-11= μ.. Let Sn = X1 +... ..... Xn. Then for all a 2 0.
Lemma F.2 (Chernoff-Hoeffie
2a² (-m
2a
P(S, 2 ny + a) S exp
and P|S. 2 nu S exp
Lemma F.3 (Confidence Intervals, Lemma in Abbasi-yadkori et al. (2011)). Assuming that the noise te is conditionally Mun--Gaussian. With probability at least 8,
(1 + mi,))
( 2og
K(1 ni,
Vi E (1,2,... K}, A 2 0, Mi 1al S
nie
Lemma F.4 (Lemma 8 in Antos et al. (2010)), Let a > O. For any 2 (2/a) - b|,, at + b > logt. Lemma F.5 (Exercise 3.7 in Lattimore Szepesvári (2020)). Let FF (Fi)ien be a filtration, and be a stopping time with respect to F. Then F, is r-algebra.
Lemma F6 (Bretagnolle-Huun Huberr inequality, Theorem 14.2 in Lattimore & Szepesvári (2020)). Let P and Q be probability measures on the same measurable space (2, and let A E F be an arbitrary event. Then
P(A)+Q(19) 2 exp (KL((,Q)),
where A = S \ is the complement of Lemma F.7 (Pinsker's inequality, Equation (14.12) in Lattimore & Szepesvári (2020)). For measures and Q on the same probability space (2, that
drv (P,Q) S KL(P,Q).
Lemma F88 (Theorem 26.13 in Lattimore & Szepesvári (2020)). Let n 0 and be Legendre and twice differentiable with positive definite Hessian in A = int(dom(f)). Then all x. E A, there exists 2 E [x, vl {(1 {(1-a)+tay : [0. 1]} such that
(x y, u)
Dy(2,8) v)
< |1112
" 2
G. Additional Experimental Results
G.1. Additional Experimental Details
For the experiments stated in Section 6.1, report the memory utilization in this section. As Ealge-2 (Li et al., 2024b) is one of the best speculative decoding methods, we adopt it as the baseline. The Normalized Memory (NM) and Normalized Memory Bandwidth (NMB) are presented in Table G.1. The result shows that the proposed methods do not incur additional memory consumption compared to the baseline method.
We further remark that this result is achieved by our superior algorithm design, where several non-parametric models (PLD, REST, Suffix Tree) to enhance parametric SOTA model (Fagle-2). Specifically,
21

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
Table The memory and memory bandwidth utilized by our method. As Eagle-2 is of the hest SD methods, we adopt it the baseline to normalize the results of other methods. NM-Normalized Memory and NMB-Normalized Memory Bandwidth.
Spec Bench
Alpaca
Code Editor
Debug Bench
Methods
NM
NMB
NM
NMB
NM
NMB
NM
NMB
-LaMA3-8B-Inssrrue
Eagle-2
1.0000
1.0000
1.0000
1.0000
1.0000
1.0000
1.0000
1.0000
EXP3SPEC
0.9981
1.0171
0.9950
1.0170
1.0200
0,9980
1.0100
0.9960
UCBSPEC
1.0059
1.0093
1.0130
1.0090
0.9990
0.9820
1.0150
1.0020
Qwen2-7B-Instruct
Eagle-2 1.0000
1.0000
1.0000
1.0000
1.0000
1.0000
1.0000
1.0000
EXP3SPEC 1.0043
1.0095
1.0050
0.9980
1.0400
0.9850
0.9890
0.9960
UCBSPEC
0.9929
1.0036
1.0080
0.9900
1.0270
0.9930
1.0320
0.9950
"Non-parametric" means that these methods do not have any parameters in GPU, and directly predict the future tokens based on the past tokens according to the data structures like Trie Tree, which are python objects and stored in CPU RAM. All these show that the storage of the draft models will not increase the GPU memory. Our model only requires approximately an additional 100MB of CPU RAM. Since CPU memory is typically much larger (ITB our server) and cheaper than GPU memory (40 GB in our server). this cost is negligible.
All the draft models share the same verifier model, which is the target model (LLaMA3-8D-Instre) (Dubey et al., 2024) and Qwen2-7B I-Innrucct (Yang et al., 2024) in our experiments). So that the storage of the verifier does not increase the GPU memory.
The reduction in memory usage comes from the fact that non-parametric models require fewer verification tokens (e.g., 40 for Suffix Tree) compared to the baseline Eagle-2 (c.g., 64). As result, when invoking these models, a slight decrease in activation memory usage may be observed. Additionally, slight differences GPU memory may be observed, arising randomly from the short-lived activation tensors rather than from the method itself.
We note that the size of SpecBench is not large enough, i.c.. the number of arms pulls is not large, to derive statistically sound result. We enable Mixture-of-AAee Afent (Wang et al., 2024) on the prompts whose responses are shorter than 100 tokens to increase the number of arm pulls.
G.2. Experiments on Larger Models
In addition to the two models in the main paper, we conduct an addtional set of experiment on a larger target model, namely LLaMA-2-13B (Touvron et al., 2023). As Table indicates Eagle-2 (Li et al., 2024b) is one the best speculative decoding methods, we adopt as the baseline. The other setups are same as the ones in Section 6.
From the result reported in Table G.2, the proposed methods, UCBSPEC and EXP3SPEC, demonstrate their efficacy on larger models.
G.3. Experiments on Different Hardwares
In the main paper, the experiments are conducted on single A100 GPU. In this section, we conduct an additional set of experiment on GeForce RTX 4090. We adopt Eagle-2 (Li al., 2024b) the baseline and Spec Bench (Xia et al., 2024) as the benchmark. The result is presented in Table G.3. We observe a similar trend the result presented in Table 1. The proposed method remains useful with different hardware setup.
22

---

BANDITSPEC: Adaptive Speculative Decoding via Bandit Algorithms
Table Empirical Comparison between the proposed algorithms and Eagle-2 (Li et al., 2024b) with LLaMA-2-13B as the target model, measured by Mean Accepted Tokens (MAT) (1) and Tokens/s (f). The best result is highlighted in bold, while the second best result underlined. The proposed algorithms remain effective on larger models.
Spec Bench
Alpaca
Code Editor
Debug Bench
Methods
MAT(1)
Tokens/s(f)
MAT(1)
Tokens/s(t)
MAT(1)
Tokens/s(1)
MAT(1)
Tokens/s(1)
LLaMA-2-13B
Eagle-2
4.35
91.94
4.32
96.59
5.19
107.57
5.16
108.45
EXP3SPEC
4.05
95.52
4.32
99.64
5.22
115.65
5.03
116.65
UCBSPEC
4.43
97.16
4.36
102.29
5.27
113.97
5.27
118.67
Table Empirical comparison between Eagle-2 and the proposed algorithms GeForce RTX 4090. We observe similar trend the result presented Table 1.
Spec Bench
Methods
MAT
Tokens/s
LLaMA3-8B-Instee
Eagle-2
4.14
97.01
EXP3SPEC
3.95
102.24
UCBSPEC
4.16
107.38
Qwen2-7B-Instruct
Eagle-2
3.65
94.16
EXP3SPEC
3.96
111.74
UCBSPEC
4.17
112.21
23


Reviews:
[{"review_id": "oQINIu6yNI", "reviewer": "Reviewer_pdNK", "summary": "This paper proposes a training-free online learning framework to adaptively choose the configuration of the hyperparameters for speculative decoding as text is being generated. Specifically, this paper first formulates this hyperparameter selection problem as a Multi-Armed Bandit problem, and proposes two bandit-based hyperparameter selection algorithms to adaptively select configurations for speculative decoding. Experiments with LLaMA3 and Qwen2 demonstrate that the proposed method is effective compared to existing methods.\n\n## update after rebuttal\n\nThanks for the authors' detailed rebuttal. I will maintain my positive score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed feedback and helpful suggestions.\n\n>**Q1**: It would be more convincing if the authors could provide more supports for the reasonableness of formulating draft model selection across different decoding steps as a over-simpliﬁed multi-armed bandit problem.\n\nWe thank the reviewer for the question.\nWe believe the reviewer thinks that it is an oversimplification to employ the **standard stochastic Multi-Armed Bandits (MAB)** to model the Speculative Decoding (SD) problem, because the rewards (the number of accepted tokens in our case) of an arm (hyperparameter configuration) in the vanilla MAB problem are i.i.d. and hence, cannot capture the correlation between different decoding steps.\n\nOn the theoretical side, we clarify that our stationary mean assumption is strictly **weaker** than the i.i.d. assumption (this is discussed in Line 209 on the second column of page 4). In particular, the number of accepted tokens **can depend on the generated tokens**. Therefore, the assumption is aligned with real-world scenarios in which different decoding steps are correlated. Furthermore, the basic MAB model can be generalized to contextual bandits and non-stationary bandits. The proposed BanditSpec framework provides a basic template to apply these more general MAB setups to SD. Our formulations under the stationary/adversarial mean assumptions are just basic setups and we leave the more general/elaborate setups as future research (please refer to Appendix B for more details).\n\nOn the experimental side, as our experimental results indicate (Table 1), the performance of UCBSpec significantly outperforms the SOTA in SD, namely, Eagle-2 (Li et al., 2024). This **corroborates** the stationary mean assumption in our formulation. \n\n>**Q2**: Evaluation of the proposed method on larger models.\n\nThank the reviewer for the advice. We further conduct the experiment with LLaMA-2-13B (Touvron ea al., 2023) as the target model. As Eagle-2 is one of the best SD methods, we adopt it as the baseline. The result is as follows:\n\n\n\n| Methods | Spec Bench |  | Alpaca |  | Code Editor |  | Debug Bench |  |\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n|  | MAT | Tokens/s | MAT | Tokens/s | MAT | Tokens/s | MAT | Tokens/s |\n| Eagle-2 | 4.35 | 91.94 | 4.32 | 96.59 | 5.19 | 107.57 | 5.16 | 108.45 |\n| EXP3Spec | 4.05 | 95.52 | 4.32 | 99.64 | 5.22 | **115.65** | 5.03 | 116.65 |\n| UCBSpec | **4.43** | **97.16** | **4.36** | **102.29** | **5.27** | 113.97 | **5.27** | **118.67** |\n\n\nThis indicates the proposed BanditSpec framework is useful on larger models. We will incorporate the results in the revised version.\n\n**We hope our responses have addressed your concerns and would greatly appreciate your kind consideration in increasing your score.**", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "**Method**\n\n1.\t(Strengths) The proposed bandit-based online hyperparameter configuration method for speculative decoding is interesting and practical in real applications.\n\n2.\t(Strengths) The authors shown that the regret performance of the proposed method is optimal up to universal constants by deriving an information-theoretic impossibility result.\n\n3.\t(Weaknesses) The authors propose to formulate the draft model selection in standard speculative decoding as a multi-armed bandit problem. However, it simplifies the correlation between different decoding steps, which can be unaligned with realistic decoding process. It would be more convincing if the authors could provide more supports for the reasonableness of formulating draft model selection across different decoding steps as a over-simplified multi-armed bandit problem. \n\n**Evaluation Criteria**\n\n1.\t(Strengths) Experiments demonstrate the proposed method outperforms existing methods in terms of inference latency.\n\n2.\t(Weaknesses) The authors use LLaMA3-8B-Instruct and Qwen2- 7B-Instruct as the target models. However, it would be more convincing to evaluate the proposed method on larger models.", "theoretical_claims": "Yes, the theoretical claims are correct. However, it would be more convincing to explain the reasonableness of the assumptions.", "experimental_designs_or_analyses": "1.\t(Strengths) Experiments demonstrate the proposed method outperforms existing methods in terms of inference latency.\n\n2.\t(Weaknesses) The authors use LLaMA3-8B-Instruct and Qwen2- 7B-Instruct as the target models. However, it would be more convincing to evaluate the proposed method on larger models.", "relation_to_broader_scientific_literature": "1.\tThe authors propose to formulate the draft model selection in speculative decoding as a multi-armed bandit problem, which is an interesting formulation.\n\n2.\tThe authors propose to leverage two widely-used multi-armed bandit methods to adaptively select the draft models, and provide theoretical guarantees of the proposed method under mild assumptions.\n\n3.\tExperiments demonstrate that the proposed method outperforms existing methods.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Please see the above comments.", "other_comments_or_suggestions": "No", "questions_for_authors": "Please see the above comments.", "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": "o5c6bLEqd9", "reviewer": "Reviewer_S6Tu", "summary": "The paper introduces BanditSpec, a training-free online learning framework to route prompts to suitable off-the-shelf specualtive decoding methods. The authors formulate the problem as a Multi-Armed Bandit (MAB) problem and propose two bandit-based algorithms, UCBSpec and EXP3Spec, to adaptively choose different draft models and speculation lengths. \n\nThe paper provides theoretical analysis, including upper bounds on the stopping time regret under both stochastic and adversarial reward settings, and demonstrates the effectiveness of the proposed algorithms through empirical experiments with LLaMA3 and Qwen2 models. The results show that the proposed algorithms achieve competitive performance compared to existing methods, with throughput close to the oracle best hyperparameter in simulated real-life LLM serving scenarios.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed reading and feedback.\n\n>**Q1**: Computational overheads introduced by the proposed method, such as memory usage and latency.\n\n\nThank the reviewer for the advice! For the memory and memory bandwidth usage, please kindly refer to this anonymous link [Table_Memory_Utilization](https://ibb.co/TBDQ0wZ2). \nFor latency introduced by the bandit algorithm, we clarify that the latency (token/s) has already taken this factor into account. According to Table 1, the speedup of the proposed method is $\\times 2.96$ with repsect to the vanilla SD method and is $\\times 1.08$  with repsect to the SOTA, namely Ealge-2. In conclusion, the benefits of BanditSpec come at a negligible cost.\n\n>**Q2**: Generalization quality of the proposed method.\n\nIn the speculative decoding literature, it is theoretically guaranteed that the distribution of the generated sequence is the **same** as that of the target model, which means the quality of the output is **maintained** and **lossless acceleration** is achieved.(see e.g., Leviathan et al., 2023, Chen et al., 2023, Yin et al., 2024). Therefore, the quality metric is often omitted in the experiments and more emphasis is put on the acceleration and latency metrics (see Leviathan et al., 2023, Chen et al., 2023, Yin et al., 2024, Li et al., 2024b, etc.).\n\nAs we have also provided the same theoretical guarantees for the quality of the generated tokens from the BanditSpec framework in Proposition 1, the quality metric is omitted, following the convention in the community (see the above references). \n\n>**Q3**: Assumption of Fixed Prompts.\n\nWe would like to clarify that the theoretical guarantees are derived to bound the latency given **any** prompt, as explained in Line 201 \"Interpretation of the Desired Result\". We completely agree that prompts are diverse and performance of the speculative decoding (SD) method can be significantly influenced by the input prompts. This is indeed why we propose the BanditSpec framework, where given an input prompt, BanditSpec gradually learns the best SD method for this specific input prompt as the decoding process proceeds. Observe that our objective in equation (1) is $\\text{Reg}(\\text{ALG}, \\text{pt}, \\nu)$, which depends explicitly on the prompt $\\text{pt}$. We show that the additional SD rounds (the stopping time regret) is sublinear in the SD rounds required by the best SD method, indicating that the best SD method is adopted \"most of the time\" under BanditSpec. This perfectly aligns with real-world scenarios. We will further highlight this setup in the revised manuscript.\n\n>**Q4**: Stationary Mean Acceptance Length Assumption and its applicability to real-world scenarios.\n\nWe understand that in real-world scenarios, there can be many factors that influence the acceptance rate, making it non-stationary or even adversarial. This is also why we relax the i.i.d. assumption that is commonly used in the standard stochastic Multi-Armed Bandits (MAB) and we do allow the acceptance rate to be dependent on the input prompts and generated tokens. In particular, the distribution of the acceptance rate can also change with the only constraint on its mean under our assumption. Additionally, our formulation under the stationary mean assumption paves the way for the application of more generalized setups, like contextual bandits and non-stationary bandits which can lead to future research.\n\nOn the experimental side, the experimental results in Table 1 show that the proposed method exhibits competitive empirical performance compared to the current methods, including Eagle-2, the SOTA in SD, which strongly corroborates the validity of the assumptions. Additionally, the applicability of our theoretical results has also been empirically verified by the experimental results.\n\n>**Q5**: Adversarial Setting.\n\nWe would like to clarify that we derive the results for the adversarial setting under the **greedy decoding** strategy. Given a draft model, a target model and an input prompt, under the greedy decoding strategy, the output tokens are **fixed** but **unknown** at the beginning of the algorithm. This is modeled precisely by our adversarial MAB setup. \n\nFurthermore, we include the adversarial setting in our paper as a means to compare it to the stochastic setting. Prior to our work, it was a prior unclear how to use MAB to improve SD. Should one employ a stochastic, adversarial or even more generalized model? We consider a range of such MAB models and do a comparison among them to provide the community with a guide on which MAB model is best suited to the SD problem. \nAs the empirical performance of UCBSpec is better than EXP3Spec, it implies that real-life scenario tends to be benign and may be more aligned with the stationary mean assumption. We will highlight this observation in the revised version.\n\n**We hope our responses have addressed your concerns and would greatly appreciate your kind consideration in increasing your score.**", "claims_and_evidence": "The claims made in the paper are generally supported by clear evidence (though I didn't have a chance to check the proof details). The authors provide a detailed theoretical analysis, including upper bounds on the stopping time regret, and demonstrate the effectiveness of their algorithms through numerical experiments.", "methods_and_evaluation_criteria": "The proposed method uses a bandit framework to adaptively select hyperparameters (e.g., base speculative decoding strategies and number of speculations), which makes sense but does not seem very practical as it requires loading all base models into the GPU, significantly increasing the computational resource overheads. \n\nThe evaluation criteria can be improved by adding the numbers of actual performance on downstream tasks, providing a comparison of the generation quality of the proposed method and compared baselines.", "theoretical_claims": "The authors derive upper bounds on the stopping time regret for both UCBSpec and EXP3Spec under stochastic and adversarial settings, and the proofs appear to be technically sound. However, the theoretical analysis relies on assumptions that may not hold in practical scenarios, which limits the practical value of the results. Specifically,\n\n- Assumption of Fixed Prompts: The theoretical analysis assumes that the input prompts are fixed, which is rarely the case in real-world applications. In practice, prompts are highly diverse and dynamic, and the performance of speculative decoding methods can vary significantly depending on the input. This raises questions about how well the theoretical bounds translate to real-world decoding latency improvements. \n\n- Stationary Mean Acceptance Length: The stochastic setting assumes that the mean number of accepted tokens for each hyperparameter is stationary (Assumption 4.1). This assumption may not hold in practice, as the acceptance rate of speculative tokens can vary depending on the context and the specific input prompt. This limits the applicability of the theoretical results to real-world scenarios where the acceptance rate is non-stationary.\n\n- Adversarial Setting: While the adversarial setting relaxes the stationarity assumption, it still assumes that the number of accepted tokens is fixed by the environment before the algorithm starts (Assumption 5.1). This is also an unrealistic assumption in practice, as the acceptance rate can depend on the interaction between the draft model, the target model, and the input prompt.\n\nI understand the intuition of such theoretical analysis is to provide insights into the behavior of the proposed algorithms under idealized conditions, but its practical implication is limited by such assumptions -- For example, a lower regret bound under such assumptions does not necessarily guarantee lower decoding latency in practical settings.", "experimental_designs_or_analyses": "The experimental designs and analyses can benefit from a more detailed analysis of the following aspects:\n\n1. Computational resource overheads: The proposed method requires loading all base model weights into the GPU, which may introduce significant computational overhead compared to baseline methods. The paper does not provide a detailed discussion of the inference overheads, such as memory usage, GPU utilization, or latency introduced by the bandit algorithm itself. This information is crucial for understanding the practical feasibility of the proposed method, especially in resource-constrained environments.\n\n2. Generation quality: It'd be great to assess and compare the generation quality of the proposed method with baselines. It is unclear whether the proposed method guarantees output parity with standard autoregressive decoding. For example, does the adaptive selection of hyperparameters introduce any degradation in the quality of the generated text?", "relation_to_broader_scientific_literature": "The paper builds on existing sepcualtive decoding methods work by introducing a badnit framework to route input prompts to different off-the-self methods.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Other  Strengths:\n\n- The second experiment is conducted in simulated real-life scenarios with diverse input prompts (however, no comparison with baseline models in this setting)\n\nOther Weaknesses:\n- The paper lacks a detailed analysis of the computational overheads introduced by the proposed method, such as memory usage and GPU utilization.\n- The evaluation does not include an assessment of generation quality, which is critical for understanding the practical utility of the method.", "other_comments_or_suggestions": "Please address the concerns mentioned in the above sections.", "questions_for_authors": "1. The proposed method requires loading all base model weights into the GPU, which may introduce significant computational overhead. Could the authors provide a detailed analysis of the inference overheads, including memory usage, GPU utilization, and latency? How do these overheads compare to existing methods?\n\n2. The paper does not evaluate the generation quality of the proposed method. Does the adaptive selection of hyperparameters introduce any degradation in the quality of the generated text? Could the authors provide an evaluation of generation quality measured by downstream task metrics such as accuracy? How does the proposed method perform compared to standard autoregressive decoding in terms of downstream task performance?", "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": {"ethical_review_concerns": "N/A"}}, {"review_id": "bcaqNuKU1t", "reviewer": "Reviewer_XwVz", "summary": "This paper introduces BANDITSPEC, which adaptively selects configurations for speculative decoding to improve inference speed. Unlike previous approaches that use fixed speculative decoding configurations regardless of context, BANDITSPEC formulates hyperparameter selection as a Multi-Armed Bandit problem, enabling dynamic adaptation to different inputs. The authors develop two key algorithms: UCBSPEC for stochastic environments and EXP3SPEC for adversarial settings, both with theoretical stopping time regret guarantees. They demonstrate BANDITSPEC's performance through extensive experimentation with LLaMA3 and Qwen2 models, showing that adaptive configuration selection outperforms existing fixed methods, approaching the performance of oracle best configurations. The framework proves particularly effective in real-world LLM serving scenarios with diverse input prompts, establishing a theoretically sound approach to minimize speculative decoding latency.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed reply and acknowledge the soundness of our theoretical results. We answer the questions regarding the experiments as follows:\n\n>**Q1**: Incomplete competitor comparison with adaptive speculative decoding algorithms like SpecDec++\n\nThank the reviewer for the pointing this good work.\n\n- Firstly, we highlight that our proposed method is **training-free** which can be deployed easily along with **existing off-the-shelf methods**.\nIn contrast, SpecDec++ focuses on **training** of an acceptance prediction head. Currently, SpecDec++ is only available when using LLaMA-2-Chat-7B as the draft model and LLaMA-2-Chat-70B as the target model (bfloat 16). It remains unclear how to integrate SpecDec++ with other potentially superior draft models/methods beyond LLaMA-2-Chat-7B. This lack of flexibility poses challenges for our implementation.\n\n- Secondly, the proposed BanditSpec framework considers the more general hyperparameter selection problem that goes beyond merely the speculation length. Therefore, it is \"orthogonal\" to SpecDec++ in the sense that any methods with (or without) SpecDec++ can also be candidates for the hyperparameter in our framework, e.g., {Eagle-2, LLaMA-2-Chat-7B} with SpecDec++ can also be regarded as arms (if they are available).\n\n- Thirdly, our ultimate goal is to devise algorithms that are competitive compared to the SOTA method, which is Eagle-2 (Li et al., 2024b). Given the currently available experimental results on the Alpaca dataset,  the speedup measured by the throughput (tokens/s) is as follows:\n\n| Methods                                   | Target Model        | Speedup                        |\n| --- | ------ | --- |\n| SpecDec++ with LLaMA-2-Chat-7B as drafter | LLaMA-2-Chat-70B    | 2.04 (Huang et al., 2024)                       |\n| Eagle-2                                   | LLaMA-2-Chat-70B    | 3.51 (Li et al., 2024b)               |\n\nGiven the superior performance of Eagle-2, we adopt it as the backbone and baseline in our experiments.\n\nWe will include the discussion about SpecDec++ in our revised version.\n\n\n>**Q2**: Resource utilization metrics\n\nThank the reviewer for the advice! We highlight that the arm set for the model selection problem consists of one parametric model (Eagle-2) and several non-parametric models (PLD, Rest, Suffix Tree). These non-parametric models hardly consume the GPU resources. We measure the memory consumption and memory bandwidth utilization of our method. As Eagle-2 is one of the best SD method, we adopt it as the baseline to normalize the results of other methods. The result is accessible via this anonymous link [Table_Memory_Utilization](https://ibb.co/TBDQ0wZ2).\n\nIn our experiments, slight differences in GPU memory usage were observed, arising randomly from the short-lived activation tensors rather than from the method itself.\n\nAlthough speculative decoding increases the reuse of I/O operations through parallel decoding, it does **not directly** affect memory bandwidth utilization. This is because memory bandwidth measures how much data can be transferred per unit time. According to our experiments, the memory bandwidth utilization of our approach is about the same as EAGLE-2.\n\nWe will incorporate these discussions in the revised version.\n\n>**Q3**: Limited hardware scenarios\n\nThank the reviewer for the question.\nWe clarify that we do indeed investigate different computational setups. In Experiment 1, the batch size is to be 1 in order to compare the performance between the proposed method and the existing ones. In Experiment 2, we model the real-life scenario with diverse inputs and various batch sizes (ranging from 1 to 50) across the sample indices. According to the results in Figure 3, the throughput improvement is greater than 1 in most cases. This indicates the application of speculative decoding is still beneficial under reasonably high batch sizes. \n\nAdditionally, we conduct our experiments on GeForce RTX 4090, whose result is accessible via this anonymous link [Table_Empirical_Comparison_on_4090](https://ibb.co/DDhL7BJh). We observe a similar trend as the result presented in Table 1 of the manuscript. The proposed method remains useful under this setup.\n\n\nIf the reviewer has any other suggestions on other hardware scenarios for us to investigate, we would be happy to conduct such experiments to improve our paper. \n\n**We hope our responses have addressed your concerns and would greatly appreciate your kind consideration in increasing your score.**", "claims_and_evidence": "The paper provides solid theoretical guarantees for their claims through rigorous mathematical analysis, including lower bounds for the regret of their algorithms (UCBSPEC and EXP3SPEC). This theoretical foundation is a strength of the work. However, the experimental evidence has several limitations:\n\n1. Incomplete competitor comparison: The paper lacks comparisons with direct competitors in adaptive speculative decoding, such as SpecDec++. This omission makes it difficult to fully evaluate BANDITSPEC's relative performance against state-of-the-art adaptive approaches.\n2. Resource utilization metrics: While the authors demonstrate speedup compared to several baseline algorithms (vanilla, PLD, Rest, Suffix Tree, Eagle-2), they do not provide crucial metrics on resource utilization - specifically memory consumption and memory bandwidth utilization. This is particularly important since BANDITSPEC utilizes multiple speculative decoding algorithms, which likely has implications for resource overhead.\n3. Limited hardware scenarios: Experiments are conducted only on a single A100 GPU with batch size 1, with limited exploration of more diverse computational environments that would be encountered in production settings. Speculative decoding is more suitable for edge applications and not as effective on data-center applications with high batch-size.", "methods_and_evaluation_criteria": "Some of the concerns regarding the evaluation criteria mentioned above. In addition, some of the numbers in the experiments do not match with the source numbers. For instance, the Llama 3.1 8B Spec Dec using Eagle 2 numbers in the respective paper is higher than the one mentioned in the paper. It might be due to different hyperparameters.", "theoretical_claims": "I examined the key theoretical proofs in the paper, particularly those related to the regret bounds in Theorems 4.3 and 5.3, which establish guarantees for UCBSPEC and EXP3SPEC respectively.\nThe proofs appear technically sound, with appropriate application of martingale theory and self-normalized concentration bounds to handle the unique challenges of the stopping time regret minimization problem.", "experimental_designs_or_analyses": "The concerns regarding the experimental designs are mentioned above.", "relation_to_broader_scientific_literature": "This proposal can be beneficial for practical use cases of speculative decoding in memory-bound settings.", "essential_references_not_discussed": "As mentioned, some direct competitors such as SpecDec++ was not discussed nor compared with.", "other_strengths_and_weaknesses": "Mentioned above.", "other_comments_or_suggestions": "NA", "questions_for_authors": "Mentioned above.", "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:

Meta-review: This paper proposes a multi armed bandit based framework to adaptively choose the configuration of the hyperparameters for speculative decoding. The reviews mention the approach as convincing and appreciated the theoretical analysis (e.g., XwVz: “solid theoretical guarantees for their claims through rigorous mathematical analysis, including lower bounds for the regret of their algorithms”). They also mentioned the empirical analysis to be convincing of the value of the method. The latter is important since the theoretical analysis relied on assumptions that may not always hold, and in fact were mentioned as unrealistic. 

A few concerns were raised by the reviewers, and were addressed during the rebuttal stage: (1) Impracticality due to memory overhead (S6Tu, “the proposed method requires loading all base model weights into the GPU, which may introduce significant computational overhead compared to baseline methods.”): The authors analyzed the memory overhead, showing it is negligible, and in CPU memory. (2) Experiments require more baselines, XwVz: “The paper lacks comparisons with direct competitors in adaptive speculative decoding, such as SpecDec++.”: Here, the authors mentioned that the mentioned baselines are compared with the Eagle-2 method that performed better, and this method was tested in the paper. (3) Additional experiments needed: More scenarios (mentioned by XwVz), additional LLMs (mentioned by pdNK). The authors provided experiments in the rebuttal. 

Given these additional experiments and clarifications provided in the rebuttal, the remaining concerns are quite mild, making this paper sufficient in quality for ICML. Furthermore, the scope of the change does not seem very large to me and can be done towards a camera-ready version.
Decision: Accept (poster)