The Leap of Sparse Attention in Transformers

The transformer’s attention layer computes a weighted average of value vectors, where the weights come from a softmax over query–key scores (Vaswani et al., 2017). Softmax has one stubborn property: it never assigns exactly zero weight to anything. Every token, however irrelevant, gets a positive slice of the average.

Around 2023–2025, as context windows exploded, the community rediscovered that this is wasteful and started bolting on sparsity: keep the top‑\(k\) scores, mask everything else, and call it sparse attention. That works, but it treats sparsity as a post‑hoc heuristic applied to a fundamentally dense model. It also quietly reinvents ideas that were worked out — with proofs — years earlier, under names like sparsemax, fusedmax, and α‑entmax.

Below, four ways to turn the same score vector into attention weights. softmax spreads mass everywhere; top‑k keeps \(k\) and renormalizes with a hard, non‑differentiable cut; sparsemax and α‑entmax produce exact zeros too — but differentiably, and they choose how many tokens to keep based on the scores themselves.

Softmax turns scores into weights by exponentiating and normalizing:

The normalizer \(\tau\) is just the log‑sum‑exp — a smooth, closed‑form quantity. This is exactly why softmax is so friendly to GPU kernels: \(\tau\) is a single stable reduction along each row, computable in one streaming pass. It is also why softmax is dense: because \(e^{x}>0\) for all \(x\), we have \(\operatorname{softmax}(s)_i>0\) for every token \(i\), no matter how small its score (Martins & Astudillo, 2016).

For short sequences this is harmless. But as \(n\) grows, that guaranteed‑positive mass has to be shared among more and more tokens. The weight on the single most relevant key shrinks, and the output average is increasingly diluted by a long tail of near‑zero — but nonzero — contributions. Drag the context length below and watch the peak weight collapse.

The first crack in the dense default came from a beautifully simple idea (Martins & Astudillo, 2016): instead of exponentiating, just take the Euclidean projection of the scores onto the probability simplex.

Projecting onto the simplex is a constrained least‑squares problem, and its solution is a soft‑threshold: subtract a constant \(\tau(z)\) from every score and clip negatives to zero,

where \(\tau(z)\) is chosen so the surviving entries sum to one. Concretely, if \(S=\{i: z_i>\tau(z)\}\) is the support, then summing the constraint gives a closed form for the threshold:

Everything below the waterline \(\tau(z)\) is set to exactly zero. Unlike top‑\(k\), the size of \(S\) is not a hyperparameter — it falls out of the geometry of \(z\). And crucially, the map is differentiable almost everywhere, with a cheap Jacobian supported only on \(S\):

Below: the same soft‑threshold picture. Raise the scores as lollipops, drop the waterline \(\tau\), and read off the survivors. This is the whole mechanism.

sparsemax also comes with its own loss for training classifiers — a piecewise‑quadratic that behaves like a multi‑class Huber loss — but the attention story only needs the map itself. The deeper question is: why this map? Is the Euclidean projection special, or one member of a family? That question is what turned a clever trick into a theory.

Here is the unification (Niculae & Blondel, 2017). Both softmax and sparsemax solve the same optimization problem — maximize expected score, minus a regularizer \(\Omega\) that pulls the solution toward the interior of the simplex:

The choice of \(\Omega\) is the whole design space:

• Shannon negative entropy \(\Omega(p)=\sum_i p_i\log p_i\) recovers softmax.
• The squared norm \(\Omega(p)=\tfrac12\lVert p\rVert_2^2\) recovers sparsemax.
• A total‑variation penalty \(\Omega(p)=\tfrac12\lVert p\rVert^2 + \lambda\sum_i|p_{i+1}-p_i|\) yields fusedmax, which produces contiguous, group‑sparse attention.

The reason softmax is dense and sparsemax is sparse is now visible: it is the behaviour of \(\Omega\)’s gradient at the boundary. Shannon entropy’s derivative, \(\log p_i + 1\), blows up to \(-\infty\) as any \(p_i\to 0\), so moving off a face always improves the objective — the optimum is pushed strictly inside the simplex (no zeros). The squared norm has a finite gradient at the boundary, so the optimum can — and does — sit on a face where some \(p_i=0\).

Blondel et al. (2020) completed the picture with Fenchel–Young losses. For any convex \(\Omega\) with convex conjugate \(\Omega^\ast\), define

which, for strictly convex \(\Omega\), is zero exactly when \(y = \Pi_\Omega(s)=\nabla\Omega^\ast(s)\). This single template generates softmax + cross‑entropy, sparsemax + sparsemax‑loss, and more, as special cases. And it delivers the pivotal theorem: for the separable entropies used throughout this essay (the whole Tsallis family), \(\Pi_\Omega\) can produce sparse outputs if and only if the induced loss has a separation margin — a finite score gap beyond which the loss is exactly zero, and for α‑entmax that margin is exactly \(1/(\alpha-1)\). Sparsity in the forward map and a margin in the loss are the same phenomenon.

If \(\Omega\) is a dial, why stop at two settings? Peters et al. (2019) instantiated the family with Tsallis \(\alpha\)-entropies, a one‑parameter generalization of Shannon entropy:

which recovers Shannon entropy as \(\alpha\to1\). Plugging \(\Omega = -H^{\mathsf T}_\alpha\) into the regularized‑argmax gives the α‑entmax transformation, with the same threshold structure as sparsemax but a tunable exponent:

with \(\tau\) again set so the weights sum to one. The knob \(\alpha\) interpolates a whole spectrum:

• \(\alpha=1\): the exponent blows up into an exponential — softmax (dense).
• \(\alpha=2\): the exponent is \(1\) — \([z_i-\tau]_+\), i.e. sparsemax.
• any \(\alpha>1\) is sparse — heavier tails as \(\alpha\downarrow 1\), harsher selection beyond \(2\). \(\alpha=1.5\) is the popular middle ground and admits an exact \(\mathcal O(n\log n)\) sorting algorithm; general \(\alpha\) uses bisection on \(\tau\).

Its Jacobian is as cheap as sparsemax’s (supported on the survivors), so it drops into a transformer with no change to backprop. Slide \(\alpha\) below and watch the support tighten from “everything” (softmax) to a handful (sparsemax).

Once sparsity is a differentiable design choice, you can ask for more than “some tokens are zero.” You can control which tokens, how many, and even make the shape learnable. Four directions, all pre‑LLM, all with closed‑form gradients.

Structured — contiguous attention (fusedmax)

Add a total‑variation penalty to sparsemax’s regularizer to reward contiguous supports (Niculae & Blondel, 2017):

The TV term glues neighboring weights into flat plateaus, so attention selects segments (phrases, spans) rather than scattered tokens — with a group‑averaging Jacobian that is still cheap. Increase the fusion strength \(\lambda\) below and watch adjacent tokens merge.

Constrained — budgeted attention

You can also cap how much any single token receives. Constrained sparsemax (Malaviya et al., 2018) solves the projection under an extra upper bound, \(p_i \le u_i\):

solvable in \(\mathcal O(n)\) by the Pardalos–Kovoor algorithm. In machine translation the bound acts as a fertility budget: it stops the model from dumping all its attention on one source word, curing the classic coverage failures (source words dropped, target words repeated) directly in the attention map.

Learnable — a different \(\alpha\) per head

Why fix \(\alpha\) at all? Adaptively Sparse Transformers (Correia et al., 2019) make \(\alpha\) a learned parameter per attention head. Because \(\partial\,\alpha\text{-entmax}/\partial\alpha\) has a closed form, each head slides along the dense↔sparse spectrum by gradient descent — some heads stay near softmax to integrate broadly, others sharpen toward argmax to point at a single token, all “at no cost in accuracy,” improving head diversity and interpretability.

Stochastic — the other road to zero

There is a parallel, older lineage that reaches sparsity by sampling rather than projection. Soft vs. hard attention appeared in image captioning (Xu et al., 2015): soft attention is the deterministic expected glimpse; hard attention samples a single location and trains with REINFORCE. Later, reparameterized discrete gates — HardKuma (Bastings et al., 2019), built on the Gumbel‑softmax (Jang et al., 2017) and \(L_0\) relaxations (Louizos et al., 2018) — place probability mass on exact 0 and 1, giving stochastic sparse selection under a Lagrangian budget on the expected number of active tokens. The same machinery resurfaces years later inside KV‑cache gating (Part II).

Everything so far decides sparsity after computing all the scores. But if most tokens will be zeroed anyway, why compute their scores at all? The pre‑LLM efficient‑attention era — the one modern long‑context papers rarely cite — is precisely about deciding connectivity from content, before the full \(QK^\top\). Two philosophies emerged.

Fixed patterns

Hard‑wire a sparse graph with \(\mathcal O(n\sqrt n)\) or \(\mathcal O(n)\) edges. Sparse Transformers (Child et al., 2019) factorize attention into strided + local patterns with \(|A_i|\!\sim\!\sqrt n\). Longformer (Beltagy et al., 2020) combines a sliding window with a few global tokens; Big Bird (Zaheer et al., 2020) adds random edges and proves the result is a universal approximator and Turing‑complete. Robust and kernel‑friendly — but the graph is chosen a priori, blind to the data.

Content‑based routing

Let the tokens cluster themselves. The Reformer (Kitaev et al., 2020) hashes queries and keys with angular LSH, \(h(x)=\arg\max\big([\,xR;\,-xR\,]\big)\) for a random projection \(R\), so a query only attends within its bucket — \(\mathcal O(n\log n)\), plus reversible layers for depth‑independent memory. The Routing Transformer (Roy et al., 2021) runs online spherical \(k\)‑means on \(\ell_2\)‑normalized queries/keys, so that “only queries and keys from the same cluster are considered for attention” — an \(\mathcal O(n^{1.5}d)\) approximation of maximum‑inner‑product search.

Below, the zoo of connectivity patterns as \(n\times n\) masks: which keys (columns) each query (rows) is allowed to see. This is the vocabulary Part II’s “native sparse attention” methods inherit.

For years, sparse attention was sold on efficiency. Then a wave of theory reframed it as a matter of capability: dense softmax does not just waste compute at long context — it is provably unable to do certain things. This is the intellectual turn that vindicated the whole 2016–2019 lineage.

Dispersion

With all logits bounded in \([m,M]\), every softmax weight is squeezed toward uniform as the sequence grows (Veličković et al., 2025):

And boundedness is not an assumption you can train away: with inputs from a finite vocabulary, every layer maps compact sets to compact sets, so the logits are bounded. The consequence is stark: “even for tasks as simple as finding the maximum key, any learned circuitry must disperse as the number of items grows at test time.” A model that must stay sharp at length cannot, if its weights are forced to be everywhere‑positive. Push the limit further and the pathology deepens: as the prompt grows, softmax attention provably converges to a linear attention operator — not just blunted, but with its nonlinearity washed out (Boursier & Boyer, 2026).

Collapse & over‑squashing

Relatedly, Barbero et al. (2024) prove representational collapse: take a long sequence and the same sequence with its final token repeated, and their final‑token representations become arbitrarily close as length grows — eventually indistinguishable at floating‑point precision, so the model provably cannot respond to them differently. The engine is softmax itself: the total‑variation distance between the two attention distributions tends to zero as the sequence grows (their Lemma 4.1); decaying positional information does the rest. They also formalize over‑squashing: the sensitivity of an output to a distant input is bounded by a sum over weighted causal attention paths, which thins out with distance and depth. Both are pathologies of a distribution that is not allowed to be zero anywhere.

There is a documented failure at the opposite end, too. Zhai et al. (2023) show that too‑sharp attention — “entropy collapse” — travels with training instability, and their remedy (σReparam) works by capping the spectral norm of the logits. Read the two results together and softmax is squeezed from both sides: its attention entropy is lower‑bounded in terms of the logit scale, so the only way it can be sharp is through extreme logits — precisely the regime that destabilizes training — while at bounded logits it must disperse. α‑entmax dissolves the dilemma: it reaches exact sharpness — even one‑hot — at finite, moderate logits, the margin of §4 made kinetic.

Patch the softmax, or replace it?

One camp keeps softmax and rescales. Scalable‑Softmax multiplies the logits by \(s\log n\) (Nakanishi, 2025); Veličković et al. tune an adaptive temperature at inference (2025); scale‑invariant attention derives a position‑dependent logit transform (Anson et al., 2025); entropy‑invariance reaches the same place from an information‑theoretic argument (Li et al., 2025). Four papers, four principles, one medicine — a length‑dependent temperature. That convergence says the diagnosis is right; it also says the treatment is palliative, because a colder softmax thins the tail without ever removing it.

The model’s own coping strategy points the same way. Attention sinks — the notorious pile of mass on the first token — turn out to be softmax’s pressure valve: heads park probability they are forced to spend somewhere on a semantically useless token, which provably slows the over‑mixing of everything else (Barbero et al., 2025). The folk remedy of adding \(1\) to softmax’s denominator so heads can attend to nothing — “softmax‑1” (Miller, 2023) — is the same escape hatch built by hand, and OpenAI’s gpt‑oss now ships learned per‑head sink logits (OpenAI, 2025). A model that invents a place to dump unwanted probability is a model asking for zeros. Some do abandon the simplex altogether — sigmoid attention drops normalization (Ramapuram et al., 2025), Softpick rectifies softmax so sinks vanish (Zuhri et al., 2025), stick‑breaking replaces it with a recency‑biased cascade (Yang et al., 2025) — a genuine “replace” camp, though one that gives up the convex‑duality toolbox of §4 along with the simplex.

Our camp keeps the simplex and fixes the support. With α‑entmax, irrelevant tokens get exact zeros, so the entropy of attention is bounded by the support size \(s\), not the sequence length — \(\mathcal O(\log s)\) instead of \(\mathcal O(\log n)\) — and the number of effective gradient paths shrinks from \(\mathcal O(n^L)\) to \(\mathcal O(s^L)\) (Vasylenko et al., 2025). The payoff is dramatic length extrapolation — up to \(1000\times\) on synthetic recall — from models trained on short contexts. One honest caveat: that result also needs position encodings that cooperate (NoPE (Kazemnejad et al., 2023) mixed with ALiBi (Press et al., 2022)) and a learnable temperature inside entmax — the positional side of length generalization is part of the cure, not a rival to it.

Here is the reframing that ties the whole essay together (Santos et al., 2026). A single attention output is a Nadaraya–Watson kernel‑regression estimate — the classic nonparametric regressor from 1964:

The choice of kernel \(K\) is the choice of normalizer. With \(\ell_2\)‑normalized queries and keys, \(\lVert q-k_i\rVert^2 = 2 - 2\,q^\top k_i\), so the Gaussian kernel \(K(u)\propto e^{-\lVert u\rVert^2/2}\) is exactly softmax over dot‑product scores — and its infinite support means every value contributes: that is dispersion, restated in the language of statistics. And α‑entmax? Setting \(\alpha = 1 + \tfrac1r\) makes it exactly Nadaraya–Watson with a compact, bounded‑support kernel:

• \(r=1\) → the Epanechnikov kernel → \(\alpha=2\) → sparsemax;
• \(r=2\) → the biweight kernel → \(\alpha=1.5\) → entmax‑1.5;
• \(r=3\) → the triweight kernel → \(\alpha=\tfrac43\).

Compact kernels are exactly the ones statisticians reach for when they want a local estimate: they ignore data outside a bandwidth. One precision matters. With a fixed bandwidth \(h\), the Epanechnikov estimator above is normalized‑ReLU attention; sparsemax and its entmax siblings correspond to an adaptive bandwidth, chosen per query so the truncated kernel weights sum to one on their own — that adaptive normalization is the threshold \(\tau\) of §3, wearing statistician’s clothes. So the exact zeros of α‑entmax are not a hack: they are the finite support of a textbook kernel, with the entmax exponent \(1/(\alpha-1)\) as the kernel’s polynomial degree. The softmax/Gaussian pair emerges in the \(r\to\infty\) limit, and even top‑\(k\) attention finds a place in the same frame — it is Nadaraya–Watson with a truncated‑Gaussian kernel, i.e. \(k\)‑nearest‑neighbour regression: a crude compact kernel rather than a lawless heuristic. This regression reading is not an isolated trick, either — it is the same lens under which attention is a modern Hopfield update and a fast‑weight memory (Ramsauer et al., 2021): attention as test‑time regression over the keys it has seen.

Part I made a single argument from five directions. A regularizer decides dense vs. sparse; a margin is the same as a zero; one knob \(\alpha\) spans the family; exact zeros cure the dispersion that provably cripples dense softmax at length; and all of it is compact‑kernel regression wearing different clothes. The math says the right object is bounded support.

The objection was always practical: “exact, data‑dependent zeros are pretty, but GPUs want dense blocks, and computing that entmax threshold needs iterative root‑finding — it will never match FlashAttention.” Part II is the story of dismantling that objection: IO‑aware exact kernels, entmax made hardware‑efficient (AdaSplash → AdaSplash‑2), sparsity trained in natively (NSA, MoBA), and the sprawling modern KV‑cache and long‑context literature — which, read through Part I, is largely a rediscovery of support recovery.