用注意力混合加速语言模型推理

语言模型 (LLMs) 正在彻底改变世界，但其巨大的计算需求使其部署成本高昂且速度缓慢。推测性解码提供了一个有前景的解决方案，它利用较小的模型更有效地预测未来令牌，然后由较大的LLM进行验证。然而，推测性解码技术面临着部分可观察性和低效训练等挑战。在这篇文章中，我们将探讨一种名为注意力混合 (Mixture of Attentions) 的新方法，它通过解决这些限制显著增强了推测性解码。

为什么要阅读此帖子？

简而言之，就是最先进的推理结果！注意力混合架构旨在优化单设备和客户端-服务器设置中的推测性解码。可观察性、训练效率和灵活性方面的改进使此架构更加健壮，并适用于各种用例。具体来说，LSA、CA和TLI的组合可实现更快、更准确的令牌生成，具有以下主要优势：

1. 与EAGLE-2相比，解码速度提高9.5%：与之前最先进的推测性解码方法 (EAGLE-2) 相比，注意力混合解码速度提高了9.5%。
2. 接受率提高25%：通过改善较小模型对大型模型状态的理解并在策略上对其进行训练，注意力混合将大型模型接受的令牌数量增加了25%。
3. 对客户端-服务器场景的适应性：该架构在客户端-服务器部署中特别有效，其中较小的模型可以在客户端设备上运行，即使服务器（托管大型模型的设备）不可用，也能继续生成令牌。
4. 我们还共享了可供您使用的检查点。您可以在这里找到模型：https://huggingface.co/huawei-noah/MOASpec-Llama-3-8B-Instruct
5. 我们还在 vLLM 中实现了该方法。您可以在这里找到代码：https://github.com/huawei-noah/HEBO/tree/mixture-of-attentions/

了解推测性解码

随着LLM规模的增长，其计算需求也随之增长。这些模型的自回归特性，即每个新令牌都基于先前的序列生成，使得它们在实时应用中部署起来特别缓慢且成本高昂。推测性解码通过引入一个可以“草拟”令牌的较小模型（然后由较大模型验证）来解决这个问题，从而提供了一个创新的解决方案。此方法有助于减轻大型模型的负担，从而加速整个令牌生成过程。

推测性解码的核心是一个两步过程，涉及一个较小、高效的模型（通常称为草稿模型）和一个较大、更强大的模型（验证模型）。草稿模型生成一系列令牌，这些令牌是关于较大模型将生成什么的推测性猜测。然后将这些令牌发送到较大模型进行验证，较大模型可以接受或拒绝它们。

1. 草拟：较小的模型根据提供的上下文建议未来的令牌。
2. 验证：较大的模型检查令牌的正确性。如果序列中的令牌与大型模型的预测不匹配，则整个序列将被丢弃，并重新开始该过程。这种草拟和验证过程持续进行，直到模型达到所需的输出长度。这里的主要优点是，在许多情况下，较小的模型可以准确地预测多个令牌，从而减少了大型模型所需的前向传播次数。

推测性解码的当前挑战

尽管推测性解码前景广阔，但它也面临挑战。其中最紧迫的两个问题是部分可观察性和离策略训练。

1. 部分可观察性：在传统的推测性解码方法中，较小的模型无法访问较大模型使用的所有信息。具体来说，它无法访问大型模型的整个状态，包括更深层次的关键激活和隐藏状态。这可能导致次优的令牌预测，因为较小的模型仅以部分视图操作系统。结果是验证期间更频繁地出现不匹配，导致序列被丢弃和效率低下。
2. 离策略训练：训练较小的模型是当前方法不足的另一个领域。较小的模型通常在理想条件下进行训练，假设它从大型模型接收完美的输入。然而，在实际使用中，较小的模型通常必须生成自己的预测，这可能并不总是准确的。训练和推理之间的这种不匹配被称为离策略训练，它可能导致模型部署时性能显著下降。较小的模型独立草拟令牌的时间越长，它偏离正确序列的可能性就越大，从而导致错误增加。

传统方法与注意力混合

现有的推测性解码模型，如EAGLE和MEDUSA，试图解决这些挑战，但存在局限性。例如，EAGLE利用大型模型的激活来指导小型模型的预测，但仍然存在部分可观察性问题，并且在离策略训练方面遇到困难。

在下一节中，我们将介绍注意力混合架构，它为这些问题提供了更扎实的解决方案。通过利用多个注意力机制，注意力混合方法增强了小型模型准确草拟令牌的能力，同时在更真实的策略设置中对其进行训练。

引入注意力混合架构

为了克服推测性解码的挑战，本文引入了一种名为注意力混合的新颖架构，它为推测性解码带来了三项主要创新：层自注意力 (LSA)、交叉注意力 (CA) 和目标层推理 (TLI)。这些组件共同解决了部分可观察性和离策略训练的问题，同时还提供了一种更灵活的方式来平衡速度和准确性。

注意力混合架构的关键组件：

层自注意力 (LSA)：传统推测性解码中最重要的限制之一是较小的模型无法完全了解较大模型的内部状态。较小的模型只能观察到大型模型最终层的激活，这导致了部分可观察性——对上下文的不完整理解。

注意力混合架构引入了层自注意力 (LSA) 来解决这个问题。这种注意力机制聚合了来自大型模型所有层（而不仅仅是最终层）的关键激活。通过汇总多个层的信息，LSA为较小的模型提供了对当前状态更丰富的理解，使其能够做出更明智的令牌预测。

工作原理

1. 大型模型在令牌生成过程中为每个层生成激活。
2. LSA将注意力应用于这些激活，从每个层提取相关信息并降低维度，使其易于较小模型处理。
3. 这种增强的视图减少了不正确令牌草稿的可能性，从而提高了解码过程的整体效率。

交叉注意力 (CA)：传统推测性解码的另一个主要限制是训练过程中缺乏策略性。在推测性解码中，较小的模型通常需要根据其先前的输出（而不仅仅是来自较大模型的完美输出）生成令牌。然而，大多数模型都是离策略训练的，这意味着它们是在假设完美输入的理想条件下训练的。训练和实际使用之间的这种差异在模型部署时导致性能显著下降。

CA通过允许较小的模型在更真实、更策略性的设置中学习来解决这个问题。CA机制使较小的模型能够一次预测多个未来令牌，同时仅依赖于来自大型模型直至当前令牌的激活。通过在训练期间模拟真实世界条件，较小的模型能够更好地处理实际推理过程中的错误和不确定性。

工作原理

1. 在草拟阶段，较小的模型使用CA来预测一系列令牌（而不是一次一个令牌）。
2. 交叉注意力层使用来自大型模型直至当前令牌的激活，但允许较小的模型生成多个令牌而无需大型模型持续反馈。
3. 这使得较小的模型T步受限，这意味着它可以在一次通过中草拟多达T个未来令牌，从而降低了计算成本并提高了训练效率。

目标层推理 (TLI)：传统推测性解码假设较小的模型应始终预测大型模型最终层的激活。然而，作者通过目标层推理 (TLI) 挑战了这一假设，允许较小的模型以大型模型的更深层为目标。

这里的直觉是，预测中间层可能比预测最终层的输出更容易，这仍然可以导致准确的令牌预测。通过针对不同的层，TLI可以在速度和准确性之间进行权衡：针对较早的层速度更快，但可能导致预测不那么准确，而针对较晚的层会提高准确性，但需要更多的计算。

工作原理

1. 该架构引入了一个超参数N，用于定义小型模型的目标层。
2. 如果N = 0，则小型模型以最终层为目标（标准方法）。如果N > 0，则以较早的层为目标。
3. 这种灵活性允许模型根据任务要求调整其行为，平衡速度和准确性。

主要贡献和结果

注意力混合架构通过解决部分可观察性和离策略训练的挑战，同时在平衡速度和准确性方面提供灵活性，显著推动了推测性解码领域的发展。在本节中，我们将重点介绍该论文的主要贡献，并回顾证明该架构有效性的实验结果。

解码加速：推测性解码的主要目标之一是在不牺牲准确性的前提下加速推理过程。注意力混合架构通过引入更智能、更高效的草稿过程来实现这一点。通过利用层自注意力 (LSA) 和交叉注意力 (CA)，较小的模型可以更准确地草拟令牌，同时减少大型模型所需的验证循环次数。

关键成果

1. 与之前最先进的模型EAGLE-2相比，注意力混合实现了9.5%的解码速度提升。
2. 这种改进在单设备设置中尤为明显，注意力混合显著缩短了生成响应所需的时间，同时保持了高准确性。

更高的接受率：另一个显著的贡献是小型模型生成的令牌的接受率提高了。得益于层自注意力，小型模型对大型模型的内部状态有了更全面的了解，这使得它能够草拟在验证过程中更有可能被接受的令牌。此外，交叉注意力提高了小型模型训练的策略性，增加了其草稿被接受的可能性。

关键成果

1. 注意力混合架构的接受率比EAGLE-2高25%。
2. 这意味着较小的模型生成的序列更有可能被较大的模型批准，从而减少了被丢弃的令牌数量并提高了整体效率。

客户端-服务器部署：注意力混合架构最令人兴奋的方面之一是其在客户端-服务器部署场景中的有效性。在此设置中，较小的模型在客户端设备（例如移动电话）上运行，而较大的模型托管在服务器上。较小的模型生成令牌并将其发送到服务器进行验证，但如果网络断开，它可以继续自主生成令牌。

此功能对于边缘计算以及无法保证持续访问强大服务器的情况至关重要。通过允许较小的模型在服务器不可访问时独立运行，注意力混合能够更健壮、更灵活地部署LLM到实际应用中。

关键结果：注意力混合在客户端-服务器部署中实现了最先进的延迟，即使在具有挑战性的网络条件下（例如，4G和5G网络）。在完全断开连接的情况下，注意力混合模型继续以比其他推测性解码方法更高的准确性生成令牌，而其他推测性解码方法在没有服务器访问的情况下将失败。

目标层推理 (TLI) 的适应性：引入目标层推理 (TLI) 为注意力混合架构增加了另一层灵活性。通过允许小型模型以大型模型的更深层为目标，该架构可以根据任务调整其行为。这种灵活性实现了速度和准确性之间的权衡

1. 如果目标是更快的令牌生成，小型模型可以以大型模型的早期层为目标，从而降低计算成本。
2. 如果准确性更重要，小型模型可以以更深层为目标，从而提高其预测质量。

关键结果：该架构表明，调整目标层（TLI机制中的N）可以根据特定应用需求在速度和准确性之间进行微调。

影响和未来方向

注意力混合架构在推测性解码方面取得了突破性进展，对大型语言模型 (LLMs) 的研究和实际应用产生了重大影响。通过解决部分可观察性和离策略训练等核心挑战，该架构为更高效、可扩展的LLM部署铺平了道路，尤其是在计算资源有限的边缘计算场景中。

对大型语言模型部署的影响 LLM在医疗保健、教育到金融和客户服务等各个行业的快速采用意味着提高这些模型的效率比以往任何时候都更加关键。注意力混合架构为在实际设置中部署LLM提供了几个关键优势：

1. 更快的推理：解码速度提高9.5%使得LLM的实时应用（例如聊天机器人和虚拟助手）响应更快。这可能转化为更流畅的最终用户交互，即使使用具有数十亿参数的模型也是如此。
2. 边缘和客户端-服务器计算：在LLM需要部署到边缘设备（如智能手机或物联网设备）的场景中，小型模型在服务器断开连接时能够继续独立生成令牌，这是一项颠覆性功能。这为在离线或低连接性环境中使用LLM开辟了新的可能性，例如偏远地区或自主系统。
3. 能源效率：随着对大型模型能源消耗日益增长的担忧，提高推测性解码的接受率，同时减少对大型模型的依赖，可以降低计算成本和能源消耗，使LLM部署更具可持续性。

对模型训练的影响：注意力混合引入的架构变化也为LLM的训练阶段带来了好处

1. 提高训练效率：交叉注意力 (CA) 层实现了更策略性的训练，其中较小的模型在与真实世界推理场景密切相关的条件下进行训练。这减少了模型从训练到部署时通常出现的性能下降。
2. 更具适应性的模型：通过允许较小的模型以大型模型的不同层为目标（通过目标层推理），该架构提供了模型训练和优化方式的灵活性。开发人员可以根据手头的任务调整推理深度，动态平衡计算成本和准确性。

未来方向：注意力混合架构为未来的研究和开发开辟了令人兴奋的途径。未来工作的一些潜在方向包括

1. 动态目标层推理：这项工作的一个潜在扩展是使模型能够根据任务的复杂性或客户端-服务器场景中的当前网络条件动态选择最佳目标层（目标层推理中的N）。这将允许更高效和适应性更强的部署，其中模型可以根据需要自动平衡速度和准确性。
2. 隐私保护推测性解码：在客户端-服务器设置中，有可能探索隐私保护方法，其中用户输入的敏感部分保留在客户端，只有非敏感数据发送到服务器。注意力混合架构可以进行调整，以确保某些激活或令牌序列在本地处理，从而允许在医疗保健或法律服务等领域进行隐私敏感的LLM部署。
3. 扩展到其他领域：虽然这项工作侧重于LLM中的令牌生成，但推测性解码和注意力混合的原理可以扩展到使用预测模型的其他领域。例如，它可以应用于机器翻译、代码生成，甚至机器人控制系统，在这些领域中，快速、准确的预测至关重要。先进的边缘计算应用：随着分散式和边缘计算趋势的增加，注意力混合可以在自动驾驶汽车、智能家居和实时翻译设备等应用中发挥关键作用。这些系统需要以最小延迟运行的能力，而注意力混合处理断开操作的能力是一个有价值的特性。

该架构降低了能耗，并增强了离线和低连接性设置的功能，也使其非常适合可持续人工智能计划。随着对人工智能环境影响的担忧日益增加，在不牺牲性能的情况下更节能地部署模型将变得越来越重要。

如果您觉得这有用，请点赞、分享和关注！