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

• URL: http://arxiv.org/abs/2609.13141v1
• Date: Fri, 11 Sep 2026 17:58:06 GMT
• Title: SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking
• Abstract summary: Post-training attention sparsification reduces the quadratic cumulative attention cost of pretrained Transformers.<n>We propose Simple Attention Sparsification (SAS), a sparse attention mechanism that optimize context ranking end-to-end with the language modeling loss.<n>SAS consistently outperforms trainable sparse attention baselines across attention budgets.
• Score: 39.77760016762439
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Sparse Attention for Dense Open-Vocabulary Prediction in CLIP [54.45759517856271]
  We study an inference-time substitution of the row-wise softmax in the final visual self-attention layers.<n>Because entmax applies a data-dependent threshold that maps low scores exactly to zero, it acts as an implicit denoiser.<n>We evaluate on open-vocabulary tasks-dense semantic segmentation (Pascal VOC, Pascal Context, ADE20K) and fine-grained retrieval (FG-OVD)
  arXiv  Detail & Related papers  (2026-07-08T08:30:48Z)
• Beyond Attention Magnitude: Leveraging Inter-layer Rank Consistency for Efficient Vision-Language-Action Models [75.27089289058613]
  We introduce textbfTIES (textbfTau-guided textbfInter-layer textbfEfficient textbfSelection), a dynamic framework guided by inter-layer token ranking consistency.<n>On the CogACT + SIMPLER benchmark, TIES improves average success rates by 6% while reducing token usage by 78%, and demonstrate strong generalization across diverse decoders and benchmarks.
  arXiv  Detail & Related papers  (2026-03-26T02:13:03Z)
• AdaSpot: Spend Resolution Where It Matters for Precise Event Spotting [59.31340724915079]
  Event Spotting is a key task for applications in sports analytics, robotics, and autonomous systems.<n>bfAdaSpot achieves state-of-the-art performance under strict evaluation metrics.
  arXiv  Detail & Related papers  (2026-02-25T16:24:48Z)
• STILL: Selecting Tokens for Intra-Layer Hybrid Attention to Linearize LLMs [23.745366354566315]
  Linearizing pretrained large language models (LLMs) primarily relies on intra-layer hybrid attention mechanisms.<n>We propose STILL, an intra-layer hybrid linearization framework for efficiently linearizing LLMs.
  arXiv  Detail & Related papers  (2026-02-02T14:49:18Z)
• Long-Context Modeling with Dynamic Hierarchical Sparse Attention for On-Device LLMs [17.499497967319332]
  We introduce Dynamic Hierarchical Sparse Attention (DHSA), a data-driven framework that dynamically predicts attention sparsity online without retraining.<n>DHSA matches dense attention in accuracy, while reducing prefill latency by 20-60% and peak memory usage by 35%.<n>Our experiments on Gemma2 with Needle-in-a-Haystack Test and LongBench show that DHSA matches dense attention in accuracy, while reducing prefill latency by 20-60% and peak memory usage by 35%.
  arXiv  Detail & Related papers  (2025-10-28T16:34:18Z)
• DELTA: Dynamic Layer-Aware Token Attention for Efficient Long-Context Reasoning [6.468843780300177]
  We present textbfDELTA, a training-free sparse attention mechanism that achieves computational efficiency without sacrificing model accuracy.<n>Our results show that selective reuse of intermediate attention maps offers a robust path toward efficient long-context reasoning.
  arXiv  Detail & Related papers  (2025-10-10T21:37:49Z)
• ESLM: Risk-Averse Selective Language Modeling for Efficient Pretraining [53.893792844055106]
  Large language model pretraining is compute-intensive, yet many tokens contribute marginally to learning, resulting in inefficiency.<n>We introduce Selective Efficient Language Modeling, a risk-aware algorithm that improves training efficiency and distributional robustness by performing online token-level batch selection.<n> Experiments on GPT-2 pretraining show that ESLM significantly reduces training FLOPs while maintaining or improving both perplexity and downstream performance compared to baselines.
  arXiv  Detail & Related papers  (2025-05-26T12:23:26Z)
• Tactic: Adaptive Sparse Attention with Clustering and Distribution Fitting for Long-Context LLMs [10.52833484759311]
  We propose Tactic, a sparsity-adaptive and calibration-free sparse attention mechanism.<n>It dynamically selects tokens based on their cumulative attention scores rather than a fixed token budget.<n>We show that Tactic outperforms existing sparse attention algorithms, achieving superior accuracy and up to 7.29x decode attention speedup.
  arXiv  Detail & Related papers  (2025-02-17T08:39:43Z)
• SeerAttention: Learning Intrinsic Sparse Attention in Your LLMs [10.702409298302547]
  SeerAttention learns the block-level attention sparsity from the Large Language Models itself.<n>Inspired by the gating mechanism in Mixture of Experts (MoE), SeerAttention augments the conventional attention with a learnable gate.<n>Our evaluation results demonstrate that SeerAttention achieves better model accuracy and lower latency for long-context pre-filling.
  arXiv  Detail & Related papers  (2024-10-17T07:07:09Z)
• Data Selection for Language Models via Importance Resampling [90.9263039747723]
  We formalize the problem of selecting a subset of a large raw unlabeled dataset to match a desired target distribution. We extend the classic importance resampling approach used in low-dimensions for LM data selection. We instantiate the DSIR framework with hashed n-gram features for efficiency, enabling the selection of 100M documents in 4.5 hours.
  arXiv  Detail & Related papers  (2023-02-06T23:57:56Z)
• PCPL: Predicate-Correlation Perception Learning for Unbiased Scene Graph Generation [58.98802062945709]
  We propose a novel Predicate-Correlation Perception Learning scheme to adaptively seek out appropriate loss weights. Our PCPL framework is further equipped with a graph encoder module to better extract context features.
  arXiv  Detail & Related papers  (2020-09-02T08:30:09Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.