Amphista: Bi-directional Multi-head Decoding for Accelerating LLM Inference

📄 arXiv: 2406.13170v2 📥 PDF

作者: Zeping Li, Xinlong Yang, Ziheng Gao, Ji Liu, Guanchen Li, Zhuang Liu, Dong Li, Jinzhang Peng, Lu Tian, Emad Barsoum

分类: cs.AI, cs.CL

发布日期: 2024-06-19 (更新: 2024-10-18)

💡 一句话要点

Amphista：一种双向多头解码方法，加速LLM推理。

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

关键词: 大语言模型 LLM推理加速 推测解码 双向注意力 并行解码 多头解码 自回归解码 非自回归解码

📋 核心要点

  1. 自回归解码是LLM推理速度的瓶颈，缺乏并行性导致效率低下。
  2. Amphista通过引入双向注意力的自嵌入块，实现不同预测位置的信息交互，提升并行解码能力。
  3. 实验表明，Amphista在Vicuna模型上实现了显著的加速，同时保持了生成质量。

📝 摘要（中文）

大型语言模型（LLM）本质上使用自回归解码，这导致推理过程中缺乏并行性，显著降低了推理速度。虽然像Medusa这样的方法构建了并行化的头部，但它们缺乏不同预测位置之间充分的信息交互。为了克服这个限制，我们引入了Amphista，这是一个增强的推测解码框架，建立在Medusa之上。具体来说，Amphista建模了一个能够进行并行推理的自嵌入块，该块结合了双向注意力，以实现不同草稿头部之间的交互。此外，Amphista集成了分阶段适应层，确保了语义信息从目标模型的自回归推理到草稿头部的非自回归推理的无缝过渡，有效地实现了范式转变和特征融合。在Vicuna模型上使用MT-Bench和Spec-Bench的实验结果表明，Amphista在保持生成质量的同时实现了显著的加速。在MT-Bench上，Amphista在实际运行时间上，相对于原始自回归解码，在Vicuna 33B上实现了高达2.75倍的加速，相对于Medusa实现了1.40倍的加速。

🔬 方法详解

问题定义：大型语言模型依赖自回归解码，推理速度慢。Medusa等方法尝试并行化，但缺乏不同预测位置间的有效信息交互，导致性能提升受限。因此，需要一种既能并行解码，又能充分利用上下文信息的方案。

核心思路：Amphista的核心在于构建一个能够并行推理的自嵌入块，并引入双向注意力机制，使得不同的草稿头部之间能够进行信息交互。通过这种方式，每个头部在生成token时，可以考虑到其他头部的信息，从而提高预测的准确性。

技术框架：Amphista框架主要包含两个关键模块：Auto-embedding Block和Staged Adaptation Layers。Auto-embedding Block负责并行生成多个token草稿，并利用双向注意力进行信息融合。Staged Adaptation Layers则负责将目标模型的自回归特征平滑过渡到草稿头部的非自回归特征空间，实现特征融合。整体流程是先通过目标模型进行初步的自回归推理，然后利用Amphista进行并行草稿生成和验证。

关键创新：Amphista的关键创新在于双向多头解码机制。与Medusa等单向并行解码方法不同，Amphista允许不同草稿头部之间进行信息交互，从而提高了预测的准确性和一致性。此外，Staged Adaptation Layers的设计也保证了特征空间的平滑过渡，避免了因范式转变带来的性能损失。

关键设计：Auto-embedding Block采用了多头注意力机制，每个头部负责预测一个token。双向注意力机制允许每个头部访问所有其他头部的信息。Staged Adaptation Layers通过一系列的线性变换和非线性激活函数，将目标模型的特征逐步映射到草稿头部的特征空间。具体的参数设置和网络结构需要根据目标模型的特点进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Amphista在MT-Bench和Spec-Bench上进行了实验验证。在Vicuna 33B模型上，Amphista相对于原始自回归解码实现了高达2.75倍的加速，相对于Medusa实现了1.40倍的加速。实验结果表明，Amphista在显著提高推理速度的同时，保持了良好的生成质量。

🎯 应用场景

Amphista具有广泛的应用前景，可以应用于各种需要加速LLM推理的场景，例如在线对话系统、文本生成、机器翻译等。通过提高推理速度，可以降低计算成本，提升用户体验，并促进LLM在资源受限设备上的部署。未来，Amphista可以进一步扩展到支持更复杂的模型结构和任务类型。

📄 摘要（原文）

Large Language Models (LLMs) inherently use autoregressive decoding, which lacks parallelism in inference and results in significantly slow inference speed. While methods such as Medusa constructs parallelized heads, they lack adequate information interaction across different prediction positions. To overcome this limitation, we introduce Amphista, an enhanced speculative decoding framework that builds upon Medusa. Specifically, Amphista models an Auto-embedding Block capable of parallel inference, incorporating bi-directional attention to enable interaction between different drafting heads. Additionally, Amphista integrates Staged Adaptation Layers, which ensure a seamless transition of semantic information from the target model's autoregressive inference to the drafting heads' non-autoregressive inference, effectively achieving paradigm shift and feature fusion. Experimental results on Vicuna models using MT-Bench and Spec-Bench demonstrate that Amphista achieves substantial acceleration while maintaining generation quality. On MT-Bench, Amphista delivers up to 2.75$\times$ speedup over vanilla autoregressive decoding and 1.40$\times$ over Medusa on Vicuna 33B in wall-clock time.