KOALA: Enhancing Speculative Decoding for LLM via Multi-Layer Draft Heads with Adversarial Learning

📄 arXiv: 2408.08146v1 📥 PDF

作者: Kaiqi Zhang, Jing Zhao, Rui Chen

分类: cs.CL

发布日期: 2024-08-15

💡 一句话要点

KOALA：通过对抗学习的多层Draft Head增强LLM的推测解码

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 推测解码 Draft Head 对抗学习 多层架构 模型加速 推理优化

📋 核心要点

  1. LLM自回归解码导致高推理延迟，推测解码通过draft head缓解，但潜力未充分挖掘。
  2. KOALA将单层draft head扩展为多层架构，并引入对抗学习，提高draft head预测token的准确性。
  3. 实验表明，KOALA在多种任务中显著提升了推测解码的速度，加速比提高了0.24x-0.41x。

📝 摘要（中文）

大型语言模型（LLM）由于其自回归解码的特性而表现出较高的推理延迟。推测解码中的draft head虽然缓解了这个问题，但其全部潜力尚未被充分挖掘。本文介绍了一种正交于draft head的方法，名为KOALA（K层优化对抗学习架构）。通过将传统的单层draft head转换为多层架构，并将对抗学习融入到传统的监督训练中，KOALA显著提高了draft head预测后续token的准确性，从而更紧密地反映了LLM的功能。虽然这种改进以略微增加drafting开销为代价，但KOALA极大地释放了draft head的潜力，从而大大增强了推测解码。我们对KOALA进行了全面的评估，包括跨各种任务的自回归和非自回归draft head，结果表明延迟加速比提高了0.24x-0.41x，比原始draft head快10.57%-14.09%。

🔬 方法详解

问题定义：大型语言模型推理速度慢，推测解码旨在加速推理过程。现有推测解码方法中，draft head的预测精度不足，成为性能瓶颈。如何提高draft head的预测精度，使其更接近LLM的预测能力，是本文要解决的核心问题。

核心思路：KOALA的核心思路是通过构建多层draft head并引入对抗学习，来提升draft head的预测能力。多层结构允许draft head学习更复杂的特征表示，而对抗学习则促使draft head的输出分布更接近目标LLM的输出分布。

技术框架：KOALA的技术框架主要包含两个部分：多层Draft Head和对抗学习模块。首先，将传统的单层draft head替换为多层Transformer结构，以增强其建模能力。然后，在训练过程中，引入一个判别器，用于区分draft head生成的token和目标LLM生成的token。通过对抗训练，draft head不断优化自身，以生成更逼真的token序列，从而提高预测精度。

关键创新：KOALA的关键创新在于将多层架构和对抗学习相结合，应用于推测解码的draft head。与传统的单层draft head相比，多层结构具有更强的表示能力，能够学习更复杂的语言模式。对抗学习则提供了一种有效的训练方法，可以促使draft head的输出分布更接近目标LLM的输出分布。

关键设计：KOALA的关键设计包括：1）多层Draft Head的具体层数和结构；2）对抗学习中判别器的设计，包括网络结构和损失函数；3）对抗训练的策略，例如如何平衡生成器（draft head）和判别器的训练。论文中可能还涉及一些超参数的调整，例如学习率、batch size等。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，KOALA在各种任务上都显著提升了推测解码的性能。具体而言，KOALA实现了0.24x-0.41x的延迟加速比提升，比原始draft head快10.57%-14.09%。这些结果验证了KOALA的有效性，并表明其具有很大的应用潜力。

🎯 应用场景

KOALA可应用于各种需要加速LLM推理的场景，例如在线对话系统、文本生成、机器翻译等。通过提高LLM的推理速度，KOALA可以降低计算成本，提升用户体验，并促进LLM在资源受限设备上的部署。未来，KOALA的思路可以推广到其他加速LLM推理的技术中。

📄 摘要（原文）

Large Language Models (LLMs) exhibit high inference latency due to their autoregressive decoding nature. While the draft head in speculative decoding mitigates this issue, its full potential remains unexplored. In this paper, we introduce KOALA (K-layer Optimized Adversarial Learning Architecture), an orthogonal approach to the draft head. By transforming the conventional single-layer draft head into a multi-layer architecture and incorporating adversarial learning into the traditional supervised training, KOALA significantly improves the accuracy of the draft head in predicting subsequent tokens, thus more closely mirroring the functionality of LLMs. Although this improvement comes at the cost of slightly increased drafting overhead, KOALA substantially unlocks the draft head's potential, greatly enhancing speculative decoding. We conducted comprehensive evaluations of KOALA, including both autoregressive and non-autoregressive draft heads across various tasks, demonstrating a latency speedup ratio improvement of 0.24x-0.41x, which is 10.57%-14.09% faster than the original draft heads.