Vol. VII · Modern Architectures Plate 8 of 10linear attention, recurrent state, state-space models, Mamba, hybrid models, quadratic cost, KV cacheAC

Beyond Quadratic Attention.

Swap exp(q · k) for φ(q) · φ(k), and the past folds into one matrix that never grows.

Softmax attention costs O(T²) work over a sequence and needs a key–value cache that grows with every token (plate VI.9). If the similarity exp(q · k) is replaced by a product of feature maps φ(q) · φ(k), the sum over the past can be reordered: Σ_j φ(q_i) · φ(k_j) v_j = φ(q_i) · (Σ_j φ(k_j) v_jᵀ). The bracket is a d × d matrix that can be updated one token at a time. The model becomes a recurrent network with a fixed-size state, and generation costs the same at token one million as at token one.

Many architectures descend from this idea: linear Transformers (Katharopoulos et al., 2020), RetNet, state-space models such as Mamba (Gu and Dao, 2023), and gated variants such as DeltaNet. The catch is memory of a different kind: a fixed-size state cannot store every earlier token exactly, so retrieval of a specific fact from far back suffers. Most production models that use these layers are therefore hybrids, keeping a few full-attention layers among many linear ones.

Below: left, linear attention computed two ways — in parallel over all pairs and recurrently with a fixed state — and the largest difference between them; and the memory each needs as the context grows, against a key–value cache. Right, held-out loss of two recorded runs of micro’s shape: softmax attention and linear attention.

What you will learn
Workshop
Left: memory per layer and head, in numbers, for a softmax key–value cache (2 T d) and a linear-attention state (d² + d), d = 32, as the context grows. Right: held-out loss of recorded runs of micro’s shape with softmax attention (blue) and linear attention with φ(x) = elu(x) + 1 (red).

One matrix instead of a growing cache

Over 32 random tokens the parallel computation — every query against every earlier key — and the recurrent one — a single d × d state updated token by token — agree to 4.4 × 10⁻¹⁶, which is rounding error: they are the same function computed in a different order. The difference is in what must be kept. A softmax head stores 2 T d numbers for its keys and values, 8,388,608 at 131 072 tokens with d = 32; the linear head stores 1,056 at any length, and each new token costs the same small update.

The price shows in the recorded runs. With everything else equal, softmax attention ends at 4.208 nats and linear attention at 4.268, 0.060 worse — after 6 000 steps on short sequences, where the cache was never a problem. Studies of recall find the gap much larger on tasks that need an exact token from far back, because a fixed state must blend everything it has seen. That trade-off explains the shape of current designs: hybrids such as Jamba interleave state-space layers with a few attention layers, keeping exact lookup where it matters and constant memory elsewhere.

Softmax attention can be seen as linear attention with an infinite-dimensional feature map: exp(q · k) is a kernel. Linear attention’s weakness is that a finite map cannot make attention as sharp, which is why retrieval suffers.

Terms

Linear attention
Attention with similarity φ(q) · φ(k), computable as a recurrence with a fixed-size state.
Feature map φ
A function applied to queries and keys so that their dot product stands in for exp(q · k).
State-space model
A sequence layer that updates a fixed-size hidden state token by token, such as Mamba.
Hybrid model
A model that mixes a few full-attention layers with many linear or state-space layers.
Replacing exp(q · k) with φ(q) · φ(k) folds the past into a fixed state — constant memory per token, at a cost in exact recall.

Sources

  1. Katharopoulos, Vyas, Pappas and Fleuret, “Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention”, ICML 2020.
  2. Gu and Dao, “Mamba: Linear-Time Sequence Modeling with Selective State Spaces”, arXiv:2312.00752, 2023.
  3. Lieber et al., “Jamba: A Hybrid Transformer-Mamba Language Model”, arXiv:2403.19887, 2024.
  4. Arora et al., “Zoology: Measuring and Improving Recall in Efficient Language Models”, arXiv:2312.04927, 2023.