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.

Advertisement

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.

Advertisement

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.