用于推测性解码的接受感知草稿模型训练

推测性解码通过使用轻量级草稿模型生成多个候选标记，并在单次前向传递中由目标模型验证，从而加速大型语言模型 (LLM) 推理。它的加速很大程度上取决于接受长度，但现有的草案模型训练方法主要以交叉熵或 Kullback-Leibler (KL) 散度作为代理进行优化。这些目标鼓励分布匹配，但不直接优化接受长度，并且贪婪解码和基于采样的解码之间的接受机制也有所不同。在这项工作中，我们提出了接受长度感知的训练损失，可以直接优化推测窗口内接受词元的预期数量。对于贪婪验证，我们得出预期接受长度（EAL）损失，该损失显式最大化预期接受长度。对于基于采样的解码，我们引入了窗口总变差（WTV）损失，它优化了温度缩放草图和目标分布之间的重叠，同时考虑了顺序接受依赖性。这两个目标都可以进一步与使用模拟接受长度作为奖励的群体相关强化学习阶段（GRPO）相结合。跨不同目标和草案模型、任务和解码设置的实验表明，与基于 KL 的训练相比，我们的损失持续提高了接受长度。 WTV 在基于采样的解码下提供了特别强大的增益，而 EAL 更好地匹配贪婪验证。这些结果表明，直接优化接受目标，并根据解码模式定制损失，比传统的分布匹配目标更有效。