AdaSPEC: Selective Knowledge Distillation for Efficient Speculative Decoders
Yuezhou Hu, Jiaxin Guo, Xinyu Feng, Tuo Zhao
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, typically enhanced through knowledge distillation (KD). However, draft models often struggle to fully assimilate the target model's knowledge due to capacity constraints. 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 overall token acceptance rates 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.