Module 18 · Part IV · The Modern Architecture

Beyond Attention: Linear Attention and Hybrids

Linear Attention as a Recurrence, the Gated Delta Rule, State-Space Models, and Hybrid Models

By the end you can explain how a fixed-size state replaces the KV cache, and what it costs.

Builds onModule 15·Module 5·
Check yourself ▸

Before you start, can you already…

  • ☐turn the 2017 block into a 2024 block and justify each change. (Module 15)
  • ☐explain why recurrence hit a wall, and how attention first appeared inside RNN translation models. (Module 5)

If any is shaky, the linked module teaches it.

Open in Colab
3 Hours·Elective·Advanced·0/4

Softmax attention compares each new token with every past key and value: exact, but its cost grows with the context. This module covers the alternative that has re-entered frontier models, a layer that keeps a fixed-size state updated once per token, like the recurrent networks of Module 5 but trainable in parallel, and the hybrid models that mix both kinds of layer.

18.1 Linear AttentionRun this section

FlashAttention reads the growing KV cache efficiently (section 16.2), GQA shrinks it (15.5), and local windows cap it (16.3). Linear attention removes it.

Vectors here are columns, as in the papers, and is a outer product. Drop the softmax (Katharopoulos et al., 2020) and the newest token's output regroups:

∑j≤t(qt⊤kj) vj=(∑j≤tvjkj⊤)qt\sum_{j \le t} (\mathbf{q}_t^\top \mathbf{k}_j)\, \mathbf{v}_j = \left(\sum_{j \le t} \mathbf{v}_j \mathbf{k}_j^\top\right) \mathbf{q}_t

since each qt⊤kj\mathbf{q}_t^\top \mathbf{k}_j is a scalar. The bracket no longer involves the query, so it can be kept as a running dv×dkd_v \times d_k matrix, the state StS_t, updated once per token:

St=St−1+vtkt⊤,ot=StqtS_t = S_{t-1} + \mathbf{v}_t \mathbf{k}_t^\top, \quad \mathbf{o}_t = S_t \mathbf{q}_t

The state is the same size at token 10 as at token 100,000, so memory and per-step cost are fixed. This is a recurrent network again (section 5.2), but its linear update lets training run in parallel across the sequence.

The weakness is recall. The state blends everything it has seen and only ever adds: nothing is forgotten, and a value written under a key similar to an old one piles on top instead of replacing it.

18.2 The Gated Delta RuleRun this section

Gated DeltaNet (Yang et al., 2024), the layer in the Qwen model below, changes the update:

St=gt St−1(I−βtktkt⊤)+βtvtkt⊤S_t = g_t \, S_{t-1}\left(I - \beta_t \mathbf{k}_t \mathbf{k}_t^\top\right) + \beta_t \mathbf{v}_t \mathbf{k}_t^\top

gtg_t and βt\beta_t are numbers between 0 and 1 computed from the current token, and keys have length 1. The gate gtg_t fades old content, so the layer can forget. The factor (I−βtktkt⊤)(I - \beta_t \mathbf{k}_t \mathbf{k}_t^\top) is the delta rule: before writing, it removes a share βt\beta_t of what the state currently returns for kt\mathbf{k}_t. With gt=1g_t = 1 the update is St−1−βt(St−1kt−vt)kt⊤S_{t-1} - \beta_t (S_{t-1}\mathbf{k}_t - \mathbf{v}_t)\mathbf{k}_t^\top: one gradient-descent step on the error between what the state recalls for kt\mathbf{k}_t and what it should.

Figure 18.2
A fixed-size memory: adding versus the delta rule

Four writes into a 4 × 4 state. Plain linear attention only adds, so an updated key recalls old and new values piled together. The delta rule first erases what the state returns for the key, then writes.

Keys a and b are perpendicular; key c overlaps a, as real keys often do. Recall for a key is S k: the bars show it against the value most recently stored under that key (dashed). Green bars are negative entries.
Predict, then try

Plain linear attention writes a second value under key a. Before recalling, predict what it returns for a: the new value, the old one, or something else.

In words. Four writes into a four-by-four state: values under keys a, b and c (which overlaps a), then a new value under a. Plain linear attention returns the old and new values for a piled together, plus some of c's. The delta rule erases before writing, so a recalls exactly the new value, though c is partly disturbed, the price of non-perpendicular keys. Lowering the gate fades older content.

18.3 State-Space ModelsRun this section

State-space models (SSMs) reach the same kind of layer from the mathematics of continuous signals, with a linearly updated hidden state:

ht=A ht−1+B xt,yt=C ht\mathbf{h}_t = A\,\mathbf{h}_{t-1} + B\,\mathbf{x}_t, \quad \mathbf{y}_t = C\,\mathbf{h}_t

with learned AA, BB and CC. A linear update, the same at every step, can be trained in parallel, as one long convolution (a sliding weighted sum, section 25.1) or a parallel scan, and run one step at a time in generation. Early SSMs (S4) used a fixed AA: fast, but poor at choosing what to remember.

Mamba (Gu & Dao, 2023) computes the decay and the amount written from the current token, so the layer decides token by token what to keep. That input-dependent decay is the gate gtg_t above, and Mamba-2 showed its layer is a form of linear attention with decay. The two lines have converged: a fixed-size state, a linear update with learned forgetting, and parallel training. They share the weakness too: exact recall of a specific earlier token, trivial for softmax attention, degrades as the context grows.

Browser lab

Linear attention: exact as a recurrence, limited as a memory

Two checks. First, for 50 random keys and values, the regrouped sum (Σ (q·kj) vj) equals the state recurrence (S = Σ vjkjT, then Sq), as section 18.1 derives. Second, write a value under each of n keys and then a new value under the same keys: section 18.2 predicts plain linear attention recalls a blend of old and new, while the delta rule recalls the new value. Third, store n distinct pairs and try to recall each: here neither rule escapes the fixed size of the state.

Random unit keys in 32 dimensions, which are not perpendicular. That is why, for many distinct pairs, the delta rule does worse than plain addition here: each erase also removes part of older items that share the key's direction, so it favours recent items. Real layers learn their keys and use a gate and a fractional β; the fixed state size remains the limit either way.

18.4 Hybrid Models

Hybrid models keep a few full-attention layers for exact lookup among many linear layers that carry the bulk of the computation at constant cost per token.

Qwen3.5-9B, a full-parameter distillation (section 20.4), is such a hybrid. The table gives its configuration, run under MLX (Apple's machine-learning library), with 4-bit weights on a 16 GB laptop. The first five rows are from its published configuration, the sixth is estimated from its parameter counts, and the last was measured on that laptop:

value
layers32
hidden size4096
intermediate size12288
vocabulary248,320
attentionhybrid: 24 gated-DeltaNet linear-attention layers, and full attention in every 4th layer
estimated prefill cost13.8–15.9 GFLOP/token (from 4.832B parameters in MLPs, 1.617B in linear attention, 0.470B in full attention, at about two FLOPs per parameter)
measured throughput224 tok/s at 2.5k context, 199 tok/s at 16k

A context 6.4×6.4\times longer, 2.5k to 16k tokens, lowers the speed only 11%, from 224 to 199 tokens per second, because only 8 of the 32 layers attend over the whole sequence, confining the cost to a quarter of the stack. FlashAttention makes full attention read memory efficiently and GQA shrinks its cache; this model uses it in fewer layers.

So the design question is no longer only how many query heads share each KV head, but whether a layer needs a KV cache at all. For linear-attention layers, the per-sequence KV cache of section 23.1 does not shrink; it no longer applies.

Exercises

Check

What does a linear-attention layer keep instead of a KV cache, and how does its size grow with the sequence length?

Apply

A linear-attention layer has dk=dv=128d_k = d_v = 128. How many numbers are in its state, and how does that compare with a KV cache for 100,000 tokens of the same head?

Apply

Show that the gated delta update with gt=1g_t = 1, St=St−1(I−βtktkt⊤)+βtvtkt⊤S_t = S_{t-1}(I - \beta_t \mathbf{k}_t\mathbf{k}_t^\top) + \beta_t\mathbf{v}_t\mathbf{k}_t^\top, equals St−1−βt(St−1kt−vt)kt⊤S_{t-1} - \beta_t(S_{t-1}\mathbf{k}_t - \mathbf{v}_t)\mathbf{k}_t^\top.

Extend

A hybrid model has 32 layers, of which every fourth uses full attention and the rest linear attention. By what factor is its KV cache smaller than an all-attention model of the same shape?

Core Academic References

  • Katharopoulos et al. (2020): Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. ICML 2020. arXiv:2006.16236
  • Schlag, Irie & Schmidhuber (2021): Linear Transformers Are Secretly Fast Weight Programmers (delta rule). ICML 2021. arXiv:2102.11174
  • Yang, Kautz & Hatamizadeh (2024): Gated Delta Networks: Improving Mamba2 with Delta Rule. ICLR 2025. arXiv:2412.06464
  • Gu, Goel & Ré (2022): Efficiently Modeling Long Sequences with Structured State Spaces (S4). ICLR 2022. arXiv:2111.00396
  • Gu & Dao (2023): Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752
  • Dao & Gu (2024): Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality (Mamba-2). ICML 2024. arXiv:2405.21060
  • Arora et al. (2024): Zoology: Measuring and Improving Recall in Efficient Language Models. ICLR 2024. arXiv:2312.04927