State space models vs attention: S4, Mamba and hybrids

What is a state space model? It is a sequence layer that reads tokens through one running state instead of looking back at all of them. Attention, the core of a transformer, scores each new token against every key it has cached, so the cache grows with every token and the total work grows with the square of the length. The film opens on that bill: eight tokens take 36 scores, and by 4,096 tokens a single layer 2,048 numbers wide has computed 8,390,656 scores and holds a 32 MiB KV cache in 16-bit numbers.

A state space model (SSM) replaces the lookup with a linear recurrence: h_t = Ā h_{t-1} + B̄ x_t, then y_t = C h_t. Each step keeps a share of the state and adds a share of the new token. In the film’s one-number toy, with Ā = 0.8 and B̄ = 0.2, the inputs 5, 0, 0 and 10 leave the state at 1, 0.8, 0.64 and 2.51. Nothing is stored per token, so the memory never grows: a Mamba layer of the same width, with Mamba’s defaults of an expansion of 2 and N = 16 state numbers per channel, holds 128 KiB at any length. Attention’s cache matches that at 16 tokens and is 256 times larger by 4,096, and each new token costs the SSM one fixed update rather than a read of every earlier key. That read is what limits decoding speed: generating a token is bound by memory bandwidth, and attention must stream its whole cache from the GPU’s memory for every token it writes, so a fixed state speeds decoding as well as fitting in memory. That is the whole argument for SSMs at long context: linear time, and constant memory per token at inference.

Ā and B̄ are not free parameters. The model is defined in continuous time, h′ = A h + B x, and a step size Δ turns it into steps: Ā = e^{ΔA} and B̄ = (ΔA)⁻¹(e^{ΔA} − 1) ΔB, the zero-order hold. With A = −1, a small Δ of 0.1 keeps 90% of the state and barely writes, while a Δ of 2 keeps 14% and nearly overwrites it. S4 learned one Δ, one B and one C for every token, which made the layer time-invariant: it could be computed as one long convolution, but it treated a filler word exactly like a name.

Mamba’s change, and the part people miss, is selectivity. It computes Δ, B and C from each token, so the layer can choose what to write and what to let pass. In the film’s toy (weights set by hand, not trained), filler gets a Δ of 0.01 and slips past almost untouched, the key gets a Δ of 5 and is written in, and 91% of the key’s value survives nine fillers. With one fixed Δ of 0.5 for every token, 0.4% of it does. The price is that an input-dependent layer is no longer a convolution. Training recovers the parallelism with a hardware-aware scan, which loads the parameters from the GPU’s large, slow HBM into fast on-chip SRAM, discretises and runs the recurrence there, and writes only the outputs back, so the expanded state never makes the round trip through main memory. The scan works because each step is affine and two affine steps compose into one, so a parallel scan finds all eight states in three rounds rather than seven steps, or 4,096 states in 12 rounds. At inference the model simply runs the recurrence, one update per token. Mamba-2 goes further by restricting A to a scalar per head, which lets most of the work run as matrix multiplications; its paper reports a core layer 2 to 8 times faster than Mamba’s.

Where SSMs struggle is exact recall. A fixed state is a summary, and a summary has a capacity. The film’s toy writes key and value pairs into a 16 by 16 state, the outer-product form a Mamba-2 state takes, and reads them back: with 16 pairs it returns the right value 91% of the time, with 64 pairs 28%, while softmax attention over the same 64 cached keys returns all of them. Jelassi et al. prove the same limit for copying: a model with a fixed-size state cannot copy strings longer than its state can hold, while a two-layer transformer can copy strings of exponential length, and they find the gap in trained models too. This is why production models are increasingly hybrids. Jamba interleaves one attention layer among every seven Mamba layers; for a stack 2,048 wide, eight layers deep, that cuts the memory at 262,144 tokens from 16 GiB to about 2 GiB, an eighth, and the Jamba paper reports a 4 GB KV cache at 256K tokens against 32 GB for Mixtral.

The practical advice: reach for SSM or hybrid layers when long context and decode memory dominate your costs, keep some attention wherever exact lookup matters (retrieval over a long document, quoting, code), and evaluate candidates on recall and copying at your real lengths, not only on perplexity, which a good summary can score well on. The stats under the film let you change the context length and the state size N; the toy’s recall of 64 pairs climbs with N because a bigger state holds more.

The honest caveats. The memory numbers are shapes, not a trained model: one layer 2,048 wide in 16-bit numbers against a Mamba layer of the same width with its default E = 2 and N = 16, leaving out Mamba’s small convolution state and any grouped-query sharing an attention model might use, which would shrink its cache several times. The recurrence, Δ and selectivity toys have one channel and one state number with hand-set weights; a real Mamba layer learns its A, its projections for Δ, B and C, and runs thousands of channels. The recall toy has no decay (A = 1), random unit keys and 40 seeded trials per point; a trained model can do better than random keys by learning where to put things, but not beyond what its state can hold. The Mamba, Mamba-2 and Jamba figures quoted are the papers’ own.