The loop that beats attention

In the transformer post I made attention sound like an unalloyed win: every token reads every other token in a single parallel step, so there's no slow recurrent baton crawling left to right. That parallelism is exactly what let attention be trained on a GPU, and it's why the old recurrent networks were left for dead.

Here's the catch I glossed over. Reading every token against every other costs you $n^2$ work for a sequence of length $n$, and at generation time you have to keep every past token's keys and values around in a cache that grows without limit. I spent a whole post on tricks to shrink that cache. State-space models ask a more radical question: what if you never re-read the past at all?

The bill attention can't stop paying

Slide the context length up on this and watch the two costs pull apart. Attention's work climbs as the square of the length while a state-space model's climbs in a straight line, and the gap is the whole reason anyone went back to the loop.

The square is structural, not an implementation detail. If token 4,000 must attend to all 3,999 tokens before it, and so must every other token, you've signed up for $n^2$ comparisons no matter how clever your kernels are. And the cache that holds those keys and values grows one slab per token, forever. For a chatbot that's annoying; for a model meant to read a whole codebase or an hour of audio, it's the binding constraint.

Going back to the baton

So we go back to the idea the transformer threw away: a running summary, updated one token at a time. Not an LSTM's tangle of gated nonlinearities, though. A plain linear recurrence with learnable matrices.

$$h_t = \mathbf{A}\, h_{t-1} + \mathbf{B}\, x_t, \qquad y_t = \mathbf{C}\, h_t$$

Copy LaTeXwith $$ delimitersCopy mathexpression only

Read it slowly, because the payoff is in the shapes. The state $h_t$ has a fixed size, set once and never growing. To process the next token you update that one state and move on, so inference is $O(n)$ overall and the memory you carry between steps is $O(1)$, completely independent of how much text you've already read. Where attention's cost grows with the conversation, this doesn't.

Which probably looks too simple to compete with attention, and honestly it should. Getting from "too simple" to "actually competitive" took until 2023, and it's worth seeing what was missing.

Where A, B and C come from

These matrices aren't picked out of the air. The cleanest way to read equation (1) is as a discretisation of an underlying continuous-time system, $x'(t) = Ax + Bu$, sampled at a step size $\Delta$. The discrete matrices are what you get when you ask "if the continuous system runs for $\Delta$ seconds, where does the state land?".

$$\overline{\mathbf{A}} = \exp(\Delta \mathbf{A}), \qquad \overline{\mathbf{B}} = (\Delta \mathbf{A})^{-1}\big(\exp(\Delta \mathbf{A}) - \mathbf{I}\big)\,\Delta \mathbf{B}$$

Copy LaTeXwith $$ delimitersCopy mathexpression only

The clever bit in S4 was choosing the structure of $A$ (the HiPPO construction) so that the state provably keeps a useful compression of a long history rather than forgetting it after a few steps. I'll wave at that rather than derive it; the point you need is that the matrices encode a principled, long memory, not an arbitrary one.

The same loop is also a convolution

Now the duality that makes the whole thing trainable. Unroll the recurrence, and the output is just the input convolved with one long, fixed kernel.

$$y = x * \mathbf{K}, \qquad \mathbf{K} = \big(\mathbf{C}\overline{\mathbf{B}},\ \mathbf{C}\,\overline{\mathbf{A}}\,\overline{\mathbf{B}},\ \mathbf{C}\,\overline{\mathbf{A}}^{2}\overline{\mathbf{B}},\ \dots\big)$$

Copy LaTeXwith $$ delimitersCopy mathexpression only

One operation, two ways to compute it. At inference you run the recurrence: cheap, streaming, $O(1)$ state. At training you treat the whole sequence as a single convolution and do it in one parallel shot, which (since a long convolution is a multiplication in the frequency domain, the Fourier trick) S4 evaluated with an FFT. That's how a "recurrent" model trained as fast as a transformer. Hold onto this duality; in a moment we lose it and have to win it back.

Why the first version underperformed

Here's the puzzle that stumped the field for a while. S4 had the linear-time recurrence and the convolution trick, yet transformers still beat it on language. Why?

Because $A$, $B$ and $C$ were time-invariant: the same filter applied to every token, regardless of what that token was. A fixed kernel has no way to say "this token is a delimiter I should ignore" or "this is the variable name I'll need later, hold onto it". It treats the word the and the one fact that answers the question with exactly the same response. Attention, by contrast, decides what to look at based on content. That missing content-awareness is the whole gap.

Selectivity: let the input steer the state

Mamba's fix is almost cheeky in its directness. Make the step size and the read/write matrices functions of the current token.

$$\Delta_t,\ \mathbf{B}_t,\ \mathbf{C}_t = \operatorname{Linear}(x_t)$$

Copy LaTeXwith $$ delimitersCopy mathexpression only

Recall from equation (2) that a big $\Delta$ takes a big step and writes the current input in strongly, while a small $\Delta$ barely moves the state. Make $\Delta$ depend on the token and you've handed the model a remember/forget gate: open it wide to overwrite the state with an important token, leave it nearly shut to let a filler word pass through untouched. Step through it here and watch the single fixed-size state get steered token by token, with attention's all-pairs wiring drawn above for contrast.

The catch, and the parallel scan that saves it

There's a price for selectivity, and it's the duality we were told to hold onto.

That older trick is the parallel scan, and it's the prettiest idea in the post.

So Mamba computes all $n$ states with a scan in $O(\log n)$ depth instead of an $O(n)$ crawl, keeps the whole thing in fast on-chip memory, and never materialises the full sequence of states. The loop trains in parallel after all.

The catch you can't optimise away

I've been selling the win. Here's the honest cost, and it's not a bug you can patch out.

Which is why nobody actually had to pick a side.

One honest aside before the recap. If equation (1) looks familiar to anyone who's done control or signals, that's because it is familiar. $h_t = A h_{t-1} + B x_t,\ y_t = C h_t$ is the standard discrete-time linear state-space model, the same $(A, B, C)$ an engineer writes down for a filter or feeds to a Kalman filter. The "hidden state" a control engineer means by state is, almost word for word, the "hidden state" the old RNN meant. What's new isn't the equation. It's making $A$, $B$, $C$ learned and input-dependent and stacking the thing fifty layers deep. Mamba quietly reunited two fields that had been writing down the same maths for different reasons.

So, which loop wins

Attention reads everything and pays for it forever. The state-space model reads once, keeps a small steerable summary, and trades exact recall for linear cost. Neither strictly wins, which is why the field, sensibly, kept both, layered together. The thing I find satisfying is that the "dead" recurrent loop didn't come back unchanged. It came back having borrowed a content-gate from attention and a parallel scan from 1990s parallel computing, and only then could it compete.

Recap

Attention's $n^2$ cost and ever-growing cache are structural. A linear recurrence with a fixed-size state fixes both, at $O(n)$ time and $O(1)$ memory per step. S4 had that already and lost, because a time-invariant filter can't choose what to ignore. Mamba makes the update input-dependent (selectivity), which costs it the convolution trick and forces a parallel scan to keep training fast. The finite state is the whole bargain: linear cost in exchange for fuzzy recall, which is why the best systems hybridise it with a little attention.

Reading further

• Gu & Dao (2023), Mamba: Linear-Time Sequence Modeling with Selective State Spaces: selectivity (S6) and the hardware-aware scan, the central reference here. arXiv:2312.00752
• Gu, Goel & Ré (2021), Efficiently Modeling Long Sequences with Structured State Spaces (S4): the continuous-time SSM, HiPPO, and the convolution view. arXiv:2111.00396
• Smith, Warrington & Linderman (2022), Simplified State Space Layers (S5): the parallel-scan formulation that bridges S4's convolution to Mamba's scan. arXiv:2208.04933
• Blelloch (1990), Prefix Sums and Their Applications: the associative scan primitive, decades older than the architecture that needed it. cs.cmu.edu
• Rush, The Annotated S4: a runnable, heavily-annotated walk through the discretisation and the kernel. srush.github.io
• Lieber et al. (2024), Jamba: A Hybrid Transformer-Mamba Language Model: the hybrid that keeps exact retrieval while paying linear cost for the bulk. arXiv:2403.19887

Next I want to take that parallel scan apart on its own, because the prefix-sum primitive quietly powers half of modern GPU code and it deserves its own bench. Watch this space.