Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads

📄 arXiv: 2401.10774v3 📥 PDF

作者: Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen, Tri Dao

分类: cs.LG, cs.CL

发布日期: 2024-01-19 (更新: 2024-06-14)

备注: The code for this implementation is available at https://github.com/FasterDecoding/Medusa

💡 一句话要点

提出Medusa框架以加速大型语言模型推理

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 推理加速 自回归解码 树形注意力 并行处理 微调策略 生成质量 模型优化

📋 核心要点

  1. 现有大型语言模型在推理时面临计算瓶颈，尤其是在自回归解码过程中，依赖于前一步输出导致效率低下。
  2. Medusa通过增加多个解码头并采用树形注意力机制，使得多个标记的预测可以并行进行，从而加速推理过程。
  3. 实验结果显示，Medusa-1实现了2.2倍的推理加速，而Medusa-2进一步提升至2.3-3.6倍，且未降低生成质量。

📝 摘要（中文）

大型语言模型（LLMs）采用自回归解码，导致每一步都依赖于前一步的输出，从而形成计算瓶颈。尽管已有一些方法如推测解码被提出，但由于维护单独草稿模型的挑战，其实施受到限制。本文提出Medusa，通过增加额外的解码头并利用树形注意力机制，使多个后续标记的预测能够并行进行，从而显著减少解码步骤。Medusa提供两种微调方案，分别为Medusa-1和Medusa-2，前者在冻结的基础LLM上进行微调，后者则与基础LLM共同微调，提升了预测准确性和加速效果。实验表明，Medusa-1可实现超过2.2倍的加速，而Medusa-2的加速效果更是达到2.3-3.6倍。

🔬 方法详解

问题定义：本文旨在解决大型语言模型推理中的计算瓶颈问题，现有方法如推测解码因维护草稿模型的复杂性而难以实施。

核心思路：Medusa的核心思路是通过增加多个解码头，利用树形注意力机制并行预测多个后续标记，从而减少解码步骤，提高推理效率。

技术框架：Medusa的整体架构包括基础LLM和多个解码头，解码过程通过树形结构生成多个候选标记，并在每一步同时验证这些候选。

关键创新：Medusa的主要创新在于其并行解码能力，通过多个解码头的设计，显著减少了推理所需的步骤，与传统的逐步解码方法形成鲜明对比。

关键设计：在微调过程中，Medusa-1在冻结的基础LLM上进行，确保无损加速；而Medusa-2则与基础LLM共同微调，采用特殊的训练方案以保持基础模型的能力。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Medusa-1在不降低生成质量的前提下，实现了超过2.2倍的推理加速，而Medusa-2的加速效果进一步提升至2.3-3.6倍，展示了其在不同模型和训练过程中的有效性。

🎯 应用场景

Medusa框架在自然语言处理、对话系统和文本生成等领域具有广泛的应用潜力。通过加速推理过程，能够提升用户体验，尤其是在实时应用场景中。此外，Medusa的设计理念也可为其他类型的模型推理提供借鉴，推动更高效的AI系统发展。

📄 摘要（原文）

Large Language Models (LLMs) employ auto-regressive decoding that requires sequential computation, with each step reliant on the previous one's output. This creates a bottleneck as each step necessitates moving the full model parameters from High-Bandwidth Memory (HBM) to the accelerator's cache. While methods such as speculative decoding have been suggested to address this issue, their implementation is impeded by the challenges associated with acquiring and maintaining a separate draft model. In this paper, we present Medusa, an efficient method that augments LLM inference by adding extra decoding heads to predict multiple subsequent tokens in parallel. Using a tree-based attention mechanism, Medusa constructs multiple candidate continuations and verifies them simultaneously in each decoding step. By leveraging parallel processing, Medusa substantially reduces the number of decoding steps required. We present two levels of fine-tuning procedures for Medusa to meet the needs of different use cases: Medusa-1: Medusa is directly fine-tuned on top of a frozen backbone LLM, enabling lossless inference acceleration. Medusa-2: Medusa is fine-tuned together with the backbone LLM, enabling better prediction accuracy of Medusa heads and higher speedup but needing a special training recipe that preserves the backbone model's capabilities. Moreover, we propose several extensions that improve or expand the utility of Medusa, including a self-distillation to handle situations where no training data is available and a typical acceptance scheme to boost the acceptance rate while maintaining generation quality. We evaluate Medusa on models of various sizes and training procedures. Our experiments demonstrate that Medusa-1 can achieve over 2.2x speedup without compromising generation quality, while Medusa-2 further improves the speedup to 2.3-3.6x.