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

Long-context attention becomes increasingly costly, while existing training-based sparse-attention methods often require architectural additions. SSA uses interleaved gist tokens for query-adaptive chunk routing and selective unfolding, consistently improving over compression and sparse-attention baselines across 8×–32× compression while achieving up to 3.37× decoding speedup over Flash-Decoding.

  • Problem

    As context length grows, attention increasingly dominates inference cost, while existing training-based sparse-attention methods often require additional architectural components or specialized mechanisms.

  • Method

    SSA trains interleaved gist tokens to summarize chunks and uses them as routing signals to selectively unfold the most relevant raw-token chunks without architectural modifications.

  • Results

    SSA consistently improves over gist-compression and sparse-attention baselines across compression ratios from 8× to 32×, with up to 3.37× end-to-end decoding speedup over Flash-Decoding.

  • Takeaways & Limitations

    Selective unfolding enables query-adaptive recovery of fine-grained detail and extends to hierarchical compression with logarithmic per-step decoding complexity.

  • Takeaways & Limitations

    Experiments use Qwen2 and Llama3.2-1B model choices selected because key baseline methods had been evaluated on them.

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

Long-context applications make attention increasingly costly as context length grows. Simplified Sparse Attention addresses this by training gist tokens to compress chunks under a restricted attention mask, enabling input-dependent sparse patterns without changing the standard next-token objective.

  • Motivation: As context length grows, attention computation becomes increasingly dominant in the overall cost of long-context modeling.Long-context capabilities support in-depth reasoning, repository-level software engineering, and multi-turn autonomous agents.
  • Motivation: Recent sparse-attention methods incorporate sparsity during pretraining to preserve or improve model quality while achieving substantial speedups.The cited methods include Native Sparse Attention, Deepseek Sparse Attention, and MoBA.
  • Method: SSA inserts gist tokens at fixed positions to partition context into chunks, restricting chunk tokens to access earlier raw tokens through preceding gist tokens.Training retains the standard next-token prediction objective while changing the attention mask.
  • Method: Under this constraint, end-to-end next-token optimization teaches the model to pack information needed for future predictions into gist tokens, making compression emerge from training.Tokens in a chunk cannot directly attend to earlier raw tokens; they can access them only through the preceding chunk’s gist tokens.

2 Problem Setup

The problem setup considers autoregressive language modeling with full attention over prior key-value pairs, whose cost and memory become prohibitive for long contexts. It seeks a compact, query-adaptive context that preserves prediction quality while remaining trainable and compatible with the standard Transformer.

  • Problem formulation: Autoregressive language modeling predicts each next token from all preceding tokens.The prediction is computed from the hidden state at the current position.
  • Problem formulation: Standard Transformer decoding computes each hidden state with full attention over all preceding key-value pairs.The output projection maps the hidden state to the vocabulary prediction.
  • Problem formulation: Long-context modeling becomes computationally prohibitive because per-step computation and stored key-value memory scale linearly with context length and quadratically over the full sequence.Both hidden-state computation and KV-cache storage grow as the context expands.
  • Design goals: The goal is a compact context containing far fewer than t key-value pairs while preserving prediction quality.At each decoding step, the method constructs a context V_t with |V_t| ≪ t.
  • Design goals: The compact context should select history dynamically per query, learn relevance end to end, and require no architectural modifications to the standard Transformer.These desiderata support straightforward integration while avoiding heuristic or non-differentiable post hoc selection.

3 Method

SSA uses interleaved, learned gist tokens both to compress local chunks and to route queries toward the most relevant chunks, which are selectively unfolded with their raw tokens. The method is trained with standard autoregressive loss and extends hierarchically to achieve log-linear complexity.

  • Gist-token compression: Interleaved gist tokens summarize local chunks under a causal information bottleneck, while later tokens can access each chunk only through its gist token.Within each chunk, tokens attend locally and to preceding gist tokens; tokens after the gist cannot attend to the chunk’s raw tokens.
  • Chunk selection: SSA scores each chunk by query–gist affinity, using learned gist representations as relevance proxies for routing attention.Unlike mean-pooled block keys, gist tokens are learned representations intended to preserve semantically important signals.
  • Selective unfolding: The method selects the top-k chunks and unfolds each selected gist by reintroducing its gist KV pair and full raw KV pairs, excluding unselected chunks.Among tested hybrid contexts, attending exclusively to selected gist–chunk pairs performed best under a fixed token budget.
  • Training: SSA uses standard teacher forcing and vanilla cross-entropy loss, with required continued pretraining teaching the model to compress context into gist tokens.Selective finetuning is optional, while continued pretraining processes compressed prefixes with the gist causal mask and generation suffixes autoregressively.
  • Hierarchical extension: The single-level framework has per-step complexity O(M + kL), and recursively applying gist-based routing yields a hierarchical design with log-linear complexity.Hierarchical retrieval progressively recalls higher-level memories before refining them to specific sections and reconstructing detail.

4 Experiments

Experiments across Qwen2-7B-Instruct and Llama3.2-1B show that SSA and H-SSA outperform compression and sparse-attention baselines across LongBench, RAG, and KV-cache reuse settings. Selective unfolding also improves efficiency, with adaptive top-k outperforming top-p selection and sparse decoding remaining nearly flat in latency.

  • LongBench: SSA averages 46.20 at 8× compression, outperforming ActivationBeacon (42.52) and UniGist (43.40) on LongBench under continued pretraining.At 16× and 32×, SSA scores 45.39 and 44.07, versus 40.64 and 38.30 for ActivationBeacon.
  • LongBench: 46.48: H-SSA surpasses SSA’s 45.26 at 16× finetuning, widening to 44.94 versus 43.35 at 32×.The hierarchical variant becomes more advantageous as compression increases.
  • RAG: 33.68: SSA surpasses KVLink (21.58), UniGist (22.53), vanilla (27.99), and Full-PT (27.14) at 8× RAG compression.This result uses Llama3.2-1B with continued pretraining only.
  • KV-cache reuse: 48.07: SSA achieves this average RAG score with KV-cache reuse at 8× compression, while H-SSA reaches 46.34 at 16×.SSA also scores 41.59 at 16×, substantially outperforming KVLink and UniGist.
  • Selection: Adaptive top-k selection consistently outperforms top-p thresholds of p ∈{0.80, 0.85, 0.90, 0.95} across six LongBench tasks.Top-p dynamically varies the number of unfolded chunks according to cumulative score mass.
  • Efficiency: 21.9 ms to 76.4 ms: Dense Flash-Decoding latency rises from 8K to 44K, whereas H-SSA remains near 25 ms while reading around 1% of the cache.At the operator level, H-SSA reaches 8.3× prefill speedup at 200K tokens and 1.24–1.62× decoding speedup over 32K–200K.

5 Related Works

Related work addresses long-context efficiency through sparse attention, context compression, alternative sequence architectures, and external memory. These approaches reduce attention costs, compress information, replace softmax attention, or extend usable context beyond the native window.

  • Sparse Attention Mechanisms: Sparse attention methods use fixed patterns, KV-cache eviction, or hardware-aware and learnable mechanisms to identify relevant blocks.Examples include Sparse Transformer, Longformer, H2O, StreamingLLM, LongLLMLingua, Quest, Native Sparse Attention, DeepSeek Sparse Attention, and MoBA.
  • Context Compression: Context compression methods encode prompts or documents into compact representations using special or learnable compression tokens.Gist Tokens and AutoCompressors use an attention-mask bottleneck, while KVLink produces document-level summaries from appended tokens.
  • Alternative Sequence Modeling Architectures: Alternative sequence architectures replace softmax attention with convolutional, feature-map, or state-space approaches targeting linear or near-linear complexity.The passage names MultiresLayer, linear attention, and S4 as examples.
  • Memory-Augmented Language Models: Memory-augmented language models add external memory modules to support contexts exceeding the model’s native attention window.LongMem separates memory encoding from retrieval, while MemGPT explicitly pages information between a limited context and external memory.

6 Conclusions

SSA unifies gist-based context compression with sparse attention through selective unfolding, using gist tokens for both compressed summaries and learned routing. It enables query-adaptive detail recovery without architectural modifications and extends hierarchically to logarithmic per-step decoding complexity, with consistent gains on LongBench and RAG benchmarks.

  • SSA unifies gist-based context compression with sparse attention through selective unfolding.
  • 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.
  • Hierarchical gist-of-gist compression extends SSA with logarithmic per-step decoding complexity through coarse-to-fine selection.
  • Experiments on LongBench and RAG benchmarks demonstrate consistent gains.

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

The adaptive top-k scheme determines each level’s selection budget from the available KV positions, effective compression, query-head grouping, and chunk size. It accounts for grouped unfolding, guarantees at least one selected chunk, and applies the same k at both hierarchical levels.

  • A More Details on Adaptive Top-k.: The selection budget uses n_kv, L_eff, G, and L to determine the per-level number of selected chunks.Here, n_kv is the total available KV positions, L_eff is the effective compression factor, G is the number of query heads per KV group, and L is the chunk size.
  • A More Details on Adaptive Top-k.: L_eff equals L for single-level SSA and L · J for hierarchical H-SSA.The effective compression factor therefore reflects whether selection uses one level or the hierarchical variant.
  • A More Details on Adaptive Top-k.: The adaptive scheme is designed to keep the total number of unfolded tokens approximately proportional to the available context and compression settings.The passage states that the total number of unfolded tokens remains approximately controlled by this adaptive budget.
  • A More Details on Adaptive Top-k.: Dividing by G accounts for grouped unfolding, where heads in one GQA group may select different chunks.Their union is unfolded into the shared KV-cache, so the effective unfolded-chunk count per group scales with G.
  • A More Details on Adaptive Top-k.: The +1 term ensures that at least one chunk is always selected.This guarantee applies to the adaptive selection budget.
  • A More Details on Adaptive Top-k.: For H-SSA, the same k is used at both the meta-gist and gist levels.Thus, hierarchical selection applies one shared per-level budget across these two levels.

B Complexity Analysis

This section analyzes the computational complexity of SSA in both its single-level and hierarchical forms, using input length n and head dimension d as notation.

  • B Complexity Analysis: The analysis covers both single-level SSA and its hierarchical variant.It provides a detailed complexity analysis for the two SSA variants.
  • B Complexity Analysis: n denotes the input length in the complexity analysis.The notation defines n as the input length.
  • B Complexity Analysis: d denotes the head dimension and is treated as a constant.The head dimension d is assumed constant for the analysis.

B.1 Single-Level Complexity

Single-level SSA uses M = ⌈n/L⌉ gist tokens and unfolds k selected chunks, yielding sub-quadratic prefill under typical settings and linear per-step decoding for fixed k and L. Across a full n-token generation, decoding remains quadratic but reduces full-attention cost by a factor of L.

  • M = ⌈n/L⌉ chunks produce the corresponding number of gist tokens, while k denotes the chunks selected for unfolding.
  • O(nL + n^2/L^2) is the overall prefill cost, with both terms becoming O(n√n) when L = O(√n).For fixed L, the gist-to-gist term O(n^2/L^2) dominates and is reduced by a factor of L^2 versus full attention.
  • O(M) relevance scoring and top-k selection plus O(kL) hybrid attention make each decoding step linear in n when k and L are fixed.The query scores all M gist keys, then attends to k gist tokens and kL raw tokens from selected chunks.
  • O(n^2/L + nkL) is the aggregate decoding cost over n steps, remaining quadratic but reduced by a factor of L compared with full attention’s O(n^2).

B.2 Hierarchical Complexity

H-SSA organizes summary tokens across hierarchy levels and achieves linear prefill, logarithmic routing, and log-linear aggregate decoding complexity. Its decoding attention cost combines selected raw-token expansion with a logarithmic number of summary tokens.

  • Hierarchy structure: H-SSA uses hierarchy depth t and grouping factor J, with Mℓ=⌈M/Jℓ−1⌉ summary tokens at level ℓ.The bottom level contains M gist tokens, while higher levels contain progressively fewer meta-gist tokens.
  • Prefill complexity: O(nL + nJ/L) is the overall prefill cost, linear in n and lower than the single-level cost O(nL + n2/L2).The hierarchical summary-attention cost is O(MJ)=O(nJ/L), independent of depth t.
  • Decoding: Routing: O(log n) is the total routing cost when t=⌈logJ M⌉, assuming fixed kmax and J.Top-down routing scores all top-level tokens, then scores only candidates exposed by selected higher-level tokens.
  • Decoding: Attention: O(k1L + log M) is the decoding attention cost after routing, combining k1L unfolded raw tokens with O(log M) selected summary tokens.Attention includes selected summary tokens at each hierarchy level and the raw tokens unfolded from selected bottom-level gist tokens.
  • Decoding: Total: O(n log n) is the aggregate cost over a full sequence of n decoding steps, making hierarchical decoding log-linear.This follows from combining logarithmic routing with the post-routing attention computation.

C Experiments on Passkey Retrieval

The passkey retrieval experiment evaluates fine-grained retrieval under compression by requiring models to locate and exactly reproduce a short passkey embedded at a random position in a long distractor context. Figure 4 measures SSA and H-SSA accuracy across context lengths and relative insertion positions on Qwen2-7B-Instruct and Llama3.2-1B.

  • Evaluation setup: Passkey retrieval tests whether models can locate and exactly reproduce a short passkey inserted randomly within a long distractor context.The benchmark evaluates fine-grained retrieval under compression.
  • Evaluation setup: Figure 4 reports retrieval accuracy as a function of context length and the passkey’s relative insertion position.Context length is plotted vertically and relative insertion position (%) horizontally.
  • Models and methods: The experiment evaluates after-continued-pretraining SSA and H-SSA on Qwen2-7B-Instruct and Llama3.2-1B.Qwen2-7B-Instruct is evaluated up to 50K words, described as 2.5× the maximum sequence length seen during training.
Loading 2604.20920v2…