PaperlayerView exact revision on arXiv

Source-linked AI summary

Simplified Sparse Attention via Gist Tokens

Yuzhen Mao, Michael Y. Li, Emily B. Fox

arXiv:2604.20920v2cs.LG

TL;DR

Sparse attention can reduce long-context inference costs, but many variants add architectural components. SSA learns gist-token routing for selective unfolding without architectural changes and reports consistent gains over gist-compression and sparse-attention baselines.

  • Problem

    Sparse attention can reduce long-context inference costs, but many existing training-based approaches require additional architectural components or specialized mechanisms.

    Source: §1 · pp. 1–2

  • Method

    SSA trains interleaved gist tokens to summarize local chunks and route queries, then selectively unfolds the top-k relevant chunks by restoring their raw tokens.

    Source: §3 · p. 4, §3.4 · p. 7

  • Results

    Across LongBench and RAG benchmarks, SSA consistently gains over gist-compression and sparse-attention baselines from 8× to 32× compression; on RAG, it scores 33.68 versus 27.14 for Full-PT at 8×.

    Source: §4.3 · p. 12, §6 · p. 18

  • Takeaways & Limitations

    The hierarchical gist-of-gist variant achieves logarithmic per-step decoding complexity through coarse-to-fine selection.

    Source: §6 · p. 18

    The experiments use Qwen2 and Llama3.2-1B because the relevant baselines had been evaluated on those models.

    Source: §4.1 · p. 10

Abstract

from arXiv · show

Sparse attention can reduce the cost of long-context inference, but most variants introduce new architectural components. We introduce Simplified Sparse Attention (SSA), a simpler approach to sparse attention that requires no architectural changes. Concretely, we first perform continued pretraining on sequences interleaved with gist tokens. We optimize the standard next-token loss as usual, but the gist tokens use an attention mask to restrict what parts of the context the language model can attend to; this teaches the model to pack each chunk's important information into the gist tokens. At inference time, SSA scores chunks via attention between the current query and the small set of gist tokens, selectively unfolding the top-k chunks by reintroducing their corresponding raw tokens. Since the query is scored only against the gist tokens, we avoid the memory-bandwidth cost associated with naive scoring against the full KV cache, without requiring the auxiliary KV cache approach used by sparse attention methods. On LongBench, SSA consistently outperforms compression and inference-time sparse-attention baselines under the same compression ratio. More strikingly, in retrieval-augmented generation, SSA can even outperform full attention after continued pretraining by over 5.7 points. We attribute this to the ability of SSA's selective unfolding, which concentrates attention on the query-relevant chunks and effectively filters out noise. SSA further extends to a hierarchical gist-of-gist variant (H-SSA) that achieves log-linear decoding complexity while maintaining or improving accuracy at high compression ratios up to 32x. The code is available at https://github.com/yuzhenmao/simplified-sparse-attention/.

1 Introduction

SSA seeks the benefits of training-based sparse attention without architectural changes: it learns gist-token summaries during continued pretraining and uses them for query-guided chunk selection at decoding.

1 Introduction

  • Prior inference-time methods sparsify by selecting or evicting KV-cache entries post hoc, while training-based sparse-attention approaches often require additional architectural components or specialized mechanisms.SSA addresses the question of whether training-based sparse-attention benefits can be obtained without modifying the model architecture.Source: §1 · pp. 1–2
  • SSA learns chunk summaries by interleaving gist tokens during continued pretraining, then scores gists against each query and unfolds the top-k corresponding raw chunks.The attention mask restricts each chunk’s access to earlier raw tokens, routing it through the preceding chunk’s gist tokens.Source: Figure 1 · p. 2, §1 · p. 3
  • Standard next-token training under a mask makes chunks pass earlier context through preceding gists, yielding compression without auxiliary loss, new parameters, or architectural additions.SSA repurposes continued pretraining as a general-purpose way to turn a full-attention architecture into a sparse-attention method.Source: §1 · p. 3
  • Under a fixed token budget, SSA significantly improves accuracy over append-only and interleaved gist compression baselines and inference-only sparse-attention techniques.The introduction also notes that selective unfolding can be incorporated into training to teach the model to adapt to sparse attention.Source: §1 · p. 3
  • H-SSA hierarchically summarizes gists into meta-gists for coarse-to-fine context selection, producing a multiresolution representation with log-linear context-length scaling.Selection proceeds from higher-level summaries toward relevant subcontexts and eventually the original raw tokens.Source: Figure 1 · p. 2, §1 · p. 3

2 Problem Setup

The problem is to preserve next-token prediction quality while replacing full attention over the growing history with a compact context. The desired selection is query-adaptive and trainable, without modifying the standard Transformer architecture.

2 Problem Setup

  • The setup is autoregressive language modeling, in which each token is predicted from all preceding tokens.The next-token prediction uses the hidden state at that position and an output projection matrix.Source: §2 · p. 3
  • Full attention makes long-context modeling prohibitive because per-step computation and KV-cache memory grow linearly with context length and quadratically over the full sequence.The standard Transformer computes each hidden state using all preceding key-value pairs.Source: §2 · p. 3
  • Context selection should adapt to each query, be end-to-end trainable, and require no architectural changes to the standard Transformer.The model should learn relevant-context selection during pretraining or finetuning rather than use heuristic or non-differentiable post hoc operations.Source: §2 · p. 4

3 Method

SSA uses interleaved gist tokens as both compressed chunk summaries and learned routing signals, selectively restoring raw tokens from relevant chunks at inference. Its method also includes optional selective finetuning, a hierarchical extension, and kernels designed to realize sparse computation and memory access.

3.1 Preliminaries

  • Interleaved gist tokens are trained under a causal mask that lets chunk tokens attend to their chunk and preceding gists, while later tokens access the chunk only through its gist.This information bottleneck trains each gist token to summarize its chunk, after which its KV pair can replace that chunk’s full KV cache during decoding.Source: §3.1 · p. 5

3.2 Gist-Based Relevance Scoring

  • Learned gist keys estimate chunk relevance from query-key dot products, jointly optimizing chunk compression and selection without adding architectural components.Unlike mean-pooled keys, gist tokens are learned representations; they reuse the standard input sequence and attention mechanism.Source: §3.2 · p. 5

3.3 Selective Unfolding and Hybrid Attention

  • SSA selects top-k chunks using query-to-gist attention, then attends to each selected chunk’s gist and raw tokens while excluding unselected chunks.This selective unfolding bounds the context to |Vt| = k · (1 + L), independent of the full sequence length, and lets decoding route through the gist KV cache rather than score against the full KV cache.Source: §3 · p. 4, §3.3 · pp. 5–6

3.4 Training

  • Required continued pretraining uses standard teacher-forced next-token loss to teach gist tokens chunk summaries; optional selective finetuning trains with query-dependent top-k sparse masks.Continued pretraining alone supports inference-time selective unfolding; finetuning can further improve performance but may require custom CUDA kernels.Source: §3.4 · pp. 6–7

3.5 Extension to Hierarchical Gist-of-Gist Compression

  • H-SSA recursively summarizes gist groups with meta-gist tokens and selects coarse-to-fine, reducing per-query decoding cost to O(log n) and total decoding cost to O(n log n).Additional hierarchy levels can be added as context grows, with fine-grained details recoverable through selective unfolding.Source: §3.5 · pp. 7–8

3.6 Kernel Design

  • For prefill, a key-column permutation groups globally visible summary columns before local columns, enabling block-sparse attention to skip empty blocks while preserving the defined attention.The permutation and block mask depend on the attention mask, so they are computed once per prefill and shared across layers.Source: §3.6 · pp. 8–9
  • For decoding, SSA compacts selected keys and computes split-K partial attention, touching only selected keys plus O(n) bytes of metadata rather than the full KV cache.On-the-fly keep predicates avoid materializing a dense attention mask; the compact operator uses three kernels for compaction, partial attention, and combination.Source: §3.6 · pp. 9–10

4 Experiments

Across LongBench, multi-document RAG, and KV-cache reuse, SSA and H-SSA outperform compression baselines at reported compression ratios; selective unfolding also delivers decoding and attention-operator speedups at long contexts. Experiments cover two model families, ablations of context selection, and end-to-end efficiency.

4.1 Experiment Setup

  • The experiments evaluate Qwen2-7B-Instruct and Llama3.2-1B using continued pretraining with interleaved gist tokens, with optional selective finetuning, across LongBench, multi-document QA, and KV-cache reuse.Long-context samples are single continuous sequences, whereas RAG samples contain multiple documents with gist tokens inserted independently within each document.Source: §4.1 · pp. 10–11, §4.4 · p. 12

4.2 Long-context Results

  • 46.20 at 8× compression: SSA exceeds ActivationBeacon (42.52) and UniGist (43.40) after continued pretraining, nearly matching Full-PT (47.78).SSA also scores 45.39 at 16× and 44.07 at 32×, while finetuned SSA scores 46.74 at 8×; finetuned H-SSA scores 46.48 versus SSA’s 45.26 at 16× and 44.94 versus 43.35 at 32×. SSA exceeds Full-FT on MF-en at 8×, scoring 54.24 versus 50.33.Source: §4.2 · pp. 11–12, Table 3 · p. 13

4.3 RAG Results

  • 33.68 at 8× compression: SSA exceeds KVLink (21.58), UniGist (22.53), vanilla (27.99), and Full-PT (27.14) on RAG after continued pretraining.After finetuning, SSA scores 53.39 at 8× and 49.29 at 16×; H-SSA scores 30.19 versus SSA’s 27.74 after continued pretraining and 50.72 versus 49.29 after finetuning at 16×. In RAG, selective unfolding concentrates attention on query-relevant chunks and filters distractors; Full-PT underperforms SSA before finetuning but surpasses it after finetuning.Source: §4.3 · p. 12, Table 4 · p. 14

4.4 Parallel Computing with KV-cache Reuse

  • 48.07 at 8× and 41.59 at 16×: SSA substantially outperforms KVLink and UniGist on RAG with KV-cache reuse; H-SSA reaches 46.34 at 16×.SSA encodes documents independently with interleaved gist tokens; link tokens enable cross-document routing, while selective unfolding restores detail for the most relevant segments without traversing full raw KV caches.Source: §4.4 · pp. 12–13, Table 5 · p. 14

4.5 Ablation Analysis

  • Figure 2 compares adaptive top-k across five LongBench tasks and NQ, where it consistently outperforms top-p; on 2WikiMQA, top-k exceeds 51% versus 47.5% for top-p at p = 0.80.On GovReport, top-k scores approximately 30 versus 28–29 for top-p settings; SG+SR averages 53.39, versus 52.76 for SR-only and 53.27 for AG+SR.Source: Figure 2 · p. 15, §4.5 · p. 15

4.6 Efficiency and Scalability

  • At 44K, decoding speedups reach 3.05× for H-SSA and 3.37× for SSA over dense Flash-Decoding, while SSA prefill reaches 0.90× dense latency.Across 32K–200K tokens, the H-SSA operator is 5.8×–8.3× faster than FlashAttention for prefill and 1.24–1.62× faster than Flash-Decoding for decoding. Figure 3 shows decoding latency nearly flat for SSA and H-SSA, while dense decoding latency grows with context length.Source: §4.6 · p. 16, Figure 3 · p. 16, Table 7 · p. 17

5 Related Works

Related work addresses long-context processing through sparse attention, context compression, alternative sequence architectures, and external memory; SSA combines chunk-level gist summaries with on-demand retrieval.

Sparse Attention Mechanisms

  • Sparse attention progresses from fixed local windows with global tokens (Child, 2019; Beltagy et al., 2020), through KV-pair eviction (Zhang et al., 2023b; Xiao et al., 2023; Jiang et al., 2024; Tang et al., 2024), to hardware-aware block selection (Yuan et al., 2025; Liu et al., 2025; Lu et al., 2025).The cited approaches reduce attention costs using fixed or heuristic patterns, cache management, or dynamic selection of relevant blocks.Source: §5 · p. 17

Context Compression

  • Compression methods use attention-mask bottlenecks, document-level summaries, or interleaved local gist tokens, but typically discard raw tokens after compression (Mu et al., 2023; Chevalier et al., 2023; Yang et al., 2025; Zhang et al., 2024; Deng et al., 2025; Petrov et al., 2025).Unlike these one-way approaches, SSA selectively unfolds raw tokens from retrieved chunks.Source: §5 · p. 17

Alternative Sequence Modeling Architectures

  • Alternative sequence architectures use convolutions, feature-map approximations, structured recurrences, or recurrent attention-like mechanisms for linear or near-linear inference, often changing the core attention mechanism (Shi et al., 2023; Katharopoulos et al., 2020; Gu et al., 2021; Gu and Dao, 2023; Peng et al., 2023; Sun et al., 2023; Yang et al., 2024c).The passage distinguishes these architectures from SSA, noting they often require pretraining from scratch on the new architecture.Source: §5 · p. 17

Memory-Augmented Language Models

6 Conclusions

SSA combines gist-based compression with sparse attention, using gist tokens for query-adaptive detail recovery without architectural modifications. Its hierarchical extension and experiments demonstrate efficient decoding and gains over established baselines.

6 Conclusions

  • SSA consistently gains over gist-compression and sparse-attention baselines across 8×–32× compression ratios, with up to 3.37× end-to-end decoding speedup over Flash-Decoding.During prefill, SSA matches and eventually outperforms FlashAttention.Source: §6 · p. 18
  • Interleaved gist tokens serve as compressed summaries and learned routing signals, enabling query-adaptive recovery of fine-grained detail without architectural modifications, external indexers, or non-differentiable gates.SSA unifies gist-based context compression with sparse attention through selective unfolding.Source: §6 · p. 18
  • Hierarchical gist-of-gist compression extends SSA with coarse-to-fine selection and logarithmic per-step decoding complexity.Experiments on LongBench and RAG benchmarks demonstrate consistent gains over the baselines across compression ratios from 8× to 32×.Source: §6 · p. 18

Supplementary Material · A More Details on Adaptive Top-k.

Adaptive top-k sets a per-level selection budget using the available KV positions, effective compression, grouped unfolding, and chunk size. The scheme preserves at least one selected chunk and keeps unfolded-token counts approximately proportional to the available KV positions across compression ratios.

A More Details on Adaptive Top-k.

  • The per-level selection budget uses n_kv, effective compression L_eff, group size G, and chunk size L.L_eff is L for single-level SSA and L · J for H-SSA.Source: A More Details on Adaptive Top-k · p. 23
  • Dividing by G accounts for each head selecting chunks independently before their union is unfolded into the shared KV cache.The +1 ensures at least one chunk is selected.Source: A More Details on Adaptive Top-k · p. 23
  • H-SSA uses the same k at both the meta-gist and gist levels, keeping unfolded-token counts approximately proportional to n_kv across compression ratios.This supports fair comparisons across settings.Source: A More Details on Adaptive Top-k · p. 23

B Complexity Analysis

The section analyzes the complexity of both single-level and hierarchical SSA, with input length n and head dimension d treated as a constant.

B Complexity Analysis

  • The complexity analysis covers both single-level and hierarchical SSA.The analysis takes input length n as a variable and treats head dimension d as a constant.Source: B Complexity Analysis · p. 23

B.1 Single-Level Complexity

Single-level SSA has subquadratic prefill when chunk size L scales as √n, and decoding has linear per-step cost for fixed k and L but quadratic aggregate cost reduced relative to full attention.

B.1 Single-Level Complexity

  • O(nL+n^2/L^2) prefill becomes O(n√n) when L=O(√n); for fixed L, the gist-to-gist term is quadratic but reduced by a factor of L^2 versus full attention.The prefill cost combines intra-chunk attention, O(nL), and gist-to-gist attention.Source: §B.1 · p. 23
  • O(M) relevance scoring and top-k selection plus O(kL) hybrid attention give a per-step cost, where M=⌈n/L⌉ gist tokens and k chunks are unfolded.Hybrid attention covers k gist tokens and k·L raw tokens from the selected chunks.Source: §B.1 · p. 23
  • O(n^2/L + nkL) is the aggregate decoding cost over n steps, quadratic in n but reduced by a factor of L versus full attention's O(n^2).For fixed k and L, the per-step decoding cost is linear in context length n.Source: §B.1 · p. 24

B.2 Hierarchical Complexity

The hierarchical variant uses grouped summary tokens to achieve linear prefill and log-linear decoding complexity. Its decoding combines coarse-to-fine routing with attention over selected summary tokens and unfolded raw tokens.

B.2 Hierarchical Complexity

  • At level ℓ, Mℓ = ⌈M/J^(ℓ−1)⌉ summary tokens are used, with gist tokens at the bottom and meta-gist tokens above.Level-ℓ summary attention costs O(Mℓ · J); summed across levels, its cost is O(MJ) = O(nJ/L), independent of depth t.Source: §B.2 · p. 24
  • Prefill costs O(nL + nJ/L), linear in n and substantially below the single-level cost O(nL + n^2/L^2).Source: §B.2 · p. 24
  • With t = ⌈log_J M⌉ and fixed kmax and J, coarse-to-fine routing costs O(log n).The top level scores all Mt tokens; each lower level scores candidates exposed by selected tokens at the level above.Source: §B.2 · p. 24
  • Attention after routing costs O(k1L + log M), from selected summary tokens and k1 gist tokens’ unfolded raw tokens.Each level contributes kℓ selected summary tokens, while the bottom level also contributes k1 · L raw tokens.Source: §B.2 · p. 25
  • Across n decoding steps, hierarchical SSA has aggregate O(n log n) log-linear cost.Its per-step decoding cost is logarithmic in context length.Source: §B.2 · p. 25

C Experiments on Passkey Retrieval

SSA and H-SSA achieve 100% passkey retrieval accuracy across tested context lengths, insertion positions, and both model families, including sequence lengths well beyond continued-pretraining lengths.

C Experiments on Passkey Retrieval

  • The benchmark inserts a short passkey at a random position in a long distractor context and requires the model to reproduce it exactly.Source: C Experiments on Passkey Retrieval · p. 25
  • Tests extend to 50K words for Qwen2-7B-Instruct and 40K words for Llama-3.2-1B, respectively 2.5× and 10× their continued-pretraining maximum sequence lengths.Perfect accuracy also holds for passkeys placed at the beginning, middle, or end of the context.Source: C Experiments on Passkey Retrieval · p. 25
  • 100% accuracy across all tested context lengths, insertion positions, and both model families demonstrates perfect passkey retrieval by SSA and H-SSA.The result includes H-SSA’s multi-level coarse-to-fine routing.Source: Figure 4 · p. 25, C Experiments on Passkey Retrieval · p. 25

Related papers

Loading 2604.20920v2…