HiLS-Attention replaces attention’s mean-pooled chunk summaries with a learnable landmark-token summary that participates in the forward softmax, so the language-modeling loss trains chunk selection end-to-end. Trained at 8K, it keeps retrieval accuracy above 90% at 512× that length.
You’re shipping a chatbot or coding agent whose context balloons past 128K tokens. Full attention gets quadratically slow and the KV cache blows up your GPU memory. The usual escape hatch is Chunk-wise sparse attention: split history into chunks, pick the top few relevant ones per query, ignore the rest. The dominant recent example is Native Sparse Attention (NSA) from DeepSeek, which scores chunks by mean-pooling their token keys. The catch: existing sparse variants consistently lose to full attention on retrieval-heavy long-context tasks. This paper argues the reason is a specific, fixable design flaw in how chunks are summarized and selected.
Start with the intuition. When you decide whether a chunk is relevant, you want to estimate the total attention mass it would get under full attention. Mean-pooling the chunk’s keys assumes token logits inside the chunk are roughly uniform. But long-context retrieval is the opposite: one or two “needle” tokens dominate and everything else is noise. Mean-pooling averages the signal away.
HiLS fixes this in two coupled moves. First, each chunk gets a special Landmark token appended, whose query vector attends over the chunk’s own keys to produce a compressed summary key plus an entropy bias term. The authors show this summary is the first-order Taylor expansion of the true LogSumExp chunk mass, so it interpolates between the mean-logit and max-logit regimes automatically. Second, the surrogate score doesn’t just pick top-K chunks and get discarded. It multiplies into the final attention weights via a hierarchical softmax: inter-chunk mass first, intra-chunk distribution second. Because the surrogate is in the forward pass, gradients from next-token prediction flow back into the landmark representation, so selection is trained end-to-end.
for each query q_i:
# score every chunk using its landmark summary + entropy bias
scores = [q_i @ k_c_prime / sqrt(d) + b_c_prime for c in chunks]
top_K = argtopk(scores, K)
# hierarchical softmax: chunk-level mass * intra-chunk weights
for c in top_K:
w_intra = softmax(q_i @ K_c / sqrt(d))
w_inter = exp(scores[c]) / Z_hat_i # surrogate mass, trainable
out += w_inter * (w_intra @ V_c)
out += sliding_window_attention(q_i)
A small Low-Rank Query Calibration (Q-Cal) module adds a residual to the query specifically for chunk-level scoring, since a token-level query isn’t ideal for scoring a compressed summary. Positional encoding is HoPE positional encoding rather than plain RoPE, which matters because rotating all dimensions of aggregated keys distorts the chunk summary.
The prevailing sparse-attention recipe treats chunk selection as a separate retrieval step: summarize chunks with mean pooling, hard-select top-K, then do attention on the survivors. Because the selection scores never appear in the loss, gradients can’t teach the model which chunks matter. HiLS does the opposite: it splices the chunk-selection score into the attention softmax itself, so “which chunk is worth reading” is trained by the same next-token loss that trains everything else. The evidence that this is the load-bearing change is the ablation where removing Proposition 3.1 (LogSumExp linearization) (the entropy bias, keeping only the raw landmark key as summary) badly degrades extrapolation, and the fact that mean-pooling variants like Native Sparse Attention (NSA) fail perfect needle retrieval even inside their training length.
The finding that makes the thesis true: on the 345M model trained at only 8K context, HiLS holds >90% needle-in-a-haystack retrieval out to 4M tokens, a 512× extrapolation. Every competing sparse method the authors ran, including Native Sparse Attention (NSA), DashAttention, and InfLLM v2, collapses to near zero on multi-key retrieval well before that, and even Naive Block Sparse Attention (which computes full attention to pick chunks, so it should be an oracle for selection) is beaten by HiLS at long context. That gap between naive-BSA and HiLS is the paper’s cleanest signal that end-to-end selection isn’t just imitating full attention’s chunk picks, it’s learning better ones by suppressing noise from irrelevant tokens.
•
On the RULER variable-tracking task at 32K, HiLS scores 68 vs 0 for full attention with RoPE, and even beats full attention with HoPE (which gets 11).
•
Converting a 7B OLMo 3 checkpoint to HiLS with 50B tokens of continued pretraining preserves short-context scores on MMLU, ARC, GSM8K, etc., while lifting RULER-128K from 0 to 94.7 and beating a YaRN-extended baseline on LongBench (33.2 vs 31.7).
•
Inference: at 512K context, HiLS is 13.5× faster prefill and 15.7× faster per decode step than full attention on an H800. Crossover with full attention is around 16K tokens.
Reach for HiLS when you’re building a long-context retrieval or agent system on your own base model and can afford a modest continued-pretraining budget. The paper’s most practical recipe is the 50B-token conversion of a full-attention 7B: replace the full-attention layers with HiLS, shrink the sliding window to 512, keep chunk size 64 and top-K 32 (so each query sees 2K retrieved tokens plus the local window). You get near-constant per-token decode cost and retrieval quality that survives well past your training length. There’s also a landmark-token-only tuning mode that freezes the base and updates <1% of parameters, useful for cheap experimentation but with weaker extrapolation.
Code is at Tencent-Hunyuan/HiLS-Attention. The paper builds directly on the OLMo 3 7B checkpoint and its long-context data mix, both released by AI2, so the full conversion pipeline is reproducible. No new benchmark is released; evaluation reuses RULER and LongBench.
If you want the LM loss to train chunk selection, put the selection score inside the softmax, not next to it. Hard top-K over a detached scoring head is a dead end for sparse attention because gradients never reach the thing that decides what the model reads.
•
The strong extrapolation numbers depend on using HoPE positional encoding positional encoding, not vanilla RoPE. HiLS on plain RoPE fails past the training length in the ablations. If your base model is deeply committed to a different positional scheme, migration cost is real.
•
Below roughly 16K context, HiLS is slower than full attention because the top-K routing overhead isn’t amortized. This is a long-context tool, not a general drop-in replacement.
•
Small-model RULER is a mostly find-and-copy task. The paper explicitly argues large models can “compensate” for retrieval errors, which is honest but also means the 512× extrapolation headline may look smaller on messier real workloads. The 7B LongBench gap over the YaRN baseline is ~1.5 points, not the order-of-magnitude gap the small-model plots suggest.