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

SAS. Simple Attention Sparsification via End-to-End Optimization of Context Ranking. By Zhiwei Li and eight others, from Tencent HY LLM Frontier, the Hong Kong University of Science and Technology, Guangzhou, and one other institution. Published on the eleventh of September. Its code repository is worked out: about forty source files, with a readme.

Post-training attention sparsification reduces the quadratic cumulative cost of attention in a pretrained Transformer, by having each query attend to a small set of context units, either tokens or blocks. Existing trainable methods usually score those units with a lightweight selector, then apply a hard Top-K selection. That hard step blocks gradients from the language modeling loss, so these methods commonly distill dense attention distributions, layer by layer, instead. The selector cannot learn from the loss the model actually cares about, so it is given a stand-in target: copy where the full, dense attention would have looked.

The trouble lies in what the distillation teaches. The selector learns to rank context units by their dense attention weights in the original model. Under a fixed attention budget, meaning a fixed number of attended units per query, that ranking is not directly aligned with how much each unit affects the predictions. Some of a limited budget can end up spent on units that help less.

Simple Attention Sparsification, or SAS, is a gated sparse attention mechanism that optimizes context ranking end-to-end, with the language modeling loss itself. The key idea is to inject the selector's continuous scores into the attention logits during training. The loss can then update the selector through standard backpropagation. Nothing exotic is needed to get a gradient through: the scores become part of the attention computation itself, so the ordinary training signal reaches them the same way it reaches every other weight in the model.

Several choices turn out to be crucial for this simple design to work well in practice. The gate is placed inside the attention softmax, in log form. The gates are normalized softmax gates, which calibrate historical context against the current block, a block that is always retained. And the selector's scores are kept continuous, so the model learns relative priorities rather than only hard selections.

To support training on long sequences, the authors implement a memory-efficient Triton kernel that integrates SAS into FlashAttention-style computation. A method aimed at long contexts has to be trainable on long contexts, and the kernel is what makes that practical.

Across reasoning, long-context understanding and agentic tasks, SAS consistently outperforms trainable sparse attention baselines across attention budgets. The gains are especially large when budgets are tight, which the authors take as evidence of more effective context ranking for downstream tasks. That pattern fits the diagnosis. When only a few units can be attended, spending one of them on a less useful unit costs the most, so a better ranking should show its value most clearly there.

The simplicity in the name is about what the selector answers to. Instead of a layer-wise distillation target, it answers to the same loss as the model. The fix is to stop asking the selector to imitate the dense model, and start letting the loss train it. Once the scores carry gradient, the ranking is shaped by what improves predictions under the budget, which is the thing a sparse model is finally judged on.