HISA:高效分层索引用于细粒度稀疏注意力
机器学习
2026-04-07 v3 人工智能
摘要
以深度优先稀疏注意力(DSA)为代表的逐标记稀疏注意力机制,通过轻量级索引器对每个查询评分所有历史键,然后仅对选中的子集计算注意力。虽然下游稀疏注意力本身具有良好的扩展性,但索引器仍需为每个查询扫描整个前缀,引入了一个随上下文长度增长而变得不可接受的每层瓶颈。我们提出 HISA(分层索引稀疏注意力),作为索引器的即插即用替代方案,将搜索路径从扁平标记扫描重写为两阶段分层过程:(1)一个块级粗过滤阶段,通过评分池化的块表示来丢弃无关区域,随后(2)一个标记级细化阶段,仅在保留的候选块内应用原始索引器。HISA 保留了下游稀疏 MLA 算子所消费的相同逐标记 top-sparse 模式,且无需额外训练。在核级基准测试上,HISA 在 64K 上下文下实现了显著加速。在 Needle-in-a-Haystack 和 LongBench 上,我们直接将 DeepSeek-V3.2 和 GLM-5 中的索引器替换为我们的 HISA 索引器,无需任何微调。HISA 在质量上紧密匹配原始 DSA,同时大幅优于块稀疏基线。
引用
@article{arxiv.2603.28458,
title = {HISA: Efficient Hierarchical Indexing for Fine-Grained Sparse Attention},
author = {Yufei Xu and Fanxu Meng and Fan Jiang and Yuxuan Wang and Ruijie Zhou and Zhaohui Wang and Jiexi Wu and Zhixin Pan and Xiaojuan Tang and Wenjie Pei and Tongxuan Liu and Di Yin and Xing Sun and Muhan Zhang},
journal= {arXiv preprint arXiv:2603.28458},
year = {2026}
}