Linear Attention and Kernel Feature Maps

Sparse attention (see Sparse Attention: Longformer and BigBird) escapes the $n^2$ cost of full attention by refusing to compute most pairs at all. Linear attention takes a different route: keep every token influencing every other token, in principle, but restructure the arithmetic so the full $n \times n$ matrix never has to be built. The trick hinges on a piece of algebra from linear algebra — matrix multiplication is associative, so $(AB)C$ and $A(BC)$ give the same answer but can cost wildly different amounts of work — combined with replacing the softmax similarity function with something that factors cleanly.

The analogy: totalling a big multiplication table two different ways

Suppose you want $\sum_j (q \cdot k_j)\, v_j$ for a fixed query $q$ against $n$ keys and values. One way: compute all $n$ similarity scores first, then weight and sum the values — this is what standard attention does, and doing it for every one of $n$ queries means building the whole $n \times n$ table. But if the similarity factors as $q \cdot k_j = \phi(q) \cdot \phi(k_j)$ for some feature map $\phi$, you can instead sum $\sum_j \phi(k_j) v_j^\top$ once, into one small matrix that doesn't depend on the query at all — then every query just does one small multiplication against that shared summary. Same total, computed in a cheaper order, the way $(2\times3)\times4$ and $2\times(3\times4)$ both equal 24 but can be organised differently.

Why softmax normally blocks this trick

Ordinary attention weights are $\text{softmax}(QK^\top)$ — the exponential inside softmax does not factor as $\phi(q)\cdot\phi(k)$ for any simple $\phi$, so the reordering trick above does not apply directly, and the full $n\times n$ matrix must be formed.

Linear attention replaces softmax similarity with a kernel function that does factor:

$$\text{sim}(q, k) = \phi(q)^\top \phi(k)$$

In words: instead of directly measuring how similar a query and key are, first map each one through a feature function $\phi$ (a fixed or learned transformation into a new space), and then just take a dot product there. Because this similarity is literally a dot product of transformed vectors, the sum over keys can be computed once, shared across all queries, exactly as in the analogy.

$$O_i = \phi(q_i)^\top \sum_{j} \phi(k_j) v_j^\top \Big/ \phi(q_i)^\top \sum_j \phi(k_j)$$

In words: each query's output is its transformed vector $\phi(q_i)$ multiplied against a running sum built once from all keys and values (numerator), divided by a similarly shared normalising sum (denominator) — no query ever needs the full set of pairwise scores, only these two small shared summaries.

Worked example 1: cost comparison by direct counting

With $n$ tokens and feature dimension $r$ for $\phi$: standard attention forms an $n\times n$ matrix, cost $O(n^2 d)$. Linear attention computes the shared sum $\sum_j \phi(k_j) v_j^\top$ once — an $r \times d$ matrix built by summing $n$ outer products, cost $O(nrd)$ — then applies it to each of the $n$ queries, cost $O(nrd)$ again. Total: $O(nrd)$, linear in $n$.

At $n=8{,}192$, $d=64$, $r=64$: standard cost $\propto 8{,}192^2 \times 64 \approx 4.3 \times 10^9$. Linear cost $\propto 8{,}192 \times 64 \times 64 \approx 3.4 \times 10^7$ — roughly 125× fewer operations at this length, and the ratio keeps growing with $n$.

Worked example 2: the recurrent-state view

Linear attention has an equivalent formulation as a running state updated one token at a time: autoregressive generation can maintain a small fixed-size state $S = \sum_j \phi(k_j) v_j^\top$, updated incrementally as $S_t = S_{t-1} + \phi(k_t) v_t^\top$, output $O_t = \phi(q_t)^\top S_t$. Starting from $S_0=0$ (a $64\times64$ matrix of zeros), after 3 tokens the state is the sum of 3 rank-one updates — a fixed 64×64 object regardless of whether 3 tokens or 3 million have been processed. This is why linear attention is described as "transformers behaving like RNNs": constant per-step memory and cost.

Set the exponent low above to trace roughly linear growth — that's how linear attention's cost scales with sequence length, against the steep exponent-2 curve of standard attention on the same axes.

What this means in practice

Linear attention is genuinely sub-quadratic, unlike FlashAttention and IO-Aware Attention, which is fast but still $O(n^2)$ in FLOPs — a real algorithmic change, not a scheduling one. The tradeoff is representational: choosing $\phi$ to approximate softmax well is an active research problem (random-feature approximations like Performer, or simpler maps like ELU+1), and linear attention variants often trail exact softmax attention on tasks needing very sharp, peaked attention, since a linear kernel cannot express an arbitrarily peaked softmax exactly.

Practice

1. If $\phi(q)^\top \phi(k) = q^\top k$ exactly (the identity feature map), what does linear attention reduce to, and why does that particular choice fail to approximate softmax attention well?
2. Explain why the recurrent-state view means linear attention needs the same fixed amount of memory to generate token 10 as it does to generate token 10,000.