Linear attention: move the nonlinearity

Softmax attention cannot be factorised because the exponential sits on top of the dot product. Move the nonlinearity to before the dot product — apply it to \(q\) and \(k\) separately — and the entire computation re-associates into a single fixed matrix. The memory stops growing. Something else starts breaking.

The substitution

Write softmax attention with the exponential named as a similarity function \(\text{sim}\):

Nothing about the structure of attention requires \(\text{sim}\) to be that particular function. It requires only that \(\text{sim}\) be non-negative, so the weights form a genuine average. Choose instead a similarity that factorises:

elu(x)+1 equals \(x+1\) for \(x>0\) and \(e^{x}\) for \(x\le 0\): smooth, strictly positive, and cheap. Any elementwise positive map works; this one is the original choice.

\(\phi\) is applied to each vector on its own, before they meet. That one change is the whole chapter, because it makes the numerator re-associable.

The re-association

Matrix multiplication is associative: \((AB)C = A(BC)\). Same result, wildly different intermediate. With \(\phi\) folded into \(q\) and \(k\) so we can stop writing it:

On the left, \(q_t k_i^{\top}\) is a scalar score, and you need one per key — an \(N\times N\) matrix over a whole sequence. On the right, \(k_i^{\top} v_i\) is an outer product: a \(d\times d\) matrix, the same size no matter how many tokens you have seen. And outer products can be summed before the query ever arrives.

The recurrence you get for free

Because the numerator is now a sum of outer products, it can be accumulated one token at a time. Same for the denominator. Two running quantities, both fixed size:

That is a recurrent neural network. Not "like" one — it is one, with a linear update and a \(d\times d\) hidden state. The paper that introduced this was titled Transformers are RNNs, and it meant it literally: given the factorised similarity, the transformer and the RNN are the same computation written in two orders.

# the whole of linear attention, prefill and decode
q = qkv[:, :, :d  ].view(b,t,h,d_head).transpose(1,2)
k = qkv[:, :, d:2*d].view(b,t,h,d_head).transpose(1,2)
v = qkv[:, :, 2*d:].view(b,t,h,d_head).transpose(1,2)

k = F.elu(k) + 1                 # φ, applied to k alone
q = F.elu(q) + 1                 # φ, applied to q alone
k = k.transpose(-1,-2)

S, z = cache if cache is not None else (0.0, 0.0)
S = S + k @ v                    # fold this token into the state
z = z + k                        # and into the normaliser

o     = q @ S                     # read: one d×d matmul, no loop over history
denom = q @ z
o     = o / denom
return self.o_proj(o.transpose(1,2).reshape(b,t,d)), (S, z)

There is no past_kv to concatenate and no history to scan. The two highlighted lines replaced the KV cache with a matrix that never changes shape.

What that costs, precisely

For GPT-2's geometry — \(d = 64\) per head, 12 heads, 12 layers — the crossover is around 32 tokens. Past that, linear attention's state is strictly smaller, and the gap opens fast:

The three-step contract

Softmax and linear attention are the same three operations with different choices at step one:

1. Make the scores non-negative. Softmax exponentiates. Linear attention applies \(\phi\) to \(q\) and \(k\) separately.
2. Divide by their sum, so the weights are a distribution.
3. Take the weighted average of the values.

Only step 1 changed, and the whole cost profile inverted. Worth noting for later: from chapter 3 onward, step 2 gets dropped entirely — modern variants normalise \(q\) and \(k\) to unit length and put an RMSNorm on the output instead of carrying a running denominator. The contract loosens as the field gains confidence that the denominator was doing less work than assumed.

What you gave up, part one: sharpness

\(\exp\) is unbounded and convex, so a single high-scoring key can capture nearly all the attention mass. \(\phi(q)\cdot\phi(k)\) with \(\phi = \mathrm{elu}+1\) is a much flatter, more bounded function of the inputs. The resulting weight distribution is smoother; the model retrieves a blur where softmax could retrieve one item.

This is a real loss but a soft one — a worse approximation to the same kernel, recoverable in part with better feature maps. The second loss is structural and no feature map fixes it.

What you gave up, part two: capacity

Here is the failure, stated exactly. After \(N\) tokens the state holds \(S_N = \sum_{i=1}^{N} k_i^{\top} v_i\). Query it with one of the keys you wrote, \(q = k_j\):

The second term vanishes only if \(k_j\) is orthogonal to every other key. In \(\mathbb{R}^d\) you can have at most \(d\) mutually orthogonal directions. So the moment \(N > d\), some pair of keys must overlap, and the overlap leaks one fact's value into another fact's readout.

You can quantify how bad. If the keys behave like random unit vectors in \(\mathbb{R}^d\), each cross term has magnitude around \(1/\sqrt d\), and the \(N\) of them accumulate incoherently, so their total scales as \(\sqrt{N/d}\). The signal is \(1\). Therefore

With \(d = 128\): perfectly usable at \(N = 128\), down to 0.18 by \(N = 4{,}096\). Learned keys do better than random ones — the model can allocate directions deliberately, and values that should be confusable are not — but no amount of learning creates a \(129\)th orthogonal direction in a 128-dimensional space. The ceiling is dimensional, not statistical.