DistillSpec: Improving Speculative Decoding via Knowledge Distillation
Authors: Yongchao Zhou
Kaifeng Lyu
Ankit Singh Rawat
Aditya Krishna Menon
Afshin Rostamizadeh
Sanjiv Kumar
Jean-François Kagy
Rishabh Agarwal
Venue: ICLR 2024 poster (2024)
Keywords: large language model
knowledge distillation
speculative decoding
Primary area: generative models
TL;DR: We propose to improve speculative decoding via knowledge distillation.

Abstract: Speculative decoding~(SD) accelerates large language model inference by employing a faster {\em draft} model for generating multiple tokens, which are then verified in parallel by the larger {\em target} model, resulting in the text generated according to the target model distribution. However, identifying a compact draft model that is well-aligned with the target model is challenging. To tackle this issue, we propose {\em DistillSpec} that uses knowledge distillation to better align the draft model with the target model, before applying SD. DistillSpec makes two key design choices, which we demonstrate via systematic study to be crucial to improve the draft and target alignment: utilizing \emph{on-policy} data generation from the draft model, and \emph{tailoring the divergence function} to the task and decoding strategy. Notably, DistillSpec yields impressive $10 - 45\%$ speedups over standard SD on a range of standard benchmarks, using both greedy and non-greedy sampling. Furthermore, we combine DistillSpec with lossy SD to achieve fine-grained control over the latency vs. task performance trade-off. Finally, in practical scenarios with models of varying sizes, first using distillation to boost the performance of the target model and then applying DistillSpec to train a well-aligned draft model can reduce decoding latency by $6 - 10\times$ with minimal performance drop, compared to standard decoding without distillation.

Published as a conference paper at ICLR 2024
DISTILLSPEC: IMPROVING SPECULATIVE DECODING VIA KNOWLEDGE DISTILLATION
Yongchao Zhou' :3: Kaifeng Lyu Ankit Singh Rawat', Aditya Krishna Menon',
Afshin Rostamizadeh', Sanjiv Kumar', Jean-François Kagy"i, Rishabh Agarwal 2.5f Gooole Research GGooge DeepMind University of Toronto *Princeton University Mila
ABSTRACT
Speculative decoding (SD) accelerates large language model inference by employing faster draft model for generating multiple tokens, which are verified in parallel by the larger target model, resulting in the text generated according to the target model dis- tribution. However, identifying a compact draft model that is well-aligned with the target model is challenging. To tackle this issue, knowledge distillation better align the draft model with the target model before applying SD. DistillSpec makes two key design choices, which we demonstrate via systematic study to be crucial to improving the draft and target alignment: utilizing on-policy data genera- tion from the draft model, and kailoring the divergence function to the task and decoding strategy. Notably, DistillSpec yields 10 45% speedups over standard SD range of benchmarks, using both greedy and non-greedy sampling. We show that the distilled model can be well transferred to various tasks with an average speedup of 26%. Furthermore, we combine DistillSpec with lossy SD to achieve fine-grained control over the latency vs. task performance trade-off. Finally, in practical scenarios with models of varying sizes, first us- ing distillation boost the performance of the target model and then applying DistillSpec to train well-aligned draft model can reduce decoding latency by 10x with minimal performance drop, compared to standard decoding without distillation.
propose DistillSpec, method that uses
INTRODUCTION
Large language models (LLMs) have revolutionized natural language understanding and genera- tion across diverse applications (OpenAT [023] Anil et al. (023). However, their autoregressive generation nature poses significant computational challenges, especially in real-time deployments with stringent latency constraints (Thoppilan al. 2022 Pope al.) (023). Conversely, smaller language models, while computationally efficient, often lack the expressive power of their larger counterparts and achieve subpar performance. While reducing the inference cost of larger models, e.g., via quantization or pruning, or improving the performance of the smaller models, c.g., via knowledge distillation (KD) (Hinton et al. 2015). constitute natural approaches enable a favor- able performance versus inference cost trade-off, these approaches frequently result in unacceptable performance gap compared to the high-quality large models. This has inspired growing literature on designing mechanisms that combine both large and small models at inference approximate the performance of the larger models without incurring their high computational cost.
Among conventional approaches, model cascading aims to identify easy instances where smaller models suffice to achieve good performance, and soliciting larger models only on a subset of hard instances (Rowley et al. 1998 Xu et al. 2014) tasks (Cai et al. 2023b). Different from such task- or instance-level cascading, speculative decoding (SD) (Leviathan et al. 2023 Chen et al., (203) exploits the token-level variability in the computation demand during LLM inference by interac- tively invoking a small "draft" model and a large "target" model. At given stage during inference, the draft model generates successive candidate tokens for multiple inference steps via autoregressive decoding. The target model then verifies the candidate tokens via parallel decoding, and employs rejection sampling accept a subset of candidate tokens contiguous positions.
The main objective of SD is to speed up text generation while guaranteeing that the decoded tokens follow the target model distribution. SD relies on the insight that the combined cost of autoregres- sive decoding with a small draft model followed by parallel verification with the target model is
'Student Researcher at Google Research. Advising contribution. <yezhou www.ooototooooduuu..... <jfkagy wom. and crishabhagarwalee gowgle..oo....mmmm....
Corresponding authors:

---

Published as a conference paper at ICLR 2024
Distill (Greedy)
Standard (Greedy)
DistillSpec Greedy)
Standard INon- Mon-Greedyl
Speculative Decading
150
+10%
Speedup (%)
+41%
+28%
100
+10%
+23%
+25%%
+19%
-16%
+18%
50
XSum
GSMBK
CNNDM
WMT BBH (Transfer)
Figure 1: Performance comparison of standard speculative decoding (SD) vs. our proposed DistillSpec, with small- and XL-sized models from the T5 v1.1 family (Raffel al. 2020 being utilized the draft and the target models, respectively. DistillSpec enhances SD speed by better aligning the draft with the target via white-box knowledge distillation, resulting in a consistent 10-45% speedup improvement over standard SD across various datascts. The distilled draft model from GSM8K transfers well to 23 unseen BigBenchHard tasks (Suzgun et al. 2022), resulting in an average speedup of 26%. See $ for additional details.
lower than the cost of autoregressive decoding with the target model alone. However, the realized inference cost reduction or latency improvement crucially depends on the acceptance rate of the draft- traft-generated tokens by the target model, which can be shown to be directly tied the alignment between the token distributions of the draft and target models. Thus, a successful application of SD hinges on identifying compact draft model that simultaneously has small autoregressive decoding cost and is closely aligned with the target model.
In this work, we propose DistillSpee, a novel approach that relies on KD (Hinton et al. 2015) to obtain an effective draft model. Unlike the standard application of KD which primarily focuses on improving the task performance of a small student model, DistillSpec aims at aligning the stu- dent (draft) model with the teacher (target) model to enhance the acceptance rate during SD. We undertake a comprehensive exploration of the distillation process for speeding up SD, considering several factors including the composition of training data, choice of divergence functions define the training objective for KD, and decoding strategies. Notably, our findings underscore that using model-generated data crucial for ensuring strong student-teacher alignment across various tasks via KD, and that the selection of the best-performing divergence function in DistillSpec is highly uusk-dependent and sensitive to the decoding strategy (i.e., greedy versus non non-greedy). Furthermore, we explore the utility of DistillSpec for lossy SD (Leviathan et al., (0233) which allows for sampling away from the target model distribution. We show that combining DistiliSpec with lossy SD enables more fne-grained control over the latency versus task performance trade-off.
Finally, we carry out systematic study of how to design an efficient inference scheme in a practical setting where one has access to multiple models of increasing size and quality. Leveraging the insights that we have laid out about KD and SD, our study concludes that the most effective strategy involves first distilling large model into smaller one the potential target model for performance optimization, followed by DistillSpec for distilling an even smaller model to be used as the draft model in SD. This approach results in remarkable 6 10x reduction in latency, compared to standalone non-distilled target model of the same size, with minimal performance degradation.
Our key contributions are:
We propose DistillSpec, method that uses KD to enhance draft model alignment with the target model ($4), and show that our method can improve SD speed by 10-45% while preserving model performance across diverse datasets under greedy and non-greedy sampling (Figure
(ii) We conduct extensive analysis of the optimal distillation recipe ($5.2) for model alignment, encompassing factors such as training data generation and different divergences, and emphasiz- ing the distinctions between standard KD and distillation tailored for SD.
(ii) We extend DistillSpec to lossy SD, enabling refined control over the quality-lateney trade-off. Moreover, we offer insights for combining KD and SD when several models are available ($5.).
2 RELATED WORK
Speculative decoding (SD). Due to the inherent sequential nature of autoregressive decoding, the primary latency bottleneck in LLM inference arises from memory read/write operations rather than arithmetic computations (Pope et al. 2023). Speculative decoding (Leviathan al. 2023
2

---

Published as a conference paper at ICLR 2024
Chen al. 2023) (SD) addresses this challenge by utilizing compact draft model to generate batch of tokens sequentially, while validating them in parallel with a larger target model. Prior to SD, various parallel computing paradigms have been explored for autoregressive models, including block parallel sampling (Stern et al., 2018), shallow aggressive decoding (Sun et al. [021], and aggressive decoding (Ge et al. 2022). However, these approaches not readily adaptable to typical language models due to potential deviations from target model's distribution, strict input constraints, or limited support for general stochastic sampling. Notably, recent variants of SD have considered different interactions between the draft and target model to reduce unnecessary computation (Kim et al.) 2023) and incorporated parallel computation along the batch axis, sometimes combined with token tree verification, as seen in SpecTrr (Sun et al. (023). SpecInfer (Miao et al. (023), and Medusa (Cai et al. 2023a). In contrast, our work focuses on enhancing SD by improving the alignment between the small draft model and the large target model through KDD, which does not require any changes to serving infrastructures already implementing SD and is complementary to the recent variants of SD. Furthermore, we conduct a systematic study of lossy SD for providing nuanced control over the trade-off between quality and latency for specific serving models.
Knowledge distillation (KD) for LLMs. KD (Buciluà et al. 2006 Hinton et al. 2015),
which trains high-quality smaller student models with the supervision of larger teacher models, has emerged vital technique for reducing inference cost while maintaining model quality across a range of domains. In the context of LLMs, prior uses of KD (Taori et al.) 2023 Fu et al., 2023) have mostly focused on black-box KD, wherein only the teacher's 's output generations, generally via APIs. are accessible during student training. However, with the proliferation of open-source LLMs (Zhang al. 2022 Touvron et al. 2023), which enable access to teacher weights and logits, there is growing interest in white-box KD. White-box KD allows student models to benefit from richer supervision signals provided by white-box teacher models, leading to enhanced language 2023 Wen et al. (023)..
abilities (Agarwal et al. 2023
Gu et al.
Unlike prior works focused on creating highly capable standalone student models, we harness KD to foster closer collaboration between smaller and larger models in SD, which may be particularly valuable when a small distilled model alone cannot meet stringent quality requirements. While Stern et al. (2018) use black-box KD approach (SeqKD) to enhance blockwise parallel decoding, their samples are generated from the large target model, which is prohibitively expensive for LLMs. Furthermore, they ignore the teacher model's logits and train their draft model using only one-hot teacher labels a reasonable choice for greedy decoding but a less effective one for non- non-greedy sampling (Figure D. Concurrently, Liu al. (2023) propose to improve SD using KD, but they assume an online setup with a changing query distribution, and focus on improving the acceptance rate rather than reducing the actual latency.
BACKGROUND: SPECULATIVE DECODING
Notation. Given an input sequence x comprising tokens from pre-defined vocabulary, lan- guage model M provides a distribution over possible output sequences y. Suppose we employ SD with compact draft model Mg to assist larger target model Mg. Let p(yo z, yee)) and q(ye x, Bee)) represent the distributions governing next-token predictions at time step t for Mc and Mg. respectively, given the context م {sr, v<t)- Given input as prefix, let Psr(v) ar) and 4sr(y] x)) represent the distributions governing the sequence sampled autoregressively from Mp and My, respectively, where the generation stops either when an uudou--e token is sampled or the maximum sequence length T' is reached. For simplicity, we use p(y:) and q(y:) as shorthands for p(yel x. Met) and g(ve Ix, Met), whenever the context p is clear. Similarly, PST (y) and qgr(y) serve as shorthands for PSr(B] z)) and 9<r(y) z). whenever the input x is clear.
Speculative sampling.
Standard SD uses a procedure called speculative sampling to generate tokens from the draft model while maintaining the same output distribution as the target model. As detailed in Algorithm A. ( each step of SD works as follows. First, a block of ? tokens, denoted as Ut...., --+ is autoregressively sampled from g(v))... q(V++-11) Next, the tokens are verified in parallel by passing them to M, as a whole block, which sequentially accepts token Me++i with probability min (1, p( bt++)// (3)++)). If any token Btos is rejected before the end of the block, the subsequent tokens are discarded and the rejected token is resampled from the adjusted distribution p' (3)++))
max(0, p()(+s) - ((v++))) otherwise, the drafted tokens are all accepted and an extra token is sampled from p()++>) and appended to the output sequence. This process guarantees that the sequence of accepted and resampled tokens follow the same output distribution
3

---

Published as a conference paper at ICLR 2024
as p(yc+s) (Leviathan et al. 2023). The procedure is repeated until an end-of-sequence token is accepted, or the maximum sequence length T has been reached.
Efficiency measure: acceptance rate. Each SD step takes constant amount of time, so the wall-clock time scales linearly with the number of steps. This number is equal to the total number of times that the target model rejects token, plus the number of blocks accepted as a whole, where the latter term is small for large "y. This motivates us to use the acceptance rate as a surrogate efficiency measure for the wall-clock time. For ideal SD process with 7 = oc, we define the sequence-level acceptance rate (x) for a given input x as follows:
lul B(x,
E [number of accepted tokens in generating y| Ey>per(y 2))
a(x)
E [number of tokens in y]
Lp(x)
where B(x, U-ee) == Ey.-)(u.) [min (1, p(ue)/ tations are taken over the randomness in SD. Since speculative sampling preserves the output dis- tribution of the target model, the denominator is simply equal to the expected length of the target output Lx)x) ;= Ey-pgr(yl acceptance rate is directly related to the expected total number of rejected tokens (1 (1-a al))())) Lp()), which lower bounds the expected number of SD steps.
is the token-level acceptance rate, and expec-
[[v]]. which is invariant to the choice of draft model. Therefore, the
Efficiency measure: block efficiency. In practice, SD is usually employed with a fixed finite block size y; thus, a more relevant efficiency measure is the block efficiency т. Given input x, we compute the block efficiency (I) as the expected number of accepted tokens per block. Note that, for a given block size 7, the maximum value of (xx) is 7 1, corresponding the case where all drafted tokens are accepted and augmented with an additional token sampled by the target model. If we assume that the loken-level rates 3(x, y<e) are i..dd. then the sequence- level acceptance rate satisfies = E [3] and r(x)= = (1 a ))((1 a:) (Leviathan et al. 2023
Wall-clock time improvement. For given block efficiency T(x), the expected speedup of SD is given by (x))/eey + 1), where the relative latency is the ratio between the times elapsed when making a single forward pass through the draft model M, and through the target model Mp.
DISTILLSPEC: KNOWLEDGE DISTILLATION FOR SPECULATIVE DECODING
As described in $ | speculative decoding (SD) leverages a small (draft) model to reduce the latency of decoding from the larger (target) model distribution without any performance drop. However, the realized speedup critically depends on how "well-aligned" the draft model is to the target model. In this work, we propose DistillSpec, general framework that improves SD by better aligning the tar- get model and draft model using white-box knowledge distillation (KD). We first present KD-based training of the draft model, and highlight how our objective of enhancing SD via KD influences our selection of training data generation method and divergence function-twe two key ingredients of DistillSpec. We then discuss how DistillSpec can be extended to lossy SD.
Let the draft model M be parameterized by DistillSpec utilizes predictions from the target model Mp as a source of supervision in training the draft model Mr We assume white-box access to both models, i.c., we can obtain their next-token distributions p(ve) and q(vt), and therefore we are able to generate samples from both models. Given a divergence function D that measures the misalignment between two distributions, KD-based training of the draft model seeks to minimize the divergence between the teacher (target) and student (draft) distributions over training set 9:
8" == arg min E(aa)> IMM
(- ). .e We note that flexibility in how 9 is constructed and the choice of divergence D gives rise to different possible KD algorithms. For instance, 9 may consist of task task-specific input-output pairs (X, Y) sequences generated from M, or M While forward KL (DFK1.) ) is the commonly used divergence for KD (Hinton et al. 2015), recent works (Agarwal et al., 2023 Wei et al. 2022) have shown the effectiveness
where ((M) ا!
alternative
divergences, including reverse KL (DRKL). Jensen Shannon divergence (Dıspja)), ). and total varia- tion distance (DryD). Further details on each divergence can be found in Appendix A. Table summarizes various distillation recipes, each being specialized instance of Algorithm A.2
4

---

Published as a conference paper at ICLR 2024
Table 1: Summary various KD algorithms terms of training data and divergence D (ef. Eq. ct al.
Wen
2023);
Agarwal et al. 2023 also consider other D) we list the most representative one.
Name
Divergence
Training Data
S (Kim & Rush 2016
Black-box KD)
FKL
Data generated by teacher Mr
Supervised KD (Hinton al. 2015
FKI.
Fixed dataset of input-output pairs
ImitKD Lin et al. 2020
FKL
Fixed dataset + data generated by My
[-Distill Wen et al. 2023
TVD
Data generated by M., and M,
On-policy GKD (Agarwal et al. 2023
FKI./ ISD
On-policy data from student M,
Our choices for 9 and D are guided by how the resulting distilled model, once employed as draft model, improves the speed of SD. Towards this, we first highlight the role that DrvD between p(v)) and (vr.) plays in dictating the acceptance rate ($ 3 key efficiency measure for SD.
TVD proxy for the acceptance rate. Leviathan et al.. (2023, Corollary 3.6) show that the token-level acceptance rate (1,, yet) satisfies (xx Mct) = DryD(p(Be) (v+)). Hence, Eq. implies that maximizing the sequence-level acceptance rate (xx) is equivalent to minimizing the expected DTVD between p(y)) and ((v)) over the output sequence distribution of My i.c.:
lyl
a(x) = - Eymper(ylas
D(vnlo(un) (z)
(3)
Choice of divergence. Based on Eq. 3 it appears that directly minimizing DIVD may be a princi- pled objective for draft model distillation. While optimizing (P, q) is theoretically inspired, our empirical study shows that such an objective may not consistently yield optimal results. We find that the choice of the most suitable divergence is highly task-dependent (S (.2).
Choice of training data. As for S, one could resort to an existing ground-truth dataset, however the teacher's output distribution may deviate from the ground-truth distribution despite the teacher having been fine-tuned on it. Moreover, ground-trute datasets are often limited in size, so training only on such data could result in overfitting. To mitigate these issues, we use gener-ggnerrted outputs for distillation. Specifically, we prompt the model with a task input sampled from ground- truth training dataset, and use the model response as data for distillation. Both the teacher and student models may be used to generate the distillation examples.
Model-generated distillation data. Eq. 3] suggests optimizing the expected DTvD over outputs generated from the teacher. Decoding from a large teacher is generally prohibitively expensive, pecially at the scale of dataset required for KD. To reduce the generation cost, we explore using on-policy data during distillation, i.e., output sequences sampled from the student itself. Besides being more computationally efficient compared to teacher generations, this approach motivated by Gu et al. 2023);; Agarwal et al. (202)), who have shown that distilling on-policy data improves student task performance. However, different from these prior works, our primary focus is on im- proving the student-teacher alignment. Thus, it may not be immediately clear whether minimizing the expected DTVD over on-policy (student-generated) data ensures an improved acceptance rate, which is computed an expectation over the teacher's output distribution (cf. Eq. Our following result shows that this is indeed the case.
Theorem 4.1. For SD, if the draft model My achieves on-policy KD loss Eg-X, [x)
x)). then the sequence- acceptance rate is at least
Ez~x [a(=)] 1-
[,)(2)
When the target output length is always T, the bound simplifies to E [a(x)] Te.
We defer the proof to Appendix A.2 Intuitively, it builds upon the following insights. If the on- policy KD loss is small, then, for any 1 < < T, the same loss evaluated only at the t-th token should also be small. Since the first token generation is independent of any other tokens, small value of on-policy KD loss ensures that the first token distributions of the draft and target models close. Then, inductive argument shows that once the draft and target are similar on the first t tokens, the distributions of the (t + 1) token should also be close. Our proof makes this argument rigorous by utilizing variational representations of DTVD, leading linear error bound in T.
5

---

Published as a conference paper at ICLR 2024
SupervisedKD
SeqKD
ImitKD
F-Distill
GKD
3.40
XSum
5.00
CNNDM
5.10
WMT
4.25
GSMBK
Block Effileecy
3.20
4.50
5,00
3.00
3.75
2.80
4.00
4.90
3.50
Witthout Dist
Without Distilation
without Distllation
3.25
WithettDis Distilation
2.60
4.80
Figure 2: Distillation enhances block efficiency (r) across diverse datasets, highlighting the superiority of model-generacee data over fixed ground truth data (SupervisedKD) and emphasizing the importance of white- box distillation in addressing SeqKD's subpar performance.
DistillSpec with lossy SD. DistillSpec enhances SD efficiency without any quality loss com- pared to the target model Mp. In practical upplications, a reduction in model quality may be jus- tified in order support even faster inference. For such scenarios, extend DistillSpec to lossy SD (Leviathan et al. 2023), which uses lenience function f(p.e) that modifies the acceptance probability from min (1, p((k) /q(()) to min (1, ())(()) (cf. 3 Here f :0 1)2 K+ increasing and decreasing in its first and second arguments, respectively, and E [0, 1] is a free parameter (ef. Algorithm A.I). In this work, evaluate multiple choices for the lenience functions: fün(p, e) = p/c, (p, e) p/e", and (exp(p, E) p². For example, when the lenience function fag e) and accepted by the target, thus enabling faster inference at the expense of potential drop in generation quality. Lenience was discussed by Leviathan al. (2023) in the context of fin and their treatment focuses solely on latency improvements, whereas we explore the use of different lenience functions as a precise control mechanism to achieve the desired quality-latency profile.
0.1, token x sampled from (3va) becomes hundred times more likely to be
5
EXPERIMENTS
5.1 ENHANCING SPECULATIVE DECODING THROUGH DISTILLATION
We evaluate the effectiveness of KD in improving the speed of speculative decoding (SD). We fol- low Leviathan et al. (2023) and investigate its impact on the acceptance rate a. block efficiency т, and actual latency speedup with a batch size of under greedy sampling (T = 0) and standard temperature sampling (7=1).
Tasks and models. Following Leviathan et al. 2023), , we evaluate two model types: 1) GPT- like decoder-only Transformer models trained on LM1B task (Chelba et al.) (013) using a standard autoregressive objective, where the target and draft models have 234M and 33M parameters, re- spectively; and 2) standard encoder-decoder T5 v1.1 models (Raffel et al. 2020 fine-tuned on four different tasks, with T5-XL (3B) and T5-Small (77M) serving as the target and draft mod- els, respectively. We utilize two datasets from the T5 paper, namely WMT EnDe (Bojar et al., 2014) and CNN/DM (Hermann et al.] (015), which deal with translation and text summarization, respectively. The remaining two tasks used test T5 models XSum (Narayan al. 2018) and GSM8K (Cobbe et al. 2021), which deal with abstractive summarization and arithmetic reasoning, respectively. See Appendix for more details.
Decoding speedup. Figure I shows that the impact of distillation on SD speed is evident, con- sistently yielding 10-46% improvement across various datasets. This effect is most pronounced when employing greedy decoding. The performance of KD for different block sizes and decoding strategies across five datasets is presented in Table C.T (Appendix). These findings demonstrate that KD significantly enhances the acceptance rate and block efficiency of SD for both decoder-only and encoder-decoder models across all datasets. Distillation algorithms utilizing model podel-generated data consistently outperform other approaches, resulting in ~20% additional speedup compared to standard SD on LMIB, XSum, CNN/DM, and GSMSK.
Block efficiency. Figure 2 presents comparison of block efficiency across different algorithms, employing temperature sampling (T = 1) with a block size 7 = 7. The figure underscores the utility of model- poodel-generated data: using ground-trutt truth data (i.e., Supervised KD) ranks lowest across all settings. In contrast, ff f-Ditilll and GKD, which only use model-generated data, significantly out- perform other KD variants. The subpar performance of SeqKD, despite being purely trained data

---

Published as a conference paper at ICLR 2024
Drafter Standard Drafter
Empinical Block Efficieccy
4.5
xsum
condm
Block afficiency
wint
4.0
0.66
ImiKD
SupKD
3
3.5
Sequi
3.0
0.64
F-Distill
0.62
GKD
20
40
250 5.00
750
1000 1250
2.5
3.0
3.5
4.0
Training Efficiency (Waltimme hrs)
Sorted Test Data Index
Theoretial Block Efficiency
(a) acceptance rate vs, wall-time
(b) instance-level block eff. (T)
(c) agreement of
Figure 3: (a) White-box KD using teacher logits and model-generated data crucial. Draft model's on on-policy data can be as effective as target model data. GKD achieves the best wall-time performance improvement on XSum, (b) Distillation improves the block efficiency for all examples on GSM8K. (c) Empirical block efficiency aligns well with its Dryn-based theoretical counterpart.
generated by the target model, suggests that white-box distillation (i.e., supervision from the target model's logits) is vital for SD. This is corroborated by Figure 3a which illustrates the evolution of the acceptance rate throughout training. Supervised KD ranks lowest, and its performance plateaus early during training due to the static nature of the dataset. In contrast, algorithms using model- generated data lead continual improvement of the acceptance rate. Despite f f-Distill being much more computationally costly than GKD due to the use of teacher- seabber-geeerteed data, both algorithms exhibit comparable performance. Notably, GKD achieves the best wall-time performance improve- ment. See Appendix C.1.2 for more visualizations of performance improvement during training.
We also investigate whether KD improves block efficiency universally or impacts limited subset of examples. Figure 3b depicts the improvement per example. We observe consistent gains in block efficiency across most examples, which can also be seen in various datasets (see Figure Fig- ure 3c illustrates strong agreement between theoretical and empirical block efficiency values for several distilled models (each model is represented as a filled cirele). Despite theoretical values oc- underestimating empirical values, possibly due to potential deviations casionally overestimating from the i.i.d. token-level assumption (ef. $3). the ranking of distilled models remains highly con- sistent. In summary, these findings largely confirm that KD effectively optimizes block efficiency.
Transferability of distilled models. We next examine the transferabilityy of distilled models on diverse datasets unseen during training. We use a draft model distilled on GSM8K and test its ability to speed up SD on zero-shot chain-of-thought (CoT)) prompting over 23 reasoning tasks from the BigBenchHard suite (Suzgun et al. 2022). The results, illustrated in Figure I indicate effective transferability to other datasets. Compared to standard SD, the distilled model significantly enhances average decoding speeds, yielding speedup improvements from 1.93x and 1.78x to 2.21x and x.02x for greedy and non-greedy decoding methods, respectively. Further analysis in Figure C.1 reveals that using our distilled T5-Small as draft model is also compatible with larger target models (T5-XXL) despite being distilled from different-sized model (T5-XL). Despite being not fully optimized, this configuration consistently outperforms standard SD by 7% 37% across various datasets. See Appendix C.I for more details.
5.2 DISTILL.PECC RECIPE
We now focus on identifying the optimal KD approach for SD. Following the training and evaluation protocols in $ S.I we explore four training data construction methods and four divergence functions on XSum and GSM8K. Specifically, we explore the following variants of training data: 1) fixed ground-truth dataset Dirain: 2) data generated only from the draft model My. 3) data generated only from target Mp. 4) data generated from both M and M, in equal proportion. We also examine the following divergences: 1) forward KL (FKL), 2) Jenson-Shannon divergence (JSD), 3) reverse KL (RKL), and 4) total variation distance (TVD).
Importance of training data and divergence in DistillSpec. Figure 4 illustrates the block efficiency improvement on XSum and GSM8K, in line with observations from $ [.11 We note that using model-generaaed data (last three rows) yields superior performance than using fixed dataset (first row). Specifically, on XSum with greedy decoding, using data generated from both MD and Mr leads the best performance, with JSD slightly outperforming the other divergences. However,
7

---

Published as a conference paper at ICLR 2024
Block Efficienty Improvement
Block Efficiency Improvement
Block Ufficinccy Im improvement
Fixed 1Drua)
0.01
8:08
8:13
0:17
Fixed (Draal
0:23
0.29
0.29
0.26
Fixed 1Dhaal
0.22
0.31
0.32
0.29
64
Student (xe)
0.48
0.49
0.45
0:46
Student (M@)
0:61
0.42
0.43
0.38
Student (Mt)
0.42
0.43
0:45
0.41
MC+
0.52
0.5
Mr Me
0.49
0.43
0.47
0:47
0.44
Teacher (#6g)
0.51
0.53
0:51
Teacher (Mp)
0.35
0.43
0.46
21
Teacher (Me))
0:44
0.49
0.43
FKL
J5D
RXL
TVD
FXL
150
AKL
TVD
FKL
150
RKL
TVD
(a) XSum: T (greedy)
(b) GSM8K: 7 (greedy)
(c) GSM8K: (non-greedy)
Figure 4: DistillSppe recipe. We report improvement in empirical block efficiency post-distillation on (a) XSum with greedy sampling, and GSM8K with (b) greedy and (c) non-greedy sampling. Which divergence function and data construction lead to the largest block efficiency improvements highly depends on the task. They should be treated as hyperparameters to be tuned on our task of interest.
on GSM8K with greedy decoding, FKL with only grantt-geneated data emerges as the best setup. In contrast, with temperature sampling (at T = 1), a different trend is observed as RKL combined with data generated by Mp is the most effective setup. See Appendix C.2. I for results on different datasets and decoding strategies. Nonetheless, using only graft-geneeated data is found to be competitive.
Impact of distillation on draft quality vs. compatibility. also study how different distillation recipes affect task performance and whether there is any one design choice that is simultaneously optimal for improving both draft model task performance and its utility for SD (cf. Figure C.23
C.24p. Similar to our earlier observations, the use of generated data is paramount for improving draft performance. Notably, utilizing data generated from M yields comparable or superior results compared to using data generated from Mp. However, which KD algorithm is optimal largely depends on the task at hand and the underlying decoding strategy. Figure FigureBaalighht haghigllgght an interesting dichotomy between block efficiency improvements and task performance gains via KD: distilled models with high task performance do not necessarily translate into more aligned drafters for SD.
Interestingly, although TVD is the objective we should aim to optimize for SD (cf. Eq. 33. its direct optimization does not yield the best performance in most of the settings explored. We generally find that the choice of divergence in KD is a hyperparameter that needs to be tuned based on the task at hand and decoding strategy used. For training data construction, pose using the draft model M for data generation as it can achieve similar or superior performance compared to the target model Mp, but at much lower cost.
Recommendation.
pro-
5.3 QUALITY VERSUS LATENCY TRADE-OFF
Lossy speculative decoding. We analyze the quality- latency trade-off using lossy SD variants, as detailed in Algorithm| A.I] As Figure 5b illustrates, employing either KD (x) or SD (x) alone does not fully bridge the performance or latency gaps, respectively. In such case, a leniency parameter (E) can help interpolate between these two approaches, point within a given group corresponds to a different value of E. As the GSM8K experiment shows, be limited: even using permissive lenience of ε 10-5, füi still
demonstrated in Figure 5b where each
the power of interpolation results in high performance but high latency, while fsq traces similar but slightly extended trade- off curve. Although fexp makes interpolation possible, yields a worse quality-latene trade-off. Interestingly, it is possible to significantly reduce latency while preserving most of the quality on GSM8K, possibly because many tokens are inconsequential for final performance, and variety of proposals can be safely accepted with minimal effect on generation quality. See Appendix C.3. for a comparison between non-distilled and distilled draft models, where we show that distilled draft model enables much better quality vs. latency trade-off.
DistillSpec meets model garden. In many practical scenarios, we have access to multiple mod- els of different sizes-a a model garden-to to design the inference pipeline. We emulate this setting by focusing on the five model sizes in the T5 model family: T5-Small (77M), T5-Base (250M), TS-Large (800M), T5-XL (3B), and T5-XXL (1IB). We study the quality-latency trade-off curves obtained from applying KD and SD as follows: 1) raw: deploying supervised fine-tuned (SFT) T5 models; 2) distilled: applying KD by distilling smaller models from the larger T5 models; 3) spec- ulative: applying SD using T5 models; and 4) DistillSpec: applying KD on T5 models and using SD with distilled models target and draft models.
8

---

Published as a conference paper at ICLR 2024
0.50
FKL
Divergence
Accuracy
30
J
25
Improvementt
RKL
0.40
TVD
fun e) = p/r
0.35
TrainData
Performan
fa p/c²
Teacher (MMg)
(P, (p. =p
0.30
Student (M()
10
No Lenience
Fixed (Dras)
Raw Draft
AÇ + Mr
Distilled Draft
0.5
1.0
1.5
2.0
2.5
3.0
3.5
Accuracy Improvement
Relative Latency
(a) quality vs. SD compatibility (GSM8K)
(b) DistillSpec with lossy SD (GSM8K)
Figure 5: (a) The improvement on speculative decoding and downstream task performance are only weakly correlated. A high quality distilled model does not imply it can be an effective draft model in speculative decoding. (b) We employ leniency precise control mechanism to achieve the desired quality-latency profile.
Performance Rouge2
22
Performance Accuracy
30
20
Target Model
Serving Strategy
Target Model
Serving Strategy
18
T5-Small
Raw
-20
T5-Smail
Raw
Distilled
Base
Distilled
TS Large
Speculative
TS-Large
Speculative
T5 XL
DistillSper
10
XX-
DistillSpec
14
T5 XXL
TS-XXL
0.0
2.5
7.5
10.0
2
10
12
14
Relative Latency
Relative Latency
Figure 6: DistillSpec excels in both quality and latency, offering a remarkable 6.4x and 10.7x latency reduction on XSum (left) and GSM8K (right), while maintaining nearly identical performance.
Figure 6 shows that SD effectively shifts the trade-off curve leftward, especially with larger model sizes. However, its efficacy diminishes with smaller model sizes when the computation time between the draft and target models is closely matched. In contrast, distillation, which optimizes the model for downstream task performance, appears to offer a superior trade-off between quality and latency, particularly for smaller models. Conversely, reverse trend observed for larger model sizes when evaluating the model with temperature sampling. Figure C.30a indicates a substantial gap between the distilled model and the larger teacher model, while the SD-based method is able to significantly reduce latency. This suggests that when stringent performance and decoding strategy constraints are in place, SD remains a valuable approach. Our DistillSpec method, which combines the benefits of distillation and SD, consistently achieves the best trade-off between quality and latency, yielding an impressive reduction in latency while maintaining nearly identical performance. Specifically, DistillSpec reduces relative lateney from 17.3 to 2.7 and from 15.0 to 1.4 on XSum and GSM8K, respectively, representing speedup improvements of 6.4x and x. In contrast, the Rouge2 score only marginally decreases, from 23.1 to 23.0 on XSum, while the model accuracy GSM8K actually improves from 33.1 to 34.8.
CONCLUSION
In this paper, we evaluate the efficacy of white-box knowledge distillation (KD) in enhancing spec- ulative decoding (SD) through improved alignment between target and draft models. A thorough analysis is conducted understand the impact of training data construction and divergence func- tions on KD performance. We underscore the significance of utilizing model-generatee data and argue that employing the draft model's on-policy data during KD a cost-efficient way of improv- ing model alignment. Additionally, we assess the trade-off between quality and latency within the scope of lenience and availability of multiple models of varying quality and size (model garden), concluding that KD enables a superior trade-off compared to standard SD. The optimal strategy in- volves first applying KD for downstream task performance, followed by SD using a distilled draft model. resulting in six ten-fold decrease in latency with negligible performance loss. Our study contributes novel insights into white-box KD algorithms for LLMs and provides guidance for strik- ing an effective balance between quality and latency using KD and SD.

---

Published as a conference paper at ICLR 2024
REFERENCES
Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, Matthieu Geist, and Olivier Bachem. On-policy distillation of language models: Learning from self generated mistakes. arXiv preprint arXiv:2306.12: 13649, 2023.
Rohan Anil, Andrew M Dai, Orhan Firat, Mclvin Johnson, Dmitry Lepikhin, Alexandre Passos, Siamak Shakeri, Emanuel Taropa, Paige Bailey, Zhifeng Chen, et al. Palm 2 technical report. arXiv preprint arXiv:2305. 10403, 2023.
Ondrej Bojar, Christian Buck, Christian Federmann, Barry Haddow, Philipp Koehn, Johannes Leveling, Christof Monz, Pavel Pecina, Matt Post, Herve Saint-Amand, Radu Soricut, Lucia Specia, and Ale s Tamchyna. Findings of the 2014 workshop on statistical machine trans- lation. In Proceedings of the Ninth Workshop Statistical Machine Translation, pp. 12- 58, Baltimore, Maryland, USA, June 2014. Association for Computational Linguistics. URL http: / wwww aclweb org/ anthologyyl.ooooooooooooooooooaa ww/pg//////////////// /wtp://////1////3333333
Cristian Buciluă, Rich Caruana, and Alexandru Mieulescu-Mizil Model compression. In Pro- ceedings of the 12th ACM SIGKDD international conference on Knowledge discovery and data mining, pp. 535-541, 2006.
Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, and Tri Dao. Medusa: Simple frame- work for accelerating Ilm generation with multiple decoding heads. httts://githhb. com/
FasterDecod /Medusa 2023a.
Tianle Cai, Xuezhi Wang, Tengyu Ma, Xinyun Chen, and Denny Zhou. Large language models as tool makers. arXiv preprint 23Xi::330..71220 2023b.
Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robinson. One billion word benchmark for measuring progress in statistical language modeling. arXiv preprint 1312.v1312.330 2013.
Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean-Baptiste Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.013/8, 2023.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021.
Yao Fu, Hao Peng, Litu Ou, Ashish Sabharwal, and Tushar Khot. Specializing smaller language models towards multi-step reasoning. arXiv preprint arXiv:2301.12726 2023.
Tao Ge, Heming Xia, Xin Sun, Si-Qing Chen, and Furu Wei. Lossless acceleration for seq2seq generation with aggressive decoding. arXiv preprint arXiv:2205. 2022.
Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang, Knowledge distillation of large language models.
arXiv preprint ar0iv:2306. 08500,0085,,.... 2023.
Karl Moritz Hermann, Tomas Kocisky, Edward Grefenstette, Lasse Espeholt, Will Kay, Mustafa Suleyman, and Phil Blunsom. Teaching machines to read and comprehend. In Advances in neural information processing systems, pp. 1693-1701, 2015.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv preprint 1xXiv:150300253 2015.
Ferenc Huszár. How (not) to train your generative model: Scheduled sampling, likelihood, adver- sary? arXiv preprint a51 1.05101, 2015.
Schoon Kim, Karttikeya Mangalam, Jitendra Malik, Michael W Mahoney, Amir Gholami, and Kurt Keutzer. Big little transformer decoder. arXiv preprint NXXv:2302.008633. 2023.
Yoon Kim and Alexander M Rush. Sequence-level knowledge distillation. arXiv preprint arXiv:1606.070 07947, 2016.
10

---

Published as a conference paper at ICLR 2024
Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. In International Conference on Machine Learning, PP. 19274- 19286. PMLR, 2023.
Alexander Lin, Jeremy Wohlwend, Howard Chen, and Tao Lei. Autoregressive knowledge distilla- tion through imitation learning, arXiv preprint arXiv:2009.0722 07255,9 2020.
Xiaoxuan Liu, Lanxiang Hu, Peter Bailis, lon Stoica, Zhijie Deng, Alvin Cheung, and Hao Zhang. Online speculative decoding. arXiv preprint arX01:2310.07117, 2023.
Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang, Rae Ying Yee Wong, Zhuoming Chen, Daiyaan Arfeen, Reyna Abhyankar, and Zhihao Jia. Specinfer: Accelerating generative Ilm serving with speculative inference and token tree verification. arXiv preprint arXiv:2305.09781, 2023.
Shashi Narayan, Shay B. Cohen, and Mirella Lapata. Don't give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization. ArXiv, abs/ 808.08745, 2018.
OpenAL. Gpt-4 technical report. ArXiv, abs/2303.08774, 2023. iettantiscchlaar. wtw.coo......ooroooooraaaaaa233.
URL sttpss//aap
Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling transformer inference. Proceedings of Machine Learning and Systems, 5, 2023.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu, Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485-5511, 2020.
H. A. Rowley, S. Baluja, and T. Kanade. Neural network-based face detection. IEEE Transactions on Pattern Analysis and Machine Intelligence, 20(1):23-38, 1998. doi: 00.110933..555447
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp. 4596-4604. PMLR, 2018.
Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit. Blockwise parallel decoding for deep autore- gressive models. Advances in Neural Information Processing Systems, 31, 2018.
Xin Sun, Tao Ge, Furu Wei, and Houfeng Wang. Instantaneous grammatical error correction with shallow aggressive decoding. arXiv preprint (oX.v:2116.04997 2021.
Ziteng Sun, Ananda Theertha Suresh, Jae Hun Ro, Ahmad Beirami, Himanshu Jain, Felix Yu, Michael Riley, and Sanjiv Kumar. Spectr: Fast speculative decoding via optimal transport. In Workshop on Efficient Systems for Foundation Models @ ICML2023, 2023. URL https: //openreview. wwt.:/et.//orie
Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung. Aakanksha Chowdhery, Quoc V Le, Ed H Chi. Denny Zhou, et al. Challenging big-bench tasks and whether of can solve them. arXiv preprint arXiv:2210.09261, 2022.
Rohan Taori, Ishaan Guirajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An instruction-follooin llama model. https://githuit com/ tat su- hts//taan stanford alpaca] 2023.
Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, Heng-Tze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, et al. Lamda: Language models for dialog applications. arXiv preprint Niti:22001002333 2022.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint 200.02307..02288, 2023.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi. Quoe V Le, Denny Zhou, et al. Chain-of thothu prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 35:24824-24837, 2022.

---

Published as a conference paper at ICLR 2024
Yuqiao Wen, Zichao Li, Wenyu Du, and Lili Mou. f-divergence minimization for sequence-level knowledge distillation. arXiv preprint 1rXiv:2307.15/90, 2023..
Zhixiang (Eddie) Xu, Matt J. Kusner, Kilian Q. Weinberger, Minmin Chen, and Olivier Chapelle. Classifier cascades and trees for minimizing feature evaluation cost, Journal of Machine Learning Research, 15162):113-2214, 2014.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christo- pher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, Opt: Open pre-trained transformer language models. arXiv preprint wrwiv:2205.010000 2022.
12

---

Published as a conference paper at ICLR 2024
Appendix
Table of Contents
Method
14
A.1 Description of divergence functions
14
A.2 Justification of using policy data
15
A.3
DistillSpec algorithms
18
B Implementation Details
19
Datasets
................................
Models
................................
Distillation
20
4 Evaluation
21
Additional Results
22
Enhancing speculative decoding through knowledge lation
22
Distillation recipe)
35
Quality versus latency trade-ofi
38
13

---

Published as a conference paper at ICLR 2024
A METHOD
A. DESCRIPTION OF DIVERGENCE FUNCTIONS
Below are some common divergence functions used in distillation, given two discrete probability distribution P(C) and Q(C).
Kullback-Leibler (KL) divergence.
P(c) Q(c)
DKL (PIQ)== P(c) P(()) log
(A.1)
OEC
We note that the KI. divergence not symmetric, i.e., DKL (PIQ) * DKI(Q]P). As such, we refer to (Pl()) as the forward KL (FKL) and DAKL (P) == (Q) P) as the reverse KL (RKL) between P and Q. Minimization of the FKL under an empirical data distribution cor- responds to maximum likelihood estimation (MLE), which is the typical loss objective used in supervised learning given a fixed dataset.
Jensen-Shannon (JS) divergence.
Dis(P))Q) = 2 (DRL(P)IM) + DKL (Q)
where M = (P
(A.2)
Generalized Jensen-Shannon divergence (DJsD(a))
DISDLa (PIIQ) = DKx.. (p| 3P (1 (1- + -1- B) 3P (1))
+
(A.3)
Interestingly, it can be proved that lim Pasp(a) DKL (P (Huszár 2015). As such, Daspie behaves similarly forward KL for small values of B. Similarly, Disp[g] has similar
behavior to reverse KL for 3 close to 1, since DasDia (Q))))
Total variation distance (TVD).
P(c) - Q(c)
DTvD(PIIQ)
Σ
(A.4)
2
14

---

Published as a conference paper at ICLR 2024
A.2 JUSTIFICATION OF USING ON-POLICY DATA
In this section, we prove Theorem 4.1 which motivates our use of on-policy data. We follow the notations in Sections 3 and 4 In addition, we write c(x) == By-agr(vl [DIvD(M,IDM))) z)]
for the distillation loss of a single input x.
To ease the analysis, let decompose a(x) and e(x) into sums of contributions from each token.
Lemma A.I. For all a(x) 1- (m) 1=1 A and e(x) 2 $ Et=1 E, where:
(x))) E. (vl=))()))) (1 (pv x. BC.) g(v a v.b))))]] (x) = E E (w)
(A.5)
[ (estvl) DTVD (plve x, v,,)e),, a(v a. ve)))
(A.6)
Proof. By Equation 3 we have:
a(x) = Eyppgr(vla)
Elel
DTvD(p(M [., a(val as wee))
(x)x)
We rewrite Σ vv between the summation and the expectation to obtain:
== DTVD (p(v)) ((r)) E. ((sil)) PTVD(((()) q())) and swap the order
a(r) = 1 Etn Expper(vl
[(sl Drvb(p(u, x, ve,. g y<t))]
Lp(x)
which proves Equation A.S
Similarly, by definition (xx) and DTVD we have:
lpl
( = E (y (2)
I vx,e)) (x,
V
lul
2 Eyrogr(U1e
Drvp(p(y x. v(ee v<1))
T
Again, we can rewrite E/r (PvD(p(u.). (v))) = Σt ( DTVD ((v.)) and then swap the order between the summation and the expectation obtain:
T
(x)
2
E
E [A(ssto)) DrvD a vec). g(v. v<)))]
T
1=1
which proves Equation A.6
Lemma A. motivates
study As(x) and (x))) instead. Below we rewrite them in variational forms that will be used later. For this, we introduce some defintiions.
Definition A.I. For any sequence = E [P,Q}7 that consists only of letters P and Q we define M(z, y, 2) as the distribution of sequences sampled as follows:
1. Initialize a sequence of tokens as y;
2. If there are tokens, sample the 1-th wken fromfrom M, if 21 = P, and from My otherwise;
3. Repeat until an end-of-sequence token is sampled, or the sequence length has reached т.
We use the shorthand p and q to represent the sequence of & consecutive letters of and Q re- spectively. Let R denote the set of all possible strings of length k and Ω - [-1/2, 1/2] be generic function that maps a sequence t tokens to real number in [- 1]. We abuse the notation and assign 5(y) for all y !
15

---

Published as a conference paper at ICLR 2024
Lemma A.2. For all and all 1 S T:
A((x) sup
-:12-11112.1/] (E [6(u)] - -y,,((,,,.-- [d()))
(A.7)
Ex(x)= =
6:0'- (T [(y) - Ey [[(v)]}
sup
(A.8)
Proof. For a fixed pair z and y, we rewrite the total variance distance between p(ye x. yee) and (Vv x. Vet) as the following variational form:
DrvD(p(yt x. Vet), (V xt vc+))
E gerpll (a,8)))) [ō(vn)
Eye (x,aeee)) ((sr)
sup 6:2-1-1/21121]
After taking the expectations:
(x) E Fymper(vla)
DrVD(p(u a. Met). a(ve 1x, [r,ye))]]
AAr ((tr 1-1-1111221/21
= (ypper(vlz)
sup
111 v. 88}
sup 8/2-111/2.1/22
E Ey-M(,,,, 6:2'-1-1/2.121
sup
((b)= A(2 u. 8(y. 1)]]
where the third equality uses the observation that for any function S. Ey>per(0 (s) ([x y<+)) can be replaced with Ey>:M(rr,.pr-)) [1 w)]. and the last equality swaps the order between the expectation and the supremum. Finally, we move the expectations in A to merge with the expectation outside and obtain:
Ax(x) sup 6:03 6-12-111/21//1
(.) [1])) Ey-MM(r,,,P--q0 [[
sup
E (yN>M(a,,)) [6()) Ey- [8(y))]
6:00 --122111/
where the last step follows from our abuse of notation that 5(y) = 0 for all y & Ω. This proves Equation A.7 and Equation can be proved similarly.
Based variational forms of Ar(x) and (x(),, obtain the following lemma for bounding Ax(2) in terms of (x).
Lemma A.3. For all x and TT
1-1
A(xx) < 2 Σ Ex(x) E(().
(A.9)
Proof. It suffices to find an upper bound on Ey-M( [8())] and a lower bound on Ey-M((.22,p-11g [ō(z)] for all d ot + [-1/2,
For all I < t, we can replace the first P in By (6()) with q by only intro- ducing an error of (x))
Ey-M(r, Qh lpc 1) [6(y]]
= Ey-M((,,oo 1P) [[ 6) 1)(00//1] 8) 16(2/1]] + Ex(x) = Ey 6) [5(v)] + (x(x),
< EyrM((x,000)))
16

---

Published as a conference paper at ICLR 2024
where the second inequality holds because hatgo
define function 8' 22
[-1/2, 1/2], suchthatö(y) = (Ivl- ((v/=A)Ey [6(v/) and then apply Equation A.8
Now taking a telescoping over S k we obtain:
Ey-M(u, [(()] < My- [o(y)] [(()) +
Σ Bx(x).
(A.10)
Similarly, for all 1 S k S 1
replace the first P in F. [6(y)] with Q and
only introduce error of (x)):
E (0- [8())] E Mc 1p)
(Ivl= / (6(y')]]
= E
- E [6('']]]
2 (Ey~M((s,
| 1 k/= By= 14) [0(1)]] + Ex)))
= Ey 1q) [ó(v)] - (x(x).
Taking a telescoping sum over S 1s St yields:
I My 10) [6(y) 2 Ey-M((,, [8(1)] Σ Fx(x).
(A.11)
k=1
Subtracting Equation A.11 from Equation A. we have the following holds for all functions 8 - 1-1/2.1/2
2 t-1
Ey-M((,,))))) [6(y)] Ey 10) 10(y) S (x(
Ex(2)
Σ + E,())
which proves the claim.
Proof of Theorem 4.1 Summing Equation A.9
over 1 < 1 S T, we have:
A1(z) Σ
(1 + (TT - (z).
t=1
Combining this with Lemma A. yields:
a(x)=1 =1
1
Σ
1
1
Σ
(1 2(T - (x)
Lp(x)
Lp(x)
1=1
2T
2
Fx(x)
Ip xx)
2T
2 1
Lx(x) e(x).
which proves the theorem statement after taking the expectation over x X.
17

---

Published as a conference paper at ICLR 2024
A.3 DISTILLSPEC ALGORIHHM
Algorithm A.I Speculative decoding step
Require: target model My. draft model M, context p {r, yet). Hyperparameters: block size (Y), lenience Function f(p, e) (c= for lossless and for lossy decoding).
1: for all i to 1 do
(+++1)) <- Mr (144)))))) M++ ~ q++(())
P Sample 7 tokens from M2 autoregressively.
3: end for
4: (pe(v)..... ..... (v)) + (M,(a, v<t). (1
> Run Mp in parallel.
5: r. + Keilohe))
VL S +77
P. Compute the rejection thresholds. D Generate 7y random values. Determine the number of accepted tokens n.
6: u: ~ Uniform (0, .)............................ 447-1 ~ (ni
7: min (fi 0 & 7,44 u
8: n 7 then
Vttn norm (max (0,pt+n(0) (v) ge+n (v)))
Sample from the adjusted distribution.
10: else
11: Vetn (V)
b Sample from Mp.
12: end if
Return { vct+n+1}
b Append n tokens from M, and one token from Mp-
Algorithm A.2 Knowledge distillation
Require: target model My. draft model My, datasct (X, Y) containing input and possibly output Hyperparameters: fixed data fraction Aι [o, 1], student data fraction A2 [0, 1]. divergence function D, learning rate 1).
I: U1 Uniform(O, 1), U2 Uniform 1)
b Generate two random values.
uı < Aı then
3: B = [(16, (b) 1om) where (I,, M:) (X,Y)
b Sample inputs and outputs from (X, Y).
4: else
5:
B' {(=x)) where 2, X
Sample batch of inputs from X.
< X2 then
7:
B {(rs, yb) Yo where n. B¹, y1 ~ My (-120.)
Sample data from Mg.
8: else
9:
B {(Ts. (b) 1o) where B', Mr (+(s))
Sample data from Mp.
10:
end if
end
Return
1/ E(aWED VoD(M,IIM))
Update 0 minimize D(M, (MM))
18

---

Published as a conference paper at ICLR 2024
B IMPLEMENTATION DETAILS
B.1 DATASETS
In this section, we present comprehensive overview of the datasets employed in this study.
XSum Narayan et al. 2018). The Extreme Summarization (XSum) dataset serves as an eval- uation benchmark for abstractive single-document summarization systems. This dataset comprises 226,711 I news articles, sourced from BBC articles spanning the years 2010 to 2017. These articles encompass a wide range of domains, including News, Politics, Sports, Weather, Business, Technol- ogy, Science, Health, Family, Education, Entertainment, and Arts. Summarization performance is evaluated using ROUGE scores on the validation dataset split of XSum and primarily emphasizes ROUGE-2, but with similar trends observed for ROUGE-LSum and ROUGE-1. A maximum input length of 1,024 and a maximum output length of 64 are employed for distillation and evaluation.
CNN/DM (Hermann et al. 2015). The CNN/Daily Mail (CNN/DM) dataset is tailored for text summarization. It comprises abstractive summary bullets generated by humans from news stories on CNN and Daily Mail websites, presented in the form of questions with entities hidden. The questions can be answered using relevant passages from the source text. Similar to XSum, ROUGE scores on the validation dataset are reported. primarily emphasizing ROUGE-2 but with similar trends observed for ROUGE- -Sum and ROUGE-I. A maximum input length of 2,048 and maximum output length of 128 are used for distillation and evaluation.
WMT EnDe (Bojar et al., 2014). The WMT14 EnDe dataset stands as standard benchmark for machine translation. The task entails translating English text to German while preserving content, semantics and style. Evaluation is based on BLEU scores, which measures the similarity of machine- translated text to high-quality reference translations. A maximum input length of 80 and a maximum output length of 80 are employed for distillation and evaluation, with performance assessed on the original test split.
GSM8K (Cobbe et al. 2021). The GSM8K dataset comprises 8.5K high-quality, linguistically diverse grade school math word problems crafted by human problem writers. The dataset is divided into 7.5K training problems and IK test problems, with solutions typically requiring 2 to 8 steps involving elementary calculations using basic arithmetic operations. To enhance reasoning abilities, explored distillation alongside zero-shot chain-of-thouge (CoT), as described in (Agarwal et al. (023). A maximum input length of 256 and a maximum output length 320 are used for distillation and evaluation.
LMIB Chelba et al. (2013). The One Billion Word dataset (LMIB) is widely recognized benchmark for language modeling. The training and held-out data are derived from the WMT 2011 News Crawl dataset using Bash shell and Perl scripts. A maximum input length of 128 and maximum output length of 128 are used for distillation and evaluation.
B.2 MODELS
In accordance Leviathan et al.) (2023), we evaluate two model types: 1) GPT-like decoder-only Transformer models trained on the LMIB task (Chelba et al. 2013) using the standard autoregres- sive objective, where the target and draft models have 234M and 33M parameters, respectively; and 2) standard encoder-decoder T5 v1.1 models (Raffel et al. 2020) supervised fine-tuned four different tasks, with T5-XL (3B) and T5-Small (77M) serving as the target and draft models, respec- tively.
The target M, in the decoder-only model experiment has hidden dimension 1,024, feed-forward dimension 4,096, layers and 16 attention heads per transformer block, for total of 234M parameters. The corresponding draft model M has hidden dimension 512, feed-forward dimen- sion 1,024, layers and 4 attention heads per transformer block, for a total of 33M parame- ters. All models utilize the T5 tokenizer with 32k tokens. As for the T5 base checkpoints, we start from LM-adapted 1 models) These LM-adapted models are initialized from 15v1.1 and trained for an additional 100K steps the LM objective discussed in the T5 paper (Raffel et al. 2020). These checkpoints are available at ht https::///// httppp:::/cooooiio /cosool.....ccocccooooocccocooo. www.eeecccccoooooooooo google. com/ sttorage/// httee//brooseeetttttt/a///rreeee
ned wwdela........oo...mm.........m.
19

---

Published as a conference paper at ICLR 2024
In encoder-decodes distillation experiments, both the student and teacher models are initialized from models supervised fine-tuned on the original training dataset. This process for each dataset is detailed as follows:
XSum: for T5-Small, -Base, -Large, -XL, and -XXL. models,
start distillation from LM-Adapted T5v1.1 models supervised fine-tuned for 100K, 50K, 30k, 20K and 8K steps, respectively.
CNN/DM: for from LM-Adapted T5v1.1 models supervised fine-tuned for 200K, 80K, 20k, 20k and 4K steps, respectively.
-XL and -XXL models,
start distillation
WMT: for --SSmalll-aseeLLaree -XL and -XXL models, warm start distillation from LM-Adapted T5v1.1 models supervised fine-tuned for 250K, 250K, 110k, 50K and 10K steps, respectively.
GSM8K: all models are supervised fine-tuned starting from FLAN-TS models the PaLM-540 generated CoT dataset for 10K steps.
B.3 DISTILLATION
Training Data for KD. We study the five KD algorithms outlined in Table For SeqKD (Kim & Rush, 2016) and f-Distill (Wen et al. 2023), we opt for an online teacher data generation approach instead of conventional fixed offline teacher-generated dataset. This approach, while computation- ally expensive, yields more diverse dataset. For GKD, we exclude the static ground truth data and solely rely on the data generated by the online student model. All data generated by the teacher or the student is based on temperature sampling with a temperature of 1.0 (see Appendix C.1.3 for an ablation study on sampling temperature).
C.1
Training Details. We employ an Adafactor optimizer (Shazeer & Stern, 2018) to train the draft student model (M²) in all our experiments, following the guidelines outlined Algorithm A.2 In the context of our knowledge distillation (KD) loss function, defined in Eq. 2 we maintain the Tem- peratures for both the target model, denoted as Tp. and the draft model, denoted as Ty. a constant value of 1.0. We ought to emphasize the significance of maintaining this uniform temperature set- ting as it plays a pivotal role in speculative decoding, by ensuring a consistent and coherent semantic interpretation for both M, and M A summary of the hyperparameters used in our knowledge dis- tillation process can be found Table B.T.
Table B.I: Hyperparameters for distillation experiments.
hyperparameter
value
training steps
300,000
batch size
32
dropout
0.0
learning rate (LR)
0.0003
LR warmup steps
5,000
LR cooldown (begin, end)
(150,000, 300,000)
warmup schedule
linear (from 0 LR)
cooldown schedule
cosine decay (from LR to 0. 1L)
20

---

Published as a conference paper at ICLR 2024
B.4 EVALUATION
To obtain scores for each task (specifically, ROUGE-2 for XSum and CNN/DM, BLEU for WMT. and accuracy for GSM8K), we employ the evaluation methodology in Agarwal al. (2023). We assess models on all examples in the test or validation sets and report their average performance. To assess empirical speculative decoding performance, i.c., the empirical acceptance rate and empirical block efficiency, we conduct evaluations on all instances in the test or validation sets and report the average value of these metrics.
To measure the actual latency, we follow Leviathan et al. 2023 and execute both our target model and draft models are on the same TPUv4 device without utilizing model parallelism. We randomly sample 500 examples from either the test validation set, and measure the wall-clock decoding time on a batch size of This measurement procedure repeated three times, and the mean performance is reported. We have observed minimal variance across different random seeds in our results.
21

---

Published as a conference paper at ICLR 2024
C ADDITIONAL RESULTS
C.I ENHANCING SPECULATIVE DECODING THROUGH KNOWLEDGE DISTILLATION
Table C.I: DistillSpec improves speculative decoding performance across various datasets and block sizes, for both greedy decoding and temperature sampling. BBH-AVG contains the evaluation of the distilled draft model from GSM8K on all BIG-Bench Hard (BBH) tasks (Suzgun et 2022);
report the average results over the 23 BBH tasks. See Figure of performance on individual BBH tasks.
C.2
and C.3 for detailed breakdown
Sampling
w/o Distillation
with Distillation
Dataset
Model
temp
7
7
SPEED
SPEED
A
Method
3
2.31
1.44
2.62
1,58
+0.14x
f-Distill
Target Mp
T=0
5
2.57
1.43x
3.08
1,62
+0.19x
f-Distill
(3B)
7
2.68
1.36x
3.31
1.57x
+0.21
f-Distill
XSum
3
2.19
1.40x
2.58
1.57
+0.17
f-Distill
Draft M,
T=1
5
2.39
1.37
3.01
1.61
+0.25
f-Distill
TS-Small (77M)
7
2.47
1.28x
3.21
1.55x
+0.27
f-Distill
2.83
1.89x
3.19
2.11x
+0.22x
f-Distill
Target Mo
T=0
5
3.46
2.07x
4.13
2,42
+0.35
GKD
TS-XL (3B)
7
3.85
2.07x
4.83
2.53x
+0.46x
f-Distill
CNNDM
3
2.49
1.71x
2.87
1,93
+0.23x
f-Distill
Draft M,
5
2.89
1.77x
3.52
2.12x
+0.35
f-Distill
Tm-Small (77M)
3.08 7
1.71x
3.92
2.12x
+0.41x
f-Distill
3.22
2.08x
3.30
2.13x
+0.05
f-Distill
Target Mo
T=0
5
4.17
2.33x
4.32
2,41x
+0.08
f-Distill
T5-XL (3B)
7
4.83
2.36x
5.06
2,46x
+0.10x
f-Distill
WMT
2.73
1.77x
2.88
1.83
+0.07
GKD
Draft M,,
5
3.29
1.83x
3.48
1,93
+0.10x
GKI)
T5-Small (77M)
3.54
1.72x
3.85
1.84x
+0.12x
f-Distill
3
2.60
1.51
2.96
1.69x
+0.18x
GKD
Target Mp
T=0
3.06
1.42x
1.65
+0.22x
GKD
T5-XL (3B)
7
3,27
1.36x
4.14
1.60x
+0.24x
GKD
GSM8K
3
2,58
1.48x
2.84
1,64x
+0.16x
f-Distill
Draft M,
T=1
3.03
1.39x
3.45
1.58x
+0.19x
f-Distill
T5-Small (77M)
7
3.23
1.33x
3.84
1.53
+0.20x
f-Distill
3
2,96
3.66x
3.13
3.97
+0.31
f-Distill
Target Mo
T=0
3.69
3.35x
3.92
3.51
+0.16x
f-Distill
GPT-Like (234M)
4.15 7
2.52x
4.55
2,72x
+0.20x
f-Distill
LMIB
3
2.51
2.34x
2.69
2.45x
+0.11x
f-Distill
Draft Mg
T=1
2.90
2.79x
3.20
3.02
+0.23x
f-Distill
GPT-Like (33M)
3.10
1.98x
3.51
2.18x
+0.20x
f-Distill
3
2.55
1.90x
2.68
2.06x
+0.16x
f-Distill
Target M₁
T=0
5
2,98
1.92
3.19
2,16x
+0.24x
f-Distill
TS-XL (3B)
3.20
1.93x
3.49
2.21x
+0.28x
f-Distill
BBH-AVG
3
2.45
1.79x
2.57
1.93
+0.14x
f-Distill
Draft M,
T=1
5
2.81
1.80x
3.03
1.98
+0.18x
f-Distill
Small (77M)
3.01
1.78x
3.28
2.02x
+0.23x
f-Distill
22

---

Published as a conference paper at ICLR 2024
DistiliSper (Greedy)
Standard SD (Greedy)
DistillSper (Non- Greedy)
Standard SD (Non-Greedy)
Speculative Decoding Speedup (%)
+9%
+36%
+-37%
+29%
7%
+27%
+30%
100
+15%
+200%
XSum
GSMBK
CNNDM
WMT
BBH (Transfer)
Figure C.I: The distilled draft T5-Small model, derived from T5-XL teacher model, is capable of generalizing to a larger target model (TS-XXL), resulting in a substantial acceleration in various scenarios. BBH-AVG contains the evaluation of the distilled draft model from GSM8K on all BIG- Bench Hard (BBH) tasks (Suzgun et al. 2022); we report the average results over the 23 BBH tasks. and C.5 for detailed breakdown of performance individual BBH tasks.
See Figure C.4
cemporal_sequences
object counting
logical deduction five vee_ojects
web_of lies
logical deduction three oojectt
logical_ pereddutioo_svee_obbiee
geometric_shapes
reasoning about colored objects
movie_ree recommendaaiion
tracking_shuufled, objects_seveeeeeeeeesseeeeeessee tavee__
tracking liipin__hfffffee_bbie_ee objects five five_object objects
penguins_ing in__iinaa__a table
snarks
disambiguation.cm qa
dyck_languages
boolean expressions
date undersddeetgg
word_sorting
navigate
causal_judgene jusgemjudgemntt
tracking shuffled objectss three _oejccts
multistep_arithmee artthmetic two
ruin_names
sports_ sports_understarii understanddng
DistillSpec Avg
salient translation error
detection
Standard Avg DistillSpec
@ormal_fallcche
hyperbaton
Standard SD
60
100
120
140
160
180
Speculative Decoding Speedup (%)
Figure C.2: Assessing DistillSpec's model transferability the BIG-Bench Hard suite: zero-shot CoT reasoning with non- non-greedy decoding. This study examines a T5-Small draft model, initially trained the GSMSK dataset, across 23 varied tasks using T5-XL the target model. DistillSpec can deliver significant speculative decoding speedups on broad spectrum of tasks.
23

---

Published as a conference paper at ICLR 2024
temporal _equeecess
logical deduction five_on liv_objeccts
web_of_lie lies
logical deduction three _eree_b
logical deduction_seeeee sevennssseeeeeeeeeeeeeeeeeeeeeee to
dyck_language
boolean eeperesooossss
geometric_shene scapetessssppessssee
reasoning about colored tbjects
movie reremmeeddeniinnn
ruin_names
tracking_shuffledd cappe_oojjeeeeeeeeeeeeeeee
sdeven_objecte
snarks
multistep arithmetic
disambiguation_ce qa
penguins_in a_table
tracking shuffled objects five objects
sports ___ddesttaddiinggggg
causal_ jjssgm_jddeeeettt
date, understanding
salient translation, ciciineerrr_ddee detection -retttttttrooooottionnnnn
word_ par__sortigg
tracking shuffled objects three three_objects
DistiliSpec Avg
navigate
Standard Avg DistillSpec
fallacle
hyperbaton
Standard
40
60
80
100
120
140
160
180
Speculative Decoding Speedup (%)
Figure C.3: Assessing DistillSpec's model transferability on the BIG-Bench Hard suite: zero-shot CoT reasoning with greedy decoding. This study examines T5-Small draft model, initially trained on the GSM8K dataset, across 23 varied tasks using TS-XL as the target model. DistillSpec can deliver significant speculative decoding speedups broad spectrum of tasks.
24

---

Published as a conference paper at ICLR 2024
web of lies
tracking_shunne wttpssshhhhddihhhheeee objecto_seeeeeeeeeee
formal_fallacie
object counting
logical_deduce deduction, five_bjjects
logical deduction
Sereectionssveen_eeee seven_ seven_objects
rrakinngshuuflled sotes_ssffldddddddttsssssssss five fivetoo__eoeoeectt
logical deduction, three three_objects
disambiguation
reasoning about colored objects
tracking_shuffled objects thhree tetrtttrreeeeee _hjeeeoeeeereeesssssesesss
word_sorting
causaljudgement
geometric_shaper
date_und unnerssdestnn
boolean expressions
dyck languages
snarks
penguins_ in_a_______ a_bbe
ruin_na sin_names
navigate
movie recommendation
DistiliSpec Avg
cultistep_arithmeet two sports_ uoodss_nndes
Standard Avg
DistillSpec
hyperbaton
Standard SD
100
150
200
250
Speculative Decoding Speedup (%)
Figure C.4: Assessing DistillSpec's 's model transferability the BIG-Bench Hard suite: zero-shot CoT reasoning with non non-greedy decoding. This study examines a T5-Small draft model, initially trained on the GSM8K dataset, across 23 varied tasks using T5-X V5-XXL-s as the target model. DistillSpec can deliver significant speculative decoding speedups on a broad spectrum of tasks.
25

---

Published as a conference paper at ICLR 2024
web of lies
formal formal_fallan fallacies
logical_ deeuutioo_seveeeeeee seven_objects
disambiguation qa
trackingsshuffee objects_seven ses_sseeeeeoee sevven_objects
object counting
logical_dedieoo dee_ddeddoooofiiie five five_objects
reasoning_bbou ceclored_objectte
geometric_shapp
causal judgement
ruin sin_nnames
logical deduction ceaincttiinttteeeeeeeee chree_objecte
dyck_languages
boolean expressions
bootean__epressions
sracking_shuffled_ objects_fivee five five_objects
date understanding
word_ sorting
Backknggshuufled objects bebjects_three__eeee jorre__bjjcts
roviee_recommendation
snarks
penguins_ in_a table
sports understanding
multistepparithmtttt
arithmetic two
DistillSpec Avg
salient www.strannnntnnnooo cecaaaiio__rrrr_detettttet detectioen
Standard Avg DistillSpec Standard SD
navigate
hyperbaton
50
150
200
250
300
350
Speculative Decoding Speedup (%)
Figure C.5: Assessing DistillSpec's 's model transferability the BIG-Bench Hard suite: zero-shot CoT reasoning with greedy decoding. This study examines a T5-Small draft model, initially trained on the GSM8K dataset, across 23 varied tasks using T5-XXL as the target model. DistillSpec can deliver significant speculative decoding speedups on a broad spectrum of tasks.
26

---

Published as a conference paper at ICLR 2024
C.1.1
EMPIRICAL BLOCK EFFICIENCY IMPROVEMENT
SupervisedKD
SeqKD
ImitKD
f-Distill
GKD
XSum
CNNDM
WMT
2.90
GSMBK
Block Efficieccy (2)
2.60
2,80
2.85
2.40
2.80
2.60
2.20
WWthout Distilation
hevt iisinn
2.75
hoth Dist llltion 2.60
Without Distilation
Figure C.6: Empirical block efficiency improvement of DistillSpec for non non-greedy sampling (T 1) and block size 3. The draft model is trained using of the distillation methods listed in Table of Section + The dashed line indicates the block efficiency of speculative decoding using a dos-diitilled draft model. DistillSpec outperforms standard speculative decoding across all of the distillation methods being considered, with f- f-Distill and GKD yielding the highest gains.
SupervisedKD
SeqKD
ImitKD
f-Distill
XSum
CNNDM
3.50
WMT
GSMBK
Block fffieencyy (R]
3.00
3.40
2.80
3.40
3.20
2.60
3.20
3.00
2.40
Without Distilation
Without Distillation 3.30
Without Distillation
Without Distilltion
2,80
Figure C.7: Empirical block efficiency improvement of DistillSpec for non-greedy sampling (T 1) and block size 5. The draft model is trained using one of the distillation methods listed in Table of Section 4 The dashed line indicates the block efficiency of speculative decoding using a non con-distilled draft model. DistillSpec outperforms standard speculative decoding across all of the distillation methods being considered, with f- D-Distill and GKD yielding the highest gains.
SupervisedKD
SeqKD
ImitKD
f-Distill
Block Efficlnccy F
3.25
XSum
4.00
CNNDM
3.90
WMT
GSMBK
3.80
3.75
3.80
3.00
3.60
3.70
3.25
3.40
2.50
Without Distilation
Without Distillation
Without Distillaionn
Without Distilation
3.00
3.50
Figure C.8: Empirical block efficiency improvement of DistillSpec for non-greedy sampling (7 1) and block size 7 7. The draft model is trained using one of the distillation methods listed in Table [ of Section 4 The dashed line indicates the block efficiency of speculative decoding using a non distiiitlleed draft model. DistillSpec outperforms standard speculative decoding across all of the distillation methods being considered, with f- E-Distill and GKD yielding the highest gains.
27

---

Published as a conference paper at ICLR 2024
SupervisedKD
SeqKD
f-Distill
XSum
CNNDM
WMT
3.00
GSMBK
3.20
3.30
Btock Efficlency
2.60
2.90
2.50
3.28
2.80
3.00
2.40
3.25
2.70
2.30
_Wthhut D Distilation
2.80
witshus Distillliot
Withoou Distilation
2.60
Without Distilation
Figure C.9: Empirical block efficiency improvement of DistillSpec for greedy sampling (T' = O) and block size 3. The draft model is trained using one of the distillation methods listed in Table [ of Section 4 The dashed line indicates the block efficiency of speculative decoding using a non-distilled draft model. DistillSpec outperforms standard speculative decoding across all of the distillation methods being considered, with GKD weakly outperforming the other methods on average.
SupervisedKD
SeqKD
ImitKD
f-Distill
GKD
4.351
Block Efficiency (x)
4.20
3.60
2.80
3.80
4.25
3.40
3.60
4.20
2.60
Without Distilation
Distiaation
Distowtion
3.40
3.00
Figure C.10: Empirical block efficiency improvement of DistillSpec for greedy sampling (T 0) and block size 5. The draft model is trained using one of the distillation methods listed in Table of Section 4 The dashed line indicates the block efficiency of speculative decoding using non-distilled draft model. DistillSpec outperforms standard speculative decoding across all of the distillation methods being considered, with GKD weakly outperforming the other methods on average.
SupervisedKD
SeqKD
ImitKD
f-Distill
GKD
3.40
XSum
5.00
CNNDM
5.10
WMT
4.25
GSMBK
Block Efficieccy
4.00
4.50
5.00
3.00
3.75
2.80
4.00
4,90
3.50
hWth Distilation
hout Distalationn
houh Distilation
Withou Dist Wistontioon
3.25
2.60
4.80
Figure C.11: Empirical block efficiency improvement of DistillSpec for greedy sampling (T 0) and block size 7. The draft model is trained using one the distillation methods listed in Table If Section 4 The dashed line indicates the block efficiency of speculative decoding using a non non-distilled draft model. DistillSpec outperforms standard speculative decoding across all of the distillation methods being considered, with GKD weakly outperforming the other methods
average.
28

---

Published as a conference paper at ICLR 2024
C.1.2
PERFORMANCE IMPROVEMENT OVER TIME
Acceptance Rate (0)
Acceptance Rate (a)
0.70
0.76
SupKD
0.74
SupKD
0.66
ImitKD
ImitKD
SeqKD
SeqKD
D-Distill
f Distill
0.62
GKD
GKD
50K
100K
150K 200K 250K 300K
50K
100K 150K 200K 250K 300K Training Steps
Training Steps
(a) XSum
(b) CNNDM
0.80
0.810
Acceptance (e)
0.805
Acceptance Rate (e)
0.78
0.800
0.795
/ SupKD
0.76
SupKD
0.790
ImitKD
ImitKD
0.785
SeqKD
0.74
SeaKD
f-Distill
f-Distill
0.780
GKD
GKD
50K 100K 150K
200K 250K 300K
50K 100K 150K 200K 250K 300K Training Steps
Training Steps
(c) WMT
(d) GSM8K
Figure C.12: Progression of the acceptance rate a of DistillSpec over the training of the draft model, measured by the number of training steps. The draft model is trained using one of the distillation methods listed in Table Section GKD and f- f-Distill yields the most consistent improvement in over training steps, while SupKD yields the least improvement and exhibits declining acceptance rates after 400 training steps on XSum and CNNDM.
29

---

Published as a conference paper at ICLR 2024
0.72
0.78
Acceptance Rake
0.70
Acceptance Rake
0.76
0.68
SupKD ImitKD
SupKD
0.74
0.66
ImitKD
SeqKD
SeqKD
0.64
0.72
0.62
GKD
0
20
40
0
50
100
150
Training Efficiency (Walltime hrs)
Training Efficiency (Walltime hrs)
(a) XSum
(b) CNNDM
0.80
0.810
Acceptance (o)
0.805
Acceptance Rate (a)
0.78
0.800
0.795
SupKD
0.76
SupKD
0.790
ImitKD
SeoKD
0.785
f-Distill
f-Distill
0.780
GKD
Training Efficiency (Walltime in hrs)
Training Efficiency (Walltime in hrs)
WMT
(d) GSM8K
Figure C.13: Progression of the acceptance rate of DistillSpec over the training of the draft model, measured by the training wall time. The draft model is trained using one of the distillation methods listed in Table of Section 4 GKD and f-Distill yield the most consistent improvement in over training wall time, while SupKD yields the least improvement and even exhibits an inflection the acceptance rate early during training.
30

---

Published as a conference paper at ICLR 2024
3.4
4.0
Block Effidienyy
Block Efficiency
3.2
3.8
3.0
SupKD
3.6
SupKD
ImitKD
ImitKD
2.8
SDq
3.4
SDq
f-Distill
f-Distill
2.6
GKD
3.2
GKD
50K 100K 150K 200K
250K 300K
50K
100K 150K 200K 250K 300K
Training Steps
Training Steps
XSum
(b) CNNDM
4.3
Block Block Efficienyy
4.2
Block Block Emfiinncy
4.0
3.8
SupKD
SupKD
ImitKD
ImitKD
3.6
4.0
KD
SD
f-Distill
f-Distill
3.4
50K 100K 150K 200K
250K 300K
50K 100K 150K 200K 250K 300K
Training Steps
Training Steps
(c) WMT
(d) GSM8K
Figure C.14: Progression of the block efficiency of DistillSpec over the training of the draft model, measured by the number of training steps. The draft model is trained using one of the distillation methods listed in Table I of Section 4 GKD and f-iistill yield the most consistent improvement in over training, while SupKD yields the least improvement and exhibits drop in block efficiency after 400 training steps on XSum and CNNDM.
31

---

Published as a conference paper at ICLR 2024
Rouge2 Score (T=0.1)
Rouge2 Score (T=0.1)
16
15
SupKD
SupKD
ImitKD
ImitKD
SeqKD
SeqKD
14
f-Distill
F-Distill
GKD
GKD
50K 100K 150K 200K 250K 300K
50K 100K 150K 200K 250K 300K 0
Training Step
Training Step
(a) XSum
(b) CNNDM
Accuracy with Calculator (T=0.1)
16
14
26.2
12
Rouge2 Score
26.0
10
25.8
Sup
SupKD
ImitKD
ImitKD
25.6
SeqKD
SeqKD
25.4
f-Distill
D-Distill
GKD
GKD
50K 100K 150K
200K
250K 300K
0 50K 100K 150K
200K 250K 300K
Training Step
Training Step
WMT
(d) GSM8K
Figure C.15: Progression of the task performance of the draft model over the course of its training, measured by the number of training steps. Task performance is based sampling temperature T 0.1 and is measured by the Rouge2 score on the XSum and CNNDM datasets, and the accuracy score with the help of calculator GSM8K. The draft model trained using one of the distillation methods listed in Table of Section E GKD and f-Distill yield the most consistent improvement in task performance over training, while SupKD yields the least improvement and exhibits declining Rouge2 scores after 400 training steps XSum and CNNDM.
32

---

Published as a conference paper at ICLR 2024
Distilled Drafter
Distilled Drafter
Standard Drafter
Standard Drafter
Block Effic Pfficlency (R
Block Effic Priiclency (+)
2
0
2500
5000
7500
10000
0
250
500
750
1000
1250
Sorted Test Data Index
Sorted Test Data Index
(a) XSum
(b) CNNDM
Distilled Drafter
Distilled Drafter
Block Efficiency (e)
Standard Drafter
Block Efficienn lencyyycyyyy (T)
Standard Drafter
3
0
1000
2000
3000
0
250
500
750
1000 1250
Sorted Test Data Index
Sorted Test Data Index
(c) WMT
(d) GSM8K
Figure C.16: Comparison of the block efficiency of draft models trained with distillation vs. without distillation. The draft model is trained using f-Distill on XSum and CNNDM, and on-policy GKD on GSM8K. For each dataset and draft model, the block efficiency is sorted and plotted from lowest to highest index. The higher position of the curve for the distilled draft model at each index indi- cates that the distribution of block efficiency of the distilled draft model first-onder stochastically dominates the block efficiency of the non- distilled draft model.
33

---

Published as a conference paper at ICLR 2024
C.1.3
SAMPLING TEMPERATURE EFFECT
0.74
0.72
Acceptance Rate (x)
Acceptance Rate (ct
0.72
0.70
0.70
0.68
0.68
0.66
0.66
JSD0.5 (T=1e-2)
JSD0.5 (T=1e-2)
0.64
JSD0. (T-1.0)
0.64
JSD0.5 (T=1.0)
TVD (T-1e-2)
TVD (T-1e-2)
0.62
TVD (T-1.0)
0.62
TVD (T-1.0)
50K 100K 150K 200K 250K 300K Training Steps
0
50K 100K 150K 200K 250K 300K Training Steps
(a) distilling on student generations (on-policy)
(b) distilling on target model generations
Figure C.17: Effect of the sampling temperature on the acceptance rate in XSum decoding. We assess how varying the sampling temperature influences speculative decoding performance in two distinct scenarios: (a) the draft model is trained on its own generated sequences (on-policy distil- lation), and (b) the draft model is trained on sequences sampled from the target model. In both instances, we find that using a high sampling temperature is paramount for attaining superior per- formance.
C.1.4 REDUCTION IN CROSS-ENTROPY
0.8
SeqKD
SeqKD
0.8
GKD
GKD
0.7
Cross Entropy
Cross Entropy
0.7
0.6
0.6
0.5
0.5
0.4
0.4
0.3
50K 100K 150K 200K 250K 300K Training Steps
0
50K 100K 150K 200K 250K 300K
Training Steps
(a) cross-entropy reduction on XSum
(b) cross-entropy reduction on CNNDM
Figure C.18: BiLD Kim et al. 2023) proposes an SD variant that uses the cross-entropy between the target and draft model as the rejection criterion and uses SeqKD) to improve the alignment. Our experiment shows that GKD can reduce the cross-entropy more effectively than SeqKD. Thus, DistillSpec can potentially be utilized to improve BiLD.
34

---

Published as a conference paper at ICLR 2024
C.2
DISTILLATION RECIPE
C.2.1
SCORE AND BLOCK EFFICIENCY IMPROVEMENT
Rouge2 Improvement
Rouge2 Improvement
Fixed (Dirain)
4.3
5.1
4.3
Fixed (Draai)
0.52
1.1
1.3
1.8
3.0
Student (M
4
5.7
6.2
5.9
Student (MV)
3.6
3.6
3.3
3.8
2.5
3
M MMp
4.1
5.7
6.3
6
Mg + Mº
3.6
3.7
3.4
3.7
2
1.0
Teacher (Mp)
4.2
5.7
6.4
6
1
Teacher (Mp)
3.6
3.7
3.3
3.5
0.5
FKL
JSD
RKL
TVD
FKL
JSD
(a) non-greedy decoding (T' 1)
(b) greedy decoding (T'
Figure C.19; Improvement in the Rouge2 score of DistillSpec over standard speculative decoding on XSum given by different combinations of divergence functions and generative distributions used in the draft model distillation (see Table T of Section using greedy (T' = 0) and non-gree non-greedy (I" 1) sampling.
Block Efficiency Improvement
Efficiency Improvement
0.5
Fixed (Drrain)
0.11
0.18
0.2
0.27
Fixed (Drrai)
0.01
0.08
0.13
0.17
0.4
Student (MV)
0.48
0.59
0.55
0.58
Student (MV)
0.48
0.49
0.45
0.46
0.3
M& + Mp
0.51
0.64
0.6
0.6
My + Mp
0.52
0.53
0.5
0.5
0.2
0.1
Teacher (Mp)
0.52
0.63
0.58
0.62
Teacher (Mp)
0.51
0.53
0.5
0.51
JSD
RKL
TVD
FKL
JSD
RKL
TVD
(a) non-greedy decoding (I' 1)
(b) greedy decoding (T'
Figure C.20: Improvement in the block efficiency of DistillSpec over standard speculative decoding on XSum given by different combinations of divergence functions and generative distributions used in the draft model distillation (see Table [ of Section 4). using greedy (T' 0) and non-greedy (T' sampling.
35

---

Published as a conference paper at ICLR 2024
Accuracy Improvement
Accuracy Improvement
10
Fixed (Drran)
4.2
4.8
5
5.4
Fixed (Draii)
3.9
5.4
Student (MG)
9.2
6.7
7.5
5.3
Student (MQ)
10
6.7
5.2
MC MM
6.9
6.3
7.3
6.8
Mg + Mp
9.2
5.8
2
Teacher (Mp)
6.2
6
Teacher (Mp)
6.4
5.5
FKL
JSD
RKL
TVD
FKL
JSD
RKL
TVD
(a) non-greedy decoding (T 1)
(b) greedy decoding (7 0)
Figure C.21: Improvement in the accuracy score of DistillSpec over standard speculative decoding on GSM8K given by different combinations of divergence functions and generative distributions used the draft model distillation (see Table of Section using greedy (T 0) and non- -regeeedy
(T = 1) sampling.
Block Efficiency Improvement
Block Efficiency Improvement
Fixed (Drrain)
0.22
0.31
0.32
0.29
Fixed (Drrai))
0.23
0.29
0.29
0.26
0.4
Student (ME)
0.42
0.43
0.45
0.41
Student (ME)
0.61
0.42
0.43
0.38
0.4
-03
Mg Mp
0.36
0.47
0.47
0.44
0.2
MC Mp
0.49
0.43
0.47
0.41
Teacher (Mp)
0.33
0.44
0.49
0.43
Teacher (Mp)
0.35
0.43
0.42
0.0
FKL
JSD
RKL
TVD
FKL
JSD
RKL
TVD
(a) non-greedy decoding (2'
(b) greedy decoding (T' = 0)
Figure C.22: Improvement in the block efficiency of DistillSpec over standard speculative decoding on GSM8K given by different combinations of divergence functions and generative distributions used in the draft model distillation (see Table of Section 4) using greedy (T 0) and non-greedy
(T = 1) sampling.
36

---

Published as a conference paper at ICLR 2024
C.2.2
IMPACT OF DISTILLATION ON DRAFT QUALITY VS. COMPATIBILITY
Divergence
Divergence
0.6
FKL
FKL
JSDO.5
X
JSD0.5
Improvement
0.5
RKL
Improvement
0.4
RKL
TVD
TVD
0.4
Data
0.3
TraanData
Teacher (Mp)
Teacher (Mp)
0.2
Student (Mg))
Student (M@)
Fixed (Praw)
Fixed (Drae)
JV MMp
MG My
0.1
0.0
5
1
Rouge2 Improvement
Rouge2 Improvement
(a) non-greedy decoding (T' 1)
(b) greedy decoding (T
Correlation plot of the improvements in the block efficiency (y-axis) and Rouge2 score (X-axis) of DistillSpec standard speculative decoding XSum given by different combinations of divergence functions and generative distributions used in the draft model distillation (see Table of Section 4 using greedy (I 0) and non-greedy (T 1) sampling.
Figure C.23:
0.50
Divergence
Divergence
0.6
FKL
0.45
JSDO.5
JSDO.5
Improvement
RKL
0.5
0.40
TVD
TVD
TrainData
TrainData
Teacher (Mp)
0.4
Teacher (Mg)
0.30
Student (MC)
Student (M))
Fixed (Diann)
0.3
Fixed (Draa)
MG + M₁
Mr MM
6
10
Accuracy Improvement
Accuracy Improvement
(a) non-greedy decoding (T
(b) greedy decoding (T
Figure C.24: Correlation plot of the improvements in the block efficiency (y-axis) and Accuracy (x- axis) of DistillSpec over standard speculative decoding on GSM8K given by different combinations of divergence functions and generative distributions used in the draft model distillation (see Table of Section 4 using greedy (T 0) and non-greedy (T 1) sampling.
37

---

Published as a conference paper at ICLR 2024
C.3
QUALITY VERSUS LATENCY TRADE-OFF
C.3.1
LOSSY SPECULATIVE DECODING
18
Performance RRuge2
16
Performance RRoge2
16
14
filp, El - ple
fn el ple
fulp, e) plc?
kulp. e) pir?
12
foolp,
foolp.
10
Lenience
Distilled Draft
Distilled Draft
1.0
2.5 Relative Latency
2.0
3.5
Relative Latency
(a) using non-distilled draft model (standard SD)
(b) using distilled draft model (DistillSpec)
Figure C.25: Tradeoff between task performance and decoding latency on XSum enabled by com- bining lossy speculative decoding with the alternative variants of DistillSpec presented Section 4 The higher position and flatter slope of the tradcoff curve for DistillSpec indicate that the method enables larger latency gains for smaller reduction in model quality than are feasible under standard SD.
12
Performance Rouge2
Performance Rouge2
*
11
fnlo, = =pe
11
1n(0, el ple
10
falp. E) ple?
fulp, e) ple²
foolp, c) p
foulp, e) - p
No Lenience
No Lenience
9
Raw Draft
Raw Draft
Distilled Draft
Distilied Draft
0.5
1.0
1.5
2.0
2.5
0.25
0.50
0.75
1.00
1.25
1.50
1.75
2.00
Relative Latency
Relative Latency
(a) using non-distilled draft model (standard SD)
(b) using distilled draft model (DistillSpec)
Figure C.26: Tradcoff between task performance and decoding latency on CNNDM enabled by combining lossy speculative decoding with the alternative variants of DistillSpec presented in Sec- tion 4 The higher position and flatter slope of the tradeoff curve for DistillSpec indicate that the method enables larger latency gains for smaller reduction in model quality than are feasible under standard SD.
38

---

Published as a conference paper at ICLR 2024
17
Performance Bieu
Performance -Bleu
16
fn él plz
felo, plc?
16
*
TelP, ple
15
foulp, p
foulp.s) = =p
14
Raw Draft
No Lenience
15
Raw Draft Distilled Draft
No Lenience
Distilled Draft
0.5
1.0
1.5
2.0
2.5
0.25
0.50
0.75
1.00
1.25 1.50 1.75 2.00 2.25 Latency
Relative Latency
(a) using non-distilled draft model (standard SD)
(b) using distilled draft model (DistillSpec)
Figure C.27: Tradeoff between task performance and decoding lateney on WMT enabled by com- bining lossy speculative decoding with the alternative variants of DistillSpec presented in Section While the similar height and slope of the tradeoff curves of DistillSpec and standard SD) indicate comparable quality-latency tradeoff between the two setups, DistillSpec exhibits lower latency overall.
30
30
Performance Accuracy
Performance Accuracy
25
20
finlp, cl p/r
firlp, el - plc
felp, 6) = plc?
Telp, e) pir?
feplp, =p"
15
Tep
10
Lenience
10
Draft
Draft
Distilled Draft
Distilled Draft
1.0
1.5
2.5
1.0
2.0
3.0
Relative Latency
Relative Latency
(a) using non- non-distilled draft model (standard SD)
(b) using distilled draft model (DistillSpee)
Figure C.28: Tradeoff between task performance and decoding latency on GSM8K enabled by com- bining lossy speculative decoding with the alternative variants of DistillSpec presented Section
While the similar height and slope of the tradcoff curves of DistillSpec and standard SD indicate a comparable quality-latency tradeoff between the two setups, DistillSpec exhibits lower latency overall.
39

---

Published as a conference paper at ICLR 2024
C.3.2
DISTILLSPEC MEETS MODEL GARDEN
Performance Rouge2
Performance: Rouge2
14
Target Model
Serving Strategy
Target Model
Serving Strategy
T5-Small
Raw
12
T5-Base
Distilled
T5-Bad TS-Base
Distilled
10
TS La-Large
Speculative
TS-Large
Speculative
TS-XL
DistilSSec
TS XL
DistillSpec
TS-XXL
TS-XXL
0.0
12.5
15.0
17.5
Relative Latency
Relative Latency
(a) non-greedy decoding (T' 1)
(b) greedy decoding (T'
Figure C.29: Quality-latency trade-off curves on XSum using greedy =T'= 0) or non-greedy (I' = 1) sampling, for single target LLM trained without distillation ("Raw"), trained by distilling from a larger LLM ("Distilled"), speculative decoding using a draft model trained without distillation ("Speculative"), and speculative decoding combining a distilled target LLM with a distilled draft model ("DistillSpec").
Performance Accuracy
30
Perfomance Accuracy
30
Target
Serving Strategy
Target Model
Serving Strategy
20
Distilled
T5-Base
Distilled
TS-Large
Speculative DistillSpec
TS-Large
Speculative
10
TS-XL
10
TS-XL
Distil Spec
T5-XXL
TS-XXL
10
12
0
12
Relative Latency
Relative Latency
(a) non-greedy decoding (T = 1)
(b) greedy decoding (T
Figure C.30: Quality-latency trade-off curves GSM8K using greedy (I' 0) non-greed (T' 1) sampling, for a single target LLM trained without distillation ("Raw"), trained by distilling from larger LLM ("Distilled"", speculative decoding using draft model trained without distillation ("Speculative"), and speculative decoding using a distilled target LLM combining a distilled draft model ("DistillSpec").
40


Reviews:
[{"review_id": "TcqbOjr3gH", "reviewer": "Reviewer_gZt5", "summary": "This paper proposes DistillSpec, which improves the efficiency of speculative decoding (SD) by aligning the distributions (at the token and sequence level) between the draft model and the target model in advance. The paper further provides valuable insights regarding the recipe of distillation data, distillation objective and sampling strategy.\n\nExperiments show that DistillSpec speedup SD by 10-45% while preserving the model performance across four diverse datasets.", "questions": "See weakness.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. SD is an important direction in accelerating LM inference. The idea of leveraging distillation to speed up SD is novel and very effective.\n2. Using acceptance rate as the efficiency measure is well-motivated. Using total variation distance (TVD) objective is simple and straightforward in maximizing this efficiency measure.\n3. The authors provide a fast alternative of using student-generated data for training, with a theoretical justification and strong experimental results.", "weaknesses": "1. The technical novelty is a bit limited. It is a direct application of existing KD techniques in SD.\n2. 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 and TVD has varied performance on different tasks.\n3. Major experiment results are based on small target models. The results on larger models are not very clear -- in Figure 6, which data points correspond to which sizes of target and draft models?\n4. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Q4ScIuaXbF", "reviewer": "Reviewer_GLMj", "summary": "This paper proposes DistillSpec, a method that uses knowledge distillation to improve the speed and quality of speculative decoding. The paper explores various factors that affect the alignment between the draft and target models, such as the training data, the divergence function, and the decoding strategy. It also extends DistillSpec to lossy speculative decoding, which allows for sampling away from the target model distribution. The paper evaluates DistillSpec on several tasks and datasets to demonstrate its speedups.", "questions": "I am wondering how will DistillSpec be compared with recent work that combines speculative decoding and distillation [1]. It would be interesting to have this discussion.\n\n[1] Liu, X., Hu, L., Bailis, P.D., Stoica, I., Deng, Z., Cheung, A., & Zhang, H. (2023). Online Speculative Decoding. ArXiv, abs/2310.07177.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- It provides a comprehensive and systematic study of different aspects of knowledge distillation for speculative decoding, such as data generation, divergence functions, and lossy sampling. The conclusion of models generated data being important makes sense.\n- It demonstrates the effectiveness of DistillSpec on several tasks and datasets, using both greedy and non-greedy sampling, and compares it with representative baselines.\n- Its lossy speculative decoding results offer novel insights and recommendations for combining knowledge distillation and speculative decoding in different scenarios.", "weaknesses": "- The presentation of the paper is a bit messy and unclear. For instance, it is difficult to find the formal definition of DistillSpec among the analysis of various existing distillation approaches; it is not clear how distillation data is generated in detail; 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. The clarity and structure of the paper should be improved.\n- While DistillSpec is effective on certain tasks, the experiments mainly focus on T5 models and simple tasks (except for GSM8K), compared with the recent advances in LLMs. This also results in the lack of discussion on the difficulty of LLM distillation. Is DistillSpec also effective for LLMs like LLaMA-7B? Can distilling LLMs for some specific tasks also be helpful for other tasks in general? The effectiveness of DistillSpec should be evaluated using more recent large models and some zero-shot benchmarks to prove its effectiveness.\n- The paper does not compare or discuss DistillSpec with other methods that combine large and small models at inference, especially under the lossy decoding cases [1].\n- The paper does not analyze the generated texts by DistillSpec in greater detail, like the diversity and coherence, which are important aspects of text quality and largely influenced by the sampling approaches. It would be helpful to provide some examples of the generated tasks.\n\n[1] Kim, S., Mangalam, K., Moon, S., Malik, J., Mahoney, M.W., Gholami, A., & Keutzer, K. (2023). Speculative Decoding with Big Little Decoder.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9VAC54OMLv", "reviewer": "Reviewer_p4vK", "summary": "The authors propose DistillSpec that uses knowledge distillation to better align the draft model with the target model. The authors have a wide exploration of different dilation methods from SeqKD to on-plicy GKD. DistillSpec yields impressive 10 − 45% speedups over standard speculative decoding on a range of standard benchmarks.", "questions": "Why do different distillation methods vary a lot in Figure 4?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The experiments are quite solid. It explores different distillation losses and different ways to collect data for distillation. All the tables and figures are well written, and it is easy to get the difference between different methods.\n2. The authors also explore the trade-off between accuracy and latency.\n3. The paper is well-written and easy to read.", "weaknesses": "1. 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.\n2. All the experiments are done in an in-domain setting, while the benefits of LLM are on zeroshot/fewshot setting on out-of-domain datasets. Thus the experiments can not show whether it can replace the general draft models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)