分层稀疏注意力正确实现：迈向无限上下文建模

摘要

扩展现代大型语言模型（LLMs）的长上下文处理能力受限于密集型注意力机制带来的二次计算成本及较差的长度外推能力。基于分块的稀疏注意力提供了有前景的替代方案，但现有方法均因不准确的分块选择而未能达到完整注意力的性能。我们提出层次化地标稀疏注意力（HiLS Attention），这是一种基于分块的稀疏注意力机制，能够在语言建模（LM）损失下端到端地学习分块选择。HiLS通过层次化方式分解注意力：每个查询独立地对每个检索到的分块进行注意力计算以提取分块特定信息，随后根据分块检索得分融合输出结果。通过将检索得分纳入前向注意力计算，HiLS直接使用LM损失优化这些得分，实现了端到端的检索学习和原生稀疏训练。实验结果表明，HiLS-Attention在域内上下文长度上取得了与完整注意力相当甚至更优的性能。同时，HiLS-Attention能够外推超过训练上下文长度64倍的范围，且检索准确率达到90%，远超完整注意力。此外，现有的完整注意力模型可通过轻量级持续预训练转换为HiLS-Attention，在保持域内性能的同时获得超长上下文外推能力。结合其稀疏的键值访问与计算，HiLS-Attention打破了通常的效率-性能权衡，使得长上下文LLM在通用长上下文任务上比其完整注意力版本更高效、更有效。

English

Scaling modern large language models (LLMs) to long contexts is limited by the quadratic computation cost, and poor length extrapolation of dense attention. Chunk-wise sparse attention offers a promising alternative, but all existing methods fall short of full attention because of their inaccurate chunk selection. We propose Hierarchical Landmark Sparse (HiLS) Attention, a chunk-wise sparse attention mechanism that learns chunk selection end-to-end under the language-modeling (LM) loss. HiLS factorizes attention hierarchically: each query performs attention independently with each retrieved chunk to extract chunk-specific information, and the resulting outputs are fused according to chunk retrieval scores. By incorporating retrieval scores into the forward attention computation, HiLS optimizes them directly with the LM loss, enabling end-to-end retrieval learning and native sparse training. Experimental results show that HiLS-Attention achieves performance comparable to, and in some cases better than, full attention at in-domain context lengths. Meanwhile, HiLS-Attention extrapolates more than 64times the training context length with 90% retrieval accuracy, far beyond full attention. Moreover, existing full-attention models can be converted to HiLS-Attention with lightweight continued pretraining, preserving in-domain performance while acquiring ultra-long-context extrapolation. Together with its sparse KV access and computation, HiLS-Attention breaks the usual efficiency-performance trade-off, enabling long-context LLMs that are both more efficient and more effective on general long-context tasks than their full-attention counterparts.