稀疏注意力再添一员，华为诺亚推出高效选择注意力架构ESA

机器之心 · 公众号 · AI · 2025-02-24 16:37

正文

AIxiv专栏是机器之心发布学术、技术内容的栏目。过去数年，机器之心AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：[email protected]；[email protected]

当 DeepSeek 的 NSA 与月之暗面的 MoBA 以稀疏注意力掀起长序列技术热潮，行业对 “效率革命” 的追逐迎来关键一跃 —— 华为诺亚方舟实验室正式发布全新 ESA 算法（Efficient Selective Attention）。

论文地址：https://arxiv.org/pdf/2502.14477

通过稀疏化注意力的创新设计，ESA 突破了大模型在长文本处理中的瓶颈。ESA 不仅实现了数倍序列长度的拓展，还引入独创的动态计算范式，结合邻域影响力有效避免了单纯选择 top-ranked token 所带来的性能损失。通过对关键 token 的精确选择，ESA 在优化长序列处理效率的同时，提升了计算性能，为大模型在长序列任务中的应用带来了新的可能性。

在大语言模型的推理过程中，长序列模型的训练需要极高的算力和海量数据支持，理想的解决方案是通过短序列的训练成果外推到长序列。然而，随着序列长度的增加，注意力计算的复杂度呈平方级增长，这使得高效且准确的长序列推理成为了一大挑战。为此，研究人员提出了多种方法，以应对这一挑战。

ESA 方案正是在这一背景下提出的创新外推解决方案。ESA 通过对 query 和 key 的低维压缩，有效减少了 token 选择的计算复杂度。该方案通过灵活高效地选择关键 token 进行注意力计算，大幅度降低了 LLMs 在处理长文本时的计算负担，且在性能上与全注意力外推方法相当，甚至在高倍外推场景下优于全注意力算法，实现了上下文长度的有效扩展。

1. 高效外推

当大模型训练长度有限，随着序列长度的增长，一方面会出现 OOD (out-of-distribution) 的问题，另一方面注意力计算量会迅速增大。现有的研究表明，注意力矩阵具有稀疏性，对于长序列而言，稀疏程度进一步扩大。选择性注意力（Selective Attention）利用了稀疏性这一特性，选择部分 token 来计算注意力，结合外推的位置编码能将短序列模型应用到长序列任务上的同时，显著降低计算量。在计算稀疏注意力时细粒度的 token 选择方法能够更加灵活、精准地定位到关键信息。然而，token 粒度选择会引入巨大的计算开销。这引出了一个核心的问题：如何在选择性注意力方法中平衡灵活性与效率。针对这一挑战，ESA 方法通过将 query 和 key 进行低维压缩，显著降低 token 选择的计算复杂度，在外推场景下实现 token 粒度动态稀疏注意力机制。

具体而言，ESA 包括以下两个核心步骤：

高效选择：ESA 引入了一种基于 query 感知的 token 粒度选择机制，基于压缩后的 query 和 key 计算 token 的重要性分数，同时考虑周围 token 的影响（邻距影响力），以避免直接选择 top-ranked token 导致的性能下降。

注意力计算：在选择关键 token 后，ESA 使用被选中的 token 的完整的 query 和 key 进行注意力计算，而非对所有前序 token 进行计算，从而大幅降低复杂度。

2.ESA：基于 token 粒度的高效选择性注意力

ESA 的主要创新点在于通过 token 粒度选择性注意力机制，在保持模型准确率的同时显著降低计算复杂度。具体来说，与现有的长序列外推方法不同，ESA 提出了一种基于 token 的细粒度选择注意力，能够在 prefilling 和 decoding 阶段动态选择最关键的少量 token，而不是固定 block 选择或者永久丢弃不重要的 token。首先，ESA 将 query 和 key 经过简单的一层 MLP 压缩到原有维度的大约 3.2%，在低维空间计算重要性分数，显著降低计算复杂度；其次，根据重要性分数选择 topk 的 token，控制 key 的长度是固定的，这样将注意力计算由原有的平方复杂度降低为线性复杂度。虽然选择 token 是平方复杂度，但是由于将 query 和 key 压缩到了更低维的空间，使得对于算力要求大大降低。

ESA 算法示意图

ESA 的具体实现方式如下：输入序列的 token 被分为 4 部分，注意力包括全局注意力和 window 的局部注意力，初始 token 和 ESA 选择的 topk 中间 token 拼接起来计算全局注意力，localtoken 用于计算 window 的注意力，两部分注意力进行融合计算最终的注意力。ESA 按照 chunked-prefill 缓存 key 和 value，即基于当前 chunk 的 query 选择重要的中间 tokens，计算 token 的重要性时兼顾当前的所有 query；在解码阶段，只需要考虑当前的一个 token 的 query 即可。如果计算中间某个 token 重要性，需要计算和当前所有 token 的重要性，其中单个 token 的重要性用 query 和 key 的点积表示：

这里 H 是 head 的数量，为了降低复杂度 ESA 整合了所有的 head。为了进一步降低计算复杂度，不要求准确计算重要性分数，而是更关注相对大小，ESA 将 query 和 key 分别通过一层 MLP 进行压缩。ESA 采取 offline 的方式学习 MLP 的权重：

ESA 使用一个小的校准数据集用模型进行推理，保存中间的 query、key 和 value，用于训练降维 MLP，只增加了极少量的降低 query 和 key 大小的网络权重，且无需对模型微调。

为了确保分数的相对大小，避免某个 token 在重要性分数中占据主导地位，ESA 对分数进行修正：

进一步的，作者发现仅选择 topk 的 token 模型在大海捞针任务中只能检索到部分信息，提出了邻距影响力的概念，即对于某个中间的 token，其重要性分数不仅取决于自身的分数，还受到周围 token 的影响，更新后的分数为：

在选择完重要 token 后，ESA 使用完整的 query、key 和 value 计算注意力，最终的注意力输出如下所示：

ESA 的计算复杂度降低主要来源于低维的 query 和 key 计算重要性分数以及选择完成以后的线性注意力计算复杂度，经过理论计算，一步 attention 计算在长序列场景下能降低为原有的：

实际实验中我们将 query 和 key 压缩为原有的 3.2%，一步 attention 计算量在输入序列足够长时理论能降低至 1.6% 左右。

3. 实验结果

论文选择开源训练集 Pile 的 2 条 Books3 样本收集用于训练降维 MLP 的 qk 样本，query 和 key 从 4096 压缩为 128，压缩比例约为 l3.2%，注意力计算的窗口长度约为 6k。为了将开源的短序列模型应用到长序列中，ESA 沿用了 Infllm 的外推位置编码设置，使用 Llama-3-8B-Instruct 和 Mistral-7B-Instruct-v0.2，在多个公开的长序列基准测试中验证了 ESA 的性能，包括 Longbench、InfiniteBench、NeedleBench 等。作者对比了 full attention 的外推方法和同类型的基于 window 的外推方法，且同类型方法的 window 长度一致。实验结果表明，ESA 通过高效灵活选择重要的 token，总体性能在外推倍数足够大时候优于 full attention 的方法，且均明显优于同类型的方法，尤其在 multi needles 检索场景下例如数星星和 NeedleBench，在其他同类型方法失效的时候，ESA 仍然有较高的准确率。

稀疏注意力再添一员，华为诺亚推出高效选择注意力架构ESA

正文

请到「今天看啥」查看全文