Hierarchical Sparse Attention Done Right: Toward Infinite Context Modeling

Abstract: Scaling modern 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.

• The paper introduces HiLS-Attention to enable efficient infinite context modeling by using trainable landmark tokens for robust chunk selection.
• It employs hierarchical factorization and low-rank query calibration to achieve up to 15× faster inference and over 90% retrieval accuracy in ultra-long contexts.
• The mechanism preserves short-context performance while offering scalable, denoising attention that outperforms dense models on variable tracking and multi-hop tasks.

Hierarchical Sparse Attention Done Right: Toward Infinite Context Modeling (2607.02980)

Motivation and Background

The paper addresses the scalability limitations of dense full attention in Transformer-based LLMs, principally the quadratic computation cost with respect to input length and the poor length extrapolation behavior beyond the training window. Although chunk-wise sparse attention alternatives such as NSA, DashAttention, and Landmark Attention exist, their inability to select relevant chunks with high fidelity, especially under ultra-long contexts, leads to suboptimal retrieval and degraded performance relative to full attention. This work identifies the expressiveness and optimization of chunk selection as the principal failure modes of prior methods.

Method: Hierarchical Landmark Sparse (HiLS) Attention

Formulation

HiLS-Attention introduces a hierarchical chunk-wise sparse attention mechanism where chunk selection and retrieval are optimized end-to-end via the language modeling objective. Unlike prior approaches that use static or non-parametric chunk summaries (e.g., mean-pooling), HiLS appends trainable landmark tokens to each chunk, whose representations serve as compressed chunk summaries. Chunk selection is based on a mathematically motivated surrogate for the LogSumExp of attention logits, which is approximated via first-order Taylor expansion. This yields a parametric chunk-wise score comprising both a relevance term and an entropy-calibrated bias, computed by:

$$\text{score}_c(q) \approx q^\top k'_c/\sqrt{d} + b_c$$

where $k'_c$ is the weighted chunk key summary, and $b_c$ is the entropy of the attention distribution over the chunk. By incorporating these scores into the forward pass, gradients from next-token prediction directly supervise the chunk selection process.

Hierarchical Factorization

Attention is factorized hierarchically:

• Inter-chunk (Routing): Attending over chunk summaries to select a fixed number of top chunks according to the surrogate scores.
• Intra-chunk (Local): Computing vanilla attention within each selected chunk and the sliding window subset.

The hierarchical decomposition enables replacing the quadratic token-to-token computations in naive block sparse attention (BSA) with much more efficient O(N) and O(KS) routines, where N is sequence length, K is the number of selected chunks, and S is chunk size.

Architectural and Implementation Details

• Positional Encoding: HiLS attains best results when coupled with HoPE (Hybrid of Position Embeddings), which limits the extrapolation decay of traditional RoPE and is especially beneficial during chunk compression.
• Query Calibration: A low-rank query calibration module (Q-Cal) adapts token-level queries for better chunk-level discriminability.
• Grouped Query Attention Compatibility: Efficient batched kernel implementations recycle overlapping retrieved chunks across adjacent tokens and heads, enhancing GPU utilization without imposing a large head-group size constraint.
• Lightweight Conversion: Existing full-attention models can be converted to HiLS-Attention via continued pretraining with only a small number of landmark and projection parameters retrained, preserving original performance while newly conferring ultra-long context capability.

Empirical Results

Short- and Long-Context Performance

• Short-Context: HiLS-Attention matches or slightly outperforms both full attention and YoRN-extended models in perplexity and downstream tasks, with negligible degradation even under sparse training.
• Length Extrapolation: At both small (345M) and moderate (1.4B, 7B) model scales, HiLS-Attention maintains high retrieval accuracy (>90%) in RULER's needle-in-a-haystack evaluations at context lengths up to 4M tokens—far surpassing dense and prior sparse attention methods (including those using mean-pooled and max-pooled chunk summaries).
• Variable Tracking and Multi-hop Tasks: HiLS achieves up to a 50% improvement over full attention in variable tracking tasks at ultra-long contexts, attributed to its ability to suppress token-level noise via compression.

Efficiency

HiLS-Attention demonstrates marked inference speedups:

• Latency: At 512K input tokens, prefill and decode are up to 13.5× and 15.7× faster, respectively, than full attention. The speedup crossover point occurs near 16K tokens.
• Hardware Utilization: Token-chunk overlap and union chunk computation enable high efficiency on modern accelerators without requiring large query group sizes.

Ablation Studies

Key findings include:

• Landmark tokens are essential for robust length extrapolation; their removal or replacement with fixed queries severely degrades performance.
• Low-rank query calibration is critical; omitting it results in large PPL and retrieval accuracy deterioration, while excessive rank also impairs generalization.
• Positional encoding choices matter: HoPE consistently delivers the best overall results, mitigating the position-induced semantic distortion that harms chunk compression and retrieval.

Large-Scale and Downstream Evaluation

Continued pretraining of a 7B LLM shows that HiLS-Attention maintains or surpasses full attention on the LongBench suite and general knowledge, mathematics, and code tasks, with especially strong resilience for sequence lengths $>$8K tokens. Notably, even minimal retraining (landmark-token tuning only) is sufficient to close the gap on short-context tasks.

Theoretical and Practical Implications

HiLS-Attention provides empirical and algorithmic evidence that native sparse attention can not only match but sometimes exceed full attention in both efficiency and retrieval-intensive long-context tasks. Its compression-based chunk selection acts as a denoising mechanism, enhancing in-context retrieval fidelity and multi-hop reasoning—mechanisms ill-posed for dense attention due to accumulated token-level noise. By natively supporting sparse training and block-level retrieval, HiLS provides a concrete architectural pathway toward infinite context modeling, where attention and memory costs no longer scale quadratically with sequence length.

Limitations and Future Directions

• Context Parallelism: Current implementations do not yet support parallel context processing during training, which may limit scaling to truly infinite context.
• Gradient Starvation: Unselected chunks receive no gradient updates; adaptive strategies for gradient routing or memory consolidation may be required.
• Q-Cal Mechanisms: The underlying mechanisms through which query calibration enhances extrapolation remain open for deeper theoretical dissection.

Conclusion

HiLS-Attention advances the state-of-the-art in chunk-wise sparse attention mechanisms by establishing rigorous mathematical underpinnings for chunk selection and enabling end-to-end, efficiently hardware-aligned retrieval. Its strong empirical performance on both in-domain and extreme out-of-domain (ultra-long) contexts breaks the trade-off between efficiency and context modeling power for LLMs. These results present HiLS-Attention as a key architectural component for the next generation of long-context LLMs and provide foundations for practical, scalable infinite context adaptation in neural sequence modeling.