Mixture of Attentions For Speculative Decoding

📄 arXiv: 2410.03804v2 📥 PDF

作者: Matthieu Zimmer, Milan Gritta, Gerasimos Lampouras, Haitham Bou Ammar, Jun Wang

分类: cs.CL, cs.AI, cs.LG

发布日期: 2024-10-04 (更新: 2025-04-03)

备注: Accepted at International Conference on Learning Representations (ICLR 2025)

💡 一句话要点

提出混合注意力机制用于推测解码，提升单设备和客户端-服务器场景下的解码速度和精度。

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

关键词: 推测解码 大型语言模型 混合注意力机制 模型加速 客户端-服务器 低延迟推理

📋 核心要点

  1. 现有推测解码模型训练缺乏on-policyness和只能部分观察，导致性能受限。
  2. 提出混合注意力机制，使小模型能更有效地利用LLM的信息，提升预测准确性。
  3. 实验表明，该方法在单设备和客户端-服务器场景下均能显著提升解码速度和精度。

📝 摘要（中文）

大型语言模型（LLM）参数量的增长导致计算需求显著增加，部署具有挑战性且成本高昂。推测解码（SD）利用较小的模型高效地提出未来token，然后由LLM并行验证。目前，利用LLM激活的小模型实现了最快的解码速度。然而，我们发现SD模型存在一些局限性，包括训练期间缺乏on-policyness和部分可观察性。为了解决这些缺点，我们提出了一种更可靠的小模型架构，引入了用于SD的混合注意力机制。我们的新架构可以应用于两种场景：传统的单设备部署和一种新颖的客户端-服务器部署，其中小模型托管在消费设备上，LLM托管在服务器上。在单设备场景中，我们展示了最先进的加速效果，EAGLE-2的加速提高了9.5%，接受长度提高了25%。在客户端-服务器设置中，我们的实验表明：1) 在不同网络条件下，以最少的服务器调用实现了最先进的延迟；2) 在完全断开连接的情况下，与其他SD方法相比，我们的方法可以保持更高的准确性，并且比LLM的API调用更具优势，否则API调用将无法继续生成过程。

🔬 方法详解

问题定义：推测解码旨在加速大型语言模型的推理过程。现有的推测解码方法，特别是那些依赖于小型模型来预测token草案的方法，存在两个主要问题：一是训练过程中的on-policyness不足，即小模型的训练目标与实际推理时LLM的策略不一致；二是部分可观察性，小模型无法完全访问LLM的内部状态，限制了其预测能力。这些问题导致推测解码的效率和准确性受到影响。

核心思路：本文的核心思路是通过引入混合注意力机制，使小模型能够更好地利用LLM的内部信息，从而提高其预测token草案的准确性。混合注意力机制允许小模型同时关注LLM的不同层次的表示，并根据上下文动态地调整注意力权重，从而更全面地理解LLM的意图。这种设计旨在解决现有方法中on-policyness不足和部分可观察性的问题。

技术框架：整体框架包括一个大型语言模型（LLM）和一个小型模型。小型模型使用混合注意力机制来预测token草案，然后LLM并行验证这些草案。具体流程如下：1) LLM生成一部分token；2) 小型模型利用LLM的激活状态，通过混合注意力机制预测后续token草案；3) LLM并行验证这些草案；4) 根据验证结果，接受或拒绝部分草案，并重复该过程。在客户端-服务器部署中，小型模型部署在客户端设备上，LLM部署在服务器上，以减少服务器的计算负担。

关键创新：最重要的技术创新点是混合注意力机制。与传统的注意力机制不同，混合注意力机制允许模型同时关注LLM的不同层次的表示，并根据上下文动态地调整注意力权重。这种设计使得小模型能够更全面地理解LLM的意图，从而提高其预测token草案的准确性。此外，客户端-服务器部署也是一个创新点，它允许在资源受限的设备上运行推测解码，从而扩展了推测解码的应用范围。

关键设计：混合注意力机制的关键设计在于如何有效地融合来自LLM不同层次的信息。具体来说，模型使用多个注意力头，每个注意力头关注LLM的不同层次的表示。然后，模型使用一个门控机制来动态地调整每个注意力头的权重，从而根据上下文选择最相关的层次信息。损失函数包括一个标准的交叉熵损失，用于训练小模型预测token草案，以及一个额外的损失项，用于鼓励小模型更好地利用LLM的信息。

🖼️ 关键图片

img_0

📊 实验亮点

在单设备场景中，该方法在EAGLE-2上实现了9.5%的加速提升，接受长度提高了25%，达到了最先进的水平。在客户端-服务器场景中，实验证明该方法在不同网络条件下均能以最少的服务器调用实现最先进的延迟。即使在完全断开连接的情况下，该方法也能保持比其他推测解码方法更高的准确性。

🎯 应用场景

该研究成果可广泛应用于各种需要加速大型语言模型推理的场景，例如智能助手、机器翻译、文本生成等。特别是在资源受限的设备上，客户端-服务器部署模式可以有效降低计算成本，提高用户体验。未来，该方法有望进一步扩展到其他模型加速技术，并促进大型语言模型在边缘设备上的应用。

📄 摘要（原文）

The growth in the number of parameters of Large Language Models (LLMs) has led to a significant surge in computational requirements, making them challenging and costly to deploy. Speculative decoding (SD) leverages smaller models to efficiently propose future tokens, which are then verified by the LLM in parallel. Small models that utilise activations from the LLM currently achieve the fastest decoding speeds. However, we identify several limitations of SD models including the lack of on-policyness during training and partial observability. To address these shortcomings, we propose a more grounded architecture for small models by introducing a Mixture of Attentions for SD. Our novel architecture can be applied in two scenarios: a conventional single device deployment and a novel client-server deployment where the small model is hosted on a consumer device and the LLM on a server. In a single-device scenario, we demonstrate state-of-the-art speedups improving EAGLE-2 by 9.5% and its acceptance length by 25%. In a client-server setting, our experiments demonstrate: 1) state-of-the-art latencies with minimal calls to the server for different network conditions, and 2) in the event of a complete disconnection, our approach can maintain higher accuracy compared to other SD methods and demonstrates advantages over API calls to LLMs, which would otherwise be unable to continue the generation process.