Linear Attention — Math + Variants

The quadratic wall, and the kernel that removes it

Standard attention is O = softmax(QK^T / √d) V with Q, K: [N, d] and V: [N, d_v]. The cost sits in the [N, N] score matrix: N^2 d multiply-adds to build it, N^2 d_v to apply it. Yet matrix products are associative, so (QK^T)V = Q(K^T V), and K^T V is only [d, d_v] — computing that first would be linear immediately. The obstruction is the softmax: a nonlinearity applied across an entire row, so every score must exist before normalization and V cannot be pushed inside.

Linear attention swaps that exponential for a similarity kernel that factorizes, which unblocks the reassociation:

sim(q, k) = φ(q) · φ(k),   φ: R^d → R^m,  all components ≥ 0

o_i = φ(q_i)^T S / ( φ(q_i)^T z ),   S = Σ_j φ(k_j) v_j^T  [m, d_v]
                                        z = Σ_j φ(k_j)        [m]

Because φ(q_i) no longer sits inside a nonlinearity with k_j, it pulls out of the sum over the sequence. Build S and z once by scanning the tokens, and every query is one small multiply against the same shared state: O(N m d_v), linear in N. In the causal case the prefix sums become a recurrence, S_i = S_{i-1} + φ(k_i) v_i^T — a constant-size RNN state instead of a growing KV cache. Everything below varies which kernel, which state, and which update.

Four different escapes from O(N^2)

Papers get lumped together under ‘efficient attention,’ but the mechanisms are genuinely distinct, and confusing them leads to wrong expectations about which will break on your workload.

Only rows one and four are ‘linear attention’ in the strict kernel sense. Linformer is linear by compressing the sequence, not by factorizing the kernel; Reformer is not linear at all, but sub-quadratic sparse attention that still uses softmax.

Linear Transformer: elu+1 and the RNN in disguise

The original construction (Transformers are RNNs, 2020) takes the simplest workable feature map: φ(x) = elu(x) + 1, elementwise. Since elu(x) > -1 everywhere, the shift makes every component strictly positive, and m = d, so nothing widens. No random projections, no extra hyperparameters.

Its lasting contribution is the second form. Written causally, the layer becomes a linear RNN whose hidden state is the [m, d_v] matrix S_i, updated by one rank-1 outer product per token. Training runs the parallel scan; decoding runs the recurrence at O(1) time and memory per step. That duality — a parallel form for the accelerator, a constant-memory sequential form for generation — is what every later variant tries to keep. The cost is quality: elu+1 only loosely resembles the exponential kernel, and the resulting attention is diffuse rather than peaky.

Performer: unbiased softmax by positive random features

Performer’s FAVOR+ asks a sharper question: instead of inventing a kernel, why not estimate the real one? It picks random features whose expected inner product equals the softmax kernel exactly:

E[ φ(q) · φ(k) ] = exp(q · k)

φ(x)_r ≈ exp( w_r · x - ½||x||^2 ) / √m,   w_r ~ N(0, I)

Two details do the heavy lifting. First, the features are positive (exponentials, not the trigonometric features of classic random Fourier methods): trig features can produce negative estimates of a quantity that must be positive, and their variance explodes exactly where softmax concentrates its mass. Second, the projections w_r are made orthogonal, which measurably reduces estimator variance for the same m. The dial is m: more features, tighter approximation, more compute. Performer is the most principled kernel variant, but an unbiased estimator with real variance is still not exact attention.

Linformer: shrinking the sequence axis instead

Linformer reaches O(N) without touching the softmax. It observes that the attention matrix is empirically close to low-rank, and projects the length dimension of the keys and values down to a fixed k ≪ N with learned matrices E, F: [k, N]:

O = softmax( Q (E K)^T / √d ) (F V)      scores: [N, k], not [N, N]

Softmax is still there, still normalizing rows — only the column count changed. Cost is O(N k d), linear in N for fixed k. The catches are structural: the projections are tied to a fixed maximum length, so the model does not extrapolate; and E, F mix all positions including future ones, which breaks causality. There is also no constant-size recurrent state, so decoding gets none of the O(1)-per-step benefit. It is an encoder-side technique.

Reformer: hashing, not linearizing

Reformer belongs in any survey of this space but is routinely miscategorized. It keeps softmax and exact similarity, and argues instead that most of the N^2 scores are wasted: softmax is dominated by a few large entries, so the rest is effort spent on numbers that round to zero.

It finds the large ones with locality-sensitive hashing. Queries and keys are tied (k_j = q_j / ||q_j||) and hashed by random rotations, so vectors with high cosine similarity land in the same bucket with high probability. Attention is computed only within buckets plus a neighboring chunk, giving O(N log N). Reformer pairs this with reversible residual layers that recompute activations in the backward pass instead of storing them — a memory win orthogonal to the attention change. The trade: hashing is probabilistic, a relevant key can miss its bucket, and irregular gather/scatter suits dense hardware poorly.

RetNet and gated linear attention: decay in the state

Plain linear attention accumulates forever: S_i = S_{i-1} + φ(k_i)v_i^T never forgets, so a fixed-size state saturates into mush. The modern fix is explicit forgetting.

RetNet’s retention uses a fixed scalar decay, S_i = γ S_{i-1} + k_i v_i^T, with a per-head γ so heads span multiple timescales, and replaces the fragile φ(q)^T z denominator with a normalization layer. Gated linear attention goes further and makes the decay data-dependent: S_i = Diag(α_i) S_{i-1} + k_i v_i^T, where α_i = σ(W x_i) is computed from the token itself. The model can now hold a value indefinitely when the input says it matters and flush it when it does not — content-based forgetting, which a constant γ cannot express, and the single largest quality jump in this lineage.