AdaSPEC: Selective Knowledge Distillation for Efficient Speculative Decoders (2510.19779v1)

Abstract: Speculative Decoding (SD) accelerates LLM 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.

• The paper introduces a novel selective distillation framework that filters out hard-to-learn tokens, thereby improving the acceptance rate of speculative decoding.
• It demonstrates significant performance gains, including up to 15% improvement in acceptance rate and enhanced prediction confidence compared to existing methods.
• AdaSPEC efficiently bridges the capacity gap between draft and target models, enabling faster inference with modest resource requirements.

AdaSPEC: Selective Knowledge Distillation for Efficient Speculative Decoders

Motivation and Background

Speculative Decoding (SD) has become a central technique for accelerating inference in LLMs by leveraging a small draft model to propose candidate tokens, which are then verified by a larger, more accurate target model. The efficiency of SD is fundamentally determined by the alignment between the draft and target models, typically measured by the token acceptance rate—the proportion of draft-generated tokens accepted by the target. Existing approaches, such as DistillSpec, employ knowledge distillation (KD) to minimize the KL divergence between the draft and target models across all tokens. However, this uniform distillation objective is misaligned with the true goal of SD: maximizing the acceptance rate. In practice, the limited capacity of the draft model means that attempting to match the target model on all tokens, including those that are inherently difficult to learn, leads to suboptimal alignment and wasted model capacity.

AdaSPEC addresses this inefficiency by introducing a selective distillation framework that filters out hard-to-learn tokens, allowing the draft model to focus its capacity on tokens where alignment is most feasible. This approach is motivated by the observation that not all tokens contribute equally to the acceptance rate, and that prioritizing learnable tokens can yield higher acceptance rates and more efficient SD.

Methodology

AdaSPEC operates in two main phases: reference model distillation and selective draft model distillation.

1. Reference Model Distillation and Token Filtering: A reference model, initialized as a copy of the draft model, is distilled from the target model using standard KD (forward KL divergence). The reference model serves as a token filter by identifying tokens that are difficult for the draft model to learn, based on the difference in token-wise KL divergence losses between the reference and draft models.
2. Selective Draft Model Distillation:

The draft model is then distilled from the target model, but only on a filtered subset of tokens—those deemed most learnable according to the reference model. Specifically, for each token $w$, the difference $$\Delta_{\mathcal{L}}(w) = \mathcal{L}_{\mathrm{draft}}(w) - \mathcal{L}_{\mathrm{ref}}(w)$$ is computed, and the top $k$ fraction of tokens with the largest $\Delta_{\mathcal{L}}(w)$ are selected for distillation. The draft model is trained to minimize the KL divergence with the target model only on this subset.

Figure 1: Overview of AdaSPEC distillation process: AdaSPEC selects the most training-effective tokens and distills on these tokens.

This selective approach enables the draft model to allocate its limited capacity more effectively, resulting in improved alignment with the target model on the most impactful tokens for SD.

Experimental Evaluation

AdaSPEC is evaluated on a diverse set of tasks—arithmetic reasoning (GSM8K), instruction following (Alpaca), code generation (MBPP), and summarization (CNN/Daily Mail, XSUM)—using both same-family (Pythia-31M/1.4B) and cross-family (CodeGen-350M/Phi-2) model pairs. Two training regimes are considered: a resource-constrained 3-epoch setting and an optimal-epoch setting where the number of epochs is tuned for maximal performance.

Across all tasks and configurations, AdaSPEC consistently outperforms DistillSpec in acceptance rate, with improvements up to 15%. The gains are particularly pronounced as the size gap between the draft and target models increases, highlighting AdaSPEC's effectiveness in scenarios with severe capacity constraints.

Figure 2: Comparative analysis of AdaSPEC and DistillSpec performance across multiple metrics on GSM8K (a, c, e) and CNN/Daily Mail (b, d, f) datasets: (a-b) Task-level acceptance rate distributions showing AdaSPEC's superior performance across tasks. (c-d) Logit margin distributions demonstrating AdaSPEC's improved prediction confidence with higher positive margins and lower negative margins. (e-f) Token-level KL divergence distributions indicating better draft-target model alignment for AdaSPEC with consistently lower divergence values. The results demonstrate AdaSPEC's more effective knowledge transfer and improved draft-target model alignment compared to DistillSpec across different evaluation metrics.

Detailed analysis reveals that AdaSPEC achieves:

• Higher task-level acceptance rates, with distributions shifted rightward relative to DistillSpec.
• Improved logit margin distributions, indicating more confident and accurate draft predictions.
• Lower token-level KL divergence, reflecting tighter alignment between draft and target models.

Case studies further show that AdaSPEC's prediction errors are nearly a subset of those made by DistillSpec, demonstrating that selective distillation reduces inference discrepancies.

Figure 3: Comparison of prediction errors between AdaSPEC and DistillSpec on GSM8K and CNN/Daily Mail Datasets: Tokens highlighted in blue represent errors made by both methods, while tokens highlighted in red indicate errors unique to the corresponding method. As can be seen, AdaSPEC’s errors form nearly a subset of DistillSpec's errors, demonstrating the effectiveness of AdaSPEC's selective training approach in reducing inference discrepancies.

Ablation and Analysis

Ablation studies confirm the importance of the token selection mechanism. Training on the top 40% of tokens (by KL-divergence margin) yields significantly higher acceptance rates than training on the bottom 40%, with the latter even underperforming the reference model. The benefits of AdaSPEC's token selection generalize beyond KD to direct fine-tuning, and forward KL divergence is found to be the most effective distillation objective compared to alternatives such as reverse KL and total variation distance.

The acceptance rate is sensitive to the token selection ratio $k$, with lower values (e.g., $k=0.2$–$0.4$) generally yielding better results. AdaSPEC also demonstrates robust wall-clock speedups (10–20%) over DistillSpec in real-world inference settings and integrates effectively with advanced SD frameworks such as EAGLE, further improving both accuracy and decoding efficiency.

Implementation Considerations

AdaSPEC is straightforward to implement, requiring only minor modifications to the standard KD pipeline. The core logic involves computing token-wise KL divergences for both the reference and draft models, selecting the top $k$ fraction of tokens by $\Delta_{\mathcal{L}}(w)$, and restricting the distillation loss to this subset. The method is compatible with both same-family and cross-family model pairs, and scales to large models (e.g., Qwen2.5-0.5B/32B) without loss of effectiveness.

Resource requirements are modest, with training times comparable to standard KD approaches. The method is orthogonal to other SD optimizations and can be combined with tree-based or multi-step verification frameworks for further gains.

Implications and Future Directions

AdaSPEC demonstrates that selective knowledge distillation, tailored to the acceptance rate objective of SD, is a principled and effective strategy for bridging the capacity gap between draft and target models. By focusing on learnable tokens, AdaSPEC enables the use of significantly smaller draft models without sacrificing alignment or generation quality, thus unlocking greater inference speedups.

Theoretically, AdaSPEC highlights the importance of task-aligned distillation objectives in multi-model inference pipelines. Practically, it provides a scalable, model-agnostic approach for efficient LLM deployment, with direct implications for reducing inference latency and computational cost.

Future work may explore more adaptive or dynamic token filtering strategies, integration with advanced SD frameworks, and extensions to multi-modal or multi-lingual settings. Additionally, the interplay between selective distillation and other model compression techniques (e.g., quantization, pruning) warrants further investigation.

Conclusion

AdaSPEC introduces a selective distillation paradigm for speculative decoding, leveraging token-level difficulty estimation to maximize draft-target model alignment within the constraints of limited model capacity. Empirical results across diverse tasks and model configurations demonstrate consistent improvements in acceptance rate, prediction confidence, and inference speed. AdaSPEC is a practical, extensible solution for efficient LLM inference, with broad applicability in real-world deployment scenarios.