An interactive tour from α‑entmax — sparse, differentiable attention that sends irrelevant tokens to exactly zero — to AdaSplash‑2, the kernel that finally makes it as fast as FlashAttention.
An attention distribution over three keys is a point on the probability simplex — the triangle of nonnegative weights that sum to one. α‑entmax turns raw scores into that point by projecting them onto the triangle under a regularizer whose strength is the entropy, set by α. Slide α and watch the whole cloud move: high entropy keeps points in the dense interior; low entropy drives them onto the edges (one key zeroed) and corners (two keys zeroed). Drag the triangle to rotate.
● Grey points are score vectors in 3‑D space (the x, y, z axes are shown); each line projects one onto its α‑entmax image on the triangle (● 3 = interior, ● 2 = edge, ● 1 = corner). Raise α and the landing points slide to the boundary. Drag to rotate in 3‑D; add or remove sampled points. The bar charts below are one such projection at a time.
Softmax gives every token a positive weight. Over a long context, that mass spreads thin across hundreds of irrelevant tokens — a dilution known as dispersion. Drag the context length and watch the peak weight collapse.
The largest attention weight shrinks like O(1/n)‑ish as context grows — the signal drowns in a sea of tiny‑but‑nonzero contributions.
α‑entmax replaces softmax with a family indexed by a single parameter α. At α = 1 it is softmax (dense). Push α up and the smallest weights snap to exact zero — the model chooses its own support, differently for every query.
The threshold is set so the weights sum to one. The positive part is what clips everything below threshold to a hard zero.
Read the closed form geometrically: raise a waterline across the (scaled) scores. Everything above the line survives and is renormalized; everything below is zero. α controls both the waterline and the curvature above it.
Softmax gets τ for free — it’s just . For α‑entmax there is no closed form: τ is the root of the equation below. That single fact is the whole efficiency story.
Because τ needs root‑finding, α‑entmax attention historically could not keep up with optimized softmax kernels. AdaSplash and AdaSplash‑2 attack exactly this inner loop. Step through the three generations on the same curve.
Bisection halves a wide bracket. Safe, but many passes over the scores — tens of iterations for tight tolerance.
A hybrid Halley–bisection solver: second‑order steps when safe, bisection as a guardrail. ~7× fewer iterations, plus Triton kernels that skip zero blocks.
A coarse histogram of scores, built on‑chip in SRAM, gives a provable bracket on τ*. The solver then lands in 1–2 steps.
As FlashAttention streams score tiles through fast SRAM, AdaSplash‑2 also drops each score into a coarse histogram. Counting how many scores exceed each bin edge turns into a lower and upper bound on τ* — a tight bracket handed to the solver before a single Newton step. Change the resolution and watch the bracket tighten.
More bins ⇒ a tighter provable bracket ⇒ fewer solver iterations — all without materializing the dense score matrix.
A fast forward pass is only half the battle — training also needs the backward pass. Here α‑entmax gives a second gift: its Jacobian is sparse. The derivative of the weights with respect to the scores vanishes off the support, so gradients only ever flow through the tokens that survived.
AdaSplash splits the backward into two kernels — one for , one for — and reuses the forward support to skip null blocks, so gradients are computed only over the nonzeros. At long context, ~3% of entries can carry ~96% of the attention, so this is most of the cost.
The α‑entmax Jacobian is nonzero only on the support × support block (teal). Push α up: the block — and the backward cost, which scales with it — shrinks with the support.
Sparsity is only useful if the hardware exploits it. When an entire 64×64 tile of the attention matrix is zero, AdaSplash‑2 never loads or computes it. Raise α and watch the causal attention map dissolve into skippable blocks — then read off the speedup.
Relative speed vs. FlashAttention‑2 (forward+backward), following Fig. 1 of the paper.
FlashAttention‑2 is flat — it can’t use sparsity. AdaSplash‑2 crosses 1× at moderate sparsity and pulls away, roughly 2× in the high‑sparsity regime typical of long context.
Knowing which blocks are zero only helps if checking is cheap. AdaSplash‑2 records the nonzero blocks
in a binary mask and bit‑packs it — 32 column‑blocks per int32, a 32× smaller
footprint. The kernel then walks only the set bits with a hardware find‑next‑set instruction, jumping
straight over runs of zeros, in both the forward and the backward pass.
Each cell is one tile pair; teal = nonzero (compute it), dark = zero (skip it), and the digit below is its bit. Press play: find‑next‑set lands only on the teal blocks, hopping over the zeros — so traversal cost scales with the number of nonzeros, not the sequence length.
Trained with exact α‑entmax attention, 1B models match softmax on short context and pull sharply ahead as it grows. All numbers below are from the paper — 1B models trained to 32K, with α‑entmax paired with the NAPE positional scheme.
Where softmax collapses to near‑zero, sparsity holds.
Average accuracy — and it keeps climbing with context.
| Model | 4K | 8K | 16K | 32K |
|---|---|---|---|---|
| Softmax (RoPE) | 51.4 | 45.3 | 37.7 | 27.3 |
| Softmax (NAPE) | 48.1 | 46.4 | 43.3 | 36.8 |
| α‑entmax (NAPE) | 54.7 | 49.9 | 45.4 | 39.4 |
The through‑line: α‑entmax gives attention a principled “off” switch; AdaSplash‑2 makes that switch cost almost nothing. Adaptive sparsity stops being a research luxury and becomes a default you can afford.
Faster Differentiable Sparse Attention.
More α‑entmax toys from the same lab: