论文
直接优化期望解码轮数的并行投机草稿模型训练
Training Parallel Speculative Draft Models by Directly Minimizing Expected Decoding Rounds
摘要
推测性解码通过使用低成本草稿模型来提出全尺寸目标模型并行验证的标记,从而加速大型语言模型推理。并行和半自回归(半 AR)绘图器通过在单次前向传递中提出整个块来提高绘图效率,但训练它们会带来新的困难:给定位置的绘图分布取决于解码轮的开始位置,而轮的开始位置取决于先前轮接受的Token数量。现有的训练目标通常依赖于忽略这种跨轮耦合的块本地代理,因此不会直接优化全局解码效率。在这项工作中,我们通过将推测性解码表示为马尔可夫奖励过程,开发了一个用于培训和评估这些起草者的理论框架。该公式产生预期解码轮数 (EDR) 目标,该目标按状态占用率对本地拒绝成本进行加权,并且恰好等于预期解码轮数。与之前的替代目标不同,EDR 不引入任何辅助目标 超参数。然后,我们导出精确的时间差异梯度,该梯度支持从目标模型推出中进行无偏随机优化。相同的框架还生成了精确的离线轮数评估器,从而可以在共享目标部署上进行配对起草者比较,而无需运行推测性解码。通过 EDR 对两个最先进的绘图器 DSpark 和 DFly 进行微调,不断提高平均接受长度,并在涵盖数学推理、代码生成和聊天的九个基准上超越现有的训练目标。