Hierarchical Sparse Attention Done Right: Toward Infinite Context Modeling
arXiv · HuggingFace · ▲68
摘要(原文)
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.
摘要(中译)
将现代大型语言模型(LLMs)扩展到长上下文受到二次计算成本和密集注意力的长度外推能力差的限制。分块稀疏注意力提供了一种有前途的替代方案,但所有现有方法由于不准确的分块选择而无法达到完全注意力。我们提出了分层地标稀疏(HiLS)注意力,这是一种分块稀疏注意力机制,它在语言建模(LM)损失下端到端地学习分块选择。HiLS分层分解注意力:每个查询独立地与每个检索到的分块进行注意力以提取分块特定信息,并且根据分块检索分数融合得到的输出。通过将检索分数纳入前向注意力计算,HiLS直接使用LM损失优化它们,实现端到端检索学习和原生稀疏训练。实验结果表明,HiLS-Attention在领域内上下文长度上实现了与完全注意力相当的性能,在某些情况下甚至更好。同时,HiLS-Attention可以外推超过训练上下文长度的64倍,检索准确率为90%,远远超过完全注意力。此外,现有的完全注意力模型可以通过轻量级的继续预训练转换为HiLS-Attention,保留领域内性能,同时获得超长上下文外推。结合其稀疏KV访问和计算,HiLS-Attention打破了通常的效率-性能权衡,使得长上下文LLMs在一般长上下文任务上比其完全注意力对应物更高效且更有效。
背景剖析
背景剖析
1. 技术背景
现代大语言模型(LLMs)需要处理越来越长的文本上下文,例如长文档理解、多轮对话或复杂推理任务。然而,传统“全注意力”机制的计算成本随文本长度平方增长,导致长文本处理效率低下,且难以准确预测超出训练长度的内容(即“长度外推”问题)。为此,研究者提出“分块稀疏注意力”,通过选择性关注相关文本块来降低计算量,但现有方法在性能上仍无法媲美全注意力。
2. 之前的问题
现有分块方法的核心缺陷在于“块选择不准确”。例如,简单的均值池化等非参数方法无法充分表达文本块的重要性,而参数化方法虽然更灵活,但其选择的块信息在训练中未被直接优化,导致模型难以区分相关与无关内容。此外,这些方法通常需要额外的计算步骤(如硬阈值筛选),无法与全注意力一样高效学习。
3. 本文的解法
本文提出“分层地标稀疏注意力”(HiLS-Attention),通过两步解决上述问题:首先,为每个文本块添加一个“地标标记”,生成更具表达力的块摘要;其次,将块选择过程与语言模型的训练目标(如预测下一个词)直接关联,使模型能自动学习哪些块更重要。这种方法避免了全注意力的高计算成本,同时通过端到端优化提高了块选择的准确性。
4. 切入角度
与传统方法不同,HiLS-Attention的关键创新在于将“块选择”转化为可学习的注意力权重,并利用全注意力的数学特性(如泰勒展开近似)来指导块摘要的设计。实验表明,这种方法不仅在常规长度任务中表现优异,还能将上下文长度扩展至训练长度的64倍以上,远超现有方法。




