Conference: NeurIPS'25 Spotlight

Github: https://github.com/yuezhouhu/adaspec

My Thoughts

这篇论文的 methods 挺简明的，感觉可以有个新 idea:
将 MoSD 和 AdaSPEC 结合起来 —— 针对不同难度的 tokens distill 出不同的 draft models，利用 router 将不同难度的 tokens 选择最适合的对应 draft model 来进行 SD。

问题在于：对于论文中提到的“hard tokens”，是否能 distill 出一个合适且有用的 draft model？
论文结果显示当前方法对这些 token 效果较差（甚至比 reference model 还差）， 因此需要新的方式来改进这一部分的蒸馏与利用机制。

1. Motivation

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.

SD 的加速性能依赖于 draft 与 target 模型的对齐度（alignment）。传统 Knowledge Distillation (KD) 的目标是最小化 KL 散度，但这与 SD 的最终目标——最大化 token acceptance rate（被 target 接受的比例）——并不一致。

因此，即使 KD 能使分布接近，也不一定提升 SD 的实际加速效果。

核心问题：

  • 容量受限：小 draft 模型无法表达 target 的完整分布。
  • 目标不一致：KD 优化 KL，而 SD 追求 acceptance rate。
  • 浪费学习资源：draft 把容量浪费在那些“无论如何都难学”的 hard tokens 上。

2. Challenge

Conventional KD methods minimize the forward KL divergence between the draft and target models across all tokens:

$$ \mathcal{L}{KD} = \mathbb{E}{x \sim D},[,\mathrm{KL}(P_\text{target}(y|x)|P_\text{draft}(y|x)),] $$

然而，这种“对所有 token 进行均匀蒸馏”的目标存在两大问题：

  1. 与 SD 目标不匹配：SD 的目标是最大化 acceptance rate，而不是全分布 KL。
  2. token 学习难度差异巨大：
    • “easy tokens” 可以被小模型很好地学习；
    • “hard tokens” 无论怎么训练都难以对齐。

因此，draft 模型若在 hard token 上浪费容量，反而可能损害 easy token 的拟合质量。

3. Contribution

AdaSPEC 的核心贡献：

  • Selective Token Filtering：利用 reference model 识别并过滤掉 draft 难以学习的 tokens，仅对“易学 token”执行 KD。
  • Draft Model 资源再分配：使有限容量专注于容易拟合的区域，从而显著提升 acceptance rate。
  • 广泛验证：在算术推理、指令跟随、代码生成、摘要生成任务上均显著优于 DistillSpec，最高提升约 15%。

4. Method

AdaSPEC 的训练过程分为两个主要步骤：

Step 1: Constructing the Reference Model

  1. 初始化：以 draft 模型结构为基础，复制其参数得到 reference model ( M_{\text{ref}} )。
  2. 蒸馏训练：利用 target 模型 ( M_{\text{tgt}} ) 作为 teacher，对 reference model 进行常规 KD： $$ \mathcal{L}{ref} = \mathbb{E}\big[\mathrm{KL}(P{\text{tgt}}(y|x)|P_{\text{ref}}(y|x))\big] $$
  3. 目标：得到一个对 target 行为较好拟合、但容量仍接近 draft 的 reference 模型，用于衡量 token 难度。

Step 2: Selective Knowledge Distillation for the Draft Model

在训练 draft 模型 ( M_q ) 时，引入 selective token filtering：

  1. 定义 token-wise loss 差值：

    对于每个 token ( w )，计算 draft 与 reference 在相对于 target 的 KL 损失： $$ L_{ref}(w) = \mathrm{KL}(P_{\text{tgt}}(w|x)|P_{\text{ref}}(w|x)) $$ $$ L_{draft}(w) = \mathrm{KL}(P_{\text{tgt}}(w|x)|P_{\text{draft}}(w|x)) $$

    定义差值： $$ \Delta L(w) = L_{draft}(w) - L_{ref}(w) $$

  2. token 选择（Filtering）：

    计算每个 token 的 ΔL，并按值排序，仅保留 top-k%（例如 40%）的 token 参与蒸馏： $$ S = { w \mid \Delta L(w) \text{ 位于前 } k% } $$

    直觉：

    • 若 ΔL 较大 → draft 对该 token 仍有改进空间，值得学习；
    • 若 ΔL 较小或负 → token 太难或 draft/ref 都学不好，无需浪费容量。

  3. 最终的 selective distillation loss：

    仅对选中的 token 求平均： $$ \mathcal{L}{AdaSPEC} = \frac{1}{|S|}\sum{w \in S} L_{draft}(w) $$

5. Evaluation

5.1 实验设置

  • Datasets / Tasks：

    • Arithmetic reasoning (GSM8K)
    • Instruction-following (Alpaca)
    • Coding (MBPP)
    • Summarization (CNN/DM, XSUM)

  • Model pairs：

    • Pythia-31M → 1.4B
    • CodeGen-350M → Phi-2 (2.7B)
    • Phi-1.5 → Phi-2
    • Mixtral-8x7B → Mixtral-8x22B

  • Baselines：

    • Direct KD
    • DistillSpec
    • RKL-based distillation
    • Token filtering ablations (不同 filter 阈值)

  • Metrics：

    • Token acceptance rate (↑)
    • Task accuracy (e.g. pass@1, rouge-L)
    • Speed-up ratio (↑)
    • Draft perplexity vs. target perplexity

5.2 主要结果

  1. Acceptance Rate 显著提升

    • 平均提升 8–15%；
    • 在 31M → 1.4B 的极大规模差下效果最突出；
    • 即使在 capacity gap 达 64× 的情况下，仍维持高 acceptance rate。

  2. Generation Quality 无明显下降

    • Rouge-L、BLEU、pass@1 等指标几乎保持；
    • selective filtering 并未导致语义偏移或句法异常。

  1. Token Filtering 的比例影响

    • 论文实验显示当过滤比例在 30–50% 时性能最优；
    • 过滤过多（<20%）导致信息不足，过滤过少（>60%）则浪费容量。

  1. Wall-Time Speed-up

    • 在 vLLM 实测下可达 1.2–1.4× 加速；
    • 与 EAGLE 等 speculative sampling 策略结合可进一步提升。

5.3 消融研究（Ablation）

  • 无过滤（Full KD）：acceptance rate 最低；
  • 随机过滤：无规律性能下降；
  • AdaSPEC Filtering：在相同 loss 下可显著提升 α；
  • 不同 KD 损失（RKL / TVD）：forward KL 表现最稳定。

6. Discussion

6.1 Size Gap Between Target and Draft Models

论文指出，在传统 SD 中，draft 与 target 的规模差一般 <10×。 但 AdaSPEC 通过 selective KD，即使在 64× 参数差距下仍保持显著性能提升。

解释：

  • 当模型容量差距扩大时，直接 KD 的“representation mismatch”问题更严重；
  • AdaSPEC 通过聚焦于易学 token，减少了 mismatch 区域；
  • 因此模型差距越大，AdaSPEC 的相对收益越明显。

6.2 与 Rho-1 的关系（Lin et al., 2024）

  • Rho-1 (Not All Tokens Are What You Need) 强调在预训练阶段 优先学习 hard tokens；

  • AdaSPEC 则相反：在 SD 中 排除 hard tokens，聚焦于易学部分；

  • 区别源于目标不同：

    • 预训练希望扩展模型能力；
    • SD 蒸馏希望最大化小模型的实用性能（acceptance rate）。

6.3 局限性与未来方向

“As a preliminary study on selective training for SD, we limit our study on simple loss-related token filter.”

  • future work:

    • 与 tree-based speculative decoding (如 MoSD) 结合

7. Conclusion

We present AdaSPEC, a novel approach for training more efficient draft models for Speculative Decoding. AdaSPEC introduces selective token filtering based on reference model perplexity gaps, enabling draft models to focus limited capacity on tokens where alignment with the target model is most achievable.

  • Higher acceptance rates across diverse tasks;
  • Significant acceleration without degradation in output quality;
  • Robustness across large model gaps.

AdaSPEC 为小模型的高效推理（尤其是在 edge / on-device 场景）提供了新的蒸馏视角。 它启发我们：在推理加速任务中，学习什么比学多少更重要。