Yuezhou Hu Jiaxin Guo Xinyu Feng Tuo Zhao

摘要

推测解码（Speculative Decoding, SD）通过使用一个小型的草稿模型生成预测结果，并由更大的目标模型进行验证，从而加速大语言模型的推理过程。SD 的有效性依赖于草稿模型与目标模型之间的对齐程度，这一对齐通常通过知识蒸馏（Knowledge Distillation, KD）来增强。然而，传统的 KD 方法旨在最小化草稿模型与目标模型在所有 token 上的 KL 散度，这一目标与 SD 的真实目标——最大化 token 接受率——并不一致。因此，由于模型容量的限制，草稿模型往往难以充分吸收目标模型的知识，导致性能不佳。为解决这一挑战，我们提出 AdaSPEC，一种将选择性 token 过滤机制引入知识蒸馏过程的新方法。AdaSPEC 利用一个参考模型识别并过滤掉难以拟合的 token，从而实现对草稿模型的知识蒸馏，使其在简单 token 上与目标模型具有更好的对齐性。该方法在不牺牲生成质量的前提下，显著提升了整体 token 接受率。我们在多种任务上评估了 AdaSPEC，包括算术推理、指令遵循、代码生成和摘要生成，采用 31M/1.4B 和 350M/2.7B 参数规模的模型配置。实验结果表明，AdaSPEC 在所有任务中均持续优于当前最先进的 DistillSpec 方法，接受率提升最高达 15%。代码已公开，可访问 https://github.com/yuezhouhu/adaspec。