PaperlayerView exact revision on arXiv

Source-linked AI summary

DistillSpec: Improving Speculative Decoding via Knowledge Distillation

Yongchao Zhou, Kaifeng Lyu, Ankit Singh Rawat, Aditya Krishna Menon, Afshin Rostamizadeh, Sanjiv Kumar, Jean-François Kagy, Rishabh Agarwal

arXiv:2310.08461v2cs.CLcs.AIcs.LG

TL;DR

Speculative decoding needs a compact draft model that remains closely aligned with a larger target, because acceptance determines its efficiency. DistillSpec uses white-box knowledge distillation, on-policy draft data, and task- and decoding-sensitive divergences to improve that alignment, yielding faster decoding and broader latency–quality control. In multi-model settings, combining target-model distillation with DistillSpec produces large latency reductions with minimal performance degradation.

  • Problem

    Speculative decoding’s efficiency depends on finding a compact draft model that is both inexpensive to decode and closely aligned with the target.

  • Method

    DistillSpec uses white-box knowledge distillation to align draft and target models, emphasizing on-policy data and divergence choices tailored to tasks and decoding strategies.

  • Results

    10−45% speedup over standard speculative decoding was achieved across diverse datasets under greedy and non-greedy sampling while preserving model performance.

  • Takeaways & Limitations

    Combining target-model distillation with DistillSpec reduced decoding latency by 6−10× with minimal performance degradation in a multi-model setting.

  • Takeaways & Limitations

    The method assumes white-box access to both target and draft models’ next-token distributions.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Speculative decoding reduces LLM latency with a compact draft model, but its benefit depends on the draft’s alignment with the target. DistillSpec uses white-box knowledge distillation to improve that alignment and extends the approach to decoding and model-selection settings.

  • Motivation: Speculative decoding combines sequential draft generation with parallel target verification to reduce autoregressive inference cost.Its realized latency benefit depends on the acceptance of draft-generated tokens.
  • DistillSpec: DistillSpec uses white-box knowledge distillation to align a draft model with its target and improve speculative-decoding acceptance rates.The method targets draft–target alignment rather than standalone student task performance.
  • Extensions: 6−10× lower latency was obtained by first distilling a target model and then DistillSpec-distilling a smaller draft, with minimal performance degradation.The comparison is against standalone non-distilled targets of the same size in a multi-model setting.
  • Results: 10−45% speedup over standard speculative decoding was achieved across diverse datasets under greedy and non-greedy sampling while preserving model performance.The contribution reports this improvement across decoding strategies and datasets.
  • Design analysis: DistillSpec’s distillation recipe studies training-data generation, divergence functions, and decoding strategies for model alignment.The analysis emphasizes model-generated data and task- and decoding-strategy-dependent divergence choices.
  • Extensions: DistillSpec extends to lossy speculative decoding, enabling finer control over the quality–latency trade-off.Lossy decoding permits sampling away from the target distribution.

2 RELATED WORK

Prior efficiency methods reduce inference cost through parallel decoding, cascading, or knowledge distillation. Speculative decoding differs by using token-level draft–target collaboration while preserving the target distribution.

  • Speculative decoding: Speculative decoding uses a compact draft model to generate token batches and a larger target model to validate them in parallel.This approach addresses autoregressive latency while retaining the target model’s distributional behavior.
  • Knowledge distillation: Knowledge distillation trains smaller students under larger teachers and can reduce inference cost while maintaining model quality.Prior LLM work includes black-box KD from teacher outputs and white-box KD using teacher weights and logits.
  • DistillSpec: DistillSpec applies distillation to draft–target collaboration rather than creating only a capable standalone student model.This setting is useful when a small distilled model alone cannot satisfy quality requirements.

3 BACKGROUND: SPECULATIVE DECODING

Speculative decoding uses a small draft model to propose token blocks and a large target model to verify them in parallel while preserving the target distribution. Its efficiency depends on how many drafted tokens are accepted, motivating acceptance and block-efficiency measures.

  • Speculative sampling: The draft model autoregressively samples a block of γ tokens, which the target model verifies in parallel using token-wise acceptance probabilities.Rejected tokens terminate acceptance within the block, and the procedure repeats until generation ends.
  • Acceptance rate: The sequence-level acceptance rate α(x) measures accepted tokens relative to the target model’s expected output length and lower-bounds the reduction in speculative-decoding steps.Because speculative sampling preserves the target distribution, the denominator is invariant to the draft model.
  • Block efficiency: For finite block size γ, block efficiency τ(x) is the expected number of accepted tokens per block, with a maximum of γ + 1.The extra token comes from the target model when all drafted tokens are accepted.
  • Wall-clock improvement: The expected speedup for block efficiency τ(x) is τ(x)/(cγ + 1), where c is the draft-to-target single-forward-pass latency ratio.This connects acceptance behavior to wall-clock improvement.

4 DISTILLSPEC: KNOWLEDGE DISTILLATION FOR SPECULATIVE DECODING

DistillSpec trains a draft model to align with the target model for speculative decoding, using white-box distillation with task-appropriate data and divergences. Its design links divergence choice and model-generated data to acceptance and block efficiency, while also extending the framework to lossy decoding.

  • DistillSpec framework: DistillSpec uses white-box knowledge distillation to align the draft and target distributions, optimizing the draft for speculative-decoding acceptance rather than standalone task performance.The method assumes access to next-token distributions from both models.
  • Distillation design: The distillation recipe is determined by the training data G and divergence D, including task pairs or model-generated sequences and alternatives to forward KL.Table 1 summarizes these algorithms by their data and divergence choices.
  • Choice of divergence: Minimizing expected total variation distance is theoretically tied to maximizing acceptance, but experiments show the best divergence is task-dependent.The token-level acceptance rate equals 1 − DTVD(p(yt), q(yt)).
  • Choice of training data: Model-generated outputs are used because ground-truth data may not match the teacher distribution and may be limited enough to encourage overfitting.The approach explores on-policy sequences sampled from the student to reduce teacher-generation cost.
  • On-policy distillation: On-policy KD provides a sequence-level acceptance bound: when target output length is always T, Ex∼X[α(x)] ≥ 1 − Tϵ.The bound connects small on-policy KD loss to draft–target alignment across sequence positions.
  • Lossy speculative decoding: DistillSpec with lossy speculative decoding varies lenience ε to trade generation quality for faster inference through modified acceptance probabilities.The evaluated functions include flin(p, ε) = p/ε, fsq(p, ε) = p/ε^2, and fexp(p, ε) = p^ε.

5 EXPERIMENTS

Experiments show that distillation improves speculative decoding across model types, datasets, and decoding strategies, while the best recipe depends on the task and sampling method. DistillSpec also improves quality-latency trade-offs, including substantial latency reductions with nearly identical performance.

  • Decoding speedup: 10−46% improvement in decoding speed is achieved across datasets, with the strongest effect under greedy decoding.Knowledge distillation improves acceptance rate and block efficiency for both decoder-only and encoder-decoder models.
  • Decoding speedup: ∼20% additional speedup over standard speculative decoding comes from distillation algorithms using model-generated data on LM1B, XSum, CNN/DM, and GSM8K.
  • Training data and divergence: Model-generated data and teacher logits are important: f-Distill and GKD outperform variants using fixed ground-truth data, while SeqKD performs worse despite target-generated data.White-box supervision from teacher logits and continual data generation support higher acceptance rates during training.
  • Block efficiency: Distillation improves block efficiency across most examples, and empirical block efficiency aligns strongly with its DTVD-based theoretical counterpart.
  • Training data and divergence: The best divergence and data construction depend on the task and decoding strategy: XSum favors mixed data with JSD, while GSM8K favors draft-only data with FKL under greedy decoding and target data with RKL at T = 1.Draft-generated data alone remains competitive, and the choices should be tuned as task-specific hyperparameters.
  • Quality versus latency: High downstream task performance does not necessarily produce an effective speculative-decoding draft model, while leniency provides a mechanism for controlling the quality-latency trade-off.
  • Quality versus latency: DistillSpec consistently achieves the best quality-latency trade-off, reducing relative latency from 17.3 to 2.7 on XSum and from 15.0 to 1.4 on GSM8K while maintaining nearly identical performance.The combined approach is especially effective when stringent performance and decoding-strategy constraints apply.

6 CONCLUSION

The paper evaluates white-box knowledge distillation for improving alignment between target and draft models in speculative decoding. It finds that model-generated data, task- and strategy-specific divergence choices, and a distillation-then-speculative-decoding pipeline improve the quality-latency trade-off, including six- to ten-fold lower latency with negligible performance loss.

  • White-box knowledge distillation improves speculative decoding by aligning target and draft models.
  • Model-generated data and task- and decoding-strategy-specific divergence choices are central design considerations for improving model alignment.
  • A six- to ten-fold latency decrease with negligible performance loss results from first distilling for downstream performance and then applying speculative decoding with a distilled draft model.

A METHOD

This section introduces divergence functions for comparing probability distributions in distillation, emphasizing that their asymmetry and parameterization affect training behavior.

  • KL divergence compares two discrete probability distributions and is not symmetric.
  • Forward KL is DKL(P∥Q), while reverse KL is DKL(Q∥P).
  • Minimizing forward KL under an empirical data distribution corresponds to maximum likelihood estimation.
  • The section also considers Jensen–Shannon divergence and generalized Jensen–Shannon divergence.
  • Generalized Jensen–Shannon divergence behaves like forward KL for small β and reverse KL when β is close to 1.
  • Total variation distance is included among the divergence functions considered for distillation.

A.2 JUSTIFICATION OF USING ON-POLICY DATA

The appendix formalizes token-level discrepancies and the mixed-model sequence distributions used to justify on-policy data, then specifies the speculative decoding and distillation procedures.

  • A.2 JUSTIFICATION OF USING ON-POLICY DATA: Token-level total variation distance is averaged over autoregressive sequences to define the discrepancy between target and draft predictions.The appendix expresses this quantity through expectations over sequence distributions.
  • A.2 JUSTIFICATION OF USING ON-POLICY DATA: The analysis defines sequence distributions by sampling each token from either the target or draft model according to a P/Q pattern.Sampling stops at an end-of-sequence token or the maximum sequence length τ.
  • A.2 JUSTIFICATION OF USING ON-POLICY DATA: The theorem bounds the target-sequence discrepancy using quantities associated with replacing target-generated prefixes by draft-generated prefixes.The proof repeatedly replaces P tokens with Q tokens while accounting for an error term.
  • A.2 JUSTIFICATION OF USING ON-POLICY DATA: Speculative decoding samples a block from the draft model, verifies it in parallel with the target model, and applies rejection thresholds.The procedure uses block size γ and a lenience function for lossless or lossy decoding.
  • A.2 JUSTIFICATION OF USING ON-POLICY DATA: Knowledge distillation mixes fixed, student-generated, and target-generated data according to sampling fractions before updating the student.The algorithm exposes fixed-data and student-data fractions λ1 and λ2, together with a divergence function and learning rate.

B.1 DATASETS

The experiments cover summarization, translation, mathematical reasoning, and language modeling using several established datasets and decoder-only or encoder-decoder model pairs.

  • B.1 DATASETS: XSum evaluates abstractive single-document summarization on 226,711 BBC news articles using ROUGE scores.
  • B.1 DATASETS: CNN/Daily Mail evaluates abstractive summarization of news stories, with validation performance reported primarily using ROUGE-2.
  • B.1 DATASETS: WMT14 EnDe evaluates English-to-German machine translation using BLEU scores.Distillation and evaluation use maximum input and output lengths of 80.
  • B.1 DATASETS: GSM8K contains 8.5K grade-school math word problems, with solutions typically requiring 2 to 8 reasoning steps.The study combines distillation with zero-shot chain-of-thought exploration.
  • B.1 DATASETS: LM1B is used for language modeling with maximum input and output lengths of 128.
  • B.1 DATASETS: Experiments use 234M-parameter target and 33M-parameter draft decoder-only models, plus 3B-parameter T5-XL and 77M-parameter T5-Small pairs.

B.3 DISTILLATION

The distillation experiments compare data-generation strategies and train the draft model with a specified optimizer, divergence-based loss, fixed temperatures, and documented hyperparameters.

  • B.3 DISTILLATION: The study evaluates five knowledge-distillation algorithms using online teacher or student data generation rather than relying uniformly on fixed datasets.For GKD, static ground-truth data is excluded and student-generated data is used exclusively.
  • B.3 DISTILLATION: The draft student model is trained with the Adafactor optimizer in all experiments.
  • B.3 DISTILLATION: The target and draft temperatures are both fixed at 1.0 in the knowledge-distillation loss.The paper emphasizes uniform temperatures for consistent semantic interpretation in speculative decoding.
  • B.3 DISTILLATION: Table B.1 summarizes the hyperparameters used in the distillation experiments.

B.4 EVALUATION

The evaluation reports average task scores, empirical speculative-decoding metrics, and measured latency using standardized test procedures. Latency is measured on a shared TPUv4 device with repeated batch-size-one runs.

  • Evaluation methodology: Task scores and empirical acceptance rate and block efficiency are averaged across all test or validation instances.Scores include ROUGE-2, BLEU, and accuracy, depending on the task.
  • Evaluation methodology: Actual latency is measured by running the target and draft models on the same TPUv4 device without model parallelism.The procedure uses batch size 1 and randomly samples 500 test or validation examples.
  • Evaluation methodology: Latency measurements are repeated three times and reported as mean performance, with minimal variance observed across random seeds.

C ADDITIONAL RESULTS

Additional results show that DistillSpec improves speculative-decoding efficiency across datasets, decoding settings, block sizes, and transferred target models. The strongest consistency over training is reported for GKD and f-Distill, while SupKD can decline on some datasets.

  • Additional evaluation: DistillSpec improves speculative-decoding performance across datasets and block sizes under both greedy decoding and temperature sampling.BBH-AVG summarizes results over 23 BIG-Bench Hard tasks.
  • Model transferability: A distilled T5-Small draft derived from a T5-XL teacher generalizes to the larger T5-XXL target model and accelerates diverse scenarios.BBH-AVG reports average results over 23 BIG-Bench Hard tasks.
  • Model transferability: Across 23 BIG-Bench Hard tasks, a GSM8K-trained T5-Small draft yields significant speedups with T5-XL or T5-XXL targets under greedy and non-greedy decoding.
  • Block efficiency: DistillSpec outperforms standard speculative decoding across all considered distillation methods for non-greedy sampling, with f-Distill and GKD producing the highest gains.The comparison is shown for block sizes γ = 3, 5, and 7.
  • Block efficiency: For greedy sampling, DistillSpec outperforms standard speculative decoding across all considered methods, with GKD weakly outperforming the others on average.This pattern is reported for block sizes γ = 3, 5, and 7.
  • Performance over training: GKD and f-Distill provide the most consistent training improvements in acceptance rate and block efficiency, whereas SupKD shows weaker gains and later declines on XSum and CNNDM.SupKD exhibits declining acceptance rates or block efficiency after approximately 40k training steps in those datasets.
Loading 2310.08461v2…