PaperlayerView exact revision on arXiv

Source-linked AI summary

AdaSPEC: Selective Knowledge Distillation for Efficient Speculative Decoders

Yuezhou Hu, Jiaxin Guo, Xinyu Feng, Tuo Zhao

arXiv:2510.19779v1cs.CLcs.AIcs.LG

TL;DR

Speculative Decoding speeds large language model inference through draft-model proposals verified by a target model, but conventional distillation objectives do not directly optimize token acceptance under draft-capacity constraints. AdaSPEC filters difficult tokens during distillation, and experiments report higher acceptance rates than DistillSpec across tasks and model setups, up to 15%.

  • Problem

    Draft models have limited capacity, and minimizing divergence across all tokens does not necessarily optimize the acceptance rate that speculative decoding relies on.

    Source: §1 · p. 2

  • Method

    AdaSPEC uses a reference model's perplexity differences to identify hard tokens, then distills the draft on a dataset with those tokens filtered out.

    Source: §1 · p. 2

  • Results

    Across all tasks and model setups, AdaSPEC achieves higher acceptance rates than DistillSpec, with gains of up to 15%.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Experiments report higher acceptance rates for AdaSPEC in arithmetic reasoning, instruction following, code generation, and summarization.

    Source: §6 · p. 10

    As a preliminary study, AdaSPEC uses a simple loss-related token filter; more adaptive filters and integration with tree-based or multi-step verification remain future work.

    Source: §5 · p. 10

Abstract

from arXiv · show

Speculative Decoding (SD) accelerates large language model inference by employing a small draft model to generate predictions, which are then verified by a larger target model. The effectiveness of SD hinges on the alignment between these models, which is typically enhanced by Knowledge Distillation (KD). However, conventional KD methods aim to minimize the KL divergence between the draft and target models across all tokens, a goal that is misaligned with the true objective of SD, which is to maximize token acceptance rate. Therefore, draft models often struggle to fully assimilate the target model's knowledge due to capacity constraints, leading to suboptimal performance. To address this challenge, we propose AdaSPEC, a novel method that incorporates selective token filtering into the KD process. AdaSPEC utilizes a reference model to identify and filter out difficult-to-fit tokens, enabling the distillation of a draft model that better aligns with the target model on simpler tokens. This approach improves the overall token acceptance rate without compromising generation quality. We evaluate AdaSPEC across diverse tasks, including arithmetic reasoning, instruction-following, coding, and summarization, using model configurations of 31M/1.4B and 350M/2.7B parameters. Our results demonstrate that AdaSPEC consistently outperforms the state-of-the-art DistillSpec method, achieving higher acceptance rates across all tasks (up to 15\%). The code is publicly available at https://github.com/yuezhouhu/adaspec.

1 Introduction

Speculative Decoding accelerates inference by having a small draft model propose tokens for a larger target model to verify, but effective acceleration depends on draft–target alignment. AdaSPEC filters difficult tokens during distillation so the draft can focus its limited capacity on more manageable tokens.

1 Introduction

  • Large language models can be slow and computationally intensive to deploy, while compression and quantization trade off efficiency against performance.The introduction motivates methods that improve inference efficiency while maintaining high model performance.Source: §1 · pp. 1–2
  • Speculative Decoding uses a lightweight draft model to propose multiple candidate tokens for verification by a larger target model.The target model verifies the proposals, preserving its predictive quality while reducing expensive forward passes.Source: §1 · p. 2
  • Conventional distillation can optimize fidelity metrics such as forward KL without ensuring high acceptance, while spending limited draft capacity on tokens that are hard to learn.The introduction also notes that these methods may encounter loss-convergence issues.Source: §1 · p. 2
  • AdaSPEC excludes hard-to-learn tokens because reducing their loss can come at the expense of higher loss on easy tokens.The method focuses the loss on tokens the draft can learn more readily, aiming to improve alignment within its capacity constraints.Source: §1 · p. 2
  • AdaSPEC identifies hard tokens with a reference model and distills the draft on a filtered dataset of remaining tokens.The reference model is initialized as a copy of the draft and compares perplexity differences between the two models on training data.Source: §1 · p. 2
  • Across all tasks and model setups, AdaSPEC achieves higher acceptance rates than DistillSpec, with gains of up to 15%.The introduction points to Table 1 for the comparison.Source: §1 · p. 3

2 Preliminaries

The preliminaries describe speculative decoding, its acceptance-rate and wall-time metrics, and the model families used to study alignment across scales.

2 Preliminaries

  • Speculative Decoding has a draft model generate tokens that a larger target model verifies, with both models producing distributions for token selection.The presented framework uses greedy decoding to select tokens for generation or verification.Source: §2 · p. 3
  • Acceptance rate α measures draft–target accuracy; higher α indicates greater alignment and facilitates faster inference.It is based on the counts of tokens accepted and rejected by the target model.Source: §2 · p. 3
  • Block efficiency τ is the expected number of accepted tokens per block, with a maximum of γ + 1 for block size γ.The passage also states that block efficiency can be expressed in terms of acceptance rate α.Source: §2 · p. 3
  • The wall-time speed-up factor uses c, the ratio of draft-model execution time to target-model execution time.The passage introduces the factor as a measure of total wall-time improvement.Source: §2 · p. 3
  • Language-model families share a core architecture across scales, enabling knowledge transfer or prediction alignment between smaller and larger models.Examples named include Llama 3, BERT, and Pythia.Source: §2 · p. 3

3 Method

AdaSPEC constructs a reference model, then selectively distills the target model’s knowledge into the draft model. It chooses tokens by comparing token-wise losses and trains on the selected learnable tokens.

3 Method

  • AdaSPEC constructs a reference model and selectively distills knowledge from the target and reference models to the draft model.These are the method’s two key steps.Source: §3 · p. 4
  • The reference model is distilled from the target on a downstream-task dataset using DistillSpec [32] to minimize forward KL divergence.The passage defines the KL objective in terms of target and reference distributions.Source: §3 · p. 4
  • AdaSPEC computes token-wise losses for the reference and draft models, then uses their difference to identify tokens that are highly learnable for the draft.Larger differences indicate a greater draft–reference performance gap relative to the target model.Source: §3 · p. 4
  • The draft is trained on the top k×100% of tokens by loss difference, with only selected tokens contributing to the distillation loss.Figure 1 overviews this filtering and selective-distillation process.Source: §3 · p. 4, Figure 1 · p. 4

4 Experiments

Across two model configurations and five tasks, AdaSPEC achieves higher acceptance rates than DistillSpec; analyses and ablations examine token alignment, selection, and training choices. Additional tests report end-to-end speed-ups and results with EAGLE and larger models.

4.1 Experimental Setup

  • Experiments use Pythia-31M/1.4B and cross-family CodeGen-350M/Phi-2 configurations across five tasks; comparisons include three-epoch and optimal-epoch settings, capped at 30 epochs for computational feasibility.The tasks are GSM8K, Alpaca, MBPP, CNN/Daily Mail, and XSUM; CNN/Daily Mail and XSUM use a smaller optimal-epoch search range.Source: §4.1 · p. 5, §4.2 · p. 5

4.2 Baseline Setup

  • For the optimal-epoch setting, the number of epochs is tuned using validation performance, and the reference with the highest validation acceptance rate trains the draft model.The three-epoch setting provides a controlled comparison under a fixed training duration.Source: §4.2 · p. 5

4.3 Main Results

  • AdaSPEC achieves higher acceptance rates than DistillSpec across all tasks and model configurations.Acceptance rate α is the proportion of draft-generated tokens validated by the target model.Source: Table 1 · p. 6, §4.3 · p. 6

4.4 Analysis

  • AdaSPEC shifts acceptance-rate distributions rightward, yields more frequent and larger positive margins and fewer and smaller negative margins, and lowers KL divergence on GSM8K and CNN/Daily Mail.Positive margins indicate more frequent and confident correct predictions; negative margins indicate errors that are less frequent and less severe.Source: §4.4 · pp. 6–7, Figure 2 · p. 7

4.5 Ablation Study

  • Models trained on top-40% tokens selected by KL-divergence margin outperform bottom-40% models, with up to 6% gains on MBPP; lower k typically improves draft acceptance, and k = 0.4 balances efficiency and performance.For GSM8K, lower k values of 0.2–0.4 yield higher acceptance rates.Source: §4.5 · pp. 7, 9, Figure 4 · p. 9

4.4 Analysis

  • AdaSPEC’s errors are nearly a subset of DistillSpec’s, and on GSM8K it preferentially selects mathematics-related tokens and improves prediction accuracy for them.The case studies cover GSM8K and CNN/Daily Mail.Source: §4.4 · p. 7, Figure 3 · p. 8

4.5 Ablation Study

  • Direct fine-tuning retains benefits from token selection, while forward KL improves MBPP acceptance by 6%; RKL and TVD degrade acceptance, motivating the choice of forward KL.Fine-tuned drafts outperform the distillation reference on both datasets and improve by up to 4% over their reference counterparts.Source: §4.5 · p. 8, Table 3 · p. 8, Table 4 · p. 8

4.6 Additional Experimental Results

  • 10–20% speed-up is reported over DistillSpec with vLLM on one A100 GPU; AdaSPEC also improves accuracy and decoding efficiency with EAGLE and reaches 86.21% acceptance versus DistillSpec’s 84.43% on larger GSM8K models.The speed test uses greedy decoding and Pythia-31M→1.4B models trained for three epochs; the larger-model comparison uses Qwen2.5-0.5B and Qwen2.5-32B.Source: Table 5 · p. 9, §4.6 · p. 9, Table 6 · p. 9

5 Dicussion

AdaSPEC’s acceptance-rate gains become more pronounced as the target–draft size gap widens, including experiments with gaps up to 64 times. The method filters challenging tokens, while the study’s current filtering strategy remains simple.

Size Gap Between Target and Draft Models

  • Up to 64 times: AdaSPEC enhances acceptance rates across a substantial target–draft size gap while helping maintain generation quality.Traditional speculative-decoding size gaps are often within 10x.Source: §5 · p. 10

Model Size Gap and Performance Gains

  • AdaSPEC’s gains over DistillSpec grow with the reference–target size gap across both 3-Epoch and Optimal-Epoch settings.The passage attributes the trend to representation mismatch under direct distillation and selective alignment of easier tokens.Source: §5 · p. 10

Connection with Lin et al. [19]

  • Unlike Lin et al. [19], which prioritizes harder-to-learn tokens during pre-training, AdaSPEC filters challenging tokens for speculative decoding.The passage frames the approaches as serving different objectives: pre-training versus speculative decoding.Source: §5 · p. 10

Limitations

  • The study is a preliminary investigation limited to a simple loss-related token filter.Future work could explore more adaptive filtering and integration with tree-based or multi-step verification frameworks.Source: §5 · p. 10

6 Conclusion

AdaSPEC uses selective token filtering to train more efficient speculative-decoding draft models. Experiments across several task domains report higher acceptance rates than baselines.

6 Conclusion

  • AdaSPEC uses selective token filtering to train more efficient speculative-decoding draft models, outperforming baselines in acceptance rates across four task domains.The conclusion names arithmetic reasoning, instruction following, code generation, and summarization.Source: §6 · p. 10

A.1 Full Algorithms for AdaSPEC

The algorithms describe speculative decoding’s draft-and-verify loop and AdaSPEC’s selective-distillation workflow. AdaSPEC trains a reference model, filters tokens, and distills the draft model on the selected set.

Algorithm 1 Greedy Speculative Decoding

  • Speculative decoding samples draft tokens autoregressively, verifies them in parallel with the target model, and accepts tokens until a mismatch or end-of-sequence.On mismatch, the algorithm appends a target-model token and continues drafting.Source: §A.1 · p. 13

Algorithm 2 AdaSPEC: Selective Distillation for Speculative Decoding

  • AdaSPEC fine-tunes the target model, initializes a reference model from the draft, and distills the reference from the fine-tuned target.The algorithm specifies forward KL as an example distillation objective.Source: §A.1 · p. 13
  • AdaSPEC computes token losses and filters tokens by selecting the top k×100% according to the loss difference.The selected set is used to distill the draft model.Source: §A.1 · p. 13
  • The selective-distillation algorithm returns the trained draft model for speculative decoding.The algorithm labels the returned model Mq.Source: §A.1 · p. 13

A.2 Implementation Details

The implementation uses separate epoch-selection procedures for the 3-Epoch and Optimal-Epoch settings. Optimal-Epoch selection uses validation performance, with a capped candidate epoch schedule.

A.2 Implementation Details

  • The 3-Epoch setting distills both reference and draft models for three epochs.In the Optimal-Epoch setting, GSM8K target epochs are selected by validation accuracy; other experiments use validation perplexity.Source: §A.2 · pp. 13–14
  • Optimal epochs are selected by validation accuracy or, for the remaining experiments, validation perplexity; among epochs 1, 3, 6, 10, 15, 20, and 30, the distilled reference model with highest validation α is chosen.For XSUM and CNN/Daily Mail, candidate epochs are 1, 3, 6, and 10 for training efficiency.Source: §A.2 · p. 14

A.3 AdaSPEC Example Tokens

AdaSPEC’s selected GSM8K training tokens are typically mathematical tokens, including digits and operators.

A.3 AdaSPEC Example Tokens

  • AdaSPEC’s selected GSM8K training tokens are typically mathematical tokens, such as digits and operators.Listing 1 provides examples of selected tokens during GSM8K training.Source: §A.3 · p. 14

A.4 Minimal Code Implementation of AdaSPEC

The appendix presents a compact AdaSPEC implementation, notes its broader impact, and points to estimated GPU-hour totals for training.

A.4 Minimal Code Implementation of AdaSPEC

  • AdaSPEC’s core implementation could be written in approximately 100 lines of code.The listing computes target and draft distributions, uses a reference-model comparison to form a mask, and sums masked losses.Source: §A.4 · pp. 14–15

A.5 Broader Impact

  • AdaSPEC is primarily intended to improve LLM generation speed, with potential to reduce serving energy use, but could also be used for models that violate regulations or ethics.The passage gives discriminatory-content generation as an example.Source: §A.5 · p. 16

A.6 Experiments compute resources

  • The appendix lists estimated GPU hours for training models on A100 GPUs in Table 10.The supplied caption identifies the table’s scope but gives no GPU-hour values.Source: §A.6 · p. 16, Table 10 · p. 16
Loading 2510.19779v1…