AIxiv专栏是机器之心发布学术、技术内容的栏目。过去数年，机器之心AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com
当 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 不对每个 head 单独选择 token，而是将所有 head 整合到一起计算重要性分数，有利于降低计算复杂度，提升效率，为了验证这一操作对算法的影响，作者做的对比实验如下所示，可以看出这样的整合对于算法影响有限。
论文研究了邻距影响力的超参数影响，结果如下所示，对不同的测评集该参数的影响不同，取值较小有利于 multi needles 类型的检索任务，取值较大则有利于 single needle 类型任务，这可能是由于单针检索任务只需要关注 ground truth 所在的片段即可，增大邻距影响力有利于 attention 集中到较长的片段上。
4. 总结
ESA 有效平衡了长序列外推场景下的选择性注意力中的灵活性和计算效率，用于在不进行模型参数增量微调的情况下扩展上下文长度。ESA 的核心思想是在每个步骤中选择固定数量的最重要 token 来计算注意力，利用注意力矩阵的稀疏性。当输入序列足够长时，ESA 通过将 query 和 key 压缩为低维表征，有效降低选择 token 的计算复杂度。实验评估表明，ESA 能够有效处理长度为训练长度 4 倍甚至 25 倍的各种长序列任务。未来的研究需要探索更准确、更高效的选择重要 token 的方法，以及软硬件协同的高效外推方案。
