Kangaroo: Lossless Self-Speculative Decoding via Double Early Exiting
作者: Fangcheng Liu, Yehui Tang, Zhenhua Liu, Yunsheng Ni, Kai Han, Yunhe Wang
分类: cs.CL, cs.LG
发布日期: 2024-04-29
🔗 代码/项目: GITHUB
💡 一句话要点
提出Kangaroo框架以实现无损自我推测解码
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 自我推测解码 早期退出机制 推理加速 深度学习 语言模型 模型压缩
📋 核心要点
- 现有的推测解码方法需要训练单独的草稿模型，导致推理成本高且效率低下。
- Kangaroo框架通过固定的浅层子网络作为自我草稿模型，并引入早期退出机制来提高token接受率。
- 在Spec-Bench上的实验表明，Kangaroo实现了1.68倍的加速，并显著减少了参数数量。
📝 摘要（中文）
自我推测解码在加速大型语言模型推理方面表现出色，但传统方法需要训练单独的草稿模型，成本高昂。本文提出Kangaroo框架，利用固定的浅层子网络作为自我草稿模型，并在其上训练轻量适配器模块，以提高模型的表示能力。为了解决推测过程中小模型的推断延迟问题，本文引入了额外的早期退出机制，在当前token的置信度低于阈值时停止后续预测。实验结果表明，Kangaroo在Spec-Bench上实现了最高1.68倍的加速，相比Medusa-1减少了88.7%的额外参数。
🔬 方法详解
问题定义：本文旨在解决传统推测解码方法中草稿模型训练成本高、推理延迟不可忽视的问题。现有方法依赖于单独的草稿模型，导致效率低下。
核心思路：Kangaroo框架采用固定的浅层子网络作为自我草稿模型，并在其上训练轻量的适配器模块，以提升模型的表示能力，同时引入早期退出机制以减少不必要的推理步骤。
技术框架：Kangaroo的整体架构包括一个固定的浅层子网络作为草稿模型、一个适配器模块用于增强表示能力，以及一个早期退出机制用于控制推理过程。推理时，当当前token的置信度低于设定阈值时，立即停止后续的预测。
关键创新：Kangaroo的主要创新在于将自我推测解码与早期退出机制结合，显著提高了token接受率，同时降低了推理延迟。这一设计与传统方法的根本区别在于不再依赖于单独的草稿模型。
关键设计：在设计中，适配器模块的参数设置经过精心调整，以确保其在提升表示能力的同时不增加过多的计算负担。早期退出机制的阈值选择也经过实验验证，以达到最佳的推理效率。
🖼️ 关键图片
📊 实验亮点
Kangaroo在Spec-Bench上实现了最高1.68倍的推理加速，相比于基线Medusa-1，减少了88.7%的额外参数（67M对比591M），显示出其在效率和性能上的显著提升。
🎯 应用场景
Kangaroo框架在大型语言模型的推理加速中具有广泛的应用潜力，特别是在需要实时响应的自然语言处理任务中。其无损的自我推测解码方法能够有效提高模型的推理效率，降低计算资源消耗，适用于智能助手、对话系统等场景。未来，该方法还可能扩展到其他深度学习模型的推理优化中。
📄 摘要（原文）
Speculative decoding has demonstrated its effectiveness in accelerating the inference of large language models while maintaining a consistent sampling distribution. However, the conventional approach of training a separate draft model to achieve a satisfactory token acceptance rate can be costly. Drawing inspiration from early exiting, we propose a novel self-speculative decoding framework \emph{Kangaroo}, which uses a fixed shallow sub-network as a self-draft model, with the remaining layers serving as the larger target model. We train a lightweight and efficient adapter module on top of the sub-network to bridge the gap between the sub-network and the full model's representation ability. It is noteworthy that the inference latency of the self-draft model may no longer be negligible compared to the large model, necessitating strategies to increase the token acceptance rate while minimizing the drafting steps of the small model. To address this challenge, we introduce an additional early exiting mechanism for generating draft tokens. Specifically, we halt the small model's subsequent prediction during the drafting phase once the confidence level for the current token falls below a certain threshold. Extensive experiments on the Spec-Bench demonstrate the effectiveness of Kangaroo. Under single-sequence verification, Kangaroo achieves speedups up to $1.68\times$ on Spec-Bench, outperforming Medusa-1 with 88.7\% fewer additional parameters (67M compared to 591M). The code for Kangaroo is available at https://github.com/Equationliu/Kangaroo.