One-sentence Summary

Tencent HY Team, ShanghaiTech University, and collaborators propose Hierarchical Landmark Sparse (HiLS) Attention, a chunk-wise sparse attention mechanism that factorizes attention hierarchically by having each query independently attend to retrieved chunks and fuse outputs via learned retrieval scores optimized end-to-end with the language modeling loss, thereby matching full attention in-domain while extrapolating over $64\times$64× the training length with $90\%$90% retrieval accuracy and breaking the typical efficiency-performance trade-off to enable efficient, ultra-long-context LLMs.

Key Contributions

• HiLS Attention is a chunk-wise sparse attention mechanism that learns chunk selection end-to-end under the language modeling loss, enabling native sparse training.
• It achieves performance comparable to or better than full attention on in-domain context lengths and extrapolates to more than 64× the training context length with 90% retrieval accuracy.
• Existing full-attention models can be converted to HiLS Attention with lightweight continued pretraining, preserving in-domain performance while acquiring ultra-long-context extrapolation, thereby breaking the efficiency–performance trade-off.

Introduction

Scaling context windows in large language models is essential for long-horizon agentic tasks, complex reasoning, and large-scale information integration, but full attention’s quadratic complexity, poor length extrapolation, and growing KV cache costs make it impractical. Chunk-wise sparse attention methods offer a promising alternative by attending to only a few relevant chunks, yet they have not matched full attention’s quality, particularly in long-context retrieval, because they rely on weak, non-parametric chunk summaries (such as mean-pooling) and lack end-to-end optimization of the selection process. The authors propose HiLS-Attention, a native sparse attention mechanism that derives expressive chunk summaries from landmark tokens using a first-order Taylor expansion of the full-attention chunk mass, and then embeds these summaries into a hierarchical softmax factorization so that chunk selection is trained end-to-end under the language modeling loss. This design yields accurate, learnable chunk retrieval that matches full attention in perplexity, surpasses it on long-context benchmarks, and enables extreme length extrapolation (e.g., preserving over 90% needle-in-a-haystack accuracy at 4M tokens after pretraining on only 8K).

Method

The authors propose HiLS-Attention, a fully differentiable sparse attention mechanism that overcomes the expressiveness limitations of conventional block sparse attention (BSA) by learning chunk-level summaries. The overall design addresses two core questions: how to approximate the log-sum-exp chunk mass tractably using compact summaries, and how to learn those summaries end-to-end alongside the language modeling objective.

In naive BSA, a sequence is partitioned into uniform chunks of size $S$S. For each query token $i$i, a local sliding window of size $W$W is aligned to chunk boundaries, and the attention is restricted to the top‑$K$K distant chunks, selected by descending chunk mass $Z_{i,c} = \sum_{j \in \mathcal{T}_c} \exp(s_{i,j})$Zi,c​=∑j∈Tc​​exp(si,j​), where $s_{i,j}$si,j​ is the standard dot‑product attention logit. While this yields exact selection, computing the true chunk mass requires evaluating all token‑level logits inside each chunk, negating the computational savings of sparsity. Practical methods that instead use mean‑pooled keys implicitly approximate the chunk mass by the mean of logits, but the log‑sum‑exp function behaves as the mean only when logits are nearly uniform, and as the maximum when one logit dominates. This mismatch can misrank chunks and cause important tokens to be missed.

HiLS‑Attention replaces the non‑parametric chunk mass with a learnable surrogate. The first step is a linearization of the log‑sum‑exp via a Taylor expansion, which yields a chunk‑level surrogate score of the form

where $\mathbf{k}'_c$kc′​ and $b'_c$bc′​ are computed from a learnable query $\mathbf{q}'_c$qc′​ associated with the chunk. Specifically, $\mathbf{q}'_c$qc′​ induces an attention distribution over the keys in chunk $c$c; $\mathbf{k}'_c$kc′​ is the weighted sum of those keys, and $b'_c$bc′​ is the entropy of that distribution. The entropy term adaptively interpolates between the uniform and dominant‑logit regimes of the log‑sum‑exp. To instantiate $\mathbf{q}'_c$qc′​, the authors append a special landmark token to the end of each chunk and use its query vector. The resulting pair $(\mathbf{k}'_c, b'_c)$(kc′​,bc′​) serves as an entropy‑calibrated compressed key. Chunk selection then proceeds by ranking the surrogate scores $\hat{s}_{i,c}$s^i,c​ and keeping the top‑$K$K chunks, with the partition function estimated as $\hat{\mathcal{Z}}_i = \sum_{c \in \mathcal{I}_i} \exp(\hat{s}_{i,c}) + Z_{i,\mathrm{swa}}$Z^i​=∑c∈Ii​​exp(s^i,c​)+Zi,swa​. Because the routing cost is $O(N/S)$O(N/S) per token, the total complexity remains $O(N^2/S)$O(N2/S).

To make the landmark queries learnable, the authors factorize the attention mass into an intra‑chunk normalized term and an inter‑chunk mass term. For a selected token $j$j, the attention weight is rewritten as

where the second factor is the learned surrogate mass. This hierarchical softmax factorization allows gradients from the language modeling loss to flow into the landmark parameters, so that chunks more useful for prediction receive larger mass. Empirically, this end‑to‑end learned surrogate outperforms naive BSA.

Several practical design choices are introduced to make HiLS‑Attention effective in modern architectures. First, the authors find that standard RoPE degrades perplexity under sparse attention; they adopt HoPE, which retains only RoPE dimensions whose rotation periods do not exceed the pre‑training context length and replaces the rest with NoPE. Second, a low‑rank query calibration module (Q‑Cal) adapts the original query $\mathbf{q}_i$qi​ for chunk‑level scoring by adding a learned correction $\Delta\mathbf{q}_i = \mathbf{W}^{\mathrm{up}}\mathbf{W}^{\mathrm{down}}\mathbf{h}_i$Δqi​=WupWdownhi​, where $\mathbf{h}_i$hi​ is the hidden state and the projections have rank $r \ll d_{\mathrm{model}}$r≪dmodel​. This significantly improves perplexity and length extrapolation. Third, for models using grouped‑query attention (GQA), query heads within the same group must share the same retrieved chunk set for kernel efficiency. To preserve head‑level flexibility, the chunk weights are computed per head, then aggregated by taking the maximum over heads in the group, and the top‑$K$K chunks are selected from the group‑level scores. An alternative, landmark‑free variant uses a single shared learnable query per layer, which achieves comparable in‑domain performance but weaker extrapolation.

On the hardware side, a naive sparse attention kernel suffers from poor Tensor Core utilization when processing one query token at a time. HiLS‑Attention instead groups $M$M adjacent query tokens and computes attention over the union of their selected chunks. Because adjacent queries often retrieve highly overlapping chunks, this batching changes the computation shape to $(M \times G, d) \times (d, S)$(M×G,d)×(d,S), where $G$G is the query‑head count per GQA group. Efficient Tensor Core utilization then only requires $M \times G \ge 16$M×G≥16, relaxing the dependency on a large $G$G, and reduces redundant memory access by reusing loaded key/value chunks.

For training, two continuous‑training strategies are provided. In landmark‑token tuning, all base model parameters are frozen and only the landmark token embeddings and the low‑rank query adapter matrices are trained (less than $1\%$1% of parameters), requiring only a few billion tokens to reach performance comparable to the base model. Full‑parameter tuning jointly updates all parameters, initializing landmark embeddings and the adapter randomly while inheriting the rest from the base model; this strategy is especially effective when combined with HoPE positional encoding, maximizing length generalization.

Experiment

Small-scale experiments on 345M-parameter models show that HiLS-Attention matches full attention perplexity and in-context retrieval while outperforming other sparse attention methods, with compression improving retrieval and landmark tokens proving essential for length extrapolation. Scaling to 1.4B-parameter training from scratch confirms that HiLS-Attention preserves short-context performance and gives stronger extrapolation, and 7B continued pre-training yields comparable or better LongBench results. Ablations demonstrate the importance of HoPE positional encoding, low-rank query calibration, and landmark tokens. Inference efficiency benchmarks reveal that HiLS-Attention becomes substantially faster than full attention beyond 16K context length, with prefill and decoding latency scaling sublinearly, and chunk overlap analysis validates the kernel’s adjacent-query reuse strategy.

In the 345M-parameter small-scale study, HiLS-Attention achieves the best or second-best perplexity at every evaluated context length, from 64 to 512K tokens, demonstrating strong long-context extrapolation. Its performance is consistently close to full attention, while other sparse methods often degrade at longer contexts. Full attention with RoPE surpasses the HoPE variant at short contexts, and HSA-UltraLong's perplexity becomes dominated by the sliding window component, limiting its long-context effectiveness. HiLS-Attention attains the best or second-best perplexity across all context lengths (64 to 512K), outperforming all other sparse attention baselines. Full attention with RoPE yields lower perplexity than full attention with HoPE at short contexts (up to 512 tokens), making it a stronger baseline for tasks with short sequences. HSA-UltraLong's perplexity becomes dominated by the sliding window component, as indicated by its performance approaching that of the standalone sliding window attention.

At the 8K training length, full attention with RoPE or HoPE achieves perfect needle retrieval, while all sparse methods except HiLS-Attention fall short. Applying HoPE position embeddings enables full attention to extrapolate to 16K and 32K, but most sparse methods fail to generalize beyond 8K. HiLS-Attention is the only native sparse approach that matches full attention's in-domain accuracy and also demonstrates strong length extrapolation when combined with HoPE. Full-Attn RoPE and Full-Attn HoPE both reach 100% single-needle retrieval (S-N) at 8K, whereas sparse baselines like NSA (49%) and SWA (18%) lag significantly. Dash-Attention-RoPE achieves 92% S-N at 8K but collapses to 0% at 16K, illustrating brittleness in extrapolation. Full-Attn HoPE retains 99% S-N and 97% multi-key multi-query (MK-MQ) at 16K, partially preserving performance at 32K before failing at 128K. SWA-RoPE and NSA-RoPE show near-zero retrieval across all extended lengths, confirming that sparsity without improved position handling severely limits long-context capability.

HiLS-Attn-HoPE maintains strong S-N and MK-MQ scores as context grows from 1M to 4M tokens, with S-N falling only from 100 to 96 and MK-MQ from 97 to 89. The VT score declines more substantially, dropping from 53 to 43 over the same range. S-N remains near 100 across all context lengths, decreasing just 4 points from 1M to 4M. MK-MQ dips to 87 at 2M before recovering to 89 at 4M, while VT steadily drops from 53 at 1M to 43 at 4M.

When extended to 256K training, HiLS-Attention with HoPE consistently achieves lower perplexity than full attention at long evaluation lengths, from 256K up to 1M tokens. The sparse method remains stable with the original RoPE base, whereas full attention with HoPE collapses at 512K. Enlarging the RoPE base helps full attention but HiLS-Attention still holds a slight advantage, using only 0.6% extra parameters. HiLS-Attention with HoPE yields lower perplexity than full attention at evaluation lengths of 256K and beyond, with the gap reaching 0.53 at 1M tokens. Full-Attention HoPE with the original RoPE base breaks down beyond 256K, exceeding 100 perplexity at 512K, while the HiLS variant remains at 7.55. Enlarging the RoPE base from 10^4 to 10^7 reduces full attention perplexity at 256K from 7.87 to 7.49, but HiLS-Attn-HoPE still edges it out (7.45). The extra parameter cost of HiLS-Attention is only 0.6% compared to full attention.

Continuing training of 345M models with a 256K context length shows that HiLS-Attn-HoPE with a large RoPE theta (1e7) maintains near-perfect needle-in-haystack retrieval across all lengths up to 1M tokens, while the full-attention RoPE baseline degrades sharply beyond 256K. HiLS-Attn-HoPE also achieves much higher variable tracking accuracy than the baseline, sustaining scores above 50 at every context length. Full-attention models with a smaller RoPE theta (1e4) fail to generalize at all, underscoring the importance of the combined hierarchical attention and positional encoding scaling. HiLS-Attn-HoPE retains near-perfect S-NIAH (needle-in-haystack) scores from 8K to 1M tokens, whereas full-attention RoPE drops to 41 at 512K and 2 at 1M. Variable tracking (VT) accuracy with HiLS-Attn-HoPE is consistently above 50 across all context lengths, compared to near-zero values for full-attention RoPE at long contexts. When the RoPE theta is reduced to 1e4, full-attention models fail even at short lengths, but HiLS-Attn-HoPE with the larger theta succeeds, highlighting the synergy between hierarchical attention and RoPE scaling. Multi-key multi-query (MK-MQ) accuracy remains low for both methods, yet HiLS-Attn-HoPE shows a modest improvement at 512K and 1M, rising from 1 and 0 to 13 and 5.

The experiments evaluate HiLS-Attention on language modeling and needle-in-haystack retrieval across context lengths up to 4M tokens. HiLS-Attention consistently matches or surpasses full attention in perplexity and retrieval accuracy, especially when combined with HoPE position embeddings and a scaled RoPE base, while other sparse methods degrade severely at long contexts. It exhibits strong length extrapolation, maintaining near-perfect retrieval up to 1M tokens and lower perplexity than full attention at 256K and beyond, with only a 0.6% parameter overhead, highlighting the synergy between hierarchical attention and positional encoding scaling.