Linear Attention Math

Softmax attention, and the one thing it forbids

Standard attention computes, for each query q_i, a normalized weighted sum of value vectors:

o_i = Σ_j  [ exp(q_i · k_j / √d) / Σ_j' exp(q_i · k_j' / √d) ] · v_j

In matrix form that is O = softmax(QK^T / √d) V, with Q, K: [N, d] and V: [N, d_v]. The attention matrix A = softmax(QK^T) is [N, N], and building it is exactly the quadratic step: N^2 dot products to fill it, N^2 again to multiply by V.

Here is the crucial point. Matrix multiplication is associative, so in principle (QK^T)V = Q(K^T V) — and K^T V is only [d, d_v], tiny and independent of N. If we could compute that first, attention would be linear. The softmax is what blocks us: it is a nonlinear function applied to the whole N×N matrix, so the exponential and the row-normalization have to see every entry before V enters. You cannot push V inside a softmax. Linear attention’s entire trick is to remove that obstruction.

Replacing the exponential with a kernel feature map

The softmax weight exp(q_i · k_j) is really a similarity kernel sim(q_i, k_j) that must stay non-negative (weights can’t be negative). Linear attention picks a different kernel — one that factorizes as an inner product of a feature map applied separately to each argument:

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

Substituting into the attention output and dropping the softmax normalizer in favor of an explicit sum of weights:

o_i = Σ_j (φ(q_i) · φ(k_j)) v_j
      -----------------------------
      Σ_j (φ(q_i) · φ(k_j))

Both the numerator and the denominator are now sums of terms that each split cleanly into a q-part and a k-part. That separability is the whole game: because φ(q_i) does not depend on j, it can be pulled out of the sum over the sequence. The exponential could never be pulled out because exp(a+b) = exp(a)exp(b) only for scalars, and the argument q_i · k_j couples the two indices inside a single nonlinearity. A factorized kernel decouples them by construction.

The associativity trick, written out

Pull φ(q_i) out of both sums. The numerator becomes a vector-matrix product and the denominator a vector-vector product:

o_i = φ(q_i)^T S / ( φ(q_i)^T z )

  where   S = Σ_j φ(k_j) v_j^T      (shape [m, d_v])
          z = Σ_j φ(k_j)           (shape [m])

Read that carefully. S is a single [m, d_v] matrix — the ‘KV summary’ of the entire sequence — and z is a single [m] normalizer vector. Neither depends on i. You compute them once by scanning the sequence, then every query just does one small [m]×[m, d_v] multiply against the same shared S. In batched matrix notation the full layer is:

O = φ(Q) ( φ(K)^T V )          numerator, associativity moved the parentheses
    ----------------------
    φ(Q) ( φ(K)^T 1 )          denominator normalizer

The parentheses are the entire idea. φ(Q)φ(K)^T would rebuild the N×N matrix; φ(K)^T V instead contracts the sequence axis first, producing the compact [m, d_v] state before φ(Q) ever touches it.

Complexity: O(N^2) becomes O(N)

Count the work in each ordering. Softmax attention forms QK^T at N·N·d cost, then multiplies by V at another N·N·d_v: overall O(N^2 d) time and O(N^2) memory for the materialized matrix.

Linear attention forms S = φ(K)^T V at N·m·d_v cost, then φ(Q) S at N·m·d_v: overall O(N m d_v). With the feature dimension m treated as a constant on the order of d, this is linear in the sequence length. The stored state is [m, d_v] — independent of N — so memory drops from O(N^2) to O(N) for the inputs plus O(m d_v) for the state.

The trade is visible in the shapes: you have replaced an object that grows with the sequence (N×N) with one that is fixed by the model width (m×d_v). At N = 100{,}000 and d = 128, the softmax matrix alone is ten billion entries per head; the linear state is about sixteen thousand. That gap is the reason linear attention exists.

The recurrent form: a running KV state

Attention must be causal for a language model — query i may only see keys j ≤ i. That seems to break the ‘compute S once’ picture, because each position needs a different partial sum. But a prefix sum is just a recurrence. Define the state after token i:

S_i = S_{i-1} + φ(k_i) v_i^T        (running [m, d_v] matrix)
z_i = z_{i-1} + φ(k_i)             (running [m] vector)
o_i = φ(q_i)^T S_i / ( φ(q_i)^T z_i )

This is a recurrent neural network. The hidden state is the [m, d_v] matrix S_i; each step adds one rank-1 outer product and reads out one query. Crucially, the state size is constant — it does not grow as you generate. So autoregressive decoding costs O(1) time and O(1) memory per token, versus the softmax KV cache that stores every past key and value and grows O(N) per step. This is the sense in which ‘transformers are secretly RNNs’: the same layer has a parallel O(N) training form (scan the whole sequence) and a sequential O(1) inference form (update the state), and you pick whichever the workload wants.

A small worked example

Take m = 2, d_v = 2, and two past tokens with φ(k_1) = [1, 0], v_1 = [1, 0] and φ(k_2) = [0, 1], v_2 = [0, 2]. The state accumulates the outer products:

S = φ(k_1)v_1^T + φ(k_2)v_2^T = [[1,0],[0,0]] + [[0,0],[0,2]] = [[1,0],[0,2]]
z = φ(k_1) + φ(k_2) = [1, 1]

Now a query with φ(q) = [3, 1]. The numerator is φ(q)^T S = [3·1, 1·2] = [3, 2]; the denominator is φ(q)^T z = 3 + 1 = 4. So o = [3/4, 2/4] = [0.75, 0.5]. Notice what happened: the query’s larger weight on the first feature (3 vs 1) pulled the output toward v_1, and the denominator normalized the total weight to 1 — exactly the role softmax plays, but computed by a plain dot product against a state that was built without ever forming a token-by-token attention matrix. Add a third token and you update S and z in place; no earlier work is redone.

Feature maps: elu+1 and positive random features

The kernel is only as good as its feature map. Two conditions matter: every component of φ(x) should be non-negative (so attention weights stay positive and the normalizer stays well-defined), and φ(q)·φ(k) should resemble a useful similarity.

elu + 1. The simplest choice, from ‘Transformers are RNNs,’ is φ(x) = elu(x) + 1, applied elementwise. Since elu(x) > -1 everywhere, adding 1 makes it strictly positive; m = d, so there is no blow-up in width. It is cheap and stable but only crudely approximates the softmax kernel.

Positive random features. Performer’s FAVOR+ instead chooses φ so that E[φ(q)·φ(k)] = exp(q·k) — an unbiased estimate of the true softmax kernel. Each feature is roughly φ(x)_r = exp(w_r · x - ½||x||^2) / √m for random projections w_r. Using positive (exponential) rather than trigonometric features keeps every entry non-negative, which avoids the variance blow-up and negative-weight instability that plagued earlier random-feature attempts. More features m means a tighter approximation but more compute — the usual accuracy/cost dial.

Why the quality gap: normalization and expressivity

Linear attention is faster but generally weaker, and the reasons are structural, not incidental. First, rank. The state S = φ(K)^T V is at most rank m, so the effective attention matrix it induces is a rank-m object. When N > m — always, for long context — it simply cannot represent an arbitrary N×N pattern. Softmax attention has no such cap; it can route any token to any token.

Second, sharpness. The exponential in softmax has unbounded dynamic range, so it can concentrate almost all weight on a single key — the ‘hard lookup’ behavior that induction heads and retrieval depend on. A bounded positive feature map produces a much flatter, more diffuse weighting; it smears attention rather than spiking it. Third, normalization. The denominator φ(q)^T z can drift toward zero or grow uncontrolled as the sequence lengthens, causing numerical fragility that softmax’s per-row normalization avoids by construction. The net effect: linear attention trades exact, peaky, full-rank routing for a cheap low-rank summary, and on tasks that need precise long-range recall it measurably underperforms.

The kernel cousins: SSMs and RetNet

Once attention is written as a linear recurrence on a fixed-size state, a whole family of architectures lines up as variations on that state update. They differ in how the state carries information forward, not in the constant-memory idea.

RetNet is the closest relative. Its ‘retention’ is linear attention plus an explicit decay: S_i = γ S_{i-1} + k_i v_i^T with a fixed scalar (or per-head multi-scale) γ < 1, and it drops the denominator in favor of a gated normalization. The decay makes recent tokens matter more and gives a clean parallel, recurrent, and chunkwise-parallel triad of forms.

State-space models (S4, and Mamba) replace the plain sum with a structured linear dynamical system: h_i = A h_{i-1} + B x_i, y_i = C h_i. The matrix A is designed (or, in Mamba, made input-dependent so the model can selectively remember or forget) to control how the state mixes over time. Kernel-based linear attention is the special case where the recurrence is a pure accumulation driven by a feature map φ. All three share the prize — O(N) training, O(1) decoding, a constant-size state — and all three are chasing the same target: recover softmax-level quality without softmax’s quadratic bill.