SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking

Abstract: Post-training attention sparsification reduces the quadratic cumulative attention cost of pretrained Transformers by selecting a small set of context units (tokens or blocks) for each query. Existing trainable methods usually use a lightweight selector to score context units, followed by hard Top-K selection that blocks gradients from the language modeling loss. Consequently, these methods commonly distill layer-wise dense attention distributions. Although this encourages the selector to rank context units by dense attention weights in the original model, the ranking is not directly aligned with their impact on predictions under a fixed attention budget (i.e., the number of attended context units per query), potentially wasting the limited budget on less useful units. To address this misalignment, we propose Simple Attention Sparsification (SAS), a gated sparse attention mechanism that optimizes context ranking end-to-end with the language modeling loss. The key idea is to inject the selector's continuous scores into attention logits during training, allowing the loss to update the selector through standard backpropagation. We identify several choices crucial for this simple design to work well in practice: placing the gate inside the attention softmax in log form, using normalized softmax gates to calibrate historical context against the always-retained current block, and preserving continuous selector scores so the model learns relative priorities rather than only hard selections. To support long-sequence training, we implement a memory-efficient Triton kernel that integrates SAS into FlashAttention-style computation. Across reasoning, long-context understanding, and agentic tasks, SAS consistently outperforms trainable sparse attention baselines across attention budgets, with especially large gains under tight budgets, demonstrating more effective context ranking for downstream tasks.

• The paper introduces 'SAS', an attention sparsification method that directly optimizes context ranking through language modeling loss, which leads to improved prediction quality.
• SAS employs a unique approach with continuous relaxed selector scores during training that are converted to discrete Top-K ranked outputs for inference.
• When compared to attention-distillation methods, SAS enhances decoding efficiency by up to 5.6 times, especially evident in extremely long contexts—256K and 512K tokens.
• Find recent papers about attention sparsification.

Problem formulation and central claim

Long-context decoding is dominated by repeated reads over the KV cache. For a sequence of length $n$, dense autoregressive attention incurs cumulative quadratic cost, whereas block-sparse attention reduces the per-query cost from all preceding tokens to a selected subset of context blocks. The practical challenge is therefore not merely to impose sparsity, but to learn a ranking of context blocks that preserves the information most relevant to next-token prediction under a fixed attention budget.

“SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking” (2609.13141) argues that existing trainable sparse-attention methods optimize the wrong target. Their selectors are commonly trained by distilling layer-wise attention distributions from the dense model. This supervision encourages recovery of where the dense model allocates attention, but not necessarily which blocks have the greatest effect on the final prediction after many blocks have been removed. The distinction is consequential: dense attention weights ignore value content, downstream nonlinearities, and cross-layer complementarity.

SAS replaces layer-wise attention distillation with direct optimization through the language-modeling loss. The selector still produces scores and inference still uses hard Top-$K$ block selection, but training introduces a continuous relaxation in which selector scores modulate attention logits. The resulting selector is optimized for prediction quality rather than dense-attention imitation.

Differentiable context ranking

Let a lightweight selector assign scores $\mathbf{s}$ to historical context blocks. The inference procedure selects the Top-$K$ blocks and retains the current block unconditionally. Hard selection itself is nondifferentiable: within a region where the ordering of scores is unchanged, the selected set remains constant, so the language-modeling loss cannot provide a useful gradient to the selector.

SAS resolves this blockage by transforming selector scores into positive gates and adding their logarithms inside the attention softmax:

$$\mathbf{o}_{\mathrm{SAS}} = \operatorname{softmax} \left( \mathbf{q}\mathbf{K}_{\mathcal{S}}^\top + \log \mathbf{g}_{\mathcal{S}} \right) \mathbf{V}_{\mathcal{S}}.$$

The historical gates are normalized with a softmax over selector scores, while the always-retained current block has unit gate. During training, the Top-$K$ set defines the sparse routing scope, but the selected blocks retain continuous gate values. During inference, the continuous ranking is discretized into Top-$K$ indices and the gates are removed from the attention computation.

This construction is deliberately simple, but the paper shows that its success depends on four coupled choices:

• Gate placement: gates must be injected inside the attention softmax, not applied afterward to the value vectors.
• Gate activation: historical scores must be normalized competitively with a softmax.
• Ranking preservation: continuous score differences must survive the forward pass rather than being replaced by binary masks.
• Training scope: sparse-scope training is sufficient after convergence and is substantially cheaper than evaluating all blocks.

Figure 1: SAS retains discrete Top-$K$ inference while using continuous log-space gates during training to transmit language-modeling gradients to the selector.

The distinction between inner and outer gating is particularly important. Outer gating rescales value contributions after attention probabilities have already been normalized. It therefore cannot directly reallocate probability mass between competing blocks. Inner log-space gating changes the normalization itself, producing a relative signal based on the difference between a block’s value contribution and the current attention output. This gives the selector a gradient that is aligned with reallocating attention toward blocks that improve the prediction.

Why normalization and continuous scores matter

The gate activation controls the calibration between historical context and the always-retained current block. With $\mathbf{g}=\operatorname{softmax}(\mathbf{s})$, the historical log-gate is

$\log g_m=s_m-\operatorname{LSE}(\mathbf{s}),$

while the current block receives zero bias. The shared normalization term is therefore not canceled: it changes the aggregate historical-to-current attention ratio. It also makes the gates invariant to global shifts in selector logits.

The ablation results show that this calibration is essential. On GPQA-Diamond with Qwen3-4B and a 2048-token budget, the normalized inner-softmax formulation reaches 54.4% after one epoch, compared with 41.6% for outer gating, 17.0% for sigmoid gates, and 18.8% for unnormalized logit injection. The latter two parameterizations tend toward degenerate solutions: sigmoid gates saturate toward one, while raw-logit injection collapses toward small, low-variance logits. Both behaviors approximate ungated attention and erase discriminative block priorities.

Figure 2: Ablations isolate the effects of gate position, gate activation, ranking preservation, and training scope.

Figure 3: Selector-logit dynamics show saturation under sigmoid gating and collapse under raw-logit injection, whereas softmax normalization preserves competition among historical blocks.

Continuous gates also outperform straight-through hard masks. Hard Top-$K$0 forward computation excludes unselected blocks from the softmax normalizer, while a straight-through estimator attempts to assign them surrogate gradients in the backward pass. The paper argues that this produces poorly conditioned signals: an omitted token can have an effectively unbounded surrogate attention weight if its raw attention logit exceeds those of the selected blocks. Such terms can yield large, noisy gradients, especially early in training when the selector ranking is unreliable.

By contrast, continuous gating keeps all selected blocks under a shared normalizer and preserves their relative priorities. The resulting gradients remain bounded and encode ranking information instead of only membership. This explains an important empirical observation: variants with lower training loss can nevertheless produce worse sparse selectors. Minimizing the relaxed loss is insufficient if the relaxation destroys the ordering information needed at inference.

Sparse-scope optimization and implementation

SAS distinguishes between full-scope and sparse-scope training. Full-scope training allows every historical block to participate in gated attention and therefore gives each block a direct content-dependent gradient. Sparse-scope training evaluates only the current block and selected historical blocks. An unselected block then receives only an indirect gradient through the softmax normalization over selector scores.

The sparse signal is less informative during early optimization. In the controlled ablation, sparse scope reaches 24.8% GPQA-Diamond accuracy after 10 steps compared with 30.8% for full scope. However, it catches up over training: after one epoch, sparse scope obtains 54.8%, slightly above the 54.4% of full scope. Across Qwen3-4B, 8B, and 14B and budgets of 1024, 2048, and 4096 tokens, the appendix reports comparable final performance. The practical implication is that full-scope training is not required for the final selector quality, although it may improve early optimization.

Figure 4: Sparse-scope gradients for unselected blocks are noisy and correlated because they lack independent content-dependent signals; full-scope gradients provide more informative updates.

To make long-context training feasible, the authors implement a fused Triton kernel modeled on FlashAttention. The kernel injects block log-gates into tiled $K$1 computation, masks nonselected blocks, performs online softmax accumulation, and collects block-level gate gradients during the backward pass. This avoids materializing a dense gated attention matrix. Top-$K$2 selection is represented through a per-query gate threshold, allowing the kernel to skip blocks without an explicit sort inside the attention computation.

The method is implemented for block size 64, with the backbone frozen in the principal post-training experiments. The same AttnGate selector architecture and inference backend are used for SAS and SeerAttention-R, making the training objective the principal controlled difference.

Reasoning performance under tight budgets

The strongest evidence for SAS comes from reasoning benchmarks, where dropping a small amount of critical context can substantially alter the generated solution. The experiments use Qwen3-4B, 8B, and 14B and evaluate MATH500, GPQA-Diamond, AIME24, and AIME25.

At a 1024-token budget, SAS substantially exceeds SeerAttention-R, despite using the same selector architecture:

Benchmark	Qwen3-4B	Qwen3-8B	Qwen3-14B
MATH500 improvement	+5.98	+7.70	+6.81
GPQA-Diamond improvement	+10.57	+13.74	+15.50

At 2048 tokens, the gains remain pronounced. On Qwen3-4B, SAS improves over SeerAttention-R by 13.02 points on AIME24, reaching 68.85% versus 55.83%, and by 11.22 points on AIME25, reaching 56.38% versus 45.16%. At 4096 tokens, SAS frequently matches or exceeds dense attention. For example, Qwen3-4B reaches 71.72% on AIME24 with SAS, compared with 71.25% under full attention; Qwen3-14B reaches 78.28%, compared with 78.91% for dense attention.

These results support the paper’s central causal interpretation: when the backbone, selector architecture, data, and inference procedure are matched, direct next-token-prediction training produces a more useful sparse ranking than attention-distribution distillation. The gains are largest when the budget is restrictive, precisely where ranking errors have the greatest effect.

Transfer to long-context and agentic tasks

The selector is trained only on OpenR1-MATH-220K, yet it is evaluated without task-specific retraining on LongBench. SAS consistently improves over SeerAttention-R. The largest reported margin occurs for Qwen3-14B at a 2048-token budget on inputs longer than 8K tokens: SAS scores 53.9 compared with 51.5 for SeerAttention-R, a 2.4-point improvement. At a 4096-token budget, SAS nearly recovers dense performance, reaching 56.2 versus 56.6 for full attention on Qwen3-14B.

On BFCL Multi-Turn, SAS improves over SeerAttention-R at every reported model scale and budget. At 2048 tokens, the Qwen3-4B score increases from 29.00 to 32.50, a 3.5-point gain. At 4096 tokens, Qwen3-14B reaches 44.00 compared with 43.88 for SeerAttention-R and 44.50 for full attention.

The VitaBench results are less uniformly favorable, but still support improved routing. At a 4096-token budget, SAS leads SeerAttention-R on most Delivery, Instore, and OTA metrics. The exception is not hidden by aggregate reporting: some individual pass metrics remain below dense attention, and the advantage varies by scenario. Thus the agentic evidence indicates robustness, not universal recovery of full-attention behavior.

Continued pretraining extension

The paper also tests SAS during continued pretraining rather than selector-only post-training. Starting from OLMo3-7B, the authors jointly train the backbone and selector for 13,000 steps on approximately 50B tokens, using 8192-token sequences, block size 64, and Top-$K$3.

SAS-RoPE obtains an average score of 43.28 across general knowledge, mathematics, and code tasks. This is close to the dense OLMo3-Base score of 43.88, slightly above the sliding-window model’s 43.24, and well above HiLS-Attn-RoPE’s 41.68. On LongBench, SAS ties HiLS-Attn-RoPE at 30.0 and exceeds dense OLMo3-Base at 29.0 and sliding-window continued pretraining at 28.0. The gains are concentrated on inputs longer than 8K tokens.

This experiment is preliminary relative to the post-training comparison: it changes the selector architecture, trains the backbone, and follows a separate continued-pretraining protocol. It therefore establishes compatibility with joint training rather than isolating the effect of the SAS objective under the same controlled conditions.

Learned routing and generation behavior

The analysis challenges the assumption that a good sparse selector should maximize per-layer dense attention-mass coverage. SAS covers less raw attention mass than SeerAttention-R, both for ordinary attention weights and for weights augmented by value-vector norms. This is expected because distillation explicitly trains toward dense attention distributions, whereas SAS optimizes the final prediction loss.

However, when selections are aggregated across layers, SAS obtains higher recall against the full-attention oracle. The result indicates that SAS learns more complementary routing patterns: individual layers may retain less of the dense model’s attention mass, but the union of their selected blocks recovers more of the context used by the dense network. This is consistent with the paper’s criticism of layer-wise distillation, which does not directly optimize cross-layer coordination.

SAS also generates shorter reasoning traces and exhibits lower truncation rates than SeerAttention-R at a 4096-token budget on Qwen3-4B. The effect is largest on AIME24 and AIME25. The appropriate interpretation is limited but meaningful: improved context selection appears to reduce the need for redundant or prolonged reasoning. The experiment does not establish that shorter generation is intrinsically better; it shows that, under this evaluation protocol, SAS reaches completed answers with fewer generated tokens.

Figure 5: SAS produces shorter reasoning generations and fewer maximum-length truncations than the distillation-trained selector across four reasoning benchmarks.

Decode efficiency

SAS reduces decode-time attention from scanning the entire KV cache to reading a fixed number of selected blocks. In a single-GPU Qwen3-4B serving experiment using SGLang, CUDA graphs, and steady-state generation, the method is nearly equivalent to dense attention at 8K context but becomes 2.4 times faster at 64K, 4.6 times faster at 256K, and 5.6 times faster at 512K for batch size one. At batch size eight, the speedup reaches approximately 13 times at 64K.

Figure 6: Sparse decode latency remains nearly context-invariant while dense attention grows with KV-cache length; the reported maximum speedup is 5.6 times at batch one and approximately 13 times at batch eight.

These measurements isolate decode: prefill remains dense. Moreover, sparsification does not eliminate all context-length-dependent work. Selector scoring scans block summaries, and Top-$K$4 selection must rank all candidate blocks. At 512K, Top-$K$5 selection accounts for approximately 90% of sparse-decode latency, compared with 21% at 8K; selector scoring also approaches the cost of the sparse attention computation. Consequently, the asymptotic advantage of sparse attention is partly transferred to routing overhead, and the reported end-to-end speedups depend on optimized selection kernels and serving configuration.

Figure 7: At very long contexts, selector scoring and Top-$K$6 ranking dominate sparse-decode latency rather than the selected-block attention computation.

Limitations and open questions

The principal limitation is long-context retrieval fidelity. On RULER, SAS improves over SeerAttention-R across many settings but remains far below full attention as context length reaches 128K. For example, at a 2048-token budget on Qwen3-14B and 128K context, SAS reaches 23.80 compared with 14.95 for SeerAttention-R, but full attention reaches 82.23. Even with a 4096-token budget, SAS reaches 29.95 versus 82.23 for dense attention. The authors attribute this degradation to pooled block summaries that may discard localized needle-like information.

This limitation qualifies the stronger reasoning and agentic results. The selector transfers effectively to the tested tasks, but its block-level representation is not sufficient for arbitrary long-context retrieval. The paper leaves open how to construct more expressive selectors without removing the computational advantage of block compression. It also leaves unresolved whether full-scope training can materially improve very-long-context retrieval, since the reported convergence equivalence is established mainly on the evaluated reasoning settings. Finally, the speed results identify Top-$K$7 ranking as the dominant bottleneck at extreme context lengths, so the practical benefit depends on whether routing can be accelerated without weakening selection quality.

Conclusion

SAS formulates sparse attention as end-to-end context ranking rather than dense-attention imitation. Its central mechanism is a continuous, normalized, log-space gate inserted inside the attention softmax during training, followed by hard Top-$K$8 routing at inference. Controlled ablations show that inner placement, softmax normalization, preservation of score differences, and sparse-scope optimization are all important to obtaining a useful selector.

Across Qwen3 reasoning, LongBench, BFCL, and VitaBench, SAS consistently improves over a matched distillation-based baseline, with the largest gains under tight budgets and several cases approaching or matching full attention. Its fused Triton implementation yields substantial decode speedups at long contexts, although selector scoring and Top-$K$9 ranking become the dominant costs. The RULER results establish a clear boundary: end-to-end prediction optimization improves sparse routing, but pooled block summaries remain inadequate for reliable fine-grained retrieval at extreme context lengths.