Kangaroo: Lossless Self-Speculative Decoding via Double Early Exiting

2024年04月29日
  • 简介
    推测解码已经证明在加速大型语言模型的推理过程中具有有效性,同时保持了一致的采样分布。然而,训练单独的草稿模型以达到令人满意的标记接受率的传统方法可能代价高昂。受早期退出的启发,我们提出了一种新颖的自我推测解码框架Kangaroo,它使用一个固定的浅层子网络作为自草稿模型,其余层作为较大的目标模型。我们在子网络的顶部训练了一个轻量级和高效的适配器模块,以弥合子网络和完整模型的表示能力之间的差距。值得注意的是,在起草模型的推理延迟可能不再可以忽略不计,与大模型相比,需要最小化小模型的起草步骤,同时增加标记接受率的策略。为了解决这个挑战,我们引入了一个额外的早期退出机制来生成草稿标记。具体来说,在草稿阶段,一旦当前标记的置信水平低于某个阈值,我们就停止小模型的后续预测。在Spec-Bench上的广泛实验证明了Kangaroo的有效性。在单序列验证下,Kangaroo在Spec-Bench上实现了高达1.68倍的加速,比Medusa-1表现更好,额外参数少了88.7%(67M与591M相比)。Kangaroo的代码可在https://github.com/Equationliu/Kangaroo上获得。
  • 图表
  • 解决问题
    论文旨在解决在大型语言模型推理时加速推理速度的问题,同时保持一致的采样分布。同时,论文还试图通过在浅层子网络上训练轻量级适配器模块来解决训练单独的草稿模型成本高的问题。
  • 关键思路
    论文提出了一种自我推测解码框架Kangaroo,它使用一个固定的浅层子网络作为自我草稿模型,其余层作为更大的目标模型。论文还引入了一个早期退出机制来生成草稿令牌,以尽可能减少草稿模型的草稿步骤并提高令牌接受率。
  • 其它亮点
    论文在Spec-Bench上进行了广泛的实验,证明了Kangaroo的有效性。在单序列验证下,Kangaroo在Spec-Bench上实现了高达1.68倍的加速,比Medusa-1少了88.7%的额外参数(67M与591M相比)。Kangaroo的代码可在https://github.com/Equationliu/Kangaroo上获得。
  • 相关研究
    在这个领域中,最近的相关研究包括:1)Medusa-1:一种基于早期退出的推理加速方法,2)Spec-Bench:一个用于评估大型语言模型的基准测试套件。
PDF
原文
点赞 收藏 评论 分享到Link

沙发等你来抢

去评论