DistillSpec: Improving Speculative Decoding via Knowledge Distillation

We thank the reviewer for taking the time to provide your thoughtful and valuable feedback on our work. Below are responses to your concerns.

All the distillation methods are based on existing works. The work is more of a comprehensive study of knowledge distillation on LLM for speculative decoding.

Indeed, the goal of our paper is not to propose a new KD algorithm but to provide a way to improve SD using KD. That said, applying KD to SD is not well-explored in the past literature. As shown in Figure 2, 3a, the previous attempt [1] using SeqKD turned out to be not very effective to improve SD. Therefore, we provide a systematic and comprehensive study on the most cost effective way to apply KD to SD.

To understand the key factors that affect the distillation performance, we categorize most commonly used KD techniques in terms of “divergence function” and “training data'' in Table 1 and investigate many other KD variants in Section 5.2 (Figure 4). Note that many combinations of the divergence function and training data have not been studied in the past literature. We believe that our findings that 1) using student generated data is most cost-effective; and 2) the best divergence function depends on tasks and decoding strategy; are meaningful and provide new insights on how we should apply distillation in practice to improve LLM latency.

Furthermore, we wish to emphasize our contributions to the study of the quality and latency trade-off, a central challenge in LLM inference. Our research explores this trade-off through two innovative perspectives: 'lossy speculative decoding' and 'model garden'. The former offers fine-grained control over the trade-off curve, while the latter presents optimal strategies for combining SD and Distillation. These contributions are novel and highly pertinent for practitioners seeking to balance quality and latency in model deployment.

[1] Stern, M., Shazeer, N., & Uszkoreit, J. (2018). Blockwise parallel decoding for deep autoregressive models. Advances in Neural Information Processing Systems, 31.

All the experiments are done in an in-domain setting, while the benefits of LLM are on zeroshot/fewshot settings on out-of-domain datasets. Thus the experiments can not show whether it can replace the general draft models.

Thanks for the suggestion! We have added a new experiment to study the transferability of the distilled draft model across different datasets. Specifically, we take the distilled draft model from GSM8K and evaluate it on all BigBenchHard tasks (23 tasks in total and all are out-of-domain) [1] with zero-shot CoT. We observe that the distilled draft model transfers well to other datasets. Compared to standard SD, it improves the average decoding speed from 1.93x and 1.78x to 2.21x and 2.02x using the greedy and non-greedy decoding, respectively. See Appendix Figure C.2 and C.3 for a detailed breakdown of the tasks in BBH.

[1] Suzgun, M., Scales, N., Schärli, N., Gehrmann, S., Tay, Y., Chung, H. W., ... & Wei, J. (2022). Challenging big-bench tasks and whether chain-of-thought can solve them. arXiv preprint arXiv:2210.09261.

Why do different distillation methods vary a lot in Figure 4?

Indeed, it is an interesting finding from our comprehensive study on divergence function in Section 5.2. Empirically, we observe that different divergence functions perform differently on different tasks. It may be due to certain specific characteristics of the tasks but we do not have a good explanation yet. While this performance variance is disappointing, we think it truthfully reflects the fact that distillation in the real world is challenging. Therefore, we recommend to view the divergence function as one of the hyperparameters in a distillation algorithm and we should tune it just like any other hyperparameters.

We hope that most of the reviewer’s concerns have been addressed and, if so, we would appreciate it if they could reconsider their assessment. We’d be happy to engage in further discussions.

We thank the reviewer for taking the time to provide your thoughtful and valuable feedback on our work. Below are responses to your concerns.

For instance, it is difficult to find the formal definition of DistillSpec among the analysis of various existing distillation approaches

DistillSpec (Algorithm A.2) refers to our general framework that improves the SD by aligning the target model and draft model using the whitebox knowledge distillation. We do not introduce new divergence functions or new data generation procedures, but instead systematically investigate which combinations of divergence functions and data generation procedures can be best leveraged to improve upon vanilla SD. We have revised the paper to include this definition.

it is not clear how distillation data is generated in detail.

To generate new data, we closely follow the previous work [1]. Specifically, we prompt the model with a task input sampled from the training dataset and ask the model to generate an output. This data generation process works for both teacher and student models.

[1] Agarwal, R., Vieillard, N., Stanczyk, P., Ramos, S., Geist, M., & Bachem, O. (2023). GKD: Generalized Knowledge Distillation for Auto-regressive Sequence Models. arXiv preprint arXiv:2306.13649.

the specific configurations of target and draft models in figure 6 are not given, especially for the size of the target model in the DistillSpec case.

We have updated the figure and summarized the target and draft model size below.

1. Raw: Deploying supervised fine-tuned (SFT) T5 models. It contains five model points: T5-Small (77M), T5-Base (250M), T5-Large (800M), T5-XL (3B), and T5-XXL (11B).

2. Distilled: Applying KD to T5 models for optimizing downstream task performance. It contains four model points: Distilled-T5-Small (77M), Distilled-T5-Base (250M), Distilled-T5-Large (800M), Distilled-T5-XL (3B). All distilled models use Raw T5-XXL (11B) as the teacher.

3. Speculative: Applying SD to T5 models. It contains four target models and all use Raw T5-Small (77M) as draft model for best performance. The target models are T5-Base (250M), T5-Large (800M), T5-XL (3B), and T5-XXL (11B).

4. DistillSpec: Applying KD to T5 models and using distilled models as target and draft. It has three target model sizes: (1) Distilled-T5-Base (250M) (2) Distilled-T5-Large (800M), and (3) Distilled-T5-XL (3B). We use Distilled-T5-Small (77M) as the draft model for all target models.

While DistillSpec is effective on certain tasks, the experiments mainly focus on T5 models and simple tasks (except for GSM8K) ... The effectiveness of DistillSpec should be evaluated using more recent large models and some zero-shot benchmarks to prove its effectiveness.

Thanks for the suggestion! We have added a new experiment to study the transferability of the distilled draft model across different datasets. Specifically, we take the distilled draft model from GSM8K and evaluate it on all BigBenchHard tasks (23 tasks in total) [1] with zero-shot CoT. We observe that the distilled draft model transfers well to other datasets. Compared to standard SD, it improves the average decoding speed from 1.93x and 1.78x to 2.21x and 2.02x using the greedy and non-greedy decoding, respectively. See Appendix Figure C.2 and C.3 for a detailed breakdown of the tasks in BBH.

As for the model we use, we choose the T5 model series because they are a set of powerful and representative models, which are carefully studied in the original SD paper [2] and a follow up study in [3]. Further, to demonstrate the effectiveness of our method for the decoder-only model, we followed [2][4] and had already included an experiment on LM1B in our original submission. As shown in Table C.1, DistillSpec consistently yields speed gain over the vanilla SD. We expect our findings from the LM1B experiment with GPT-like decoder-only models to extend to LLaMA sized models.

[1] Suzgun, M., Scales, N., Schärli, N., Gehrmann, S., Tay, Y., Chung, H. W., ... & Wei, J. (2022). Challenging big-bench tasks and whether chain-of-thought can solve them. arXiv preprint arXiv:2210.09261.

[2] Leviathan, Y., Kalman, M., & Matias, Y. (2023, July). Fast inference from transformers via speculative decoding. In International Conference on Machine Learning (pp. 19274-19286). PMLR.

[3] Kim, S., Mangalam, K., Moon, S., Malik, J., Mahoney, M.W., Gholami, A., & Keutzer, K. (2023). Speculative Decoding with Big Little Decoder.

[4] Sun, Z., Suresh, A. T., Ro, J. H., Beirami, A., Jain, H., & Yu, F. (2023). Spectr: Fast speculative decoding via optimal transport.

Major experiment results are based on small target models.

Our main results (Figure 1) are based on the second large T5 model (T5-XL, 3B) as target and T5 (T5-small, 77M) as draft. Besides, we include additional results using the largest T5 (T5-XXL, 11B) as target in Appendix (Figure C.1). Compared to recent literature on SD, where BiLD [1] only considers T5-Large (770M) as the target model, SpecTr [2] considers a 97M decoder only model, Spector & Re [3] consider a 762M GPT2 model, our target model choices (3B, 11B) are not small, which closely follows the original SD paper [4].

[1] Kim, S., Mangalam, K., Moon, S., Malik, J., Mahoney, M.W., Gholami, A., & Keutzer, K. (2023). Speculative Decoding with Big Little Decoder.

[2] Sun, Z., Suresh, A. T., Ro, J. H., Beirami, A., Jain, H., & Yu, F. (2023). Spectr: Fast speculative decoding via optimal transport.

[3] Spector, B., & Re, C. (2023). Accelerating LLM inference with staged speculative decoding. arXiv preprint arXiv:2308.04623.

[4] Leviathan, Y., Kalman, M., & Matias, Y. (2023, July). Fast inference from transformers via speculative decoding. In International Conference on Machine Learning (pp. 19274-19286). PMLR.

The results on larger models are not very clear -- in Figure 6, which data points correspond to which sizes of target and draft models?

We agree and have updated the figure. The target and draft model size are detailed below.

1. Raw: Deploying supervised fine-tuned (SFT) T5 models. It contains five model points: T5-Small (77M), T5-Base (250M), T5-Large (800M), T5-XL (3B), and T5-XXL (11B).

2. Distilled: Applying KD to T5 models for optimizing downstream task performance. It contains four model points: Distilled-T5-Small (77M), Distilled-T5-Base (250M), Distilled-T5-Large (800M), Distilled-T5-XL (3B). All distilled models use Raw T5-XXL (11B) as the teacher.

3. Speculative: Applying SD to T5 models. It contains four target models and all use Raw T5-Small (77M) as draft model for best performance. The target models are T5-Base (250M), T5-Large (800M), T5-XL (3B), and T5-XXL (11B).

4. DistillSpec: Applying KD to T5 models and using distilled models as target and draft. It has three target model sizes: (1) Distilled-T5-Base (250M) (2) Distilled-T5-Large (800M), and (3) Distilled-T5-XL (3B). We use Distilled-T5-Small (77M) as the draft model for all target models.

The method is target-model-dependent, meaning that we need to distill a new draft model for each new target model. Such a distillation cost can be quite expensive, especially when using online data generation.

SD’s efficiency depends on the alignment between the target model and draft model, which is “model dependent”. When we have a new target model to deploy, if it is very different from the previous target model, then, it is reasonable to distill a new draft model tailored to the new target model. However, if the new target model is very similar to the old target model (e.g., they belong to the same model class, but are of different sizes), then it is sensible to reuse the existing distilled draft model. As we demonstrated in Appendix Figure C.1, the distilled draft model from T5-XL (3B) as target also works well for T5-XXL (11B) as target. Though it is a suboptimal choice, it still achieves a consistent improvement of 7%- 37% over original SD across four different datasets.

In contrast, if the new target model is very different, then the extra distillation cost is inevitable to realize significant gains via SD. Note that the distillation process only needs to be done once, which can be amortized across many serving queries and the real cost is serving to a large number of users. As for the online data generation, Figure 3a suggests that using the student model for online data generation is actually the most cost-effective way to do distillation. In contrast, when you use a static ground truth dataset (i.e., SupKD), your model performance will soon plateau.

Additionally, we have added a new experiment to study the transferability of the distilled draft model across different datasets. Specifically, we take the distilled draft model from GSM8K and evaluate it on all BigBenchHard tasks (23 tasks in total) [1] with zero-shot CoT. We observe that the distilled draft model transfers well to other datasets. Compared to standard SD, it improves the average decoding speed from 1.93x and 1.78x to 2.21x and 2.02x using the greedy and non-greedy decoding, respectively. See Appendix Figure C.2 and C.3 for a detailed breakdown of the tasks in BBH.

[1] Suzgun, M., Scales, N., Schärli, N., Gehrmann, S., Tay, Y., Chung, H. W., ... & Wei, J. (2022). Challenging big-bench tasks and whether chain-of-thought can solve them. arXiv preprint arXiv:2210.09261.

We hope that your concerns/questions have been addressed and, if so, we would appreciate it if you could reconsider your assessment. We’d be happy to engage in further discussions.

Thank you for taking the time to provide your thoughtful and valuable feedback on our work. Below are responses to your concerns.

The technical novelty is a bit limited. It is a direct application of existing KD techniques in SD.

We argue that the application of KD in SD is not as straightforward or well-explored. Our approach diverges from the standard KD, which primarily aims at enhancing the student model's performance. Our focus is on improving alignment to enhance decoding efficiency, a distinct objective as evidenced in Figure 5a.

Additionally, our work critically examines the effectiveness of applying KD to SD. Figures 2 and 3a demonstrate that previous attempts, such as the use of SeqKD [1], were inadequate in enhancing SD efficiency. Our research contributes to this field by providing a systematic and comprehensive study (Section 5.2) on the most cost-effective methods of applying KD in SD. The insights we offer, particularly regarding the cost-effectiveness of using student-generated data and the task and strategy-dependent nature of the best divergence function, are valuable for practical applications.

Furthermore, we wish to emphasize our study of the quality and latency trade-off, a central challenge in LLM inference. Our research explores this trade-off through two perspectives: 'lossy speculative decoding' and 'model garden' (Section 5.3). The former offers fine-grained control over the trade-off curve, while the latter presents optimal strategies for combining SD and KD. These contributions are novel and highly pertinent for practitioners seeking to balance quality and latency in model deployment.

[1] Stern, M., Shazeer, N., & Uszkoreit, J. (2018). Blockwise parallel decoding for deep autoregressive models. Advances in Neural Information Processing Systems, 31.

The performance of the method is task-dependent, posing concerns for using the method in practice. For example, the speedup on WMT En-De is marginal.

Indeed, the additional speedup provided by DistillSpec is task-dependent. There are two reasons for this:

1. First, different tasks have different amounts of headroom. As shown in Figure 1 and Table D.1, the original SD already shows a speedup of 2.36x on WMT, in contrast to a 1.44 speedup on XSum. It suggests that target model and draft model exhibit different levels of alignment when trained on different tasks and the model trained on WMT is already much more aligned than XSum. Thus, the headroom for improvement is smaller for WMT than XSum. The difference in task nature was also observed in previous SD literature [1][2].

2. Secondly, we did not attempt to optimize the hyper-parameters for each task: we use the same hyperparameters for all distillation experiments, so they may be suboptimal for each task. Figure C.8 and C.9 show the progression of the acceptance rate tested on the sequence generated by the target model. The unsaturated learning curves suggest that we can achieve much better performance if we train the model longer, especially in the WMT case. However, our current training strategy can already show the effectiveness of the DistillSpec and training longer to optimize the absolute performance gain might not provide any new insights.

[1] Leviathan, Y., Kalman, M., & Matias, Y. (2023, July). Fast inference from transformers via speculative decoding. In International Conference on Machine Learning (pp. 19274-19286). PMLR.

[2] Chen, C., Borgeaud, S., Irving, G., Lespiau, J. B., Sifre, L., & Jumper, J. (2023). Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.01318.

TVD has varied performance on different tasks.

We argue that “TVD has varied performance on different tasks” is itself a finding of our paper. Although TVD is the theoretically motivated objective we should optimize, our systematic study in Section 5.2 suggests that it does not consistently yield the best empirical performance. While this gap between theory and practice is disappointing, we think it truthfully reflects the fact that distillation in the real world is challenging.

Just like ML algorithms that have some hyperparameters to tune, we think the divergence function is one of the hyperparameters in a distillation algorithm and you should tune it for your task. Note that this process only needs to be done once and can be amortized across many serving queries.

Overall, we view treating the divergence function as a hyperparameter as an interesting take-away from our study.