State Space Models and Selective SSMs

Attention's cost is quadratic in sequence length during training and its decoding memory grows linearly, because the KV cache stores every past token forever. Recurrent networks have the opposite profile, constant state and linear cost, and were abandoned because they train sequentially and forget. State space models are the attempt to get the recurrent profile without either drawback: a linear recurrence, expressible as a convolution so it trains in parallel, with a state matrix structured so that long-range information survives.

From a continuous system to a trainable layer

The starting point is a continuous linear system mapping a scalar input signal \(u(t)\) to output \(y(t)\) through a hidden state:

Discretising with step size \(\Delta\) gives \(h_t = \bar{A} h_{t-1} + \bar{B} u_t\), \(y_t = C h_t\), a linear recurrence. Because it is linear and time-invariant, unrolling it yields a convolution with a kernel \(\bar{K} = (C\bar{B}, C\bar{A}\bar{B}, C\bar{A}^2\bar{B}, \dots)\), which an FFT computes in \(O(L \log L)\). Training is parallel, inference is recurrent, and memory during decoding is the state, not the history.

The part that took work was \(A\). A random \(A\) forgets almost immediately. HiPPO derives matrices whose state maintains an optimal polynomial approximation of the input history, and S4 made computing the resulting kernel numerically tractable, which is what first made these models competitive on long-range benchmarks (Gu et al., 2021, Efficiently Modeling Long Sequences with Structured State Spaces, arXiv:2111.00396).

Selectivity: making the recurrence input-dependent

Linear time invariance is what makes the convolution work and is also the weakness. A fixed kernel treats every token identically, so the model cannot decide to remember this token and forget that one, and it fails at content-based tasks like selective copying and induction.

Mamba makes \(B\), \(C\), and \(\Delta\) functions of the current input. The system stops being time-invariant, so the convolution shortcut is lost; the paper recovers parallel training with a hardware-aware associative scan that keeps the expanded state in SRAM rather than materialising it in HBM. The result was reported as matching or beating transformers twice its size on language modelling with 5× higher generation throughput (Gu and Dao, 2023, Mamba: Linear-Time Sequence Modeling with Selective State Spaces, arXiv:2312.00752). Mamba-2 then showed that this class of models and a form of linear attention are two views of the same structured-matrix decomposition, which let it reuse matmul-heavy kernels and run substantially faster (Dao and Gu, 2024, arXiv:2405.21060).

The copying gap

The honest limitation is a consequence of the architecture, not of insufficient training. A fixed-size state can hold only a bounded amount of information about the past, so a task that requires reproducing arbitrary spans verbatim must eventually fail as the span grows, while attention can look up any past token exactly. This was demonstrated directly: transformers learn to copy long strings and generalise beyond training length, while state space models of comparable size do not (Jelassi et al., 2024, Repeat After Me: Transformers are Better than State Space Models at Copying, arXiv:2402.01032).

That result explains why almost every shipped model in this family is a hybrid rather than pure. Interleaving a small number of attention layers among many SSM layers restores retrieval and in-context copying while keeping most of the memory and throughput advantage, which is the design used in Jamba and in later hybrid releases (Lieber et al., 2024, arXiv:2403.19887).

When it breaks

• In-context retrieval degrades before perplexity does. A hybrid or pure SSM can match a transformer's loss curve and still be worse at pulling a specific string out of a long prompt. Test retrieval explicitly; language modelling loss will not surface it.
• The ecosystem is thinner. Quantisation kernels, speculative decoding integrations, LoRA variants, and serving-stack support are all written for attention first. Budget for the parts of your inference stack that assume a KV cache.
• State size is a real hyperparameter. Expanding the state improves recall and costs memory and bandwidth in the scan; there is no scaling law here as well established as the ones for transformer width and depth.
• Long-context claims need task-level evidence. "Linear scaling to a million tokens" is a statement about compute, not about whether the model still uses token 900,000.