Hierarchical Sparse Attention Done Right: Toward Infinite Context Modeling

2026-07-03

AI Summary

Large language models are hindered by the quadratic complexity of dense attention and poor length extrapolation. This work proposes HiLS (Hierarchical Learned Sparsity) attention, an end-to-end trainable block-sparse attention mechanism that enables efficient long-context modeling by integrating hierarchical attention with retrieval scores and optimizing language modeling loss to guide sparse key-value access. HiLS achieves the first fully learnable block selection strategy, matching or surpassing full attention in in-domain performance while extrapolating up to 64× beyond the training sequence length—all while maintaining over 90% retrieval accuracy. This approach substantially advances the trade-off between computational efficiency and model performance in long-context scenarios.

Abstract

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.

Problem

• long-context modeling
• sparse attention
• length extrapolation
• chunk selection
• quadratic complexity

Innovation

• Hierarchical Sparse Attention
• End-to-End Retrieval Learning
• Long-Context Extrapolation
• Chunk-wise Attention
• Sparse Language Models