Beyond Attention

It is June 2020. The transformer has, in three years, become the default architecture for sequence modeling, and its central operation — every position attending to every other position — carries a cost that grows with the square of sequence length. A small paper out of Idiap and EPFL asks a narrow question: can attention be rewritten so that generating the next token costs the same regardless of how long the sequence already is? Answering that question well takes four years and produces, in sequence, a kernel trick, a rehabilitated piece of 1960s control theory, a model that learns what to remember, and a production system that concludes attention should be kept, not replaced.

Where We Were

The transformer’s self-attention, introduced by Vaswani et al. in 2017 and covered in depth there, computes a compatibility score between every pair of positions in a sequence before combining values by those scores. For a sequence of length $n$, that is $n^2$ pairwise comparisons — cheap when $n$ is a few hundred tokens, expensive when $n$ reaches the tens or hundreds of thousands that long-document, long-context, and agentic workloads started demanding. Two prior chapters chip away at the consequences of that cost without touching its shape. Chapter 12 shows how FlashAttention, arriving in May 2022, restructures the same $O(n^2)$ computation to avoid materializing the full attention matrix in slow memory — a large constant-factor speedup, but the number of pairwise comparisons a longer sequence demands does not change. Chapter 17 shows how multi-query and grouped-query attention shrink the key-value cache that autoregressive decoding must keep resident, and how PagedAttention manages that cache efficiently — again a constant-factor and memory-layout fix, not a change to the fact that generating token $n+1$ still requires attending back across all $n$ prior tokens. Language-models Chapter 16 covers a third lever, extending the positional range a transformer can address — RoPE and its successors — without touching the per-token attention cost at all.

None of these three levers change the complexity class. A transformer processing a sequence twice as long still does roughly four times the attention work, and its per-step generation cost still scales with how much of the sequence it has already produced. By 2020, a parallel research thread had formed around a different question: rather than making quadratic attention cheaper to compute, could a sequence model be built whose cost is linear in sequence length from the start, ideally with a recurrent form whose per-step generation cost does not grow with context length at all — the property RNNs had always had, and had lost to transformers on quality grounds earlier in the decade, as Chapter 3 of the language-models book covers.

The Breakthrough

The first answer came from reformulating attention itself. Katharopoulos, Vyas, Pappas, and Fleuret observed that softmax attention’s cost comes specifically from the softmax: replace the similarity function between queries and keys with a positive-valued kernel feature map, and the associativity of matrix multiplication lets the computation be reorganized so that attention can be computed as a running sum rather than an all-pairs comparison (Katharopoulos, Vyas, Pappas, and Fleuret, 2020). Written this way, attention becomes mathematically equivalent to a linear RNN with a matrix-valued hidden state — hence the paper’s title, “Transformers are RNNs.” The practical payoff is autoregressive generation that costs $O(1)$ per step instead of $O(n)$: the paper reports linear transformers up to 4000 times faster than a standard transformer on long-sequence autoregressive prediction, at similar task performance.

Linear attention answered the recurrence question by editing attention. A second, largely independent thread edited nothing about attention and instead revived a much older tool: the state-space model from control theory, in which a hidden state evolves continuously over time under a linear dynamical system and is read out through a separate output map. Prior attempts to use state-space models (SSMs) for sequence learning existed, but their computation and memory requirements made them impractical at any useful scale. Gu, Goel, and Ré fixed this by constraining the state matrix to a structured form — derived from the HiPPO framework for optimal continuous-time memory — that can be diagonalized stably and reduced to a well-studied Cauchy-kernel computation, making the model both trainable via a global convolution and, at inference, runnable as an exact recurrence (Gu, Goel, and Ré, 2021). Their model, S4, was the first sequence model to solve every task in the Long Range Arena benchmark, including Path-X — a synthetic task with 16,000-step dependencies that every prior approach, transformers included, had failed outright — while generating 60 times faster than a comparable transformer.

S4 proved that structured state-space models could match or beat transformers on long-range, largely continuous signals — image pixels, audio, and the synthetic Long Range Arena tasks. It underperformed transformers specifically on language, where the state transition matrices $A$, $B$, and $C$ are fixed once training finishes and cannot adapt to what a specific token actually is. Two lines of work closed that gap from different directions in 2023. Peng and over thirty collaborators built RWKV, which keeps the linear-attention idea of a recurrent, constant-memory formulation but engineers it specifically for language, and scaled it to 14 billion parameters — at the time, the largest dense RNN ever trained — showing it performs on par with similarly sized transformers (Peng et al., 2023). Gu and Dao, meanwhile, diagnosed the language gap in S4-style SSMs directly: the fixed transition matrices meant the model could not perform content-based reasoning, deciding token by token what to keep and what to discard. Their fix, Mamba, makes the SSM’s discretization step and its input and output matrices functions of the current token — a selection mechanism — and, because that change breaks the fixed-kernel trick S4 relied on for fast training, pairs it with a hardware-aware parallel scan that keeps the recurrent computation in fast GPU memory rather than materializing the full sequence of hidden states (Gu and Dao, 2023). Mamba integrates the selective SSM into a network with no attention and no MLP blocks at all, reports five times higher inference throughput than a comparably sized transformer, and its 3-billion-parameter model matches transformers twice its size on language modeling — the first attention-free architecture to do so credibly at that scale.

By 2024, the question had shifted from “can a linear-time architecture match attention” to “should it replace attention entirely.” Lieber, Lenz, and their AI21 Labs collaborators answered with a hybrid: Jamba interleaves transformer attention layers with Mamba layers at a ratio of one attention layer for every seven Mamba layers, and adds mixture-of-experts routing — covered in Chapter 13 — to selected layers, so the model’s 52 billion total parameters activate only about 12 billion of them per token (Lieber et al., 2024). The premise was explicit: Mamba’s constant-memory recurrence solves the long-context memory problem, but a small amount of full attention still adds retrieval-style capability that a purely recurrent model was found to lack. The resulting model supports a 256,000-token context window and fits up to 140,000 tokens of that context on a single 80 GB GPU — a footprint a comparable transformer-only model at that scale could not match.

The Mathematics

Linear attention’s contribution is a rewrite of the same attention computation Chapter 5 of the language-models book defines, replacing the softmax similarity with a feature map $\phi$ so the sum over positions can be accumulated incrementally rather than recomputed for every query:

$$\begin{aligned} o_i &= \frac{\phi(q_i)^{\top} \sum_{j \le i} \phi(k_j)\, v_j^{\top}}{\phi(q_i)^{\top} \sum_{j \le i} \phi(k_j)} \\[6pt] S_i &= S_{i-1} + \phi(k_i)\, v_i^{\top}, \qquad z_i = z_{i-1} + \phi(k_i), \qquad o_i = \frac{\phi(q_i)^{\top} S_i}{\phi(q_i)^{\top} z_i} \end{aligned}$$

Symbol	Meaning
$q_i, k_i, v_i$	query, key, and value vectors at position $i$
$\phi$	positive-valued feature map standing in for the softmax kernel
$S_i$	running outer-product state summarizing all keys and values up to position $i$
$z_i$	running normalizing sum, replacing the softmax denominator
$o_i$	attention output at position $i$

S4 starts from the classical continuous-time state-space model and derives a discrete recurrence from it:

$$\begin{aligned} x'(t) &= A x(t) + B u(t), \qquad y(t) = C x(t) + D u(t) \\[6pt] x_k &= \bar{A}\, x_{k-1} + \bar{B}\, u_k, \qquad y_k = C x_k \end{aligned}$$

Symbol	Meaning
$u(t), u_k$	continuous and discretized input signal
$x(t), x_k$	continuous and discretized hidden state, dimension $N$
$A$	state matrix, constrained by S4 to a HiPPO-structured form for long-range memory
$B, C, D$	input, output, and skip matrices of the linear system
$\bar A, \bar B$	$A$ and $B$ discretized at a fixed step size
$y(t), y_k$	continuous and discretized output

Mamba’s selection mechanism modifies S4’s recurrence in exactly one place: the matrices that were fixed for the whole sequence become functions of the current token.

$$\begin{aligned} \Delta_t,\ B_t,\ C_t &= s_\Delta(x_t),\ s_B(x_t),\ s_C(x_t) \\[6pt] h_t &= \bar A_t\, h_{t-1} + \bar B_t\, x_t, \qquad y_t = C_t h_t \end{aligned}$$

Symbol	Meaning
$x_t$	input token representation at step $t$
$\Delta_t$	learned, input-dependent discretization step size
$B_t, C_t$	input-dependent input and output matrices, from learned projections $s_B$, $s_C$
$\bar A_t, \bar B_t$	discretized transition matrices at step $t$, themselves functions of the input
$h_t$	hidden state at step $t$
$y_t$	output at step $t$

What It Enabled

Architecture	Training compute vs. sequence length $n$	Inference memory per step	Reported speed / throughput	Headline result
Transformer, baseline	$O(n^2)$	grows with the KV cache (Chapter 17)	reference point	reference point
Linear Transformers, 2020	$O(n)$	$O(1)$, fixed-size state	up to 4000x faster autoregressive generation	matches vanilla transformer accuracy
S4, 2021	$O(n)$ via global convolution	$O(1)$ recurrent at inference	60x faster generation	SotA on every Long Range Arena task; solves Path-X at 16,000 steps
RWKV, 2023	$O(n)$	$O(1)$	scaled to 14B parameters	largest dense RNN trained; matches similarly sized transformers
Mamba, 2023	$O(n)$	$O(1)$	5x higher throughput than transformers	Mamba-3B matches transformers twice its size
Jamba, 2024	mostly $O(n)$, one attention layer in eight	reduced KV cache; 140K tokens on one 80 GB GPU	higher throughput, smaller footprint than vanilla transformers	256K-token context; 52B total / 12B active parameters

Property	Attention	Linear attention	S4	Mamba
Parallel training	yes, full $O(n^2)$ matrix	yes, via chunked recurrence	yes, via global convolution kernel	yes, via hardware-aware parallel scan
Constant-memory recurrent inference	no — KV cache grows with $n$	yes	yes	yes
Transition parameters depend on input content	yes, via softmax weights	no, fixed feature map	no, $A, B, C$ fixed after training	yes, $\Delta, B, C$ are functions of the input
Demonstrated language-modeling parity with transformers	reference point	approximate, not at scale	falls short specifically on language	yes, at up to 3B parameters

The direct beneficiary of this arc is inference cost. Chapter 17 shows that a transformer’s key-value cache grows linearly with context length, which is exactly the cost a constant-memory recurrence eliminates; that is why Jamba can hold 140,000 tokens of context on hardware a comparable dense transformer cannot fit on. Chapter 19 treats inference cost as an economic variable that determines what a model can be deployed to do — a state-space or hybrid layer changes the shape of that cost curve at long context lengths, independent of the model’s raw parameter count. And the pattern itself — a hardware constraint (attention’s cost at long context) provoking an architectural response, which then gets partially rolled back once its own limitation (content-based reasoning, and Jamba’s evidence that a little attention still helps) becomes visible — is the recurring shape Chapter 20 generalizes across the whole book.

Two years after Mamba, no frontier-scale production language model had gone fully attention-free; the architectures that reached the largest deployed scale by 2025 kept some ratio of attention, exactly as Jamba did, rather than eliminating it. That outcome is itself the chapter’s finding: state-space models did not replace attention, they became one more component a hybrid architecture mixes in, alongside the mixture-of-experts routing Chapter 13 covers and the long-context positional tricks language-models Chapter 16 covers.

Key Papers

Paper	Authors	Year	Link	Why it matters
Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention	Katharopoulos, Vyas, Pappas, and Fleuret	2020	arXiv:2006.16236	Rewrote attention via a kernel feature map, exposing an equivalent linear-cost recurrent form; up to 4000x faster autoregressive generation
Efficiently Modeling Long Sequences with Structured State Spaces	Gu, Goel, and Ré	2021	arXiv:2111.00396	Made state-space models practical via a HiPPO-structured, stably diagonalizable state matrix; first model to solve all of Long Range Arena, including Path-X
RWKV: Reinventing RNNs for the Transformer Era	Peng et al.	2023	arXiv:2305.13048	Scaled a linear-attention RNN to 14B parameters, the largest dense RNN trained at the time, matching similarly sized transformers on language
Mamba: Linear-Time Sequence Modeling with Selective State Spaces	Gu and Dao	2023	arXiv:2312.00752	Made SSM transition matrices input-dependent (selection) plus a hardware-aware parallel scan; first attention-free architecture to match transformers on language at scale
Jamba: A Hybrid Transformer-Mamba Language Model	Lieber et al.	2024	arXiv:2403.19887	Production-scale hybrid interleaving attention and Mamba layers with MoE; 256K context, evidence that a little attention still earns its keep

PM Lens