12  State Space Models

You now own two kinds of memory. The transformer of Chapter 11 keeps everything: its key–value cache is an archive that grows with every token, answers exact-recall questions by lookup, and presents a bill that Section 11.3 measured directly — a bill that in production sets the economics of long-context serving. The recurrent networks of Chapter 8 keep one thing: a fixed-size state, updated in place, whose cost per token does not grow with the length of the context and whose memory of the past is whatever its finitely many bits managed to squeeze in. This chapter asks the question that sits between those designs: how far can a fixed state actually go? Its answer, which is the field’s answer as of this writing, comes as five verbs and a truce. Gate the state, so that learned sigmoids decide what is written and what is erased. Linearize it, so that training parallelizes across the sequence. Select, so that the dynamics read the data as it flows past. Edit, so that a write can correct the memory instead of piling on top of it. Learn, the reading under which all of these turn out to be one algorithm fitting a regressor at test time. And then the truce that production systems have settled on: hybridize, keeping a few layers of genuine attention inside a mostly recurrent stack.

The first three verbs are the chapter’s classical spine. Section 12.1 builds the gate once, carefully: starting from the observation that only an additive state passes gradients unattenuated, it derives the LSTM’s three gates and the GRU’s two, trains both against the vanilla RNN under one fixed recipe, and keeps the scoreboard — at a ten-epoch budget the cheaper GRU wins while the LSTM merely matches the baseline until its initialization is repaired, because architecture and optimization cannot be judged separately. Section 12.2 deletes the nonlinearity from the state path: what remains of the GRU is an affine recurrence that an associative scan evaluates in logarithmic depth, and the state space view rebuilds the same object from continuous time, where discretization derives the gate rather than positing it and the HiPPO theory supplies dynamics under which a fixed state’s memory of the past is provably good. The section trains an S4D classifier on 784-pixel image sequences, then runs the trained model both ways (all pixels at once through the scan, one pixel at a time through a carried state, verifying that the two agree) and prints the punchline: the model’s entire memory of an arbitrarily long history is a kilobyte-scale state, the same at token one hundred as at token one hundred thousand, where a transformer’s cache would have grown a thousandfold. Section 12.3 restores what linearization gave up. A selective-copying task built to defeat time-invariant dynamics motivates making the step size a function of the input (the forget gate derived a third time), and the resulting Mamba block solves the copy task, tops the chapter’s three-answers scoreboard in our PyTorch run (in the JAX run the far cheaper minGRU edges it out), and generates text at a per-token cost that does not grow with the prefix, through its own stepped path.

The next three sections carry the story to the present. Section 12.4 is where this chapter’s road meets the one from Chapter 10: linear attention’s matrix state and the selective recurrence are one template that varies only in how it forgets. The section measures the memory’s capacity law: after \(n\) independent random unit-norm writes into key width \(d_k\) the expected squared read error is \((n-1)/d_k\), and the measured curves sit on that prediction across three widths. It then climbs the decay ladder from RetNet to Mamba-2 to GLA, and derives the promised state-space duality — a gated linear recurrence and masked attention are the same matrix computed in two contraction orders, with the chunked third order being how these models train at scale. Section 12.5 changes the write rule. A memory that can only add fails when a key must be re-bound: in the section’s flagship experiment, recall of the latest value roughly halves by two writes per key and approaches chance by eight, a collapse that end-to-end training does not escape within the section’s deliberately restricted memory class, while the delta rule (read first, then write only the correction) holds recall essentially perfect throughout and turns out to be one step of gradient descent on a recall loss, running inside the forward pass. The section makes it trainable with a triangular solve, gates it into the Gated DeltaNet cell that several production models now ship, and shows that the new transition genuinely computes: a single eigenvalue explains why letting the write strength exceed one makes parity representable at any length. Section 12.6 then supplies the theory the instances have been hinting at: every memory in this chapter can be read as maintaining itself by solving a weighted regression of values on keys at test time. Softmax attention is the Nadaraya–Watson estimator (closing a loop opened in Section 10.1.3, whose one learnable bandwidth the section finally trains); linear attention is least squares with the key covariance deleted; the delta rule is one explicit gradient step; and a measured spectrum from a single online pass to the batch solve confirms that more solving buys a better memory. Two models fall out of the view rather than being designed: Longhorn, whose gate is the closed form of an implicit update, and Titans, a memory that is itself a small network trained inside the forward pass; a drifting-target experiment ends the section with the statistical reason forgetting exists at all.

Section 12.7 closes the chapter where production begins. A fixed state loses the exact-recall fight (the section quantifies what it cannot copy, and why language-modeling loss hides the deficit), but only attention layers pay a growing cache, so shipped systems interleave a few of them into a mostly recurrent stack. The centerpiece experiment trains three matched models, a pure recurrent stack, a pure attention stack, and a hybrid with a single attention layer mid-stack, and watches that one layer buy back most to all of the recall the recurrent stack loses (roughly 0.92 to 1.00 across the sweep in our runs) while perplexity barely moves; measured design rules for how much attention to keep and where to put it, and a recipe table of shipped hybrids from Jamba to Kimi Linear, turn the trade into engineering. One recipe threads all of these experiments together: every trained language model in the chapter runs on the Time Machine text of Section 8.5 — the classical spine and the Gated DeltaNet row on one shared scoreboard, the hybrid stacks on their own matched panel — and the mechanistic experiments (capacity, overwrite, the regression spectrum) run in seconds on a CPU.

The history here is a pendulum. The LSTM (Hochreiter and Schmidhuber 1997) made recurrence trainable and carried speech recognition and translation through the 2010s; the transformer displaced it almost completely, and for a few years recurrence looked finished. It returned through an unexpected door: S4 (Gu et al. 2022) arrived from continuous-time modeling rather than the RNN lineage, Mamba (Gu and Dao 2023) made the dynamics selective and competitive with transformers on language, and the state-space duality of Mamba-2 (Dao and Gu 2024) collapsed the wall between the returning recurrences and the linear attention that transformer researchers had been simplifying toward. After that the lineages merged outright: delta-rule cells and attention–recurrence hybrids now ship inside production language models from many labs at once. Whether the pendulum swings all the way back is a question with a date on it: a public wager between Jonathan Frankle (yes) and Sasha Rush (no) resolves on January 1, 2027, on whether a transformer-like model still holds the state of the art in most benchmarked language tasks (isattentionallyouneed.com). This chapter takes no side; it teaches what each side is counting on.

A word on the name, and on what this chapter is not. We use state space models the way the field now uses it: as the umbrella term for the whole fixed-state family — gated RNNs, linear recurrences, selective SSMs, matrix memories, test-time learners, and their hybrids — and not only for the continuous-time linear systems from which Section 12.2 takes the term (that section also notes what the phrase means to a statistician, which is different again). The chapter teaches algorithms, not kernels: the chunked forms here are twenty-line teaching implementations, and the Triton kernels and memory hierarchies that make them fast belong to Chapter 13. It trains no large models: pretraining recipes, data pipelines, the serving stacks that turn a trained model into a service, and everything downstream of a base model belong to the Language Models part. The efficient-attention taxonomy stays in Chapter 10, which already implemented the surviving variants; applications of state space models to vision, audio, and genomics are out of scope; and the fast-moving family of test-time-training architectures beyond Titans is fenced off at a pointer in the resources below. What remains is one adversary, met six ways: the fixed-size state, and the measured question of how much attention a model must keep when the state is not enough.

Two maps are worth carrying into the chapter. The first pins down its most overloaded word. State names five related but distinct objects in the sections ahead:

What “state” means Where Typical shape At autoregressive inference
RNN hidden vector \(\mathbf{H}_t\) (plus the LSTM cell \(\mathbf{C}_t\)) Section 12.1 \(h\) numbers per layer carried, updated in place
Continuous-time latent \(\mathbf{x}(t)\) Section 12.2 \(N\) numbers per channel analysis object; only its discretization runs
Discretized SSM state \(\mathbf{x}_t\) Section 12.2, Section 12.3 \((H, N)\) block per layer carried, updated in place
Matrix fast weight \(\mathbf{S}_t\) Section 12.4, Section 12.5 \(d_k \times d_v\) per head carried, updated in place
Inner-loop parameters of a memory network Section 12.6 a small MLP’s weights carried, updated by inner gradient steps

All but the second are one idea at different granularities: the numbers a fixed-memory model carries from token to token. The KV cache of Section 11.3 is the contrast class, per-token storage that grows with the context; “state” in this chapter never means that.

The second map is for reading the experiments. Each probes one property, and each has a confounder worth knowing before its conclusion arrives:

Experiment Probes Main confounder
Sequential-image classification (Section 12.2) long-range mixing (mean-pool readout) vs. state retention (final-step readout) the readout decides which is measured; the LSTM baseline is initialization-sensitive
Selective copying (Section 12.3) content-dependent selection a deep network around LTI mixers earns partial credit without selectivity
Random-key capacity (Section 12.4) additive-memory interference vs. key width assumes independent isotropic keys; learned keys are neither
Overwrite task (Section 12.5) key re-binding: additive vs. delta writes the trained baseline is a deliberately restricted memory class
Parity vs. length (Section 12.5) representability vs. trainability of sign-flipping transitions optimization noise across seeds and lengths
Solver spectrum (Section 12.6) value of more inner-solver compute the estimators optimize related, not identical, objectives
Hybrid recall sweep (Section 12.7) exact recall vs. attention budget position handling and parameter matching
LM scoreboards (Section 12.1, Section 12.3, Section 12.7) end-to-end quality at teaching scale one seeded run each; optimizer and parameter-count asymmetries

When a section’s conclusion reads stronger than its table, this map is the antidote.

Resources and Further Reading

Grouped by the chapter’s arc: implementations to build from, the explanations behind the ideas, the papers that organize the field, and course counterparts. All are freely accessible online.

Annotated implementations

The ideas, explained

Papers that organize the field

Course counterparts