# Hierarchical Sparse Attention Done Right: Toward Infinite Context Modeling

> HiLS-Attention achieves 512x length extrapolation and 15x speedup via hierarchical softmax and learnable chunk summaries optimized end-to-end.

- **Source:** [arXiv](https://arxiv.org/abs/2607.02980)
- **Published:** 2026-07-09
- **Permalink:** https://picx.dev/p/i8mmMJ
- **Whiteboard:** https://picx.dev/p/i8mmMJ/image

## Summary

## 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**
| Model | 64 | 128 | 512 | 8K | 32K | 128K | 512K |
|---|---|---|---|---|---|---|---|
| Full-Attn RoPE | 33.92 | 26.89 | 18.68 | 4.96 | >10² | >10² | >10² |
| NSA-RoPE | 34.15 | 27.11 | 18.85 | 5.01 | 7.62 | 11.75 | 19.14 |
| HiLS-Attn-HoPE | **33.97** | **26.91** | **18.65** | **4.94** | **4.34** | **4.71** | **5.95** |

**Table 2: RULER average exact match (%)**
| Model | 8K | 16K | 32K | 128K | 512K | 1M | 2M | 4M |
|---|---|---|---|---|---|---|---|---|
| Full-Attn HoPE | 78 | 72 | 31 | 0 | 0 | – | – | – |
| HiLS-Attn-HoPE | **89** | **86** | **88** | **85** | **85** | **83** | **78** | **76** |

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)**
| Model | 8K | 32K | 128K | 256K | 512K | 1M |
|---|---|---|---|---|---|---|
| Full-Attn RoPE (θ=1e7) | 9.11 | 8.86 | 8.11 | 7.49 | 7.54 | 8.61 |
| HiLS-Attn-HoPE (θ=1e7) | **9.08** | **8.88** | **8.15** | **7.45** | **7.37** | **8.08** |

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)**
| Benchmark | Olmo3-Base | HiLS-Attn-HoPE |
|---|---|---|
| RULER (avg) | 3.75 | **97.42** |
| MMLU | 59.90 | 56.58 |
| GPQA | 29.29 | **34.34** |
| Average (all tasks) | 43.88 | **43.35** |

**Table 11: LongBench-v1 overall score**
| Method | Score |
|---|---|
| Olmo3-Base | 29.0 |
| Olmo3-512swa-CPT + YaRN | 31.7 |
| HiLS-Attn-HoPE | **33.2** |

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.

---

_Markdown view of https://picx.dev/p/i8mmMJ, served by PicX — AI-generated visual whiteboard summaries of research papers._