华为诺亚推出高效选择注意力架构ESA，稀疏注意力迎来新成员

在 DeepSeek 的 NSA 与月之暗面的 MoBA 共同引领稀疏注意力技术热潮之际，行业对于“效率革命”的追求迎来了重大的突破——华为诺亚方舟实验室正式推出全新的 ESA 算法（Efficient Selective Attention）。论文地址：https://arxiv.org/pdf/2502.14477

通过在注意力机制中引入稀疏化的设计，ESA 成功克服了大型模型处理长文本时的瓶颈。该算法不仅实现了数倍的序列长度扩展，还采用创新的动态计算范式，结合邻域影响力，有效避免了单纯依赖选择排名前列的 token 所导致的性能损失。通过精确选择关键 token，ESA 在提升长序列处理效率的同时，也增强了计算性能，为大型模型在长序列任务中的应用提供了新的机遇。
在大规模语言模型的推理过程中，长序列模型的训练通常需要极高的计算能力与大量的数据支持，理想的解决方案是通过短序列的训练成果推广至长序列。然而，随着序列长度的增加，注意力计算的复杂度呈现平方级增长，使得高效且准确的长序列推理成为一项重大挑战。为此，研究者们提出了多种应对策略。
ESA 方案正是在这一背景下提出的有效解决方案。通过对查询（query）和键（key）的低维压缩，ESA 显著降低了 token 选择的计算复杂度。该方案灵活且高效地选择关键 token 进行注意力计算，大幅度减轻了大型语言模型在处理长文本时的计算负担，在性能上与全注意力推理方法相当，甚至在高倍率推理情况下优于全注意力算法，有效扩展了上下文长度。
1. 高效外推
在大模型的训练时，序列长度的限制导致循环外推可能出现 OOD（out-of-distribution）问题，并且随着序列长度的增加，注意力计算的需求迅速上升。目前的研究显示，注意力矩阵往往具有稀疏性，而对于长序列而言，这种稀疏性会进一步增强。选择性注意力机制（Selective Attention）正是基于这一特性，采用部分 token 进行注意力计算，结合外推位置编码将短序列模型有效地应用于长序列任务，并显著降低了计算负担。通过对 token 进行细粒度选择，可以更灵活、精准地获取关键信息。然而，token 粒度选择的过程可能带来巨大的计算开销，这就引发了一个核心问题：如何在选择性注意力方案中平衡灵活性与效率。为应对此挑战，ESA 方法通过对查询和键进行低维压缩，显著降低了 token 选择的计算复杂度，并在外推场景中实现了动态稀疏注意力机制。
具体而言，ESA 方案由以下两个核心步骤构成：
高效选择：ESA 引入了一种基于查询的感知机制，用以动态选择 token 的粒度，基于压缩后的查询和键计算 token 的重要性分数，同时顾及周围 token 的影响（邻域影响力），以避免直接选择排名前列的 token 导致的性能下降。
注意力计算：在关键 token 被选择后，ESA 使用所选 token 的完整查询和键进行注意力计算，而不是对所有前序 token 进行计算，从而显著降低计算复杂度。
2. ESA：基于 token 粒度的高效选择性注意力
ESA 的主要创新在于通过 token 粒度选择的注意力机制，在保持模型准确性的同时显著降低计算复杂度。具体而言，ESA 提出了与现有长序列推理方法不同的基于 token 的细粒度选择机制，能够在 prefilling 和解码阶段动态选择关键的少量 token，而不是固定选择特定块或者永远丢弃不必要的 token。ESA 通过一层 MLP 将查询和键压缩至原有维度的约 3.2% 的低维空间中计算重要性分数，显著减少了计算复杂度；随后，根据重要性分数选择排序前 k 的 token，固定键的长度，从而将注意力计算的复杂度从平方降低至线性。尽管 token 的选择过程是平方复杂度，由于查询和键被压缩至较低维度，因而对计算资源的需求也大幅降低。力机制包括全局注意力与局部窗口注意力。首先，将初始 token 与 ESA 选择的 top-k 中间 token 拼接，以计算全局注意力；而 local token 则用于计算窗口的局部注意力。两种注意力的计算结果将融合以获得最终的注意力。ESA 通过 chunked-prefill 机制对 key 和 value 进行缓存，根据当前 chunk 的 query 来选择重要的中间 token 算出其重要性，同时也考虑了所有相关的 query；在解码阶段，仅需考虑当前 token 的 query。如果要计算某个中间 token 的重要性，必须综合当前所有 token 的重要性，其中单个 token 的重要性通过 query 和 key 的点积表示：其中 H 表示头的数量。为降低计算复杂度，ESA 整合了所有头的计算。为进一步减少计算开销，ESA 不强求准确计算重要性分数，而是更注重相对值的比较，故采用一层 MLP 对 query 和 key 进行压缩。ESA 以离线方式学习 MLP 的权重，如下所示：在使用小型校准数据集进行推理时，ESA 保存了中间的 query、key 和 value，以此训练降维 MLP，只需添加少量的网络权重以减小 query 和 key 的维度，无需对模型进行微调。
为了确保分数的相对大小，并防止某个 token 在重要性分数中占主导位置，ESA 对分数进行了修正：进一步分析表明，仅选择 top-k token 的模型在信息检索任务中会遗漏部分重要信息，因此提出了邻距影响力的概念。即某个中间 token 的重要性不仅取决于其自身的分数，还受到周围 token 的影响。更新后的分数为：选定重要 token 后，ESA 将运用完整的 query、key 和 value 计算最终注意力，其输出结果如下：ESA 的计算复杂度降低的主要原因在于低维 query 和 key 的重要性分数计算以及选择完成后的线性注意力计算。理论分析显示，在长序列场景下，一步注意力的计算能够显著降低至原有的： 在实际实验中，我们将 query 和 key 压缩至原始维度的 3.2%，在输入序列足够长时，一步注意力的计算量理论上能降至约 1.6%。
3. 实验结果
本研究采用开源训练集 Pile 中的2条 Books3 样本进行降维 MLP 的 qk 样本训练，将 query 和 key 从 4096 压缩至 128，压缩比例约为 3.2%，注意力计算的窗口长度接近 6k。为将开源的短序列模型适配于长序列，ESA 继承了 Infllm 的外推位置编码设置，使用 Llama-3-8B-Instruct 和 Mistral-7B-Instruct-v0.2，并在多个公开的长序列基准测试，如 Longbench、InfiniteBench、NeedleBench 等中验证了 ESA 的性能。作者比较了基于完整注意力的外推方法与同类基于窗口的外推机制，且在同类方法中保持相同的窗口长度。实验结果表明，ESA 在外推倍数显著增大的情况下，凭借其高效而灵活的关键 token 选择策略，总体性能超越了完整注意力方法，且在各类同类算法中表现突出，尤其是在 multi needles 检索场景中，如数星星与 NeedleBench，在其余同类方法失效时，ESA 依然保持较高的准确率。ESA 不会对每个 head 分别选择 token，而是将所有 head 的结果综合计算重要性分数，这一方式有助于降低计算复杂度，从而提升效率。为了验证此操作对算法的影响，作者进行了对比实验，如下所示，结果表明这种整合方式对算法的影响较小。此外，论文还研究了邻距影响力超参数的影响，如下所示。不同的评测集对该参数的反应各异，较小的取值有利于 multi needles 类型的检索任务，较大的取值则对 single needle 类型任务有利。这可能是由于单针检索任务只需关注目标片段，因此增加邻距影响力能使 attention 更集中于更长的片段。4. 总结
ESA 有效地平衡了在长序列外推场景中选择性注意力的灵活性与计算效率，实现了在不增量微调模型参数的前提下扩展上下文长度的目标。ESA 的核心思想是在每一步选择固定数量的最重要 token 来计算注意力，利用注意力矩阵的稀疏特性。在输入序列长度足够时，ESA 通过将 query 和 key 压缩为低维表示，显著降低了选择 token 的计算复杂度。实验评估表明，ESA 具备有效处理训练长度 4 倍甚至 25 倍task 的能力。未来的工作将着重探索更加准确与高效的关键 token 选择方法，以及实现软硬件协同的高效外推方案。