SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking
Summary
SAS introduces a gated sparse attention mechanism that optimizes context ranking end-to-end with language modeling loss, improving performance in reasoning and long-context tasks under tight attention budgets.
View Cached Full Text
Cached at: 09/14/26, 02:33 AM
Paper page - SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking
Source: https://huggingface.co/papers/2609.13141
Abstract
SAS improves sparse attention by training a selector end-to-end with language modeling loss via continuous gating inside attention softmax, yielding better context ranking under tight budgets.
Post-trainingattention sparsificationreduces 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 hardTop-K selectionthat 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 SimpleAttention Sparsification(SAS), agated sparse attentionmechanism that optimizes context ranking end-to-end with the language modeling loss. The key idea is to inject the selector’s continuous scores intoattention logitsduring 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 normalizedsoftmax gatesto 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-efficientTriton kernelthat integrates SAS intoFlashAttention-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.
View arXiv pageView PDFGitHub3Add to collection
Get this paper in your agent:
hf papers read 2609\.13141
Don’t have the latest CLI?curl \-LsSf https://hf\.co/cli/install\.sh \| bash
Models citing this paper0
No model linking this paper
Cite arxiv.org/abs/2609.13141 in a model README.md to link it from this page.
Datasets citing this paper0
No dataset linking this paper
Cite arxiv.org/abs/2609.13141 in a dataset README.md to link it from this page.
Spaces citing this paper0
No Space linking this paper
Cite arxiv.org/abs/2609.13141 in a Space README.md to link it from this page.
Collections including this paper0
No Collection including this paper
Add this paper to acollectionto link it from this page.
Similar Articles
Hierarchical Sparse Attention Done Right: Toward Infinite Context Modeling
Introduces HiLS Attention, a chunk-wise sparse attention mechanism for LLMs that learns chunk selection end-to-end via LM loss, achieving performance comparable to full attention while enabling ultra-long-context extrapolation and faster inference.
Simplified Sparse Attention via Gist Tokens
This paper introduces Simplified Sparse Attention (SSA), a method that uses gist tokens during continued pretraining to enable efficient chunk selection at inference without architectural changes, achieving high compression ratios and outperforming baselines on long-context tasks like LongBench and retrieval-augmented generation.
CoSA: Accelerating Long-Context Inference via Proxy-Kernel Co-Designed Sparse Attention
CoSA proposes a training-free sparse attention method that co-designs a kernel-aware proxy and an ordered-skipping kernel to accelerate long-context inference, achieving up to 4.93× attention speedup and 2.53× end-to-end TTFT reduction with negligible performance degradation.
Elastic Threshold Attention: Learned Contextual Sparsity for Long-Context Decoding
Elastic Threshold Attention (ETA) is a trainable sparse attention architecture that improves long-context decoding speed without quality degradation by using dynamic thresholds predicted from query representations.
Subquadratic AI introduces SubQ-1.1-Small, a new model using Smart Sparse Attention
Subquadratic AI introduces SubQ-1.1-Small, a model leveraging Smart Sparse Attention to achieve near-perfect long-context retrieval up to 12M tokens with up to 1,000x attention compute reduction. It balances long-context optimization with strong general reasoning, outperforming baselines on benchmarks like NIAH and RULER.