AdaSPEC: Selective Knowledge Distillation for Efficient Speculative Decoders

TL;DR

AdaSPEC addresses the inefficiency of speculative decoding caused by capacity gaps between draft and target models by introducing selective token filtering during distillation. It builds a reference model through DistillSpec and uses token-wise loss gaps to filter hard tokens, enabling the draft to learn only the easily learnable tokens and better align with the target model. Across arithmetic reasoning, instruction-following, coding, and summarization tasks, AdaSPEC consistently increases token acceptance rates and reduces divergence between draft and target predictions, achieving up to 15 percentage points gains and demonstrable speed-ups. The approach is robust across model scales and can integrate with advanced SD methods, offering practical improvements for efficient LLM inference.

Abstract

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.

Figures (4)

• Figure 1: Overview of AdaSPEC distillation process: AdaSPEC selects the most training-effective tokens and distills on these tokens.
• 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.
• 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.
• Figure 4: Impact of token selection ratio $k$ on acceptance rate for GSM8K: Results show a general trend where lower $k$ values (0.2-0.4) yield higher acceptance rates.