发表机构
New York University; University of Waterloo（纽约大学; 滑铁卢大学）
机构由 AI 辅助整理，请以论文原文为准。AI 中文总结
针对投机解码，提出接受长度感知的EAL和WTV损失直接优化接受长度，结合GRPO，实验证明优于KL训练，显著提升不同解码模式下的接受长度。
AI 中文摘要
投机解码通过使用轻量级草稿模型生成多个候选令牌，并由目标模型在单次前向传播中验证这些令牌，从而加速大型语言模型（LLM）的推理。其加速效果在很大程度上取决于接受长度，然而现有的草稿模型训练方法主要将交叉熵或Kullback-Leibler（KL）散度作为代理目标进行优化。这些目标鼓励分布匹配，但并未直接优化接受长度，而且贪婪解码与基于采样的解码之间的接受机制也有所不同。在这项工作中，我们提出了接受长度感知的训练损失，直接优化投机窗口内预期接受令牌的数量。对于贪婪验证，我们推导出预期接受长度（EAL）损失，该损失显式地最大化预期接受长度。对于基于采样的解码，我们引入了窗口全变差（WTV）损失，该损失在考虑顺序接受依赖性的同时，优化温度缩放后的草稿分布与目标分布之间的重叠。这两种目标还可以进一步与使用模拟接受长度作为奖励的组相对强化学习阶段（GRPO）相结合。在不同目标模型和草稿模型、任务以及解码设置上的实验表明，与基于KL的训练相比，我们的损失持续改善了接受长度。WTV在基于采样的解码下提供了特别显著的增益，而EAL则更好地匹配贪婪验证。这些结果表明，直接优化接受目标，并采用针对解码模式定制的损失，比传统的分布匹配目标更为有效。
英文摘要
Speculative decoding accelerates large language model (LLM) inference by using a lightweight draft model to generate multiple candidate tokens that are verified by the target model in a single forward pass. Its speedup is largely determined by the acceptance length, yet existing draft-model training methods mainly optimize cross-entropy or Kullback-Leibler (KL) divergence as proxies. These objectives encourage distribution matching but do not directly optimize acceptance length, and the acceptance mechanism also differs between greedy and sampling-based decoding. In this work, we propose acceptance-length-aware training losses that directly optimize the expected number of accepted tokens within a speculative window. For greedy verification, we derive an expected accepted length (EAL) loss that explicitly maximizes expected acceptance length. For sampling-based decoding, we introduce a window total variation (WTV) loss that optimizes the overlap between temperature-scaled draft and target distributions while accounting for sequential acceptance dependencies. Both objectives can be further combined with a group-relative reinforcement learning stage (GRPO) using simulated acceptance length as the reward. Experiments across different target and draft models, tasks, and decoding settings show that our losses consistently improve acceptance length over KL-based training. WTV provides particularly strong gains under sampling-based decoding, while EAL better matches greedy verification. These results show that directly optimizing the acceptance objective, with losses tailored to the decoding mode, is more effective than conventional distribution-matching objectives.
