Module 22: Linear & Recurrent Attention

Introduction

Two modules ago (m19) you built the leading alternative to attention — the state-space model, a linear recurrence over a fixed-size state. This module closes the loop by showing that the “alternative” was hiding inside attention all along.

Linear attention is self-attention with one thing removed: the softmax. That single deletion turns attention from an \(O(L^2)\) all-pairs comparison into a linear RNN — a recurrence over a fixed-size matrix state, exactly the shape of the SSM. The slogan from the paper that started this line of work says it plainly: transformers are RNNs.

From there, one small addition — a decay on the state — gives retention (the RetNet architecture), which comes with three provably-identical faces: a parallel form for training, a recurrent form for cheap decoding, and a chunkwise form that gets both at once. And at the end, the payoff: linear attention, retention, and the SSMs of m19 are literally the same recurrence, differing only in one gate. Knowing that one recurrence is how you read the whole zoo of sub-quadratic sequence models.

Why it matters for LLMs:

  • Constant-memory decoding. Softmax attention’s KV cache (m09) grows with every token; linear attention keeps a fixed-size state, so each generated token costs the same no matter how long the context.
  • Linear-time training and inference. No \(L \times L\) score matrix — compute scales with \(L\), not \(L^2\).
  • One mental model for many architectures. RetNet, RWKV, gated linear attention, and Mamba are one idea with different gates. This module gives you that idea.

What You’ll Learn

After this module, you can:

  • Explain why the softmax is the only thing forcing attention to be \(O(L^2)\), and how a kernel feature map \(\varphi\) removes it.
  • Run linear attention as a masked matmul (parallel) and as a running-state recurrence, and see they are the same function.
  • Build retention (RetNet): a decayed state with a \(\gamma^{n-m}\) decay mask, in its parallel, recurrent, and chunkwise forms — all provably equal.
  • Write the one unifying recurrence \(S_t = A_t S_{t-1} + k_t^\top v_t\) that specializes to linear attention (\(A_t=1\)), retention (\(A_t=\gamma\)), and a diagonal SSM / gated linear attention (\(A_t\) input-dependent).
  • Quantify the win: linear vs quadratic compute, constant vs growing decode memory.

Prerequisites

This module requires familiarity with:

Intuition: Attention Without the Softmax

Recall causal self-attention for query position \(i\):

\[ y_i = \frac{\sum_{j \le i} \exp(q_i \cdot k_j)\, v_j}{\sum_{j \le i} \exp(q_i \cdot k_j)} . \]

The trouble is the \(\exp(q_i \cdot k_j)\): it entangles \(q_i\) and \(k_j\) inside one nonlinearity, so you cannot separate them. You are forced to compute a score for every pair \((i, j)\) — the \(L \times L\) matrix — before you can sum. That matrix is the whole \(O(L^2)\) cost.

Now suppose the similarity factored — suppose we could write \(\text{sim}(q_i, k_j) = \varphi(q_i) \cdot \varphi(k_j)\) for some feature map \(\varphi\). Then the numerator regroups by associativity:

\[ \sum_{j \le i} \big(\varphi(q_i) \cdot \varphi(k_j)\big)\, v_j = \varphi(q_i) \underbrace{\sum_{j \le i} \varphi(k_j)\, v_j^\top}_{S_i} . \]

The sum \(S_i\) no longer depends on \(i\) except through its upper limit — it is a running total you can carry forward one token at a time. No \(L \times L\) matrix ever forms. The same regrouping is the difference between multiplying \((QK^\top)V\) (build the big matrix first) and \(Q(K^\top V)\) (build a small state first). Step through both association orders and watch the intermediate shape:

NoteKey Insight

The softmax is the only thing that forces the \(L \times L\) matrix. Remove it and replace the pairwise score with a factored kernel \(\varphi(q)\cdot\varphi(k)\), and the associativity of matrix multiplication lets you keep a fixed-size state instead — the same move that makes an SSM linear-time.

The Math: Linear Attention

We need a \(\varphi\) that keeps similarities non-negative (so the normalizer never turns negative) and is cheap. Katharopoulos et al. (2020) use

\[ \varphi(x) = \operatorname{elu}(x) + 1 , \]

which is \(x+1\) for \(x>0\) and \(\exp(x)\in(0,1]\) for \(x\le 0\) — positive for any realistic activation. With \(\text{sim}(q,k)=\varphi(q)\cdot\varphi(k)\), causal attention becomes a pair of running states:

\[ S_i = \sum_{j\le i} \varphi(k_j)\, v_j^\top, \qquad Z_i = \sum_{j\le i} \varphi(k_j), \qquad y_i = \frac{\varphi(q_i)^\top S_i}{\varphi(q_i)^\top Z_i} . \]

\(S_i\) is the \((m \times d_v)\) state matrix — a running, weighted memory of every value seen — and \(Z_i\) is the running normalizer. Both are fixed-size and updated by one addition per token. linear_attention.py implements exactly this:

import torch
from linear_attention import feature_map, linear_attention_recurrent

torch.manual_seed(0)
Q = torch.randn(6, 8)   # (length, d)
K = torch.randn(6, 8)
V = torch.randn(6, 8)   # (length, d_v)

y = linear_attention_recurrent(Q, K, V)
print("phi positive:", bool((feature_map(K) > 0).all()))
print("output shape:", tuple(y.shape))          # (length, d_v), no L×L matrix built
phi positive: True
output shape: (6, 8)

Code: The Two Faces of Linear Attention

Just like the SSM in m19 had a recurrent face and a convolutional face, linear attention has a recurrent face (the running state above, \(O(L)\), cheap to decode) and a parallel face that keeps the familiar attention layout — build the masked score matrix \(A_{ij}=\varphi(q_i)\cdot\varphi(k_j)\) for \(j\le i\), normalize each row, and read out \(V\):

\[ y_i = \frac{\sum_{j\le i} A_{ij}\, v_j}{\sum_{j\le i} A_{ij}} . \]

The parallel face is one masked matmul (great for training on a GPU); the recurrent face never forms a matrix (great for decoding). For the same inputs they compute bit-for-bit the same function — the linear-attention analogue of m19’s recurrent≡convolutional identity:

from linear_attention import linear_attention_parallel

y_par = linear_attention_parallel(Q, K, V)
y_rec = linear_attention_recurrent(Q, K, V)
print("parallel == recurrent:", torch.allclose(y_par, y_rec, atol=1e-5))
parallel == recurrent: True

Bridge both to a plot: the recurrent dots land exactly on the parallel line, because they are the same function computed two ways.

For decoding, wrap the running state in a small object and fold in one token at a time — the constant-work-per-token generation that a growing KV cache cannot match. LinearAttentionState does this:

from linear_attention import LinearAttentionState

state = LinearAttentionState(m=8, d_v=8)
ys = torch.stack([state.step(Q[i], K[i], V[i]) for i in range(6)])
print("incremental == full recompute:", torch.allclose(ys, y_rec, atol=1e-5))
incremental == full recompute: True
NoteKey Insight

Train like attention, decode like an RNN. The parallel face gives fast, parallelizable training; the recurrent face gives constant-memory generation. The price is a matrix state of size \(m \times d_v\) (typically \(d \times d\)) instead of a growing cache — constant in \(L\), but quadratic in the model width.

Retention: Add a Decay to the State

Linear attention keeps every past token with equal weight — its state \(S_i\) is an undamped running sum, which can saturate and blur. Retention (the RetNet architecture, Sun et al., 2023) adds one ingredient: a scalar decay \(\gamma \in (0,1]\) that shrinks the old state before each update:

\[ S_n = \gamma\, S_{n-1} + k_n^\top v_n, \qquad o_n = q_n\, S_n . \]

Two things are different from linear attention. First, older tokens fade geometrically — the state prefers recent context. Second, retention drops the normalizer entirely (no softmax, no division); RetNet stabilizes the output with a GroupNorm and a swish gate instead, but the core operation is this bare decayed state. Unrolling the recurrence, \(S_n = \sum_{m\le n} \gamma^{\,n-m} k_m^\top v_m\), so the output is

\[ o_n = \sum_{m\le n} \gamma^{\,n-m} (q_n\cdot k_m)\, v_m . \]

That is ordinary \(QK^\top\) scores, causally masked, weighted by a decay \(\gamma^{n-m}\) — a fixed lower-triangular decay matrix \(D\). Drive \(\gamma\) and watch \(D\) interpolate between two familiar limits:

TipTry It!
  1. Slide γ → 1. Every causal entry becomes 1: the decay vanishes and retention is uniform causal attention (unnormalized). This is exactly linear attention’s state with \(\varphi=\text{identity}\) and no forgetting.
  2. Slide γ → 0. Only the diagonal survives (\(0^0=1\)): each token sees only itself. The decay is the dial between “remember everything” and “remember only now”.
  3. Mid-range γ. A soft, fixed locality window — and note there is no positional encoding anywhere. The decay is the position information.

The Three Faces of Retention

Retention’s real selling point is that the same function has three computational forms, and you pick whichever the situation wants:

  • Parallel — \(\text{Retention}(X) = (QK^\top \odot D)\,V\). One masked matmul, fully parallel. Best for training short-to-medium sequences. \(O(L^2)\).
  • Recurrent — \(S_n = \gamma S_{n-1} + k_n^\top v_n\), \(o_n = q_n S_n\). Constant state, one step at a time. Best for decoding. \(O(L)\), \(O(1)\) memory per token.
  • Chunkwise — split into chunks of size \(B\); compute each chunk internally by the parallel form, and carry a fixed-size state \(R\) between chunks by the recurrent form. Best for training on long sequences: parallelism inside a chunk, linear scaling across chunks. \(O(LB)\).

The chunkwise form is the one production systems use. For a token at chunk-local index \(t\) it adds an intra-chunk term (parallel) to a cross-chunk term that reads the carried state \(R\), decayed by \(\gamma^{t+1}\):

\[ o_t = \underbrace{\sum_{s\le t}\gamma^{t-s}(q_t\cdot k_s)v_s}_{\text{intra-chunk (parallel)}} \; + \; \underbrace{\gamma^{\,t+1}\,(q_t R)}_{\text{cross-chunk (recurrent)}}, \qquad R \leftarrow \gamma^{B} R + \sum_s \gamma^{\,B-1-s} k_s^\top v_s . \]

All three are the same function. linear_attention.py proves it numerically:

from linear_attention import (retention_parallel, retention_recurrent,
                              retention_chunkwise)

torch.manual_seed(0)
Qr = torch.randn(12, 6); Kr = torch.randn(12, 6); Vr = torch.randn(12, 6)
gamma = 0.9

y_par = retention_parallel(Qr, Kr, Vr, gamma)
y_rec = retention_recurrent(Qr, Kr, Vr, gamma)
y_chk = retention_chunkwise(Qr, Kr, Vr, gamma, chunk_size=4)

print("parallel == recurrent:", torch.allclose(y_par, y_rec, atol=1e-5))
print("parallel == chunkwise:", torch.allclose(y_par, y_chk, atol=1e-5))
parallel == recurrent: True
parallel == chunkwise: True

Real RetNet uses multi-scale retention: each head gets its own \(\gamma_h = 1 - 2^{-5-h}\), so some heads keep long memory and others stay local — the decay analogue of multi-head attention’s different subspaces.

from linear_attention import multiscale_decays
print("per-head γ:", [round(v, 4) for v in multiscale_decays(4).tolist()])
per-head γ: [0.9688, 0.9844, 0.9922, 0.9961]

Interactive Exploration: The One Recurrence Behind Them All

Here is the idea this whole module builds toward. Linear attention, retention, and the diagonal SSM of m19 are the same recurrence

\[ S_t = A_t\, S_{t-1} + k_t^\top v_t, \qquad y_t = q_t\, S_t , \]

differing only in the gate \(A_t\) that scales the old state:

Model Gate \(A_t\) Behavior
Linear attention \(1\) (identity) Keep everything, equal weight
Retention (RetNet) \(\gamma\) (constant) Forget geometrically
SSM / gated linear attn / Mamba \(A_t = f(x_t)\) (input-dependent) Choose what to keep, per token

The unified_recurrence function is literally one loop with a swappable \(A_t\). Below, you are the gate: check which tokens are “important”. Important tokens get a high gate (\(A_t \approx 0.95\), the state holds); the rest get a low gate (\(A_t \approx 0.2\), the state forgets). Watch the memory integrate only what you select — this is selectivity, the input-dependent gate that lifts a fixed decay into Mamba’s content-based reasoning.

Now compare the three regimes on one fixed input stream — linear attention’s state grows without bound (nothing forgets), retention’s saturates (a constant decay), and the gated state rises and falls as the gates open and close:

NoteKey Insight

This is the m19↔︎attention duality (formalized by Mamba-2) made runnable: the gate \(A_t\) here is exactly the SSM’s discrete decay \(\bar{A}\). The difference between a “transformer” and a “state-space model” collapses to a single question — is the gate constant or input-dependent?

Why It Matters: Cost

The payoff, made quantitative. Softmax attention does two \(O(L^2 d)\) matmuls and stores a KV cache that grows linearly with the context. Linear attention does an \(O(d^2)\) state update per token over \(L\) tokens — \(O(L d^2)\), linear in \(L\) — and carries a fixed \(d \times d\) state that never grows. Drive the context length and watch the compute curves diverge:

NoteKey Insight

Linear attention trades a growing KV cache for a fixed matrix state. That state is larger than an SSM’s (\(d \times d\) vs \(d \times N\) with a small \(N\)), which is why pure linear attention has historically traded a little quality for that memory constant — and why the frontier (gated linear attention, Mamba, hybrids) works to recover the quality while keeping the constant-memory decode.

Common Pitfalls

When building linear/recurrent attention, watch out for:

  1. A non-positive feature map. If \(\varphi\) can go negative, the normalizer \(\varphi(q_i)^\top Z_i\) can hit zero or flip sign and the output explodes. Keep \(\varphi\) non-negative (\(\text{elu}+1\), or \(\text{ReLU}\), or a softmax-feature map).
  2. Forgetting the normalizer in linear attention. Linear attention divides by \(\varphi(q_i)^\top Z_i\); retention does not (it uses a GroupNorm downstream instead). Mixing them up gives the wrong scale.
  3. A decay \(\gamma \ge 1\). Retention needs \(\gamma \in (0,1]\); \(\gamma>1\) makes \(\gamma^{n-m}\) grow with distance and the state diverges. (Same stability rule as the SSM’s \(\bar{A}\in(0,1)\) in m19.)
  4. Chunk-boundary bookkeeping. The cross-chunk term must be scaled by \(\gamma^{t+1}\) and the carried state updated as \(R \leftarrow \gamma^B R + \sum_s \gamma^{B-1-s}k_s^\top v_s\). Off-by-one in these exponents silently breaks the equivalence — test chunkwise against parallel.
  5. Expecting linear attention to match softmax exactly. It is a different function, not an approximation of softmax. It trades some expressivity (the softmax’s sharp, data-dependent focus) for linear cost; the research since 2020 is largely about closing that quality gap.
  6. Confusing the state size with the SSM’s. Linear attention’s state is \(d \times d\); an SSM’s is \(d \times N\) with \(N\) small (16). Both are constant in \(L\), but linear attention’s is bigger — the price of using full key/value vectors as the “state coordinates”.

Exercises

Exercise 1: The two faces agree

from linear_attention import linear_attention_parallel, linear_attention_recurrent
import torch

# Generate random Q, K, V of your choice. Confirm the parallel and recurrent
# forms match, then explain in one sentence WHY they must: what property of the
# sum over j lets you regroup it into a running state?

# Your implementation here:

Exercise 2: Decay recovers linear attention

from linear_attention import retention_parallel
import torch

# Retention with gamma=1 (no decay) and the identity feature map is exactly
# unnormalized causal attention: (Q Kᵀ ⊙ tril) V. Build that expression by hand
# with torch.tril and confirm it equals retention_parallel(Q, K, V, gamma=1.0).
# Then set gamma=0.0 and describe, in words, what the output becomes.

# Your implementation here:

Exercise 3: You are the gate

from linear_attention import unified_recurrence
import torch

# Build Q, K, V (length 8). Using mode="gated", design a `gate` vector that keeps
# a HIGH gate (~0.95) only at positions 2 and 5 and a LOW gate (~0.1) elsewhere.
# Show the output differs from mode="retention" with a single fixed gamma — i.e.
# input-dependent selectivity buys something a constant decay cannot.

# Your implementation here:

Summary

Key takeaways:

  1. Linear attention is softmax attention minus the softmax. Replace \(\exp(q\cdot k)\) with a factored kernel \(\varphi(q)\cdot\varphi(k)\) and associativity turns the \(O(L^2)\) all-pairs sum into a running state \(S_i = \sum_{j\le i}\varphi(k_j)v_j^\top\) — a linear RNN hiding inside attention.
  2. It has two equal faces. A parallel masked-matmul form for training and a constant-memory recurrent form for decoding compute the same function.
  3. Retention adds a decay. \(S_n=\gamma S_{n-1}+k_n^\top v_n\) gives a \(\gamma^{n-m}\) decay mask (RetNet), which is the positional signal — no separate encoding — and comes in parallel, recurrent, and chunkwise forms, all provably identical.
  4. Chunkwise is the practical form. Parallel inside chunks, recurrent between them: GPU-parallel training that still scales linearly in sequence length.
  5. It is all one recurrence. \(S_t=A_tS_{t-1}+k_t^\top v_t\) specializes to linear attention (\(A_t=1\)), retention (\(A_t=\gamma\)), and a diagonal SSM / Mamba (\(A_t=f(x_t)\)). Transformer vs SSM collapses to is the gate input-dependent?
  6. The cost win is long-context. \(O(Ld^2)\) vs \(O(L^2d)\) and a fixed \(d\times d\) state vs a growing KV cache; the crossover is at \(L\approx d\), so linear attention pays off exactly where softmax’s quadratic wall bites.

What’s Next

You have now seen both sides of the sub-quadratic coin — m19 approached it from control theory (SSMs), this module from attention (drop the softmax) — and watched them meet in one recurrence. The live frontier from here is gated linear attention and DeltaNet (input-dependent gates and a delta-rule state update that recover more of softmax’s expressivity) and hybrid stacks (Jamba, Griffin) that interleave a few full-attention layers among many linear/SSM layers to get global recall and linear cost. The throughline is the one this module makes concrete: the sequence-mixing layer is a design space, and the axis is how the state forgets.

Going Deeper

Core Papers:

Practical Resources: