Summary

HiLS-Attention is a novel native sparse attention mechanism that enables ultra-long context modeling through end-to-end learnable chunk selection. Key contributions:

• Hierarchical factorization: Attention is decomposed into inter-chunk (chunk-level) and intra-chunk (token-level) softmax operations, making chunk selection differentiable under the language-modeling loss.
• Expressiveness-driven design: A learnable chunk summary derived from the first-order Taylor expansion of the full-attention chunk mass (Proposition 3.1), using landmark tokens and an entropy-calibrated bias term.
• Strong length extrapolation: Trained on only 8K context, HiLS extrapolates to 4M tokens (512×) with over 90% retrieval accuracy on needle-in-a-haystack tasks.
• Efficient inference: Constant per-token decoding latency and 13.5–15.7× speedup over full attention at 512K context on a single H800 GPU.
• Lightweight conversion: Full-attention models can be adapted with ≤50B continued-training tokens, preserving short-context performance while surpassing full attention on long-context benchmarks.

Introduction and Theoretical Foundation

Modern LLMs require long-context modeling but full attention suffers from quadratic complexity, poor length extrapolation, and growing KV cache costs. Chunk-wise sparse attention offers a constant cost per token, but existing methods (e.g., NSA, DashAttention, InfLLMv2) fail to match full attention because of inaccurate chunk selection. The core problem lies in using non-parametric chunk summaries (e.g., mean pooling) that have limited expressiveness and are not optimized end-to-end. The authors show that mean-pooled keys approximate the mean of token-level logits, but the true chunk mass $Z_{i,c} = \sum_{j \in T_c} \exp(s_{i,j})$ is a LogSumExp function whose behavior depends on the logit distribution:

$$\log Z_{i,c} \approx \begin{cases} \text{mean}_{j \in T_c}(s_{i,j}) + \log S, & \text{if logits are nearly uniform},\\ \text{max}_{j \in T_c}(s_{i,j}), & \text{if one logit dominates}. \end{cases} (\text{derived in Appendix A})$$

This mismatch causes important chunks to be missed, especially in long-context retrieval tasks. The paper motivates two desiderata: trainable chunk summaries that are both expressive and end-to-end optimized.

Methodology

1. Surrogate Chunk Mass (Proposition 3.1)

The LogSumExp chunk mass is linearized via a first-order Taylor expansion around a learned query $q'_c$:

$$\log \sum_{j \in T_c} \exp\left(\frac{q^\top k_j}{\sqrt{d}}\right) \approx \frac{q^\top k'_c}{\sqrt{d}} + b'_c \qquad (\text{Eq.~7})$$

where

$$k'_c = \sum_{j \in T_c} p_j k_j, \quad b'_c = -\sum_{j \in T_c} p_j \log p_j, \quad p_j = \frac{\exp((q'_c)^\top k_j / \sqrt{d})}{\sum_{k \in T_c} \exp((q'_c)^\top k_k / \sqrt{d})}. \qquad (\text{Eq.~8})$$

$k'_c$ is an attention-weighted key summary, and $b'_c$ is the entropy of the distribution, which interpolates between the two regimes of Eq. (5). $q'_c$ is obtained by appending a landmark token to each chunk.

2. Hierarchical Softmax

To make $q'_c$ learnable end-to-end, the attention weight is factorized as:

$$w_{i,j} = \frac{\exp(s_{i,j})}{Z_{i,c(j)}} \times \frac{\hat{Z}_{i,c(j)}}{\hat{Z}_i} \qquad (\text{Eq.~10})$$

where $\hat{Z}_{i,c} = \exp(\hat{s}_{i,c})$ with $\hat{s}_{i,c} = \frac{\hat{q}_i^\top k'_c}{\sqrt{d}} + b'_c$, and $\hat{Z}_i$ normalizes over selected chunks and the sliding window. This allows gradients from the LM loss to directly supervise chunk retrieval.

3. Practical Enhancements

• Low-Rank Query Calibration (Q-Cal): $\hat{q}_i = q_i + W_{\text{up}}W_{\text{down}}h_i$ to better align queries for chunk-level scoring.
• GQA adaptation: Maximum of chunk scores over heads in a group to select shared chunks.
• Hardware kernel: Groups $M$ adjacent queries and computes attention over the union of their selected chunks, achieving Tensor Core utilization with $M \times G \geq 16$.

Empirical Validation / Results

Small-scale (345M, 8K training)

Table 1: Perplexity across context lengths

Table 2: RULER average exact match (%)

HiLS is the only native sparse method achieving in-domain NIAH performance comparable to full attention, and it extrapolates to 4M with 90%+ retrieval accuracy.

Long-context scaling (345M, 256K training)

Table 3: Perplexity (256K training)

HiLS outperforms full attention even at 1M context, and achieves substantially better variable tracking (VT) scores.

7B continued pretraining (from Olmo3-7B)

Table 9: Downstream tasks (selected)

Table 11: LongBench-v1 overall score

HiLS-Attention surpasses strong baselines (including YaRN-extended full attention) on LongBench while maintaining comparable general performance.

Inference Efficiency

Figure 6 shows HiLS-Attention achieves parity with full attention at ~16K context and is 13.5×/15.7× faster (prefill/decode) at 512K context.

Theoretical and Practical Implications

• Theoretical: The paper establishes a formal connection between chunk-level attention mass (LogSumExp) and a linearizable surrogate (Eq. 7), showing that an effective chunk summary must capture both the attention-weighted key mean and the entropy of the intra-chunk distribution.
• Practical: HiLS-Attention breaks the usual efficiency-performance trade-off: it is both more efficient and more effective than full attention on long-context tasks. Its strong extrapolation ability (512× training length) opens a practical path toward infinite-context training by learning retrieval patterns in short contexts and generalizing them.
• Conversion: Full-attention models can be cost-effectively converted with ≤50B tokens and frozen base parameters, enabling seamless adoption.
• Hardware alignment: The query-packing kernel design makes HiLS efficient even with small GQA group sizes, overcoming limitations of prior sparse attention kernels.

Conclusion

HiLS-Attention introduces a hierarchical softmax + expressiveness-enhanced chunk summaries that enables native sparse training with end-to-end retrieval learning. Comprehensive experiments (345M to 7B) show it matches or surpasses full attention in in-domain performance, provides unprecedented length extrapolation (up to 512×), and yields substantial inference speedups. The approach is compatible with continued pretraining of existing models. Limitations include lack of context parallelism support and incomplete understanding of the Q-Cal mechanism. Future work includes scaling training to ultra-long contexts and improving theoretical understanding of the extrapolation source.