# Module 18: Beyond Attention: Linear Attention and Hybrids

> Linear attention as a recurrence, the gated delta rule, state-space models such as Mamba, and hybrids that swap the KV cache for a fixed-size state.

- Part IV: The Modern Architecture · elective · Advanced · 3 hours
- By the end you can explain how a fixed-size state replaces the KV cache, and what it costs.
- Builds on: [Module 15: The LLaMA Block](https://transformers-and-math.drmhse.com/md/15-the-llama-block.md), [Module 5: Sequences Before Transformers](https://transformers-and-math.drmhse.com/md/05-sequences-before-transformers.md)
- Web page (interactive figures, labs and equation explainers): https://transformers-and-math.drmhse.com/modules/18-beyond-attention
- From *Transformers & Math* by Mike Chumba: https://transformers-and-math.drmhse.com

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 Attention

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 $\mathbf{v}_j\mathbf{k}_j^\top$ is a $d_v \times d_k$ outer product. Drop the softmax (Katharopoulos et al., 2020) and the newest token's output regroups:

$$
\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 $\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 $d_v \times d_k$ matrix, the *state* $S_t$, updated once per token:

$$
S_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 Rule

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

$$
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
$$

$g_t$ and $\beta_t$ are numbers between 0 and 1 computed from the current token, and keys have length 1. The *gate* $g_t$ fades old content, so the layer can forget. The factor $(I - \beta_t \mathbf{k}_t \mathbf{k}_t^\top)$ is the *delta rule*: before writing, it removes a share $\beta_t$ of what the state currently returns for $\mathbf{k}_t$. With $g_t = 1$ the update is $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 $\mathbf{k}_t$ and what it should.

> **Interactive figure** `delta-rule`, on the web page: https://transformers-and-math.drmhse.com/modules/18-beyond-attention#figure-delta-rule
>
> **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.
>
> **What happens:** The sum of both, plus a little of every overlapping key's value. Plain linear attention only ever adds, so writes pile up. The delta rule first subtracts what the state currently returns for *a*, so recalling *a* afterwards gives exactly the new value.

> **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 Models

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

$$
\mathbf{h}_t = A\,\mathbf{h}_{t-1} + B\,\mathbf{x}_t, \quad \mathbf{y}_t = C\,\mathbf{h}_t
$$

with learned $A$, $B$ and $C$. 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 $A$: 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 $g_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-recall`: an experiment that runs in the reader's browser on the web page: https://transformers-and-math.drmhse.com/modules/18-beyond-attention#lab-linear-recall

## 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 |
|---|---|
| layers | 32 |
| hidden size | 4096 |
| intermediate size | 12288 |
| vocabulary | 248,320 |
| attention | **hybrid**: 24 gated-DeltaNet linear-attention layers, and full attention in every 4th layer |
| estimated prefill cost | 13.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 throughput | 224 tok/s at 2.5k context, 199 tok/s at 16k |

A context $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 $N^2$ 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

**Exercise (check).** What does a linear-attention layer keep instead of a KV cache, and how does its size grow with the sequence length?

*Solution:* A fixed-size state matrix, $d_v \times d_k$, updated once per token. Its size does not grow with the sequence length at all.

**Exercise (apply).** A linear-attention layer has $d_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?

*Answer:* 16384 (within 0)

*Solution:* The state is $d_v \times d_k = 128 \times 128 = 16{,}384$ numbers, whatever the length. A KV cache for one head stores $2 \times 128 = 256$ numbers per token, so 25.6 million for 100,000 tokens.

**Exercise (apply).** Show that the gated delta update with $g_t = 1$, $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 $S_{t-1} - \beta_t(S_{t-1}\mathbf{k}_t - \mathbf{v}_t)\mathbf{k}_t^\top$.

*Solution:* Expand: $S_{t-1} - \beta_t S_{t-1}\mathbf{k}_t\mathbf{k}_t^\top + \beta_t\mathbf{v}_t\mathbf{k}_t^\top = S_{t-1} - \beta_t(S_{t-1}\mathbf{k}_t - \mathbf{v}_t)\mathbf{k}_t^\top$. The bracket is the error between what the state recalls for $\mathbf{k}_t$ and the value it should recall, so this is one gradient step on the squared recall error with learning rate $\beta_t$.

**Exercise (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?

*Answer:* 4 (within 0)

*Solution:* Only 8 of the 32 layers keep a KV cache, so the cache is a quarter of the all-attention cache (the linear layers add only their small fixed states).

### Core Academic References
- **Katharopoulos et al. (2020)**: *Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention*. ICML 2020. [arXiv:2006.16236](https://arxiv.org/abs/2006.16236)
- **Schlag, Irie & Schmidhuber (2021)**: *Linear Transformers Are Secretly Fast Weight Programmers* (delta rule). ICML 2021. [arXiv:2102.11174](https://arxiv.org/abs/2102.11174)
- **Yang, Kautz & Hatamizadeh (2024)**: *Gated Delta Networks: Improving Mamba2 with Delta Rule*. ICLR 2025. [arXiv:2412.06464](https://arxiv.org/abs/2412.06464)
- **Gu, Goel & Ré (2022)**: *Efficiently Modeling Long Sequences with Structured State Spaces* (S4). ICLR 2022. [arXiv:2111.00396](https://arxiv.org/abs/2111.00396)
- **Gu & Dao (2023)**: *Mamba: Linear-Time Sequence Modeling with Selective State Spaces*. [arXiv:2312.00752](https://arxiv.org/abs/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](https://arxiv.org/abs/2405.21060)
- **Arora et al. (2024)**: *Zoology: Measuring and Improving Recall in Efficient Language Models*. ICLR 2024. [arXiv:2312.04927](https://arxiv.org/abs/2312.04927)