AdaSPEC: Селективное дистиллирование знаний для эффективных спекулятивных декодеров

AdaSPEC: Selective Knowledge Distillation for Efficient Speculative Decoders

October 22, 2025

Авторы: Yuezhou Hu, Jiaxin Guo, Xinyu Feng, Tuo Zhao

cs.AI

Аннотация

Спекулятивное декодирование (SD) ускоряет вывод больших языковых моделей за счет использования небольшой черновой модели для генерации предсказаний, которые затем проверяются более крупной целевой моделью. Эффективность SD зависит от согласованности этих моделей, что обычно улучшается с помощью дистилляции знаний (KD). Однако традиционные методы KD направлены на минимизацию расхождения Кульбака-Лейблера между черновой и целевой моделями для всех токенов, что не соответствует истинной цели SD — максимизации коэффициента принятия токенов. Следовательно, черновые модели часто не могут полностью усвоить знания целевой модели из-за ограничений по емкости, что приводит к неоптимальной производительности. Для решения этой проблемы мы предлагаем AdaSPEC — новый метод, который включает селективную фильтрацию токенов в процесс KD. AdaSPEC использует референсную модель для идентификации и отсева токенов, которые сложно аппроксимировать, что позволяет дистиллировать черновую модель, лучше согласованную с целевой моделью на более простых токенах. Этот подход повышает общий коэффициент принятия токенов без ущерба для качества генерации. Мы оцениваем AdaSPEC на различных задачах, включая арифметические рассуждения, следование инструкциям, программирование и суммаризацию, используя конфигурации моделей с 31M/1.4B и 350M/2.7B параметров. Наши результаты показывают, что AdaSPEC стабильно превосходит современный метод DistillSpec, достигая более высоких коэффициентов принятия на всех задачах (до 15%). Код доступен по адресу https://github.com/yuezhouhu/adaspec.

English

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.