SpecDec++: Boosting Speculative Decoding via Adaptive Candidate Lengths

We sincerely appreciate the time and effort you have invested in reviewing our paper and providing valuable feedback. We would like to provide detailed responses below:

W1:The paper aims to accelerate inference by reducing the number of forward passes of the target model and decreasing the number of discarded tokens, but the experiments show only a modest improvement over the baseline. Moreover, the training cost for the head is substantial. .... This makes me somewhat skeptical about the effectiveness of this approach.

A1: We would like to re-emphasize the practical contribution of our method:

• (1) The speedup is achieved through algorithmic improvements as opposed to system-level improvements. This means that our method can be plugged in whenever the inference method adopts a form of draft models (e.g., early-exiting). Furthermore, it can be readily combined with other system-level and hardware-level improvements, e.g., fused cuda kernels, smarter KV cache management, or simply faster chips.
• (2) Our experimental results are hardware-independent: Figure 4 shows that our proposed method has better Pareto front than the baseline SpecDec method, which means SpecDec++ will still out-perform the baseline SpecDec under other hardware configurations. One may argue that the key bottlenecks of speculative sampling is the system-level optimization tricks. Nevertheless, when those bottlenecks are solved, SpecDec++ will still achieve improvement over baseline SpecDec.
• (3) Besides, achieving an additional 7-11% speed improvement will save 7%-11% of the GPU time and the financial cost for those entities that need to serve their LLMs, which could be a huge number for LLM applications with high traffic.

About training cost: The training cost for the prediction head is marginal as we only tune a scalar output head on top of the last layer of hidden features of the draft model.

• (1) As a comparison, Medusa [1] involves tuning multiple token prediction heads and the total network size of Medusa heads is larger than ours.
• (2) In Appendix C.1, we have a rough estimate of the compute resources used in the paper. Training the acceptance head only takes 16.7% of the total compute. In the real-world setting where we host the model for inference service, the one-time training cost will be essentially negligible compared to the vast amount of the inference cost.

W2: The paper introduces two evaluation metrics to demonstrate the effect of dynamically adjusting the candidate token length, but the length of accepted tokens is also a crucial metric. I am not sure if this method has any impact on the length of accepted tokens.

A2: Thank you for indicating that “the length of accepted tokens” is also a crucial metric. Indeed, our method has a substantial impact on the average length of the consecutive accepted tokens. Actually, the “verification rate” in our paper is defined as the average number of the target model forward pass per accepted token. You can see that 1/verification rate is the average length of the accepted tokens per round. By adaptively determining the draft token length, our method achieves a larger length of accepted tokens while keeping the number of discarded tokens controlled.

W3: The paper only presents results on the llama-2-chat 7B and llama-2-chat 70B models. I suspect this is due to the high training costs, but a single set of results is less convincing.

A3: Thank you for pointing the issue out. We repeat all our experiments on a new pair of models: Gemma-1.1-2B-it and Gemma-2-27B-it. We are also able to show performance gains of our proposed SpecDec++ compared to the baseline SpecDec. Please see the general response for details!

Besides, the training cost is not high; please see the complete discussion in A1.

Questions:

Q1: Are there other methods to reduce the training cost, or are there alternative approaches to determine whether to stop the draft model's generation?

A1: The training cost is not high; please see the complete discussion in answer A1 to W1.

Q2: The paper claims that the method can be seamlessly integrated with other improvements. Could at least one example be provided for experimentation?

A2: Our experimental results utilize the half-precision (bf16) models to mitigate the memory I/O overhead for both the draft model and the target model, which is a widely adopted technique. Our method can be easily integrated with a more quantized draft model.

Q3: It would be desirable to conduct experiments on a broader range of models.

A3: Please see the general response for the new experiments.

We sincerely hope that our responses can address your concerns, and we would greatly appreciate it if you would like to re-evaluate our work given the responses.

We sincerely appreciate the time and effort you have invested in reviewing our paper and providing valuable feedback. We would like to provide detailed responses below:

W1: Though the paper solves an important problem by a learning-based approach the overall gains seem small (7% - 11%) over vanilla speculative decoding.

We would like to re-emphasize the practical contribution of our method:

• (1) The speedup is achieved through algorithmic improvements as opposed to system-level improvements. This means that our method can be plugged in whenever the inference method adopts a form of draft models (e.g., early-exiting). Furthermore, it can be readily combined with other system-level and hardware-level improvements, e.g., fused cuda kernels, smarter KV cache management, or simply faster chips.
• (2) Our experimental results are hardware-independent: Figure 4 shows that our proposed method has better Pareto front than the baseline SpecDec method, which means SpecDec++ will still out-perform the baseline SpecDec under other hardware configurations. One may argue that the key bottlenecks of speculative sampling is the system-level optimization tricks. Nevertheless, when those bottlenecks are solved, SpecDec++ will still achieve improvement over baseline SpecDec.
• (3) Besides, achieving an additional 7-11% speed improvement will save 7%-11% of the GPU time and the financial cost for those entities that need to serve their LLMs, which could be a huge number for LLM applications with high traffic.

W2: The paper has not compared with existing heuristic based adaptive draft length methods which use either draft entropy or other confidence scores for stopping or continuing the draft generation

A2:

We indeed tried using entropy as an indicator during our early investigations but ruled it out as the performance was limited. The specific reasons are:

• (1) Conceptually, the entropy of the predicted token combines two uncertainties: (a) the uncertainty of the next token that is inherent in the natural languages, e.g. the words “good” and “great” have similar meanings and can often be used interchangeably. (b) the uncertainty related to the confidence of the model’s prediction. We don’t want to stop the speculation when only (a) is high.
• (2) According to the speculative decoding algorithm, when the draft model and the target model align very well, the acceptance probabilities are close to 1, no matter whether the current token has a high entropy or not.
• (3) In the extreme case when the draft model and the target model always give the same distribution, the optimal strategy is to always use the cheaper draft model till the end. However, using entropy of the draft distribution will likely to stop the speculation when it encounters the context where the next token has a high inherent entropy.

To address the problem and distinguish the confidence of the draft model out of the inherent entropy of the natural languages, we work out the theory behind it and train a prediction head to directly predict the acceptance probabilities.

W3: The hyper-parameter search required for finding the w_{reject} and probability threshold is an additional set of hassle also present in vanilla SpD, in this regard it seems like it could be better to simply try vanilla SpD with multiple draft lengths on a small subset of downstream task and then choose one with best performance for entirety of the task.

A3: Thank you for pointing out the issue of hyperparameter searches. Please note that we indeed performed a search over the draft lengths for vanilla SpD, and showed that our method even achieves improvement over the best fixed draft length for vanilla SpD:

• (1) In Figure 1. the reported speed-ups for vanilla SpD (termed SpecDec in our paper) are the best performance for the task over different draft lengths. Our proposed SpecDec++ can achieve performance improvement over the best vanilla SpD.
• (2) In Figure 4, the Pareto front of the vanilla SpD (termed SpecDec in our paper) is obtained by trying multiple draft lengths. You can see that our proposed SpecDec++ has better Pareto fronts.

About the additionally introduced hyperparameters:

• (1) We found that the performance of our methods is rather insensitive to the parameters $w_{rej}$ and the threshold. Specifically, for $w_rej \in [3, 6]$, setting the threshold from 0.5 to 0.9 achieves nearly the same performance (See the full results in Table 3).
• (2) Besides, although the inference speeds are different, they all outperform the baseline SpD method with a static $K$.

In summary, naively tuning the draft lengths for vanilla SpD has limited performance, and our work utilizes a learned acceptance prediction head to adaptively determine the draft length based on the current prefix, achieving further speedup compared to the optimally tuned vanilla SpD method.

We sincerely appreciate the time and effort you have invested in reviewing our paper and providing valuable feedback. We would like to provide detailed responses below:

W1: The reported improvement over vanilla speculative decoding appears minimal, with only a 10% speedup for a large model. Moreover, vanilla speculative decoding is significantly slower than state-of-the-art methods like Medusa [1] and Eagle [2], raising questions about the method's practicality.

A1: We would like to re-emphasize the contribution of our method:

• (1) The speedup is achieved through algorithmic improvements as opposed to system-level improvements. This means that our method can be plugged in whenever the inference method adopts a form of draft models (e.g., early-exiting). Furthermore, it can be readily combined with other system-level and hardware-level improvements, e.g., fused cuda kernels, smarter KV cache management, or simply faster chips.
• (2) Our experimental results are hardware-independent: Figure 4 shows that our proposed method has better Pareto front than the baseline SpecDec method, which means SpecDec++ will still out-perform the baseline SpecDec under other hardware configurations. One may argue that the key bottlenecks of speculative sampling is the system-level optimization tricks. Nevertheless, when those bottlenecks are solved, SpecDec++ will still achieve improvement over baseline SpecDec.
• (3) Besides, achieving an additional 7-11% speed improvement will save 7%-11% of the GPU time and the financial cost for those entities that need to serve their LLMs, which could be a huge number for LLM applications with high traffic.

W2: the additional complexity of training a candidate length prediction model.

A2: The training cost for the prediction head is marginal as we only tune a scalar output head on top of the last layer of hidden features of the draft model.

• (1) As a comparison, Medusa [1] involves tuning multiple token prediction heads and the total network size of Medusa heads is larger than ours.
• (2) In Appendix C.1, we have a rough estimate of the compute resources used in the paper. Training the acceptance head only takes 16.7% of the total compute. In the real-world setting where we host the model for inference service, the one-time training cost will be essentially negligible compared to the vast amount of the inference cost.

W3: In terms of theoretical analysis, while the optimal strategy is presented as a threshold function, the threshold’s dependency on both target and source models makes it unclear how a fundamentally residual model can capture this relationship, as the input only includes the context.

A3: Yes the acceptance probability of the current token will depend on both the target model and the draft model. The prediction head is built on top of the draft model and it only includes the context of the previous tokens.

However, such an acceptance prediction head is still able to capture the token acceptance probabilities, because:

• (1) the training dataset is generated using both the target model and the draft model. Especially, the ground-truth acceptance probabilities are computed by calculating the log probability ratios of the target model and the draft model. Therefore, the training dataset contains the information of the target distribution.
• (2) We train the acceptance prediction head so it can predict the ground-truth acceptance probabilities. Therefore, the knowledge of the acceptance probabilities of the target model is trained into the acceptance prediction head.
• (3) Besides the above (abstract) theoretical argument, the effectiveness of the prediction head model in capturing the relationship between the target model and draft model is demonstrated by our empirical experiments.

W4: Insufficient Experimentation: The experiments are limited to a single pair of models and model sizes, which leaves the method’s generalizability unaddressed. It remains unclear if the speedup improvement would increase, decrease, or remain stable with varying model sizes, making it difficult to assess its applicability across different model configurations.

A4: We repeat all our experiments on a new pair of models: Gemma-1.1-2B-it and Gemma-2-27B-it. We are also able to show performance gains of our proposed SpecDec++ compared to the baseline SpecDec. Please see the general response for details!