让AI读百万字:稀疏注意力突破长文本瓶颈
大模型处理长文本时,传统注意力机制的计算量随文本长度平方增长,导致无法处理超长文档。这篇论文提出一种分层稀疏注意力方法:模型先快速扫描全文,选出最相关的几个段落,然后只对这些段落做精细计算。关键创新是,段落选择过程直接通过语言模型训练来优化,而不是靠外部规则。实验显示,该方法在训练长度内性能与全注意力相当,但能外推到训练长度64倍以上的文本,且检索准确率达90%。这意味着,未来AI可以一次性读完整本《三体》三部曲(约90万字)并准确回答细节问题,而现有模型只能处理几万字。虽然目前仍是研究阶段,但它指明了让AI真正理解长文档的可行路径。
📄 原文摘要(英文)
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.