PaperlayerView exact revision on arXiv

Source-linked AI summary

Hierarchical Sparse Attention Done Right: Toward Infinite Context Modeling

Xiang Hu, Xinyu Wei, Hao Gu, Minshen Zhang, Tian Liang, Huayang Li, Lei Zhu, Yan Wang, Sirui Han, Yushi Bai, Kewei Tu, Haitao Mi, Leo Liang

arXiv:2607.02980v1cs.CLcs.AI

TL;DR

Long-context LLMs face quadratic full-attention costs and poor length extrapolation, while existing sparse methods struggle with inaccurate chunk selection. HiLS-Attention learns sparse retrieval end-to-end and achieves strong in-domain performance with 512× context-length extrapolation at over 90% retrieval accuracy.

  • Problem

    Full attention makes long-context modeling difficult because its computation and KV-cache costs grow quadratically, while existing sparse methods select chunks inaccurately.

  • Method

    HiLS-Attention uses hierarchical softmax to enable end-to-end trainable sparse retrieval.

  • Results

    HiLS-Attention matches full-attention perplexity, improves in-domain RULER performance, and extrapolates from 8K to 4M contexts with over 90% retrieval accuracy.

  • Takeaways & Limitations

    HiLS-Attention offers a promising alternative to full attention, combining strong in-domain performance, superior length extrapolation, and more efficient long-context inference.

  • Takeaways & Limitations

    Existing chunk-selection methods can misrepresent chunk mass because mean and max logits fail across varying logit distributions, potentially missing important chunks.

Abstract

from arXiv · show

Scaling modern large language models (LLMs) to long contexts is limited by the quadratic computation cost, and poor length extrapolation of dense attention. Chunk-wise sparse attention offers a promising alternative, but all existing methods fall short of full attention because of their inaccurate chunk selection. We propose Hierarchical Landmark Sparse (HiLS) Attention, a chunk-wise sparse attention mechanism that learns chunk selection end-to-end under the language-modeling (LM) loss. HiLS factorizes attention hierarchically: each query performs attention independently with each retrieved chunk to extract chunk-specific information, and the resulting outputs are fused according to chunk retrieval scores. By incorporating retrieval scores into the forward attention computation, HiLS optimizes them directly with the LM loss, enabling end-to-end retrieval learning and native sparse training. Experimental results show that HiLS-Attention achieves performance comparable to, and in some cases better than, full attention at in-domain context lengths. Meanwhile, HiLS-Attention extrapolates more than $64\times$ the training context length with 90% retrieval accuracy, far beyond full attention. Moreover, existing full-attention models can be converted to HiLS-Attention with lightweight continued pretraining, preserving in-domain performance while acquiring ultra-long-context extrapolation. Together with its sparse KV access and computation, HiLS-Attention breaks the usual efficiency-performance trade-off, enabling long-context LLMs that are both more efficient and more effective on general long-context tasks than their full-attention counterparts.

1 Introduction

HiLS-Attention addresses the quadratic cost and poor length extrapolation of full attention by enabling end-to-end trainable chunk retrieval through hierarchical sparse attention. Experiments indicate comparable short-context performance, stronger long-context results, and more efficient inference than full-attention alternatives.

  • Motivation: Full attention remains difficult to scale because its computation grows quadratically with context length and extrapolates poorly.
  • Motivation: Existing chunk-wise sparse methods reduce computation and memory demands but have not matched full-attention performance because of inaccurate chunk selection.These methods selectively attend to relevant chunks and dynamically swap their KV caches into fast memory.
  • Method: HiLS-Attention uses hierarchical softmax to make sparse retrieval trainable end-to-end under the LM loss.Its design is derived from naive BSA and uses landmark-token summaries aligned with the first-order Taylor expansion of full-attention-induced chunk mass.
  • Results: Experiments across 345M-to-7B models show comparable perplexity to full attention at 345M and better in-domain RULER performance.
  • Results: Full-attention models can be cost-effectively converted to HiLS-Attention while preserving short-context performance and surpassing full attention on in-domain and ultra-long-context tasks.The authors position this combination of performance, length extrapolation, and efficient inference as evidence for native sparse attention as an alternative to full attention.

2 Preliminary

Naive Block Sparse Attention combines a local sliding window with K selected distant chunks, choosing chunks by exact attention mass. However, exact mass computation removes sparsity benefits, while mean- or max-based approximations can distort chunk rankings and miss important chunks.

  • BSA lets each query attend to its local sliding window and K globally selected distant chunks.Chunks are uniformly partitioned, and the local window is aligned to chunk boundaries to avoid overlap with retrieved chunks.
  • Naive BSA selects the top-K chunks by exact chunk mass, then normalizes attention only over selected tokens.Unselected tokens receive zero attention weight.
  • Exact chunk selection requires evaluating all token-level logits within every chunk, diminishing sparse-attention advantages during training.An ideal sparse mechanism would estimate chunk mass from chunk-level representations without computing all token-to-token scores.
  • Neither mean nor max logits consistently represent the LogSumExp-based chunk mass, potentially changing rankings and missing important chunks.Max logits are accurate only when a single token dominates, whereas mean logits work only for nearly uniform logits.
  • Mean-pooled keys produce mean token-level logits, which represent chunk mass faithfully only when logits are nearly uniform.This approximation is therefore query-, head-, and data-dependent.

3 Methodology

HiLS-Attention replaces non-parametric chunk summaries with learnable landmark-based summaries and integrates their surrogate masses directly into hierarchical attention. This enables end-to-end chunk selection under the language-modeling loss while retaining sparse computation.

  • Method overview: HiLS-Attention uses learnable chunk summaries to replace the non-parametric chunk mass Z_i,c in differentiable sparse attention.The surrogate mass is designed to make chunk selection trainable under the language-modeling objective.
  • Linear surrogate for chunk mass: The LogSumExp chunk mass is approximated by an affine surrogate combining query–summary relevance with an entropy-based bias term.The learned query induces a distribution over chunk keys, whose weighted key sum forms the compact summary; the entropy bias approaches log S for uniform scores and 0 when one score dominates.
  • Linear surrogate for chunk mass: Computing the surrogate costs O(S) per chunk and O(N) across a sequence of length N partitioned into N/S chunks.A landmark token appended to each chunk supplies the learnable query used to derive its compressed key representation.
  • Sparse training: Native sparse training routes each query over all N/S chunk summaries, retains the top-K chunks, and attends only to constant selected tokens.Routing costs O(N/S) per token and O(N 2/S) per sequence, leaving routing as the only quadratic term.
  • Hierarchical softmax: Hierarchical softmax factorizes attention into intra-chunk normalization and inter-chunk surrogate mass, allowing LM-loss gradients to train landmark representations and chunk retrieval.The model aggregates information within selected chunks before fusing chunk-level outputs according to learned surrogate masses, which empirically outperform naive BSA.

4 Practical Design Choices

HiLS design choices span positional encoding, query calibration, GQA-compatible retrieval, hardware-software co-design, and continued-training strategies. These choices target perplexity, length extrapolation, sparse-kernel efficiency, and preservation of base-model capabilities.

  • Architecture: HoPE positional encoding lets HiLS outperform full attention in perplexity at 8K training context, whereas standard RoPE performs worse.HoPE retains RoPE dimensions whose rotation periods do not exceed the pre-training context length and replaces the remaining dimensions.
  • Architecture: Q-Cal uses lightweight low-rank query adaptation to calibrate chunk-level surrogate scores, significantly improving perplexity and length extrapolation.The projection weights satisfy r ≪ dmodel.
  • Architecture: In GQA, query heads within each group share one retrieved chunk set, allowing selected tokens to be gathered once and processed efficiently.Unlike MHA, where each query head independently selects top-K chunks, GQA requires group-consistent chunk selection.
  • Kernel design: 80% top-k overlap among adjacent query tokens motivates batching M tokens over their union of chunks, changing Tensor Core operations to (M × G, d) × (d, S).This requires M × G ≥16 for efficient Tensor Core utilization and enables reuse of loaded K/V chunks, reducing redundant memory access.
  • Continued training: Less than 1% of parameters are trainable in landmark-token tuning, and no more than 5B training tokens suffice for performance comparable to the base model.Full-parameter tuning jointly updates all parameters and is particularly effective when replacing positional encoding with HoPE.

5 Small-scale Studies

Small-scale studies show that HiLS-Attention matches or exceeds full and sparse baselines in perplexity and retrieval, especially at long contexts, while ablations identify HoPE, query augmentation/calibration, and landmark tokens as important components. Its gains are linked to accurate sparse retrieval and compression that improves representation quality and extrapolation.

  • Experimental setup: Models use 345M parameters, train from scratch at 8K context, and give sparse variants a 2K activation budget plus a 512-token local window.Sparse baselines share 512-token windows, 64-token chunks, and top-K 32 for comparable attention budgets.
  • Retrieval and perplexity results: Only Naive BSA and HiLS-Attention remain comparable to full attention on in-domain Single-NIAH, while mean-pooled sparse methods degrade.The result supports the claim that mean-pooled chunk summaries have limited expressiveness for retrieval.
  • Retrieval and perplexity results: 4.94 v.s. 4.94 at 8K: HiLS-Attention matches naive BSA at training length but performs better through 256K interpolation and beyond 256K extrapolation.HiLS-Attention also significantly leads naive BSA on RULER within and beyond the training context length.
  • Retrieval and perplexity results: HiLS-Attention surpasses naive BSA and full attention on Variable Tracking, suggesting compression improves multi-hop retrieval by canceling noise while preserving shared semantic signals.The proposed explanation is that aggregating keys reduces less-aligned noise and retains common semantic content.
  • Long-context training: At 256K evaluation after 256K training, HiLS-Attention consistently achieves lower perplexity than full attention across settings and retains stronger retrieval despite activating at most 2K tokens.HiLS-Attn-HoPE also outperforms Full-Attn HoPE/RoPE under the same 256K training setting.
  • Ablations: HoPE positional encoding, query augmentation, low-rank query calibration, and landmark tokens contribute to performance, while removing calibration or landmark tokens harms extrapolation.Overly expanding calibration rank to r=128 also impairs length extrapolation, and HoPE outperforms RoPE and NoPE.

6 Large-scale Experiments

Large-scale experiments show that HiLS-Attention preserves full-attention short-context modeling while maintaining stable length extrapolation, and scales to 7B models with comparable or slightly better general-task performance. It also provides substantial long-context inference speedups at 345M scale.

  • 1.4B-scale experiments: A 1.4B-parameter model trained from scratch on 300B tokens compares native HiLS-Attention against full attention with RoPE across perplexity, extrapolation, and downstream tasks.Full attention with RoPE is used as the stronger baseline because it achieves lower 512-token perplexity than full attention with HoPE.
  • 1.4B-scale experiments: HiLS-Attention and full attention achieve almost identical perplexity across context lengths and training stages, especially near 8K, while HiLS extrapolation remains stable throughout training.The result indicates that native sparse training preserves short-context modeling ability comparable to full attention without extrapolation decay.
  • 7B-scale experiments: The 7B setup replaces OLMo3-1025-7B’s full-attention layers with HiLS-Attn, reduces sliding windows from 4K to 512 tokens, and preserves the 3:1 layer pattern.Long-range modeling is shifted to the HiLS retrieval branch.
  • Inference efficiency: 13.5×/15.7× faster prefill/decode at 512K: HiLS-Attention reaches parity near 16K and then substantially outpaces full attention.This measurement uses a 345M model on one NVIDIA H800 with batch size 1, bf16, chunk size 64, top-k = 32, and a 512-token sliding window.

7 Analysis

HiLS-Attention’s latency depends on a fixed retrieval budget rather than full context length, yielding more favorable scaling than full attention. Strong chunk sharing among adjacent queries further supports its one-load-multiple-compute strategy.

  • Inference latency: HiLS-Attention’s prefill latency grows near-linearly with context length, while full attention grows quadratically.Its cost is governed by K ×chunk = 2048 tokens plus the local sliding window, rather than full context length.
  • Inference latency: O(1) per-token decoding latency remains effectively constant for HiLS-Attention, while full attention grows linearly with KV-cache length.The latency curves cross over at roughly 16K tokens.
  • Adjacent-query chunk sharing: 92.8% of chunks in an adjacent-query block already appeared in the previous block on average, supporting one-load-multiple-compute.The evaluation used Olmo3-7B with chunk size 64, top-K = 32, and M = 16 for tail-block loading from 4K to 64K contexts.
  • Adjacent-query chunk sharing: 76.4% to 9.9%: the loaded fraction fell as context increased from 4K to 64K for the final M = 16 queries.The visible historical pool grew from 57 to 1032 chunks, while the loaded union grew only from 43.6 to 102.1 chunks.

8 Related Work

Related block-wise sparse-attention methods differ in how they summarize chunks for selection, but existing approaches face inaccurate retrieval, dense-training requirements, or architectural and perplexity limitations. HiLS addresses these issues through native sparse training and more accurate chunk-mass estimation while supporting ultra-long-context extrapolation.

  • Block-wise sparse attention: Block-wise sparse attention is emphasized because it can support native sparse training, and methods are distinguished by whether chunk summaries are parameterized.Chunk summaries act as proxies for chunk relevance, so their fidelity affects selection quality.
  • Non-parametric summaries: Mean-pooled summaries in NSA and MoBA, along with DashAttention’s learnable gate, fail to achieve perfect needle retrieval even at training length.These methods therefore exhibit inaccurate chunk selection in the reported comparison.
  • Parametric summaries: Landmark Attention learns chunk summaries but requires dense attention during training, sparsifies only at inference, and has substantially higher perplexity than full attention.Its landmark-token key summary makes it difficult to match full attention on downstream tasks.
  • Parametric summaries: HSA enables sparse training through specialized kernels but adds substantial parameter overhead through SWA cross-attention, whose perplexity eventually dominates after prolonged training.The reported limitation is illustrated by Table 1 in the supplied passage.
  • Proposed method: Compared with LMK-Attn, the proposed method enables native sparse training, estimates chunk mass more accurately, and achieves better perplexity and extrapolation.It adopts HoPE to preserve in-domain positional awareness while enabling ultra-long-context extrapolation.

9 Discussion & Conclusion … D Hyper-parameters

The discussion attributes HiLS-Attention’s strong length extrapolation to more accurate in-context retrieval and outlines its sparse path toward ultra-long-context training, while noting unresolved limitations. The appendices provide mathematical justifications, a GQA formulation, and hyper-parameter material.

  • 9 Discussion & Conclusion: HiLS-Attention combines end-to-end retrieval learning with native sparse training and exhibits strong length extrapolation ability.
  • 9 Discussion & Conclusion: More accurate in-context retrieval is hypothesized to drive extrapolation, aided by compression that suppresses token noise and preserves shared semantic signals.
  • 9 Discussion & Conclusion: HiLS enables ultra-long-context training by compressing chunks into lightweight summaries before retrieval, avoiding the quadratic token-pair scoring of full attention and distillation methods.Full-attention and distillation-based methods score token-level pairs with quadratic complexity O(L^2).
  • 9 Discussion & Conclusion: O(LKS) total attention cost remains after retrieval, while approximate nearest-neighbor search can reduce retrieval to approximately O(L log(L/S)).Selected chunks and a local window restrict token-level attention to per-token cost O(KS).
  • 9 Discussion & Conclusion: HiLS-Attention lacks context parallelism, leaving effectiveness at larger training context lengths insufficiently validated, and unselected chunks receive no gradient updates.
  • 9 Discussion & Conclusion: The mechanism by which Q-Cal improves extrapolation and in-domain performance remains incompletely understood, motivating further validation and improvement.
  • A Justification of Equation 5: The appendices justify Equation 5 through near-uniform and dominant-logit cases, and derive the Log-Sum-Exp approximation using a first-order Taylor expansion.
  • B Proof of Proposition 3.1; C HiLS-Atention in GQA; D Hyper-parameters: Appendix B states that the final objective matches Equation 7, while Appendix C addresses HiLS-Attention in GQA and Appendix D contains hyper-parameter material.

E Training recipes

HiLS models use a standardized small-scale training setup based on OLMo/Dolma data augmented with RULER-style examples, while 7B adaptation uses continued pretraining on the OLMo3 corpus. The recipes specify context lengths, batch sizes, token budgets, optimizer schedules, and evaluation criteria.

  • Small-scale models: Small-scale models mix OLMo/Dolma pretraining corpora with 5% RULER-style synthetic examples and train at 8K context length.The synthetic data supports instruction following for RULER-style retrieval queries.
  • Small-scale models: Small-scale models use global batch size 128, 30K optimizer steps, approximately 30B training tokens, and AdamW with a constant 3 × 10−4 learning rate.Unless otherwise specified, models share tokenizer, data, training length, and optimization settings.
  • Small-scale models: Small-scale evaluation measures language-modeling ability with perplexity and long-context random access with RULER.The passage introduces these as the two evaluation perspectives for small-scale models.
  • 7B continue pre-training recipe: 7B continued pretraining samples approximately 50B tokens from the tokenized OLMo3 500B-token pool using 8K sequences, with maximum sequence length 8192.Adaptation draws sequences from the OLMo3 pretraining pool.
  • 7B continue pre-training recipe: The 7B recipe uses global batch size 512 for 13K optimizer steps, warms up to 2 × 10−4 for 1K steps, then cosine-declays to 2 × 10−5.The schedule specifies both the warmup duration and the final learning rate.

F Downstream Evaluation Details · G Perplexity at Different CPT Steps

The downstream evaluation covers a broad suite of language-understanding, reasoning, and code-generation benchmarks. Across CPT checkpoints, HiLS-Attn’s perplexity gap versus OLMo3-512swa-CPT narrows with training and remains nearly unchanged at 8K context.

  • F Downstream Evaluation Details: The evaluation includes LAMBDA, HellaSwag, PIQA, WinoGrande, OpenBookQA, ARC-challenge, ARC-esay, MMLU, GPQA, BoolQ, RACE, CMATH, GSM8K, and CRUX.
  • F Downstream Evaluation Details: Code-generation evaluation uses HumanEval+ and MBPP+.
  • G Perplexity at Different CPT Steps: At 128 tokens and 6K CPT steps, the signed perplexity gap is +0.23, indicating higher perplexity for HiLS-Attn than OLMo3-512swa-CPT.
  • G Perplexity at Different CPT Steps: By 10K–13K CPT steps, the 128-token perplexity gap shrinks to +0.02–+0.06.
  • G Perplexity at Different CPT Steps: By 10K–13K CPT steps, the 512-token perplexity gap narrows to +0.02–+0.03.
  • G Perplexity at Different CPT Steps: At 8K context, the perplexity gap stays within ±0.01 throughout CPT.

H Per-Step Results of the 1.4B Model

This section reports the 1.4B model’s full per-step perplexity and RULER results across training steps during 300B-token training with an 8K context length.

  • Perplexity: The section provides full per-step perplexity results for the 1.4B model across training steps.These results are presented in Table 15 after 300B-token training with an 8K context length.
  • RULER: The section provides RULER results for the 1.4B model at different training steps during 300B-token training.These results are presented in Table 16.
Loading 2607.02980v1…