让小模型只学大模型最易模仿的词,提升推理速度
AdaSPEC: Selective Knowledge Distillation for Efficient Speculative Decoders

- 只精选容易模仿的词进行知识迁移,避免无效学习
- 在多种任务上接受率最高提升15%,生成质量不下降
- 适合需要快速推理的场景,如实时对话系统
推测解码(SD)通过小型草稿模型预测文本,并由大型目标模型验证来加速大语言模型推理。其效果依赖于两模型间的对齐,通常通过知识蒸馏(KD)增强。但传统KD方法要求草稿模型在所有词上最小化与目标模型的KL散度,这与SD真正目标——最大化词接受率——不符。因此,受限于容量,草稿模型难以完全吸收目标模型知识,导致性能不佳。为此,我们提出AdaSPEC,将选择性词过滤引入蒸馏过程:利用参考模型识别难拟合词并过滤,使草稿模型专注于学习简单词,从而在保持生成质量前提下显著提升整体词接受率。我们在31M/1.4B和350M/2.7B等多组模型配置下评估,涵盖算术推理、指令遵循、编码和摘要任务,结果表明AdaSPEC持续优于当前最优的DistillSpec方法,在所有任务上接受率最高提升15%。代码已开源。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。