HiLS-Attention: Mathematically Correct Sparse Attention with 13‑15× Faster Inference and 512× Extrapolation

Background

Large language models need to process much longer contexts for agents, deep reasoning, and massive data integration, but standard full attention scales quadratically with sequence length. This causes three major problems: quadratic compute cost, severe length‑extrapolation drop‑off, and linear KV‑cache growth that quickly exhausts memory.

Why existing chunk‑wise sparse attention fails

Chunk‑wise sparse attention selects the top‑K most relevant chunks for each query, but current methods estimate chunk importance using mean or max logits of token scores. Mean pooling is accurate only when attention is uniformly distributed across a chunk; max logits are accurate only when a single token dominates. In realistic scenarios the logit distribution varies with query, head, and data, leading to systematic mis‑ranking of chunks. Needle‑in‑a‑haystack experiments show that mean‑pooling methods (NSA, DashAttention, InfLLM v2) lose performance when important tokens are concentrated, while max‑logit methods lose performance when attention is spread.

Research questions

Develop a mathematically expressive chunk‑importance estimator.

Make the chunk summary trainable end‑to‑end with the language‑model loss.

HiLS‑Attention design

HiLS‑Attention (Hierarchical Landmark Sparse Attention) solves the two questions by:

Deriving a scoring function from a first‑order Taylor expansion of LogSumExp. The score combines (i) an attention‑weighted sum of keys (the “summary key”) and (ii) an entropy‑based bias term that automatically interpolates between uniform and peaked distributions.

Appending a learnable “landmark” token to the end of each chunk; this token aggregates information via intra‑chunk attention and serves as the summary key.

Factorising attention weights into two softmax layers: a chunk‑level softmax that allocates mass to selected chunks based on the differentiable scores, and an intra‑chunk softmax that distributes mass among tokens inside each selected chunk.

Because the chunk scores appear directly in the forward attention weights, gradients from the language‑model loss flow back to the summary keys and landmark tokens, eliminating the non‑differentiable top‑K bottleneck.

Experimental validation

Systematic experiments on models ranging from 345 M to 7 B parameters show:

Short‑text language‑model perplexity matches full attention across all context lengths, with a slight gain at 8 K tokens.

Training on 8 K tokens and extrapolating to 4 M tokens (512× longer) retains >90 % needle‑in‑a‑haystack accuracy, far surpassing full attention.

At 512 K context, pre‑fill is 13.5× faster and single‑step decode 15.7× faster than full attention.

Converting a full‑attention model (e.g., OLMo 3‑7 B) to HiLS with only 50 B additional tokens preserves short‑task performance; on LongBench the long‑context benchmarks match or exceed full‑attention baselines, and out‑of‑domain length tests beat YaRN and other baselines.

On long‑context retrieval tasks HiLS outperforms full attention, likely because compressing keys into summary tokens reduces noise from irrelevant tokens while preserving shared semantic signals.

Conclusion

HiLS‑Attention demonstrates that sparse attention can simultaneously improve efficiency and effectiveness. The entropy‑biased scoring unifies mean and max logits, and the hierarchical softmax makes chunk selection end‑to‑end differentiable, leading to faster inference, massive context extrapolation, and superior retrieval performance.

Paper: https://arxiv.org/pdf/2607.02980

Code: https://github.com/Tencent-Hunyuan/HiLS-Attention