Simple Attention Sparsification

• Simple Attention Sparsification (SAS) is a post-training, block-sparse attention mechanism that ranks historical context blocks based on their impact on language model predictions under a given attention budget.
• SAS optimizes context ranking during training by directly minimizing the language-modeling loss rather than focusing on matching dense attention, demonstrating superior performance in reasoning, long-context understanding, and other tasks compared to dense attention.
• Empirical results show that SAS achieves higher reasoning accuracy, outperforming other methods by up to 15.5 points on GPQA-Diamond and 71.72 on AIME24 at 4096 tokens, and maintains comparable performance to full attention in agentic tasks and long-context understanding.

Simple Attention Sparsification (SAS) is a post-training, block-sparse attention mechanism that learns to rank historical context blocks according to their effect on language-model predictions under a fixed attention budget. Introduced in “SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking” (Li et al., 11 Sep 2026), it distinguishes the always-retained current block from selector-ranked historical blocks, injects normalized selector scores into attention logits during training, and uses hard Top-$K$ block selection at inference. Its defining objective is to optimize context ranking through the language-modeling loss rather than to imitate the dense model’s layer-wise attention distribution.

1. Problem formulation and motivation

In autoregressive decoding, a query attends to an increasingly large key–value cache. For query $\mathbf q\in\mathbb R^d$, keys $\mathbf K\in\mathbb R^{n\times d}$, and values $\mathbf V\in\mathbb R^{n\times d}$, dense attention is

$$\mathbf{o} = \operatorname{softmax}(\mathbf q\mathbf K^\top)\mathbf V.$$

A single query requires $O(n)$ attention computation. Across a length-$n$ generation, cumulative decoding cost is $O(n^2)$. If only a retained token set $\mathcal S$ is attended, the per-query cost becomes $O(|\mathcal S|)$, and cumulative decoding cost becomes $\mathbf q\in\mathbb R^d$0.

SAS operates on contiguous context blocks. Let the historical context be partitioned into

$\mathbf q\in\mathbb R^d$1

where $\mathbf q\in\mathbb R^d$2 has block size $\mathbf q\in\mathbb R^d$3 and $\mathbf q\in\mathbb R^d$4. A selector $\mathbf q\in\mathbb R^d$5 produces scores

$\mathbf q\in\mathbb R^d$6

The historical blocks are ranked according to $\mathbf q\in\mathbb R^d$7, and a fixed number $\mathbf q\in\mathbb R^d$8 is retained:

$\mathbf q\in\mathbb R^d$9

The current block $\mathbf K\in\mathbb R^{n\times d}$0 is always retained. The resulting sparse context is therefore

$\mathbf K\in\mathbb R^{n\times d}$1

The central methodological distinction is between dense attention mass and predictive utility under a fixed budget. A block that receives substantial attention in the dense model may contribute little to the final prediction once only $\mathbf K\in\mathbb R^{n\times d}$2 blocks are available. Conversely, a block with relatively small dense attention may contain value vectors that materially affect the output. Layer-wise attention matching also treats layers independently and does not directly optimize complementary context selection across layers. SAS therefore trains the selector using the language-modeling loss rather than a dense-attention imitation objective (Li et al., 11 Sep 2026).

SAS belongs to a broader family of attention-sparsification methods. SAC learns a discrete task-adaptive graph with an LSTM edge predictor and REINFORCE (Li et al., 2020). Sparsefinder predicts the support of exact entmax attention before computing it (Treviso et al., 2021). S2-Attention assigns different context shards to different heads and enforces collective context coverage (Lin et al., 2024). Saap uses asymmetric key partitions and query classifiers for inference-time long-context retrieval (Mazaré et al., 12 Feb 2025). These methods differ in selector construction, training regime, sparsity granularity, and execution strategy; SAS is specifically characterized by end-to-end optimization of a fixed-budget context ranking.

2. Architecture and gated attention

SAS separates the current block from historical candidates. The current block $\mathbf K\in\mathbb R^{n\times d}$3 has a fixed unit gate,

$\mathbf K\in\mathbb R^{n\times d}$4

while historical blocks receive selector-derived gates. The current block is not scored by the selector, preserving local and current causal attention and providing a calibration reference for historical context.

For query $\mathbf K\in\mathbb R^{n\times d}$5, the selector produces continuous scores

$\mathbf K\in\mathbb R^{n\times d}$6

These scores are converted to positive gates through a softmax:

$\mathbf K\in\mathbb R^{n\times d}$7

For historical block $\mathbf K\in\mathbb R^{n\times d}$8,

$\mathbf K\in\mathbb R^{n\times d}$9

where

$\mathbf V\in\mathbb R^{n\times d}$0

The shared normalization term does not cancel, because historical blocks are normalized while the current block remains unit-gated:

$\mathbf V\in\mathbb R^{n\times d}$1

For token $\mathbf V\in\mathbb R^{n\times d}$2 in block $\mathbf V\in\mathbb R^{n\times d}$3, the gated attention output is

$\mathbf V\in\mathbb R^{n\times d}$4

or equivalently,

$\mathbf V\in\mathbb R^{n\times d}$5

The block gate is broadcast to all tokens in that block.

Inner versus outer gating

SAS places the gate inside the attention softmax. This differs from an outer value gate,

$\mathbf V\in\mathbb R^{n\times d}$6

Outer gating rescales value contributions after attention probabilities have already been normalized. It does not reallocate normalized attention mass among blocks. Inner gating changes the normalization itself and enables the selector to learn relative block importance.

The inner-gate gradient contains a comparison between block values and the current attention output:

$\mathbf V\in\mathbb R^{n\times d}$7

whereas the outer-gate gradient is

$\mathbf V\in\mathbb R^{n\times d}$8

The former is value-aware in a relative sense: it evaluates whether increasing a block’s normalized attention would improve the output relative to the current output.

Continuous gates and ranking preservation

During training, SAS preserves continuous gates rather than replacing them with binary indicators. Hard gating would use

$\mathbf V\in\mathbb R^{n\times d}$9

The paper also evaluates a straight-through estimator,

$$\mathbf{o} = \operatorname{softmax}(\mathbf q\mathbf K^\top)\mathbf V.$$0

However, binary gating discards relative priority information among selected blocks. It also removes unselected blocks from the forward softmax normalizer, potentially producing unstable hypothetical gradients for high-scoring omitted blocks. Continuous gates keep all historical blocks in the training relaxation and allow the selector to learn both membership and relative ordering.

3. Training and inference

SAS is primarily a post-training method. In the main experiments, the Transformer backbone is frozen and only the selector is trained. The selector architecture is matched to SeerAttention-R, while the sparse attention mechanism is optimized through the standard autoregressive language-modeling objective rather than layer-wise dense-attention distillation.

Training procedure

For each query, training proceeds conceptually as follows:

1. Compute selector scores:

$$\mathbf{o} = \operatorname{softmax}(\mathbf q\mathbf K^\top)\mathbf V.$$1

2. Normalize them:

$$\mathbf{o} = \operatorname{softmax}(\mathbf q\mathbf K^\top)\mathbf V.$$2

3. Select the Top-$$\mathbf{o} = \operatorname{softmax}(\mathbf q\mathbf K^\top)\mathbf V.$$3 historical blocks:

$$\mathbf{o} = \operatorname{softmax}(\mathbf q\mathbf K^\top)\mathbf V.$$4

4. Form the retained context:

$$\mathbf{o} = \operatorname{softmax}(\mathbf q\mathbf K^\top)\mathbf V.$$5

5. Compute gated attention over the retained blocks:

$$\mathbf{o} = \operatorname{softmax}(\mathbf q\mathbf K^\top)\mathbf V.$$6

6. Continue the Transformer and compute the autoregressive loss:

$$\mathbf{o} = \operatorname{softmax}(\mathbf q\mathbf K^\top)\mathbf V.$$7

7. Update selector parameters through ordinary backpropagation:

$$\mathbf{o} = \operatorname{softmax}(\mathbf q\mathbf K^\top)\mathbf V.$$8

The gradient pathway is

$$\mathbf{o} = \operatorname{softmax}(\mathbf q\mathbf K^\top)\mathbf V.$$9

This differentiable pathway is the defining difference from selector training based on dense-attention distillation.

Sparse-scope and full-scope training

The principal implementation uses sparse-scope training, in which only the current block and selected historical blocks participate in the attention computation. Unselected blocks do not receive direct gate gradients through their own values, although they can be updated indirectly through normalization over selector scores.

The paper compares this with full-scope training, in which all blocks participate in gated attention. Full-scope training gives stronger early gradients to unselected blocks and initially converges faster. Sparse-scope training is substantially cheaper and reaches comparable final performance after sufficient training across the reported model sizes and budgets (Li et al., 11 Sep 2026).

Inference procedure

At inference time, the continuous gates are used only to produce a ranking. The attention computation itself uses hard routing:

$O(n)$0

$O(n)$1

$O(n)$2

Thus, continuous gating is a training relaxation; deployment uses ordinary sparse attention over the selected context.

Budget specification

The main implementation uses block size $O(n)$3. A budget of $O(n)$4 tokens corresponds to $O(n)$5 historical blocks, and a budget of $O(n)$6 tokens corresponds to $O(n)$7 historical blocks, excluding or separately accounting for the always-retained current block.

Tight budgets make ranking errors more consequential because fewer historical blocks can be retained. With relaxed budgets, useful context is more likely to be included even when rankings are imperfect, reducing the gap between SAS and distillation-based selectors.

4. Efficient implementation and complexity

A naive implementation would materialize the gated attention matrix and add blockwise log-gates to attention logits. SAS instead uses a custom Triton kernel integrated with FlashAttention-style tiled computation.

The kernel:

1. traverses query and key–value tiles;
2. maps key tokens to block indices;
3. loads the query-specific block log-gate;
4. adds the log-gate to the query–key scores;
5. masks nonselected historical blocks to $O(n)$8;
6. leaves the current block unbiased;
7. applies the causal mask;
8. performs online softmax without materializing the full attention matrix.

Top-$O(n)$9 selection is represented inside the fused attention kernel through a per-query threshold $n$0. A block is active when

$n$1

This avoids requiring an explicit sort inside the fused attention kernel.

During the backward pass, the kernel recomputes attention probabilities from saved row-wise log-sum-exp values, computes gradients for $n$2, and aggregates gate gradients over tokens within each block.

Computational scaling

If each query retains $n$3 blocks of size $n$4, the sparse attention component processes approximately $n$5 tokens per query. The attention computation therefore scales with the retained context rather than the full sequence length. For fixed $n$6 and $n$7, sparse attention is independent of total context length in the principal attention operation.

The selector still scans block summaries and performs Top-$n$8 ranking over candidate blocks. This produces an important systems distinction: sparse attention may become nearly context-invariant while selector scoring and ranking continue to grow with context length.

In SGLang, SAS uses paged KV caches and FlashInfer sparse decode kernels. Prefill remains dense, while decoding performs selector scoring, block gathering, and sparse attention. Reported Qwen3-4B serving measurements indicate approximately:

• $n$9, $O(n^2)$0, and $O(n^2)$1 speedups over dense attention at $O(n^2)$2, $O(n^2)$3, and $O(n^2)$4 contexts for batch size $O(n^2)$5;
• approximately $O(n^2)$6 speedup at $O(n^2)$7 for batch size $O(n^2)$8;
• at $O(n^2)$9, Top-$\mathcal S$0 selection accounts for about $\mathcal S$1 of the sparse decode step.

These results show that selector overhead can dominate once sparse attention itself has been reduced sufficiently.

The broader sparse-attention literature reports the same separation between theoretical sparsity and realized efficiency. Sparse Flash Attention supports dynamic key/query dropping and hashing through tiled kernels (Pagliardini et al., 2023). S2-Attention shows that heterogeneous head-level sharding must be co-designed with GPU execution to achieve wall-clock gains (Lin et al., 2024). SparseSAM similarly uses deterministic structured masks and fused kernels because dynamic mask construction can eliminate theoretical savings (Tran et al., 17 May 2026). Sol-Attn integrates routing, approximation, and online softmax to avoid materializing block-score maps (Li et al., 27 Jul 2026).

5. Empirical evidence

The principal experiments use Qwen3-4B, Qwen3-8B, and Qwen3-14B. The post-training configuration consists of:

• a frozen Transformer backbone;
• a trainable selector;
• block size $\mathcal S$2;
• maximum sequence length $\mathcal S$3;
• one epoch on $\mathcal S$4 OpenR1-MATH-220K examples;
• AdamW with learning rate $\mathcal S$5;
• cosine decay;
• global batch size $\mathcal S$6.

SAS is compared with full attention, SeerAttention-R, Sliding Window Attention, StreamingLLM, and Quest. Evaluation covers reasoning, long-context understanding, and agentic tasks.

Reasoning tasks

At a $\mathcal S$7-token budget, SAS substantially outperforms SeerAttention-R. On MATH500, SAS reaches $\mathcal S$8, $\mathcal S$9, and $O(|\mathcal S|)$0 for Qwen3-4B, Qwen3-8B, and Qwen3-14B, compared with $O(|\mathcal S|)$1, $O(|\mathcal S|)$2, and $O(|\mathcal S|)$3 for SeerAttention-R.

On GPQA-Diamond, SAS reaches $O(|\mathcal S|)$4, $O(|\mathcal S|)$5, and $O(|\mathcal S|)$6, compared with $O(|\mathcal S|)$7, $O(|\mathcal S|)$8, and $O(|\mathcal S|)$9. The reported improvements at the tightest budget are approximately $\mathbf q\in\mathbb R^d$00–$\mathbf q\in\mathbb R^d$01 points on MATH500 and $\mathbf q\in\mathbb R^d$02–$\mathbf q\in\mathbb R^d$03 points on GPQA-Diamond.

At $\mathbf q\in\mathbb R^d$04 tokens, SAS remains ahead on difficult reasoning tasks. For Qwen3-4B, AIME24 is $\mathbf q\in\mathbb R^d$05 for SAS versus $\mathbf q\in\mathbb R^d$06 for SeerAttention-R, while AIME25 is $\mathbf q\in\mathbb R^d$07 versus $\mathbf q\in\mathbb R^d$08.

At $\mathbf q\in\mathbb R^d$09 tokens, SAS often approaches or exceeds dense attention. Qwen3-4B achieves $\mathbf q\in\mathbb R^d$10 on AIME24 with SAS versus $\mathbf q\in\mathbb R^d$11 with full attention. Training-free baselines degrade sharply under tight budgets; Quest scores zero on AIME24 and AIME25 for several $\mathbf q\in\mathbb R^d$12-token settings.

Long-context understanding

At a $\mathbf q\in\mathbb R^d$13-token budget, SAS improves over SeerAttention-R particularly on longer inputs:

• Qwen3-4B at $\mathbf q\in\mathbb R^d$14: $\mathbf q\in\mathbb R^d$15 versus $\mathbf q\in\mathbb R^d$16;
• Qwen3-8B at $\mathbf q\in\mathbb R^d$17: $\mathbf q\in\mathbb R^d$18 versus $\mathbf q\in\mathbb R^d$19;
• Qwen3-14B at $\mathbf q\in\mathbb R^d$20: $\mathbf q\in\mathbb R^d$21 versus $\mathbf q\in\mathbb R^d$22.

At a $\mathbf q\in\mathbb R^d$23-token budget, SAS nearly recovers full-attention performance. Qwen3-14B achieves $\mathbf q\in\mathbb R^d$24 with SAS versus $\mathbf q\in\mathbb R^d$25 with full attention.

Agentic tasks

On BFCL Multi-Turn, SAS improves over SeerAttention-R by up to $\mathbf q\in\mathbb R^d$26 points at a $\mathbf q\in\mathbb R^d$27-token budget. For Qwen3-4B, the reported score is $\mathbf q\in\mathbb R^d$28 versus $\mathbf q\in\mathbb R^d$29. At $\mathbf q\in\mathbb R^d$30 tokens, Qwen3-14B reaches $\mathbf q\in\mathbb R^d$31, compared with $\mathbf q\in\mathbb R^d$32 for SeerAttention-R and $\mathbf q\in\mathbb R^d$33 for full attention.

On VitaBench at $\mathbf q\in\mathbb R^d$34 tokens, SAS is generally ahead of SeerAttention-R across Delivery, Instore, and OTA metrics and approaches full-attention performance.

Continued pretraining

An additional experiment jointly trains the backbone and selector for $\mathbf q\in\mathbb R^d$35 steps on approximately $\mathbf q\in\mathbb R^d$36 billion tokens using OLMo3-7B. SAS-RoPE obtains an average general-task score of $\mathbf q\in\mathbb R^d$37, compared with $\mathbf q\in\mathbb R^d$38 for HiLS-Attn-RoPE and $\mathbf q\in\mathbb R^d$39 for dense OLMo3-Base.

On LongBench, SAS obtains an average of $\mathbf q\in\mathbb R^d$40, tying the best reported result and outperforming the dense base’s $\mathbf q\in\mathbb R^d$41 and sliding-window continued pretraining’s $\mathbf q\in\mathbb R^d$42. These results support the extension of end-to-end selector training beyond frozen-backbone post-training, although the evidence remains limited to the reported setup.

Ablation findings

Four design choices are identified as crucial:

1. Inner softmax gating: substantially outperforms outer value rescaling because it changes normalized attention allocation.
2. Softmax-normalized gates: outperform sigmoid and raw-score injection by calibrating historical blocks against the unit-gated current block.
3. Continuous gates: outperform straight-through hard gates by preserving relative priorities and avoiding unstable gradients from omitted blocks.
4. Sparse-scope training: converges more slowly initially than full-scope training but reaches comparable final performance at substantially lower cost.

In the reported GPQA-Diamond ablation, inner softmax gating reaches approximately $\mathbf q\in\mathbb R^d$43 after one epoch, whereas outer gating reaches approximately $\mathbf q\in\mathbb R^d$44. Sigmoid and raw-logit variants remain near $\mathbf q\in\mathbb R^d$45–$\mathbf q\in\mathbb R^d$46, and hard/straight-through gating reaches approximately $\mathbf q\in\mathbb R^d$47. Sparse-scope training reaches approximately $\mathbf q\in\mathbb R^d$48, comparable to full-scope training at approximately $\mathbf q\in\mathbb R^d$49.

SAS often covers less dense attention mass per layer than SeerAttention-R because it is not trained to reproduce dense attention. However, when selected blocks are unioned across layers, SAS achieves higher overlap recall with blocks used by a full-attention oracle. This supports the interpretation that SAS learns more complementary cross-layer routing.

6. Comparisons, limitations, and significance

Relation to other sparsification strategies

SAS differs from several related approaches:

• Dense-attention distillation: trains a selector to reproduce dense attention weights. SAS instead optimizes the final language-modeling loss.
• Fixed sparse patterns: sliding windows, global tokens, strides, and manually specified masks do not learn query-dependent predictive rankings.
• SAC: constructs a discrete adaptive graph using an LSTM edge predictor and policy-gradient training (Li et al., 2020).
• Sparsefinder: predicts high-recall support supersets for exact entmax attention and evaluates sparsity–recall Pareto curves (Treviso et al., 2021).
• S2-Attention: shards context heterogeneously across heads while requiring collective context coverage (Lin et al., 2024).
• Saap: uses asymmetric key partitions and learned query bucket assignment for long-context inference (Mazaré et al., 12 Feb 2025).
• Attention Condensation: trains models so that top-$\mathbf q\in\mathbb R^d$50 attention entries capture nearly all probability mass (Sason et al., 3 Mar 2025).
• SSA: trains sparse and full attention jointly with bidirectional output alignment (Shen et al., 25 Nov 2025).
• Sol-Attn: uses on-the-fly threshold routing and approximate contributions from unselected blocks (Li et al., 27 Jul 2026).
• Counterfactual sparse-attention auditing: evaluates how route selection changes causal content influence, rather than proposing a selector (Ren et al., 3 Aug 2026).

SAS’s distinctive contribution is the direct optimization of context ranking for downstream prediction under a fixed budget.

Limitations

SAS retains several practical and conceptual limitations.

Selector overhead: at very long contexts, selector scoring and Top-$\mathbf q\in\mathbb R^d$51 ranking can dominate sparse decoding. The reported $\mathbf q\in\mathbb R^d$52 measurements show Top-$\mathbf q\in\mathbb R^d$53 selection consuming about $\mathbf q\in\mathbb R^d$54 of the sparse decode step.

Block-summary compression: the selector ranks blocks using compressed or summarized information. On RULER, SAS improves over SeerAttention-R at shorter contexts but remains below full attention at $\mathbf q\in\mathbb R^d$55 and $\mathbf q\in\mathbb R^d$56. The reported explanation is that pooled block representations can miss localized needle signals.

Fixed budgets: a single $\mathbf q\in\mathbb R^d$57 is applied despite query-dependent information requirements. Some queries may need more context, while others may require less.

Training distribution: the principal selector is trained on math data and evaluated on reasoning, long-context, and agentic tasks. Transfer to unrelated domains remains empirical rather than guaranteed.

Sparse training gradients: under sparse-scope training, unselected blocks do not receive direct gate gradients through their own values. Full-scope training provides stronger early supervision but is more expensive.

Interpretability: SAS intentionally learns predictive utility rather than reproducing dense attention. Consequently, its routes may be less directly interpretable as approximations to the original attention pattern.

Inference mismatch: continuous gates are used during training, whereas hard Top-$\mathbf q\in\mathbb R^d$58 routing is used during inference. The quality of this relaxation depends on whether continuous ranking produces a robust ordering under deployment conditions.

Causal integration: selecting salient blocks does not guarantee that their information remains integrated with the rest of the context. Counterfactual evaluation of sparse routing shows that signal concentration and integration loss can coexist: a selector may retain Gold and Poison evidence at similar rates while still altering their causal influence through the removal of cross-block pathways (Ren et al., 3 Aug 2026).

Hardware dependence: efficient deployment requires fused sparse kernels, paged KV caches, block-aware routing, and careful treatment of selector overhead. Theoretical reductions in attention work do not automatically imply proportional wall-clock gains.

SAS is therefore best characterized as an end-to-end, fixed-budget context-ranking method whose inference mechanism is simple but whose training design addresses a central weakness of prior trainable sparsification: hard Top-$\mathbf q\in\mathbb R^d$59 selection prevents the language-modeling loss from directly optimizing context ranking.

Its central principle is to learn which blocks matter under the sparse budget actually used at inference, rather than assuming that the dense model’s attention distribution is the correct target. The most transferable implementation choices are log-space gate injection inside softmax, normalized softmax gates, preservation of continuous selector scores during training, a permanently retained current block, and FlashAttention-style fused execution. The principal unresolved issue is systems-level: once selector scoring and routing are included, the achievable end-to-end benefit depends as much on index construction and kernel efficiency as on the nominal number of retained context units.