PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Zhiwei Li, Lei Zhu, Hao Gu, Xiang Hu, Yan Wang, Haitao Mi, Sirui Han, Leo Liang, Zhijiang Guo

arXiv:2609.13141v1cs.CL

TL;DR

Post-training sparse attention must rank context units by their effect on predictions under a limited budget, but hard Top-K selection blocks language-modeling gradients and encourages indirect dense-attention imitation. SAS injects continuous selector scores into attention logits and optimizes the ranking end-to-end, with design choices for calibrated gates and efficient long-context training. Across reasoning, long-context, and agentic tasks, SAS consistently outperforms trainable sparse-attention baselines, especially at tight budgets, while performance still degrades at very long contexts.

  • Problem

    Existing trainable sparse-attention selectors use hard Top-K selection and often distill dense attention, which may misalign context ranking with prediction impact under a fixed attention budget.

  • Method

    SAS injects continuous selector scores into attention logits during training and uses log-space normalized gates, preserved score differences, and a fused Triton kernel.

  • Results

    SAS consistently outperforms existing post-training sparsification baselines across reasoning, long-context understanding, and agentic tasks, with gains exceeding 10% under low attention budgets.

  • Takeaways & Limitations

    End-to-end language-modeling optimization provides more effective sparse context routing than layer-wise attention distillation across the evaluated tasks and budgets.

  • Takeaways & Limitations

    At increasing context lengths, SAS performance degrades substantially and remains considerably below full attention, likely because pooled block summaries lose fine-grained localized information.

Abstract

from arXiv · show

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.

1 Introduction

SAS addresses the mismatch between dense-attention imitation and prediction-aware context selection by training selectors end-to-end with the language modeling loss. Its gated design preserves differentiable ranking and delivers broad gains, especially under tight attention budgets.

  • Motivation: Post-training sparsification targets the quadratic cumulative attention cost of dense pretrained Transformers without requiring architectural changes or retraining.It selects a small set of context units for each query after dense pretraining.
  • Approach: The method injects continuous selector scores into attention logits, using log-space inner gating, normalized activations, and preserved score differences.These choices make gradients informative for relative context priorities while retaining discrete Top-K selection.
  • Efficiency: A fused Triton kernel adds log-space gates during FlashAttention-style tile computation, avoiding materialization of the full attention matrix for long-context training.SAS can therefore be implemented by replacing the attention kernel with the gated-attention kernel while optimizing standard language modeling loss.
  • Results: 6.0–7.7 points on MATH500 and 10.6–15.5 points on GPQA-Diamond at a 1024-token budget versus SeerAttention-R.These gains are reported across Qwen3-4B, 8B, and 14B models.
  • Approach: SAS trains context selectors directly with language modeling loss, replacing layer-wise attention distillation and avoiding teacher attention or auxiliary distillation.Controlled comparisons use matched backbones, selector architectures, and training data.
  • Results: SAS consistently improves performance across reasoning, long-context understanding, and agentic tasks, with gains exceeding 10% under low attention budgets.The reported task families include MATH500, GPQA-Diamond, AIME24, AIME25, LongBench, BFCL, and VitaBench.

2 Related Work

Related work reduces attention cost through inference-time selection or trainable sparse routing. SAS differs by using continuous log-space gates inside attention to optimize selectors directly with the language modeling loss.

  • Inference-time selection: Inference-time sparsification and KV-cache methods select relevant context or cached tokens to reduce long-context attention cost without updating the backbone LLM.KV-cache selection can avoid irreversible eviction during long reasoning generation.
  • Trainable sparse attention: Trainable sparse-attention methods learn routing or gating patterns, while earlier sparse Transformers commonly impose fixed layouts and require adaptation around those patterns.SAS is positioned within this trainable-routing line of work.
  • SAS: SAS relaxes hard block selection into continuous log-space gates inside attention, enabling end-to-end selector optimization with the language modeling loss.This distinguishes its training signal from methods that rely on indirect selection supervision.

3 Preliminaries

Block sparse attention reduces quadratic attention cost by restricting each query to selected context blocks. A lightweight selector scores blocks, but hard Top-K selection prevents the language modeling loss from training that selector directly.

  • Standard attention: Standard autoregressive attention attends to all previous keys and values, giving cumulative sequence cost O(n^2).The number of preceding tokens grows linearly with each decoding step.
  • Block sparse attention: Block sparse attention restricts each query to a selected subset of context blocks, reducing single-query cost from O(n) to O(|S|) and sequence cost to O(n|S|).Context positions are partitioned into contiguous blocks before selection.
  • Block selection: A lightweight selector scores C context blocks and chooses the K most important blocks, whose tokens form the attended set.The blocks contain b tokens, with C = n/b.
  • Training challenge: Under hard Top-K, the selected index set is piecewise constant in selector scores, so language modeling loss provides no useful gradient through selection.Making selector scores affect attention differentiably is therefore the core training problem.

4 Method

SAS reframes sparse attention as differentiable context ranking, letting language-modeling loss train selectors through gated attention. Its effectiveness depends on where and how gates are applied, whether ranking information is preserved, and the training scope.

  • From Sparse Selection to Context Ranking: SAS trains a lightweight selector to rank context blocks from the language modeling loss, replacing nondifferentiable hard Top-K supervision with differentiable ranking.The selector learns a continuous ordering before the final sparse block selection.
  • From Sparse Selection to Context Ranking: Historical-block scores become positive gates while the always-retained current block keeps a unit gate, allowing attention gradients to reach the selector.Each historical-block gate is broadcast to its tokens and modulates attention during training.
  • Components of Learnable Ranking: Inner softmax gating, normalized softmax activation, continuous scores, and sparse-scope training define the practical design choices evaluated for learnable ranking.The ablation varies gate position, activation, ranking preservation, and training scope under controlled Qwen3-4B experiments on GPQA-Diamond.
  • Ablation Analysis: Inner softmax gating outperforms outer gating, while normalized gates and continuous scores improve ranking quality; sparse scope converges more slowly but reaches comparable final performance at lower cost.Normalized gates calibrate historical context against the unit-gated current block, whereas continuous scores preserve fine-grained ranking information.
  • Ablation Analysis: Sigmoid gates saturate toward one and raw-logit injection collapses selector logits toward zero, making attention effectively ungated and weakening discriminative ranking signals.The competitive normalization of softmax avoids these trivial solutions.
  • Ablation Analysis: Hard gating can produce unbounded gradients when high-score blocks fall outside Top-K, whereas sparse-scope training eventually matches full-scope ranking quality with lower training cost.The hard-gating instability arises because dropped-block contributions can accumulate under the restricted forward normalizer.

5 Experiments

SAS is evaluated against dense attention and sparse baselines across reasoning, long-context understanding, agentic tasks, and continued pretraining. It consistently improves sparse-attention performance, especially under tight budgets, while often approaching or matching dense attention.

  • 5.1 Experimental Setup: Experiments primarily freeze the backbone and train only the selector, using Qwen3-4B/8B/14B models and a selector trained on OpenR1-MATH-220K for one epoch.The same selector is evaluated across reasoning, long-context, and agentic tasks; continued pretraining jointly updates the backbone and selector.
  • 5.2 Post-Training Results: SAS outperforms sparse baselines across nearly all reasoning budgets, backbones, and tasks, with especially large gains over SeerAttention-R on harder problems.At budget 2048 on AIME24 with Qwen3-4B, SAS improves over SeerAttention-R by 13.0 points; Quest reaches 0 on AIME24/25 while SAS remains strong.
  • 5.2 Post-Training Results: At budget 4096, SAS matches or exceeds full attention on several reasoning cases, including 71.72 versus 71.25 on AIME24 with Qwen3-4B.This result is achieved while attending to only a small fraction of the KV blocks.
  • 5.2 Post-Training Results: SAS transfers from math-only selector training to LongBench, outperforming SeerAttention-R at nearly every budget and backbone and improving the Qwen3-14B 8K+ bucket by 2.4 points at budget 2048.At budget 4096, its Qwen3-14B LongBench average is 56.2 versus 56.6 for full attention.
  • 5.2 Post-Training Results: On agentic evaluations, SAS beats SeerAttention-R across BFCL backbones and budgets and leads on most VitaBench metrics at budget 4096 while approaching full attention.On BFCL, the Qwen3-4B gain is 3.5 points at budget 2048; on Qwen3-14B at budget 4096, SAS scores 44.00 versus 44.50 for full attention.
  • 5.3 Extension to Continued Pretraining: In continued pretraining, SAS achieves the best average score among sparse continued-pretrained models and tied-best LongBench performance, with gains concentrated on inputs longer than 8K.Its average downstream score is 43.28, compared with 41.68 for HiLS-Attn-RoPE and 43.88 for dense OLMo3-Base; its LongBench average is 30.0 versus 29.0 for the dense base.

6 Analysis

SAS selects blocks for downstream prediction rather than dense-attention coverage, producing more complementary cross-layer selections, shorter reasoning traces, and large decode speedups.

  • 6.2 Generation Behavior: SAS generates shorter reasoning traces and truncates less often than SeerAttention-R at a 4096-token budget, especially on harder AIME tasks.These outcomes are consistent with more effective context selection reaching answers with fewer generated tokens.
  • 6.1 Learned Block Selection: SAS achieves higher cross-layer overlap recall than distillation across context lengths and budgets, despite covering less attention mass within individual layers.End-to-end optimization makes selections more complementary across layers, improving recovery of blocks attended by the full-attention oracle.
  • 6.3 End-to-End Decode Efficiency: At batch 8, SAS reaches approximately 13× decode speedup at 64K context, while tighter budgets provide only modest additional gains.The main benefit comes from avoiding full KV reads rather than from the exact token budget.
  • 6.3 End-to-End Decode Efficiency: At batch 1, SAS achieves up to 5.6× lower latency at 512K context, while dense-attention latency grows linearly and SAS latency remains nearly constant.These measurements use steady-state decode on Qwen3-4B with a single GPU and CUDA graphs.
  • 6.3 End-to-End Decode Efficiency: Top-K selection becomes the dominant decode bottleneck, rising from 21% of step cost at 8K to 90% at 512K.Selector scoring also grows with context length, whereas attention compute remains context-invariant under a fixed token budget.

7 Conclusion

SAS replaces layer-wise attention distillation with end-to-end language-modeling optimization, using calibrated continuous gates and fused kernels for practical long-context sparsification.

  • 7 Conclusion: SAS shifts context selection from layer-wise attention distillation to direct language-modeling-loss optimization, capturing cross-layer dependencies more effectively.Its effective gradient flow relies on log-space gate injection, normalized gate activation, and continuous rankings.
  • 7 Conclusion: SAS outperforms existing sparse-attention baselines across reasoning, long-context understanding, and agentic tasks, particularly under tight attention budgets.Fused Triton kernels make the approach computationally practical for long-context training.

A Training Dynamics

Training dynamics show that gate placement and score parameterization strongly affect optimization, while some lower-loss variants fail to learn discriminative block rankings.

  • A Training Dynamics: Inner-softmax gating achieves lower training loss than outer gating, while independent sigmoid gates and raw-logit injection achieve even lower loss but fail to improve sparsification.Those variants effectively approach ungated attention and learn little discriminative block ranking.
  • A Training Dynamics: Hard gating reduces training loss but substantially increases gradient norm, indicating a less stable optimization regime.The comparison is part of the controlled loss and gradient-norm ablations.
  • A Training Dynamics: Sparse-scope training begins with higher loss than full-scope training and gradually approaches it, while adding noise produces little difference.These dynamics distinguish scope effects from noise effects during training.

B Limitation

On RULER, SAS improves over SeerAttention-R at shorter contexts but degrades substantially as context grows, leaving a considerable gap from full attention.

  • B Limitation: SAS consistently improves over SeerAttention-R at shorter context lengths on RULER, but its performance degrades substantially as context length increases.The evaluation trains selectors on 64K sequences and extends evaluation contexts to 128K with YaRN.
  • B Limitation: Pooling-based block summaries may lose fine-grained localized signals as context grows, motivating more expressive but efficient selectors.The authors identify this as a primary reason for the remaining gap from full attention.

C Final Performance of Sparse vs. Full Training Scope

Sparse training scope reaches the same final accuracy as full training scope across model scales and token budgets, while reducing training cost.

  • Sparse scope matches full scope at convergence across Qwen3-4B/8B/14B and 1024, 2048, and 4096 token budgets, while incurring lower training cost.Table 9 reports final accuracy comparisons across all listed model scales and budgets.

D Gradient Derivation

The gradient derivation contrasts outer and inner gate injection and analyzes how soft or hard gating and routing scope affect selector gradients.

  • Inner gating injects the gate into attention logits, whereas outer gating applies it after attention probabilities are computed.The derivation distinguishes the two placements through their resulting gradient paths.
  • Soft gating keeps all blocks active, so each block remains in the normalizer and gate gradients stay bounded.The full-context normalizer includes every block, including the block receiving the gradient.
  • Hard gating with STE excludes unselected blocks from the forward normalizer, allowing their ranking signal to grow exponentially when they outscore selected blocks.The unbounded behavior begins when a dropped block's score exceeds the maximum selected-block score.
  • The derivation uses the softmax backward formulation to express gate gradients under both sparse and full routing scopes.The analysis first defines the softmax activation and then applies its backward identity to each scope.
  • Under sparse routing scope, unselected blocks have no output dependence and therefore receive no gate gradients, while full routing scope provides gradients for every block.Sparse scope updates only selected blocks directly; full scope retains gradient paths for unselected blocks.

E Kernel Implementation

The fused SAS kernel integrates log-gate injection, threshold-based Top-K masking, online softmax, and gradient accumulation into tiled FlashAttention-style computation.

  • The kernel adds per-query block log gates inside attention scores, masks blocks below a threshold, and skips tiles selected by no query.The current block uses a unit gate and causal masking, while earlier blocks are gated and thresholded.
  • The presented pseudocode omits grouped query attention and variable-length packing, which the full implementation supports.The simplification focuses on gate injection and sparse scope.
  • The forward pass traverses causal key-value tiles, computes gated scores, and updates attention online with row-wise maxima, exponentials, and accumulators.This tiled procedure avoids materializing the dense attention matrix.
  • The backward pass recomputes gated probabilities from saved log-sum-exp values and returns gradients for Q, K, V, and block log gates.Gate gradients are accumulated over tokens within each selected block and propagated to the selector.
  • The implementation uses separate tiled passes for query and gate gradients and for key and value gradients.The backward pseudocode loops over query tiles in the first pass and key-value tiles with attending queries in the second.
Loading 2609.13141v1…