Paper Walkthrough: SAS — Training a Sparse-Attention Block Selector End-to-End With the Language Modeling Loss

The selector that decides which context blocks a query reads has traditionally been trained by imitating the dense model's attention distribution. SAS injects the selector's scores into the softmax as log-space gates, so the gradient that Top-K used to block flows straight from next-token prediction into the selector. We follow the paper through its four-axis ablation, the gradient formulas behind each choice, and the 10.6–15.5 point GPQA-Diamond gain at a 1024-token budget.

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

The paper's claim in one paragraph

This walkthrough covers "SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking" (arXiv:2609.13141, published 11 September 2026) from Tencent HY LLM Frontier and the Hong Kong University of Science and Technology. The abstract argues as follows. Post-training attention sparsification cuts the quadratic attention cost of a pretrained Transformer by having a lightweight selector pick a small set of context blocks per query, but the hard Top-K step blocks gradients from the language modeling loss, so existing trainable methods fall back on distilling the dense model's layer-wise attention distributions. That ranking is not aligned with how much each block actually helps the prediction under a fixed attention budget, so the budget gets wasted. SAS injects the selector's continuous scores into the attention logits during training, letting standard backpropagation update the selector. Three choices turned out to matter: putting the gate inside the softmax in log form, normalizing the gates with a softmax so historical blocks are calibrated against the always-retained current block, and keeping the scores continuous so the model learns relative priorities. A FlashAttention-style Triton kernel makes long-sequence training feasible, and across reasoning, long-context and agentic tasks SAS consistently beats trainable sparse-attention baselines, with the largest gains at tight budgets.

Below we go through why the gradient stops, how a log-space gate lets it through, and how each of the four design decisions is backed by the experiments.

An analogy: how do you train the person who decides who gets heard?

In Attention from scratch we described self-attention as note-taking in a meeting: compare your question (Query) against everyone's name tag (Key), and listen longer to the people whose tags match, mixing their statements (Values) accordingly.

When the meeting runs long and there are tens of thousands of participants, comparing every tag every time is out of the question. So you appoint a gatekeeper who says "for this question, only listen to those few rows of seats." That is the selector in block-sparse attention: it scores groups of seats, called blocks, and only the top K blocks get attended.

The question is how to train this gatekeeper. The standard approach has been to have them imitate where the teacher (the original dense attention) was looking: record the teacher's gaze at every layer and train the gatekeeper's scores to match. That is attention distillation.

The paper's point is that this training signal is misaligned. Where the teacher looked is not the same as which blocks are actually needed to get the final answer right when you can only afford K of them. Ideally, if the gatekeeper picks badly, the answer comes out wrong and the gatekeeper is held directly responsible. That is exactly what SAS does: it folds the gatekeeper's scores into the attention computation so the final next-token loss can push gradients all the way back into the selector.

Background: block-sparse attention, and why Top-K stops the gradient

First the notation from §3. With $\mathbf{q}$ the query of the current token and $\mathbf{K},\mathbf{V}$ the keys and values of the $n$ previous tokens, plain attention is (scale factor omitted, as in the paper):

In words, this is the plain attention recipe: take the dot product of the question with every name tag, turn those scores into proportions, and mix everyone's statements in those proportions.

In words: take the dot product of the question with every name tag, turn the scores into proportions, and mix everyone's statements in those proportions. Since generation does this for every token, the cumulative cost over a sequence of length $n$ grows as $n^2$.

Block-sparse attention splits the $n$ context tokens into $C=n/b$ blocks of $b$ tokens each, has a selector $\mathcal{R}_\theta$ produce a score per block $\mathbf{s}\in\mathbb{R}^C$, and attends only to the union $\mathcal{S}$ of the top $K$ blocks.

In words, pick the K highest-scoring blocks, gather their tokens into one set, and run attention over only those tokens.

That is, pick the K highest-scoring blocks and run attention over their tokens only. Per query the cost drops from $n$ to $|\mathcal{S}|$, and over the whole sequence it becomes proportional to $n|\mathcal{S}|$ (§3).

Here is the crux: Top-K is piecewise constant in the scores. Nudge a score slightly and, unless the membership of the top K changes, the output does not move at all. The derivative is zero, so the loss carries no information about how the selector should change its scores. The paper draws this as Figure 1(a), "Gradient Blockage in Discrete Selection". Distillation and heuristic proxy signals exist precisely to route around this wall (§1, §2).

What exactly is wrong with distillation

Section 1 names two limitations of distilling attention distributions.

First, it ignores complementarity across layers. A block missed at one layer could be picked up at another, but when each layer is trained to mimic its own teacher, every layer independently chases "the same place the teacher looked."

Second, it ignores the values. Even if the attention weights match, how the mixed $\mathbf{V}$ contents affect the final prediction is not part of the distillation objective. In the paper's Figure 1, the $\mathbf{V}$ matrix is simply absent from the training signal.

SAS resolves both with the most naive possible move: train the selector on the language modeling loss itself. Cross-layer division of labor and the contents of the values are both baked into the final loss. All that remains is to open a path for the gradient.

The mechanism: scores become log-space gates inside the softmax

Section 4.1 first reframes sparsification as context ranking rather than selection. Top-K is determined by the ordering of scores, so what the selector really needs to learn is a continuous ordering of "how important is each block." During training that ordering participates in attention as a continuous quantity; only at inference is it cut down to K.

Concretely, let $B_0$ be the always-retained current block and $\mathcal{H}=\{B_1,\ldots,B_C\}$ the historical blocks the selector ranks. The historical scores $\mathbf{s}$ pass through a differentiable activation $\phi$ to become positive gates.

In words, each historical block gets a positive weight derived from its score, while the current block's weight is pinned to 1 so it carries no bias.

Which says: historical blocks get a positive weight derived from their score, while the current block's weight is pinned to 1 so it carries no bias. Each gate $g_m$ is broadcast to every token in its block.

Then the gate's logarithm is added to the attention logits. The final form in §4.3 (equation 9) is:

In words, add the log of the gatekeeper's weight to the dot-product scores before the softmax. Because addition inside an exponential is multiplication outside it, this is the same as multiplying each post-softmax attention probability by its block's gate, so the loss gradient can reach the selector.

$\mathbf{g}_{\mathcal{S}}$ is the selected blocks' gates expanded to token level, and $\mathcal{S}=B_0\cup\bigcup_{m\in\mathcal{I}}B_m$ is the current block plus the selected historical blocks. Read it as: add the log of the gatekeeper's weight to the dot-product score table, then softmax. Addition inside an exponential is multiplication outside it, $\exp(z_i+\log g_m)=g_m\exp(z_i)$, so this is the same as multiplying the post-softmax attention probability by the gate. The scores $\mathbf{s}$ now enter the attention probabilities continuously, and gradients from the loss reach $\theta$ through $\mathbf{g}$.

At inference (Algorithm 2), Top-K is taken on $\mathbf{s}$ and ordinary sparse attention runs without any gates. In other words, inference is computationally identical to SeerAttention-R; only the way the selector was trained differs (§5.1).