通过直接最小化期望解码轮数训练并行推测草稿模型

推测解码使用低成本的草稿模型提出候选词元，再由完整规模的目标模型并行验证，从而加速大语言模型推理。并行和半自回归（semi-AR）草稿模型通过单次前向传播提出一整个词元块，提高了草稿生成效率，但其训练也面临新的困难：给定位置的草稿分布取决于解码轮次的起始位置，而各轮的起始位置又取决于此前轮次接受了多少词元。现有训练目标通常依赖块内局部的替代目标，忽略了这种跨轮次耦合，因此无法直接优化全局解码效率。在本研究中，我们将推测解码表示为马尔可夫奖励过程，从而构建了用于训练和评估此类草稿模型的理论框架。基于这一表述，我们提出了期望解码轮数（EDR）目标，该目标按状态占用概率对局部拒绝代价进行加权，其值恰好等于解码轮数的期望值。与以往的替代目标不同，EDR 不引入任何辅助超参数。随后，我们推导出精确的时序差分梯度，支持利用目标模型生成的轨迹进行无偏随机优化。同一框架还提供了精确的解码轮数离线评估方法，无需运行推测解码，即可在共享的目标模型轨迹上对草稿模型进行配对比较。使用 EDR 对 DSpark 和 DFly 这两个当前最先进的草稿模型进行微调后，平均接受长度得到了一致提升，并且在涵盖数学推理、代码生成和聊天的九项基准测试中，表现优于现有训练目标。