Linear Attention and Kernel Feature Maps

Attention is quadratic because of one operator. Write unnormalised attention as \(\sum_j \mathrm{sim}(q_i, k_j) v_j\) with \(\mathrm{sim}(q,k) = \exp(q^\top k / \sqrt{d})\), and the exponential is what forces you to form the full \(n \times n\) score matrix: it cannot be factored into something computed per-query and per-key independently. Every trick that makes attention subquadratic is, at bottom, a decision about what to do with that exponential.

Linear attention makes the most direct choice available: replace it with an inner product of feature maps, \(\mathrm{sim}(q,k) = \phi(q)^\top \phi(k)\) for some elementwise non-negative \(\phi\) (Katharopoulos et al., 2020, Transformers are RNNs, arXiv:2006.16236, ICML 2020).

Associativity is the whole trick

Once similarity factorises, the sum reassociates:

The inner sum \(S\) does not depend on \(i\). Compute it once as a \(d_\phi \times d_v\) matrix and every query is a single matrix-vector product against it. Cost falls from \(O(n^2 d)\) to \(O(n d_\phi d_v)\), and crucially the memory needed to answer any query is a fixed-size matrix rather than a growing cache of keys and values.

For causal models the sum is prefix-restricted, which turns it into a recurrence:

This is the paper's title made literal: a causal linear transformer is an RNN with a matrix-valued hidden state, and autoregressive generation runs in constant time and constant memory per token. Katharopoulos et al. report up to 4000x speedup on autoregressive prediction of very long sequences.

Which feature map, and why it matters

The original paper uses \(\phi(x) = \mathrm{elu}(x) + 1\), chosen mainly to be positive and cheap. Performer instead approximates the softmax kernel itself with positive random features, so its attention converges to real softmax attention as the number of features grows (Choromanski et al., 2021, Rethinking Attention with Performers, arXiv:2009.14794). These are opposite philosophies: one accepts a different attention function and trains with it, the other tries to be softmax attention cheaply.

The distinction that matters: fixed state versus growing cache

Softmax attention keeps every past key and value, so information from position 3 is exactly as available at position 100,000 as it was at position 4. Linear attention compresses all history into \(S\), whose size never changes. That is the source of both properties: constant-memory decode, and an inability to store more than \(d_\phi \times d_v\) numbers about the past no matter how long the sequence. Every practical improvement since 2020 — gating, decay, delta rules — is a scheme for deciding what to forget from a state that is full, which is the same design pressure that shapes selective state space models (see state space models and selective SSM).

When it breaks

• Associative recall degrades. The task linear attention is measurably worst at is retrieving a specific earlier token given a key seen once — exactly the capability induction heads provide in softmax attention (see induction heads and in-context learning). Fixed state and precise recall are in direct tension.
• The recurrent form is fast to decode and awkward to train. The sequential scan does not parallelise across positions the way a matmul does, so real implementations use chunked parallel forms that recover throughput at the cost of considerable kernel complexity.
• Normalisation is fragile. The denominator \(\phi(q_i)^\top z_i\) can approach zero, and numerical stability here is a recurring source of divergence at scale.
• Speedups are length-dependent and often absent where you deploy. Linear attention only wins past the crossover length where \(n\) exceeds the feature and state dimensions. Under a few thousand tokens, a well-tuned FlashAttention kernel is usually faster in wall-clock terms despite worse asymptotics (see flash attention).
• Pure linear attention has not displaced softmax at the frontier. The models that ship treat it as one layer type among several, interleaving linear or state-space layers with a minority of full-attention layers to keep exact recall where it matters. Read claims of "linear-time transformers" with that hybridisation in mind.