階層式稀疏注意力的正確實踐：邁向無限上下文建模

摘要

擴展現代大型語言模型（LLMs）至長上下文的能力受到兩大限制：二次計算成本與密集注意力的弱長度外推能力。分塊稀疏注意力提供了有前景的替代方案，但現有方法因不精確的區塊選擇而無法達到全注意力的效果。我們提出分層地標稀疏注意力（HiLS Attention），這是一種可學習區塊選擇的分塊稀疏注意力機制，能在語言建模損失下進行端到端的區塊選擇學習。HiLS 將注意力分解為分層結構：每個查詢獨立對每個檢索到的區塊進行注意力計算，以提取區塊特定信息，並根據區塊檢索分數融合輸出結果。透過將檢索分數納入前向注意力計算，HiLS 能直接以語言建模損失優化這些分數，實現端到端的檢索學習與原生稀疏訓練。實驗結果顯示，HiLS-Attention 在域內上下文長度下可達到與全注意力相當甚至更優的表現。同時，HiLS-Attention 能以 90% 的檢索準確率，外推至訓練上下文長度的 64 倍以上，遠超全注意力。此外，現有的全注意力模型可透過輕量級持續預訓練轉換為 HiLS-Attention，在保持域內表現的同時獲得超長上下文外推能力。結合其稀疏的 KV 存取與計算，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.