提出新型稀疏注意力机制,实现超长文本建模且效率更高。
Hierarchical Sparse Attention Done Right: Toward Infinite Context Modeling

- 通过分层检索学习动态选块,端到端优化注意力块选择。
- 在训练长度64倍外推时保持90%检索准确率,性能超越全注意力。
- 可轻量迁移现有模型,兼具高效与长程建模能力,适合长文本任务。
将现代大语言模型扩展至长上下文受限于二次计算成本及密集注意力的长度外推能力差。分块稀疏注意力虽有潜力,但现有方法因块选择不准而无法媲美全注意力。本文提出分层地标稀疏(HiLS)注意力机制,在语言建模损失下端到端学习块选择。该机制分层分解注意力:每个查询独立与检索到的块进行注意力交互以提取块特定信息,并根据检索得分融合输出。通过将检索得分融入前向计算,直接以语言建模损失优化,实现端到端检索学习与原生稀疏训练。实验表明,HiLS-Attention在域内上下文长度上达到甚至超过全注意力性能;同时可在训练长度64倍以上实现90%检索准确率,显著优于全注意力。此外,现有全注意力模型可通过轻量持续预训练转为HiLS-Attention,保留域内性能并获得超长上下文外推能力。结合稀疏键值访问与计算,HiLS-Attention突破了效率-性能权衡,使长上下文大模型在通用长文本任务中更高效、更有效。
原文摘要 · Abstract (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 $64\times$ 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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。