Post-training attention sparsification is a common way to cut the quadratic cost of attention in pretrained Transformers: instead of attending to every context unit, the model selects a small set of tokens or blocks for each query. Existing trainable approaches typically rely on a lightweight selector that scores context units, followed by a hard Top-K selection step.
That hard selection blocks gradients from the language modeling loss, so these methods usually fall back on distilling layer-wise dense attention distributions. As the paper notes, this encourages the selector to imitate dense attention weights rather than to rank units by how much they actually affect predictions under a fixed attention budget — a misalignment that can waste the limited budget on less useful context.
To address this, the authors propose Simple Attention Sparsification (SAS), a gated sparse attention mechanism that optimizes context ranking end-to-end with the language modeling loss. The core idea is to inject the selector's continuous scores into the attention logits during training, letting standard backpropagation update the selector directly.
The paper highlights several design choices it says are crucial for this simple approach to work well: 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 instead of only hard selections.
To support long-sequence training, the authors implement a memory-efficient Triton kernel that integrates SAS into FlashAttention-style computation. Across reasoning, long-context understanding, and agentic tasks, SAS is reported to consistently outperform trainable sparse attention baselines across attention budgets, with especially large gains under tight budgets — evidence, the authors argue, of more effective context ranking for downstream tasks.