HISA：新稀疏注意力机制突破64K上下文瓶颈，提速2-4倍，即插即用无需重训

一项新的研究提出了一种名为HISA（分层索引稀疏注意力） 的稀疏注意力机制。该方法旨在突破长上下文模型中的索引瓶颈，在64K上下文长度下，相比DeepSeek-V3等模型使用的DSA（DeepSeek Sparse Attention）索引器，可实现2至4倍的速度提升。

该方法在显著提升速度的同时，几乎不损失模型精度，并且具备即插即用的特性，无需对现有模型进行重新训练。

研究团队直接在DeepSeek-V3.2和GLM-5模型上替换了原有的索引器，未进行任何微调。在关键信息检索、长文本理解等任务上的测试表明，其精度与原方法基本持平。

两步消除上下文索引瓶颈

该研究旨在为大型模型的稀疏注意力机制设计一个更高效的“检索器”。

当前主流的DSA等token级稀疏注意力机制，其核心思想是通过仅计算关键token之间的注意力来降低计算成本。然而，该设计存在一个关键瓶颈：为了筛选出相关token，需要一个“索引器”，该索引器需要将每个待查询的token与之前的所有token逐一进行相关性打分，再选取分数最高的部分。

随着文本长度L的增加，这个打分过程的计算量呈L的平方级增长。例如，当文本长度翻倍时，工作量将变为原来的四倍。在处理超长文本时，索引器的平方级计算成本反而会成为速度的拖累，其耗时甚至可能超过注意力计算本身。

基于此，研究团队提出了核心问题：能否在不改变最终稀疏注意力结果的前提下，降低索引器的搜索成本？

为此，他们提出了HISA（分层索引稀疏注意力）。其核心思路是：既然对每个token逐一打分耗时过高，可以先在粗粒度上进行筛选，排除大量无关内容，再在缩小的范围内进行精细选择。

HISA在功能逻辑上实现了对原有索引模块的等价替换，无需修改后续的注意力计算逻辑，相当于“更换了一个更高效的筛子，而筛选出的结果几乎不变”。

具体过程分为两步，全程复用原模型的打分规则，无需额外学习：

第一步：块级粗过滤
* 将长文本切分为固定大小的“块”（例如，每块包含128个token），并为每个块计算一个“整体特征向量”（相当于块的摘要表示）。
* 使用原索引器的打分方式，仅对这些块级别的特征向量进行打分。
* 筛选出分数最高的m个块（例如64个），直接丢弃其余所有块。由于块的数量远少于token数量，这一步可以省去绝大部分计算量。

第二步：块内精挑token
仅在第一步筛选出的m个块内部，使用原索引器的规则对单个token进行打分，并最终选出所需的k个相关token。

此外，该方法还加入了一个小优化：强制保留文本的第一个块和最后一个块。这确保了开头的背景信息和结尾的最新上下文不会被误筛，同时也能更好地处理文本拼接时的边界问题。

HISA的关键优势在于：计算复杂度显著降低，且能实现“无缝替换”。

HISA将原索引器每层O(L²)的计算复杂度，降低至O(L²/B + L×m×B)（其中B是块大小，m是选取的块数）。文本越长，块筛选越精准，其提速效果就越明显。

更重要的是其工程友好性：
* 输出格式与原索引器完全一致，下游的注意力计算模块无需任何改动。
* 无需重新训练模型，也无需调整KV缓存结构，可直接替换原索引器。
* 在处理短文本时，HISA会自动“退化”为原方法；仅在处理超长文本时才触发分层筛选机制，整个过程是自适应的。

实测显著提速，精度基本无损

论文在DeepSeek-V3.2和GLM-5两大主流模型上进行了全面测试，结果表现突出：

在速度方面，在64K长度的文本下，HISA相比原DSA索引器最高提速达3.75倍，在常规设置下也能实现2倍以上的提速。

上下文长度越长，HISA的提速效果越显著，这完全契合超长上下文（如128K、1M） 的实际应用需求。

在精度方面，HISA几乎完全保持了原DSA的精度，并且显著优于纯块稀疏的基线方法。

研究进行了“大海捞针”测试，该测试用于衡量模型在超长无关文本中精准检索指定位置关键信息的能力。结果显示，HISA与DSA的检索精度几乎相同，在各种文本长度和插入深度下，其精度均接近DSA的近乎满分水平。

在长文本理解任务（LongBench基准） 上，HISA取得的分数也与DSA基本持平。

甚至在部分对token筛选精度要求较高的场景中，如合成检索、少样本学习等任务，HISA的表现实现了小幅超越。

在超参数测试中，无论块大小和选取块数量如何变化，HISA的表现都相当稳定，其分数与DSA高度接近，未出现显著的性能差异。

这表明HISA对超参数的选择不敏感，鲁棒性强，在工程落地时无需进行精细的参数调整。

当然，目前的HISA方法仍存在可改进之处，作者也提出了后续的优化思路：
1. 动态分块：当前采用固定大小的块，若一个块内同时包含相关与无关内容，其“整体特征”可能不够准确。未来可探索自适应分块、重叠分块或更优的块特征计算方式。
2. 联合训练：目前仅在推理阶段直接应用。未来可以考虑将块筛选机制与模型一同进行训练，使筛选过程更加精准。
3. 端到端评估：当前测试主要集中于索引器本身的速度。未来需要将其整合到完整的大模型服务框架中，评估端到端的吞吐量和延迟表现。

团队背景

本论文由北京大学的张牧涵团队完成。

张牧涵是北京大学人工智能研究院的Tenure-track助理教授、博士生导师。在回国前，他曾于Facebook AI（现Meta AI）担任研究员，从事大规模图学习系统及相关问题的研究。

其在Google Scholar的总引用量超过13000次，其中两篇一作文章的引用量分别达到3100+和2400+次，并连续多年入选Elsevier发布的全球前2%顶尖科学家榜单（生涯影响力）。

论文的共同第一作者是Yufei Xu（徐宇飞）和Fanxu Meng（孟繁续）。

参考链接：
[1] https://arxiv.org/abs/2603.28458

关注“鲸栖”小程序，掌握最新AI资讯