DistillSpec: Improving Speculative Decoding via Knowledge Distillation

Yongchao ZhouKaifeng LyuAnkit Singh RawatAditya Krishna MenonAfshin RostamizadehSanjiv KumarJean-François KagyRishabh Agarwal

article2024ICLR138 citations

Develops DistillSpec, an on-policy knowledge distillation method that aligns draft models with target language models to achieve a 10% to 45% inference speedup over standard speculative decoding.

Listen

Deploying large language models in production environments presents a severe operational bottleneck due to high inference latency and substantial computational costs. While speculative decoding has emerged as a promising remedy—using a small, fast draft model to generate candidate text that a large target model validates in parallel—its real-world acceleration depends heavily on how closely the draft model mimics the target model. If the draft model generates tokens that the target model frequently rejects, inference efficiency drops sharply. Developing compact draft models that align closely with large target architectures has thus become a critical priority for practical AI deployment.

The article evaluates a framework called DistillSpec, which uses white-box knowledge distillation to systematically align a small draft model to a larger target model. Across comprehensive empirical benchmarks spanning summarization, translation, arithmetic reasoning, and general language generation, the authors evaluated both greedy and non-greedy decoding strategies. The evaluation compared fixed training datasets against model-generated data, tested multiple mathematical divergence functions, and analyzed model scaling across architectures ranging from 33 million to 11 billion parameters.

The findings establish that DistillSpec consistently enhances inference speed, delivering 10% to 45% speedups over standard speculative decoding across all evaluated benchmarks while fully preserving generation quality. Training draft models on their own generated outputs, rather than fixed ground-truth data, was found to be essential for achieving strong alignment and continuous efficiency gains. Furthermore, draft models distilled on a single task demonstrated strong cross-task transferability, achieving an average 26% speedup across 23 unseen reasoning tasks. In comprehensive model garden evaluations, first distilling a large teacher model into an intermediate target model and then applying DistillSpec to create the draft model achieved an overall 6-fold to 10-fold reduction in decoding latency with virtually no loss in performance.

These results have direct strategic implications for reducing the infrastructure overhead, energy consumption, and serving costs of large language model applications. DistillSpec demonstrates that a small model's standalone benchmark score does not predict how well it will perform as a speculative drafter; instead, direct distribution alignment is what drives inference efficiency. The article also demonstrates that combining DistillSpec with controlled lossy speculative decoding allows organizations to finely tune latency against acceptable quality boundaries depending on operational requirements.

Engineering and deployment teams should adopt DistillSpec as a drop-in enhancement for speculative decoding pipelines. Practitioners should use the draft model's own on-policy generations to minimize training compute costs and treat the choice of divergence function as a task-specific hyperparameter to tune. While these conclusions are supported with high confidence across diverse model sizes and benchmarks, testing was conducted in single-device, batch-size-one environments. Organizations should validate these gains in high-concurrency, distributed serving environments to verify real-world throughput improvements.

Google Mila Princeton University University of Toronto hierarchical model inference knowledge distillation DistillSpec speculative decoding knowledge distillation

arXiv: 2310.08461

Paper: Fast Inference from Transformers via Speculative Decoding, Yaniv Leviathan et al. (2023). Introduces the foundational framework of speculative decoding that DistillSpec specifically aims to accelerate by improving draft-target model alignment.
Paper: On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes, Rishabh Agarwal et al. (2024). Pioneers Generalized Knowledge Distillation using on-policy sampling and flexible divergence metrics, providing the core distillation paradigms adopted in DistillSpec.
Paper: MiniLLM: On-Policy Distillation of Large Language Models, Yuxian Gu et al. (2024). Establishes reverse Kullback-Leibler divergence and on-policy optimization principles for text generation distillation that DistillSpec tailors to speculative drafting.
Paper: Distilling the Knowledge in a Neural Network, Geoffrey Hinton et al. (2015). Supplies the foundational theory and mechanism of knowledge distillation that underlies the training of compact draft models.
Paper: Sequence-Level Knowledge Distillation, Yoon Kim et al. (2016). Develops sequence-level knowledge distillation methods for autoregressive generation that underpin modern LLM student-teacher alignment.

Paper: EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees, Yuhui Li et al. (2024). Extends speculative decoding acceleration by dynamically constructing context-aware draft trees to further boost verification acceptance rates.
Paper: EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test, Yuhui Li et al. (2025). Advances draft model training by incorporating multi-layer target features and simulated multi-step test dynamics during drafting.
Paper: SpecInfer: Accelerating Generative Large Language Model Serving with Tree-based Speculative Inference and Verification, Xupeng Miao et al. (2024). Applies tree-based speculative candidate generation and parallel batch verification to optimize speculative serving systems.
Paper: Unlocking Lossless Speedups in LLMs via Discrete Diffusion, Subham Sekhar Sahoo et al. (2026). Explores alternative drafting architectures by utilizing distilled discrete diffusion pathways for parallel proposal generation in speculative verification.
Paper: OPRD: On-Policy Representation Distillation, Shenzhi Yang et al. (2026). Extends on-policy distillation by shifting supervision from token output distributions to intermediate representation matching.
Paper: Trajectory-Refined Distillation, Li Jiang et al. (2026). Refines on-policy distillation methods by addressing multi-step rollout errors through trajectory-level teacher corrections.

Cover for DistillSpec: Improving Speculative Decoding via Knowledge Distillation

Abstract

Speculative decoding (SD) accelerates large language model inference by employing a faster draft model for generating multiple tokens, which are then verified in parallel by the larger 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 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 improving the draft and target alignment: utilizing on-policy data generation from the draft model, and 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-10x with minimal performance drop, compared to standard decoding without distillation.

Topics

hierarchical model inference knowledge distillation language model large language models target model trade probabilistic inference

Keywords

DistillSpec speculative decoding knowledge distillation draft model distillation autoregressive decoding parallel decoding large target model

1 Introduction
2 Related Work
3 Background: Speculative Decoding
4 DistillSpec: Knowledge Distillation for Speculative Decoding
5 Experiments
5.1 Enhancing speculative decoding through distillation
5.2 DistillSpec recipe
5.3 Quality versus latency trade-off
6 Conclusion
References
A Method
A.1 Description of divergence functions
A.2 Justification of using on-policy data
A.3 DistillSpec algorithms
B Implementation Details
B.1 Datasets
B.2 Models
B.3 Distillation
B.4 Evaluation
C Additional Results
C.1 Enhancing speculative decoding through knowledge distillation
C.1.1 Empirical block efficiency improvement
C.1.2 Performance improvement over time
C.1.3 Sampling temperature effect
C.1.4 Reduction in cross-entropy
C.2 Distillation recipe
C.2.1 Score and block efficiency improvement
C.2.2 Impact of distillation on draft quality vs. compatibility
C.3 Quality versus latency trade-off
C.3.1 Lossy speculative decoding
C.3.2 DistillSpec meets model garden

Knowls

Knowl 1 — DistillSpec Framework for Aligning Draft Models via White-Box Knowledge Distillation
model/method
DistillSpec is a framework designed to accelerate speculative decoding (SD) by improving the distributional alignment between a compact draft model $M_q^\theta$ (parameterized by weights $\theta$ ) and a large target model $M_p$ . In speculative decoding, $M_q^\theta$ autoregressively generates a block of candidate tokens, which are subsequently verified in parallel by $M_p$ . The realized wall-clock speedup depends directly on the acceptance rate of these candidate tokens, which in turn is determined by how closely the draft model's token distribution matches the target model's distribution.

DistillSpec trains $M_q^\theta$ via white-box knowledge distillation (KD) with full access to next-token probability distributions $p(\cdot \mid x, y_{<t})$ from $M_p$ and $q^\theta(\cdot \mid x, y_{<t})$ from $M_q^\theta$ . Given a divergence function $D$ and a training distribution $\mathcal{G}$ over input-output sequences $(x, y)$ , DistillSpec optimizes the draft parameters by minimizing sequence-averaged divergence:

$\theta^* := \arg\min_\theta \mathbb{E}_{(x,y) \sim \mathcal{G}} [D(M_p \parallel M_q^\theta)(y \mid x)]$

where

$D(M_p \parallel M_q^\theta)(y \mid x) = \frac{1}{|y|} \sum_{t=1}^{|y|} D\left(p(\cdot \mid x, y_{<t}) \parallel q^\theta(\cdot \mid x, y_{<t})\right)$

Unlike traditional knowledge distillation aimed solely at maximizing the standalone accuracy of a small student model, DistillSpec explicitly tunes the draft model to maximize token-level acceptance under the target distribution during speculative verification.
Knowl 2 — Total Variation Distance as an Exact Proxy for Speculative Decoding Acceptance Rate
theoretical result
In speculative decoding with target model $M_p$ and draft model $M_q$ , let $p(y_t) = p(y_t \mid x, y_{<t})$ and $q(y_t) = q(y_t \mid x, y_{<t})$ denote the next-token probability distributions over vocabulary $\mathcal{V}$ at step $t$ given prefix context $\rho = \{x, y_{<t}\}$ . Under speculative sampling, the target model accepts a drafted token $y_t \sim q(y_t)$ with probability $\min(1, p(y_t)/q(y_t))$ . The expected token-level acceptance probability $\beta(x, y_{<t})$ is exactly related to the Total Variation Distance ( $D_{\text{TVD}}$ ) between the two distributions:

$\beta(x, y_{<t}) := \mathbb{E}_{y_t \sim q(y_t)} \left[ \min\left(1, \frac{p(y_t)}{q(y_t)}\right) \right] = 1 - D_{\text{TVD}}(p(y_t), q(y_t))$

where $D_{\text{TVD}}(P \parallel Q) = \frac{1}{2}\sum_{c \in \mathcal{V}} |P(c) - Q(c)|$ .

For an ideal speculative decoding process over an input prompt $x$ with maximum length $T$ , the sequence-level acceptance rate $\alpha(x)$ , defined as the expected ratio of accepted tokens to the expected target sequence length $L_p(x) = \mathbb{E}_{y \sim p_{\le T}(y \mid x)}[|y|]$ , satisfies:

$\alpha(x) = 1 - \frac{\mathbb{E}_{y \sim p_{\le T}(y \mid x)} \left[ \sum_{t=1}^{|y|} D_{\text{TVD}}(p(y_t \mid x, y_{<t}), q(y_t \mid x, y_{<t})) \right]}{L_p(x)}$

where $p_{\le T}(y \mid x)$ is the autoregressive sequence distribution generated by $M_p$ . Consequently, maximizing the sequence-level acceptance rate $\alpha(x)$ is mathematically equivalent to minimizing the expected token-level total variation distance over output sequences sampled from the target model.
Knowl 3 — Lower Bound on Speculative Acceptance Rate from On-Policy Distillation Loss
theoretical result
Let $M_p$ be the target language model, $M_q^\theta$ be the draft language model parameterized by $\theta$ , $\mathcal{X}$ be the distribution of input prompts $x$ , and $T$ be the maximum sequence length. Let $q_{\le T}(y \mid x)$ denote the autoregressive distribution over output sequences $y$ generated on-policy by the draft model $M_q^\theta$ .

If the draft model achieves an expected on-policy knowledge distillation loss $\epsilon$ under total variation distance:

$\epsilon = \mathbb{E}_{x \sim \mathcal{X},\, y \sim q_{\le T}(y \mid x)} [D_{\text{TVD}}(M_p \parallel M_q^\theta)(y \mid x)]$

where $D_{\text{TVD}}(M_p \parallel M_q^\theta)(y \mid x) = \frac{1}{|y|} \sum_{t=1}^{|y|} D_{\text{TVD}}(p(\cdot \mid x, y_{<t}), q^\theta(\cdot \mid x, y_{<t}))$ , then the expected sequence-level speculative decoding acceptance rate $\mathbb{E}_{x \sim \mathcal{X}}[\alpha(x)]$ satisfies:

$\mathbb{E}_{x \sim \mathcal{X}}[\alpha(x)] \ge 1 - T \cdot \mathbb{E}_{x \sim \mathcal{X}} \left[ \frac{T}{L_p(x)} \right] \epsilon$

where $L_p(x) = \mathbb{E}_{y \sim p_{\le T}(y \mid x)}[|y|]$ is the expected target generation length given input $x$ . When the target output sequence length is constant ( $L_p(x) = T$ ), the bound simplifies to:

$\mathbb{E}_{x \sim \mathcal{X}}[\alpha(x)] \ge 1 - T\epsilon$

This guarantees that minimizing the distillation loss on student-generated (on-policy) tokens directly translates to a high acceptance rate under the teacher's output distribution.
Knowl 4 — DistillSpec Knowledge Distillation Training Algorithm
algorithm
DistillSpec optimizes the parameters $\theta$ of a draft model $M_q^\theta$ against a target model $M_p$ using batched stochastic gradient descent on token divergences over data drawn from ground-truth pairs, student on-policy rollouts, or teacher rollouts.
Input: Target model $M_p$ , draft model $M_q^\theta$ , prompt dataset $(X, Y)$ , fixed data fraction $\lambda_1 \in [0, 1]$ , student data fraction $\lambda_2 \in [0, 1]$ , divergence function $D$ , learning rate $\eta$ , batch size $B$
Output: Updated draft model parameters $\theta$
Sample $u_1 \sim \text{Uniform}(0, 1)$ and $u_2 \sim \text{Uniform}(0, 1)$
if $u_1 \le \lambda_1$ then
    Sample batch $\mathcal{B} = \{(x_b, y_b)\}_{b=1}^B \sim (X, Y)$
else
    Sample prompt batch $\mathcal{B}' = \{x_b\}_{b=1}^B \sim X$
    if $u_2 \le \lambda_2$ then
        For each $x_b \in \mathcal{B}'$ , sample $y_b \sim M_q^\theta(\cdot \mid x_b)$
        Set $\mathcal{B} = \{(x_b, y_b)\}_{b=1}^B$
    else
        For each $x_b \in \mathcal{B}'$ , sample $y_b \sim M_p(\cdot \mid x_b)$
        Set $\mathcal{B} = \{(x_b, y_b)\}_{b=1}^B$
    end if
end if
Compute gradient: $g \leftarrow \frac{1}{B} \sum_{(x, y) \in \mathcal{B}} \nabla_\theta D(M_p \parallel M_q^\theta)(y \mid x)$
Update parameters: $\theta \leftarrow \theta - \eta g$
return $\theta$
In standard implementations, an Adafactor optimizer is used with a base learning rate of $3 \times 10^{-4}$ , 5,000 linear warmup steps, cosine decay cooldown to $0.1\times$ base learning rate over 300,000 steps, batch size 32, and fixed generation temperatures $T_p = T_q = 1.0$ .
Knowl 5 — Speculative Decoding with Generalized Lenience Functions
algorithm
Speculative decoding verifies candidate tokens proposed by draft model $M_q$ using target model $M_p$ . When lossy speculative decoding is desired for higher throughput at the cost of bounded quality degradation, a lenience function $f(p, \epsilon)$ scales the acceptance probability, parameterized by lenience hyperparameter $\epsilon \in (0, 1]$ (where $\epsilon = 1$ recovers standard lossless speculative sampling).
Input: Target model $M_p$ , draft model $M_q$ , context $\rho = \{x, y_{<t}\}$ , block size $\gamma$ , lenience function $f(p, \epsilon)$ with parameter $\epsilon \in (0, 1]$
Output: Extended context $\{x, y_{<t+n+1}\}$
for $i = 0$ to $\gamma - 1$ do
    Evaluate draft distribution: $q_{t+i}(y) \leftarrow M_q(x, y_{<t+i})$
    Sample candidate token: $y_{t+i} \sim q_{t+i}(y)$
end for
Evaluate target distributions in parallel: $(p_t(y), \dots, p_{t+\gamma}(y)) \leftarrow (M_p(x, y_{<t}), \dots, M_p(x, y_{<t+\gamma}))$
for $i = 0$ to $\gamma - 1$ do
    Compute rejection ratio: $r_{t+i} \leftarrow \frac{f(p_{t+i}(y_{t+i}), \epsilon)}{q_{t+i}(y_{t+i})}$
    Sample $u_{t+i} \sim \text{Uniform}(0, 1)$
end for
Determine accepted count: $n \leftarrow \min(\{i \mid 0 \le i < \gamma \text{ and } u_{t+i} > r_{t+i}\} \cup \{\gamma\})$
if $n < \gamma$ then
    Sample corrected replacement token: $y_{t+n} \sim \text{normalize}(\max(0, p_{t+n}(y) - q_{t+n}(y)))$
else
    Sample bonus token: $y_{t+n} \sim p_{t+n}(y)$
end if
return $\{x, y_{<t+n+1}\}$
The evaluated lenience functions are:
1. Linear lenience: $f_{\text{lin}}(p, \epsilon) = p / \epsilon$
2. Quadratic lenience: $f_{\text{sq}}(p, \epsilon) = p / \epsilon^2$
3. Exponential lenience: $f_{\text{exp}}(p, \epsilon) = p^\epsilon$
Knowl 6 — Divergence Objectives for Draft Model Knowledge Distillation
definition
Let $P$ and $Q$ denote discrete probability distributions over a vocabulary $\mathcal{V}$ , corresponding to next-token predictions of target model $M_p$ and draft model $M_q^\theta$ , respectively. DistillSpec considers four statistical divergence functions for training:
1. Forward Kullback-Leibler Divergence (FKL): $D_{\text{FKL}}(P \parallel Q) = D_{\text{KL}}(P \parallel Q) = \sum_{c \in \mathcal{V}} P(c) \log \frac{P(c)}{Q(c)}$ FKL is mean-seeking and heavily penalizes $Q$ for assigning zero probability to tokens where $P$ has non-zero probability.
2. Reverse Kullback-Leibler Divergence (RKL): $D_{\text{RKL}}(P \parallel Q) = D_{\text{KL}}(Q \parallel P) = \sum_{c \in \mathcal{V}} Q(c) \log \frac{Q(c)}{P(c)}$ RKL is mode-seeking and focuses on high-precision alignment on major modes of the target.
3. Generalized Jensen-Shannon Divergence ( $D_{\text{JSD}[\beta]}$ ): $D_{\text{JSD}[\beta]}(P \parallel Q) = \beta D_{\text{KL}}(P \parallel \beta P + (1 - \beta)Q) + (1 - \beta) D_{\text{KL}}(Q \parallel \beta P + (1 - \beta)Q)$ where $\beta \in (0, 1)$ . For symmetric JS divergence, $\beta = 0.5$ .
4. Total Variation Distance (TVD): $D_{\text{TVD}}(P \parallel Q) = \frac{1}{2} \sum_{c \in \mathcal{V}} |P(c) - Q(c)|$
Knowl 7 — Distillation Recipe: On-Policy Data Generation and Divergence Selection
empirical result
Systematic empirical comparisons across training data generation sources and divergence functions reveal key design rules for training draft models in DistillSpec:
1. Superiority of Model-Generated Data over Fixed Datasets: Training draft models on static supervised ground truth data (Supervised KD) yields the lowest acceptance rates and plateaus early (around 40k steps). In contrast, training on dynamically generated sequences (student on-policy generations $M_q^\theta$ or teacher generations $M_p$ ) yields continuous improvement in acceptance rate and block efficiency.
2. Efficiency of Student On-Policy Data: Distilling exclusively on on-policy student generations (GKD) achieves block efficiency gains comparable to or exceeding those obtained from teacher-generated rollouts (f-Distill), while requiring substantially lower computational wall-clock time during training because running the compact draft model is much cheaper than generating from the large teacher model.
3. Task- and Strategy-Dependent Divergence Optimality: Although minimizing $D_{\text{TVD}}$ is the direct theoretical surrogate for sequence acceptance rate, optimizing $D_{\text{TVD}}$ directly with gradient descent does not consistently produce the best empirical block efficiency. Under greedy decoding ( $T = 0$ ), Forward KL or Jensen-Shannon Divergence often achieves the largest block efficiency improvements, whereas Reverse KL often performs best under temperature sampling ( $T = 1.0$ ).
4. Decoupling of Standalone Draft Performance and SD Alignment: High task performance of a distilled model evaluated standalone (e.g., accuracy or ROUGE) is only weakly correlated with its block efficiency in speculative decoding. Distillation choices must therefore be tuned specifically for token acceptance rather than standalone student accuracy.
Knowl 8 — Cross-Task and Target-Scale Generalization of Distilled Draft Models
empirical result
Draft models trained with DistillSpec generalize robustly across unseen tasks and larger target model architectures:
1. Unseen Task Transfer: A T5-Small (77M) draft model distilled exclusively on the GSM8K mathematical reasoning dataset using zero-shot Chain-of-Thought (CoT) prompting transfers directly to all 23 unseen reasoning tasks in the BIG-Bench Hard (BBH) benchmark when paired with a T5-XL (3B) target model. DistillSpec improves average decoding speedup on BBH from $1.93\times$ to $2.21\times$ under greedy decoding ( $T = 0$ ) and from $1.78\times$ to $2.02\times$ under temperature sampling ( $T = 1.0$ ), delivering an average 26% speedup increase over standard speculative decoding without distillation.
2. Target Model Scale Generalization: The same T5-Small draft model distilled from a T5-XL (3B) teacher generalizes zero-shot to assist a much larger 11B parameter target model (T5-XXL). Across XSum, GSM8K, CNN/DM, WMT, and BBH, this cross-scale pairing consistently outperforms standard non-distilled speculative decoding by 7% to 37% in decoding speedup.
Knowl 9 — Model Garden Serving Strategy via Distillation and Speculative Decoding Cascades
empirical result
When a suite of models of increasing sizes is available (evaluated across T5-Small 77M, T5-Base 250M, T5-Large 800M, T5-XL 3B, and T5-XXL 11B), four serving strategies present different quality-latency trade-offs:
1. Raw: Standalone supervised fine-tuned T5 models.
2. Distilled: Compact models trained via knowledge distillation from larger T5 models.
3. Speculative: Standard speculative decoding pairing a raw target model with a raw draft model.
4. DistillSpec: Pairing a distilled target model with a draft model distilled specifically to align with that target.
DistillSpec achieves the best Pareto frontier on both summarization and arithmetic reasoning:
- On XSum, DistillSpec reduces relative decoding latency from 17.3 to 2.7 ( $6.4\times$ latency reduction) compared to the standalone raw target model, with essentially no change in ROUGE-2 quality (23.1 raw vs. 23.0 DistillSpec).
- On GSM8K, DistillSpec reduces relative decoding latency from 15.0 to 1.4 ( $10.7\times$ latency reduction) compared to the raw target model, while task accuracy improves from 33.1% to 34.8% due to the distilled target's enhanced reasoning capability.

Knowl 10 — Empirical Block Efficiency and Speedup Improvements Across Benchmark Tasks

data/table

DistillSpec was evaluated across multiple NLP tasks using encoder-decoder models (T5-XL 3B target paired with T5-Small 77M draft) and decoder-only models on LM1B (GPT-like 234M target paired with GPT-like 33M draft). The table below compares empirical block efficiency $\tau$ (expected accepted tokens per block) and wall-clock speedup (measured on TPUv4 at batch size 1) for standard speculative decoding without distillation vs. DistillSpec across block sizes $\gamma \in \{3, 5, 7\}$ under greedy ( $T = 0$ ) and temperature ( $T = 1$ ) sampling.

Dataset	Temp	$\gamma$	w/o Distillation		with Distillation (DistillSpec)			Best KD
			$\tau$	Speedup	$\tau$	Speedup	$\Delta$ Speedup	Method
XSum	$T=0$	3	2.31	$1.44\times$	2.62	$1.58\times$	$+0.14\times$	f-Distill
(T5-XL / Small)	$T=0$	5	2.57	$1.43\times$	3.08	$1.62\times$	$+0.19\times$	f-Distill
	$T=0$	7	2.68	$1.36\times$	3.31	$1.57\times$	$+0.21\times$	f-Distill
	$T=1$	3	2.19	$1.40\times$	2.58	$1.57\times$	$+0.17\times$	f-Distill
	$T=1$	5	2.39	$1.37\times$	3.01	$1.61\times$	$+0.25\times$	f-Distill
	$T=1$	7	2.47	$1.28\times$	3.21	$1.55\times$	$+0.27\times$	f-Distill
CNNDM	$T=0$	3	2.83	$1.89\times$	3.19	$2.11\times$	$+0.22\times$	f-Distill
(T5-XL / Small)	$T=0$	5	3.46	$2.07\times$	4.13	$2.42\times$	$+0.35\times$	GKD
	$T=0$	7	3.85	$2.07\times$	4.83	$2.53\times$	$+0.46\times$	f-Distill
	$T=1$	3	2.49	$1.71\times$	2.87	$1.93\times$	$+0.23\times$	f-Distill
	$T=1$	5	2.89	$1.77\times$	3.52	$2.12\times$	$+0.35\times$	f-Distill
	$T=1$	7	3.08	$1.71\times$	3.92	$2.12\times$	$+0.41\times$	f-Distill
GSM8K	$T=0$	3	2.60	$1.51\times$	2.96	$1.69\times$	$+0.18\times$	GKD
(T5-XL / Small)	$T=0$	5	3.06	$1.42\times$	3.68	$1.65\times$	$+0.22\times$	GKD
	$T=0$	7	3.27	$1.36\times$	4.14	$1.60\times$	$+0.24\times$	GKD
	$T=1$	3	2.58	$1.48\times$	2.84	$1.64\times$	$+0.16\times$	f-Distill
	$T=1$	5	3.03	$1.39\times$	3.45	$1.58\times$	$+0.19\times$	f-Distill
	$T=1$	7	3.23	$1.33\times$	3.84	$1.53\times$	$+0.20\times$	f-Distill
LM1B	$T=0$	3	2.96	$3.66\times$	3.13	$3.97\times$	$+0.31\times$	f-Distill
(GPT 234M / 33M)	$T=0$	5	3.69	$3.35\times$	3.92	$3.51\times$	$+0.16\times$	f-Distill
	$T=0$	7	4.15	$2.52\times$	4.55	$2.72\times$	$+0.20\times$	f-Distill
	$T=1$	3	2.51	$2.34\times$	2.69	$2.45\times$	$+0.11\times$	f-Distill
	$T=1$	5	2.90	$2.79\times$	3.20	$3.02\times$	$+0.23\times$	f-Distill
	$T=1$	7	3.10	$1.98\times$	3.51	$2.18\times$	$+0.20\times$	f-Distill
BBH-AVG	$T=0$	7	3.20	$1.93\times$	3.49	$2.21\times$	$+0.28\times$	f-Distill
(T5-XL / Small)	$T=1$	7	3.01	$1.78\times$	3.28	$2.02\times$	$+0.23\times$	f-Distill

These results demonstrate consistent gains across all tasks and model architectures, with DistillSpec providing 10% to 46% additional speedup over standard speculative decoding without distillation.

Coverage note — None was omitted; all key theoretical bounds, algorithms, ablation findings, transferability results, and benchmark data were captured.

References

1.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.13649, 2023.
2.Rohan Anil, Andrew M Dai, Orhan Firat, Melvin 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.
3.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 translation. In Proceedings of the Ninth Workshop on Statistical Machine Translation, pp. 12–58, Baltimore, Maryland, USA, June 2014. Association for Computational Linguistics. URL http://www.aclweb.org/anthology/W/W14/W14-3302.
4.Cristian Buciluǎ, Rich Caruana, and Alexandru Niculescu-Mizil. Model compression. In Proceedings of the 12th ACM SIGKDD international conference on Knowledge discovery and data mining, pp. 535–541, 2006.
5.Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, and Tri Dao. Medusa: Simple framework for accelerating llm generation with multiple decoding heads. https://github.com/FasterDecoding/Medusa, 2023a.
6.Tianle Cai, Xuezhi Wang, Tengyu Ma, Xinyun Chen, and Denny Zhou. Large language models as tool makers. arXiv preprint arXiv:2305.17126, 2023b.
7.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 arXiv:1312.3005, 2013.
8.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.01318, 2023.
9.Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021.
10.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.
11.Tao Ge, Heming Xia, Xin Sun, Si-Qing Chen, and Furu Wei. Lossless acceleration for seq2seq generation with aggressive decoding. arXiv preprint arXiv:2205.10350, 2022.
12.Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang. Knowledge distillation of large language models. arXiv preprint arXiv:2306.08543, 2023.
13.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.
14.Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531, 2015.
15.Ferenc Huszár. How (not) to train your generative model: Scheduled sampling, likelihood, adversary? arXiv preprint arXiv:1511.05101, 2015.
16.Sehoon Kim, Karttikeya Mangalam, Jitendra Malik, Michael W Mahoney, Amir Gholami, and Kurt Keutzer. Big little transformer decoder. arXiv preprint arXiv:2302.07863, 2023.
17.Yoon Kim and Alexander M Rush. Sequence-level knowledge distillation. arXiv preprint arXiv:1606.07947, 2016.
18.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.
19.Alexander Lin, Jeremy Wohlwend, Howard Chen, and Tao Lei. Autoregressive knowledge distillation through imitation learning. arXiv preprint arXiv:2009.07253, 2020.
20.Xiaoxuan Liu, Lanxiang Hu, Peter Bailis, Ion Stoica, Zhijie Deng, Alvin Cheung, and Hao Zhang. Online speculative decoding. arXiv preprint arXiv:2310.07177, 2023.
21.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 llm serving with speculative inference and token tree verification. arXiv preprint arXiv:2305.09781, 2023.
22.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/1808.08745, 2018.
23.OpenAI. Gpt-4 technical report. ArXiv, abs/2303.08774, 2023. URL https://api.semanticscholar.org/CorpusID:257532815.
24.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.
25.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–5551, 2020.
26.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: 10.1109/34.655647.
27.Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp. 4596–4604. PMLR, 2018.
28.Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit. Blockwise parallel decoding for deep autoregressive models. Advances in Neural Information Processing Systems, 31, 2018.
29.Xin Sun, Tao Ge, Furu Wei, and Houfeng Wang. Instantaneous grammatical error correction with shallow aggressive decoding. arXiv preprint arXiv:2106.04970, 2021.
30.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.net/forum?id=d0mGsaheuT.
31.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 chain-of-thought can solve them. arXiv preprint arXiv:2210.09261, 2022.
32.Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An instruction-following llama model. https://github.com/tatsu-lab/stanford_alpaca, 2023.
33.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 arXiv:2201.08239, 2022.
34.Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023.
35.Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 35:24824–24837, 2022.
36.Yuqiao Wen, Zichao Li, Wenyu Du, and Lili Mou. f-divergence minimization for sequence-level knowledge distillation. arXiv preprint arXiv:2307.15190, 2023.
37.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, 15(62):2113–2144, 2014.
38.Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068, 2022.

Citation

MLA

Zhou, Y., et al. “DistillSpec: Improving Speculative Decoding via Knowledge Distillation”. arXiv, 2023, http://arxiv.org/abs/2310.08461v2.

APA

Zhou, Y., Lyu, K., Rawat, A. S., Menon, A. K., Rostamizadeh, A., Kumar, S., Kagy, J.-F., & Agarwal, R. (2023). DistillSpec: Improving Speculative Decoding via Knowledge Distillation. arXiv. http://arxiv.org/abs/2310.08461v2

Chicago

Zhou, Y., K. Lyu, A. S. Rawat, et al. 2023. “DistillSpec: Improving Speculative Decoding via Knowledge Distillation”. arXiv. http://arxiv.org/abs/2310.08461v2.

Harvard

Zhou, Y. et al. (2023) “DistillSpec: Improving Speculative Decoding via Knowledge Distillation”, arXiv [Preprint]. Available at: http://arxiv.org/abs/2310.08461v2.

Vancouver

1. Zhou Y, Lyu K, Rawat AS, Menon AK, Rostamizadeh A, Kumar S, Kagy J-F, Agarwal R (2023) DistillSpec: Improving Speculative Decoding via Knowledge Distillation. arXiv

BibTeX

@article{zhou2023distillspec,
  title = {DistillSpec: Improving Speculative Decoding via Knowledge Distillation},
  author = {Zhou, Yongchao and Lyu, Kaifeng and Rawat, Ankit Singh and Menon, Aditya Krishna and Rostamizadeh, Afshin and Kumar, Sanjiv and Kagy, Jean-François and Agarwal, Rishabh},
  year = {2023},
  journal = {arXiv},
  url = {http://arxiv.org/abs/2310.08461v2},
  eprint = {2310.08461}
}

Metadata:arXiv

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: Published with permission

DistillSpec: Improving Speculative Decoding via Knowledge Distillation

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — DistillSpec Framework for Aligning Draft Models via White-Box Knowledge Distillation

Knowl 2 — Total Variation Distance as an Exact Proxy for Speculative Decoding Acceptance Rate

Knowl 3 — Lower Bound on Speculative Acceptance Rate from On-Policy Distillation Loss

Knowl 4 — DistillSpec Knowledge Distillation Training Algorithm

Knowl 5 — Speculative Decoding with Generalized Lenience Functions

Knowl 6 — Divergence Objectives for Draft Model Knowledge Distillation

Knowl 7 — Distillation Recipe: On-Policy Data Generation and Divergence Selection

Knowl 8 — Cross-Task and Target-Scale Generalization of Distilled Draft Models

Knowl 9 — Model Garden Serving Strategy via Distillation and Speculative Decoding Cascades

Knowl 10 — Empirical Block Efficiency and Speedup Improvements Across Benchmark Tasks

References

Citation

BibTeX

Access the Paper