BanditSpec: Adaptive Speculative Decoding via Bandit Algorithms

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:

Q1: Incomplete competitor comparison with adaptive speculative decoding algorithms like SpecDec++

Thank the reviewer for the pointing this good work.

• Firstly, we highlight that our proposed method is 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 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.

• 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).

• 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:

Given the superior performance of Eagle-2, we adopt it as the backbone and baseline in our experiments.

We will include the discussion about SpecDec++ in our revised version.

Q2: Resource utilization metrics

Thank 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.

In our experiments, slight differences in GPU memory usage were observed, arising randomly from the short-lived activation tensors rather than from the method itself.

Although 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.

We will incorporate these discussions in the revised version.

Q3: Limited hardware scenarios

Thank the reviewer for the question. We 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.

Additionally, we conduct our experiments on GeForce RTX 4090, whose result is accessible via this anonymous link Table_Empirical_Comparison_on_4090. We observe a similar trend as the result presented in Table 1 of the manuscript. The proposed method remains useful under this setup.

If 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.

We hope our responses have addressed your concerns and would greatly appreciate your kind consideration in increasing your score.

We thank the reviewer for the detailed reading and feedback.

Q1: Computational overheads introduced by the proposed method, such as memory usage and latency.

Thank the reviewer for the advice! For the memory and memory bandwidth usage, please kindly refer to this anonymous link Table_Memory_Utilization. For 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.

Q2: Generalization quality of the proposed method.

In 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.).

As 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).

Q3: Assumption of Fixed Prompts.

We 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.

Q4: Stationary Mean Acceptance Length Assumption and its applicability to real-world scenarios.

We 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.

On 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.

Q5: Adversarial Setting.

We 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.

Furthermore, 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. As 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.

We hope our responses have addressed your concerns and would greatly appreciate your kind consideration in increasing your score.

We thank the reviewer for the detailed feedback and helpful suggestions.

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-simplified multi-armed bandit problem.

We thank the reviewer for the question. We 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.

On 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).

On 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.

Q2: Evaluation of the proposed method on larger models.

Thank 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:

This indicates the proposed BanditSpec framework is useful on larger models. We will incorporate the results in the revised version.

We hope our responses have addressed your concerns and would greatly appreciate your kind consideration in increasing your score.