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

Researchers introduced Simple Attention Sparsification (SAS), a gated sparse attention mechanism that optimizes context ranking end-to-end with the language modeling loss by injecting the selector's continuous scores into attention logits during training. SAS outperforms trainable sparse attention baselines across attention budgets on reasoning, long-context understanding, and agentic tasks, with especially large gains under tight budgets, according to the paper. The method uses a memory-efficient Triton kernel that integrates SAS into FlashAttention-style computation to support long-sequence training.

Text Generation • Updated • 6 /tencent/Simple-Attention-Sparsification Papers /papers SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking 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. thinkingmachines/Inkling-Small /thinkingmachines/Inkling-Small Post-training attention sparsification /papers?q=attention%20sparsification 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 /papers?q=Top-K%20selection 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 /papers?q=Attention%20Sparsification SAS , a gated sparse attention /papers?q=gated%20sparse%20attention 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 /papers?q=attention%20logits 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 /papers?q=softmax%20gates 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 /papers?q=Triton%20kernel that integrates SAS into FlashAttention /papers?q=FlashAttention -style computation. Across reasoning, long-context understanding /papers?q=long-context%20understanding , 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. 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 paper 1 Datasets citing this paper 0 No dataset linking this paper Spaces citing this paper 0 No Space linking this paper