State Space Models (Mamba)
Attention's bottleneck, and an old idea revived
Self-attention gives transformers excellent modeling power but at $O(n^2)$ cost in sequence length $n$ — every token attends to every other token. State space models (SSMs), and specifically Mamba (Gu & Dao, 2023), offer an alternative sequence-modeling backbone with linear-time complexity in sequence length, borrowing from classical control theory's state space representation rather than from attention at all.
The state space formulation
A continuous-time state space model describes how a hidden state $h(t)$ evolves and produces an output $y(t)$ from an input $x(t)$:
$$h'(t) = A h(t) + B x(t), \qquad y(t) = C h(t) + D x(t)$$
Discretized for sequence data, this becomes a simple recurrence: at each timestep, the new hidden state is a learned linear combination of the previous hidden state and the current input, and the output is a learned linear readout of the hidden state. This recurrence is structurally similar to an RNN, but the key difference is that $A$, $B$, $C$ can be chosen (or made input-dependent, as in Mamba) such that the recurrence can also be computed via a parallel scan or convolution — giving it RNN-like linear-time inference *and* the parallelizable training of a convolution, rather than forcing a choice between the two the way plain RNNs do.
import numpy as np
def ssm_scan(x_sequence, A, B, C, h0):
"""Linear-time sequential recurrence version of an SSM (illustrative, not the
parallel-scan form actually used for fast training)."""
h = h0
outputs = []
for x_t in x_sequence:
h = A @ h + B @ x_t # state update: linear, no quadratic attention cost
y_t = C @ h # readout
outputs.append(y_t)
return np.stack(outputs)
# Cost per step is O(state_dim^2), independent of how many tokens came before --
# contrast with self-attention, where each new token's cost grows with sequence length.Mamba's key addition: input-dependent (selective) parameters
Self-attention vs selective SSM (Mamba)
| Property | Self-attention (transformer) | Selective SSM (Mamba) |
|---|---|---|
| Compute cost (training) | O(n^2) in sequence length | O(n) in sequence length |
| Inference memory | Grows with KV cache (all past tokens) | Fixed-size hidden state, independent of length |
| Content-based retrieval | Explicit, any-to-any attention | Implicit, via learned state updates |
| Long-context scaling | Expensive without approximations | Naturally linear, but state is finite capacity |
The real tradeoff: finite state vs explicit memory
The practical tension Mamba introduces is that its hidden state has fixed size, no matter how long the input sequence gets — unlike attention, which can in principle look back at any token in the full context explicitly (at quadratic cost). This makes SSMs extremely memory-efficient for long sequences but means they must compress everything relevant from the past into a fixed-size summary, which can lose precise access to specific far-back details that attention would retrieve exactly. This is why many production architectures now hybridize — interleaving SSM layers (for efficient long-range summarization) with a smaller number of attention layers (for precise retrieval) rather than committing fully to either.
Hardware-aware implementation, not just the algorithm, was essential
What's Next?
State space models are one of several answers to the same question that motivates sparse and linear attention mechanisms — how to model long sequences without paying attention's quadratic bill — and the design tradeoffs between fixed-size recurrent state and explicit retrieval keep reappearing across that whole family of efficient-sequence-modeling approaches.
I build these systems professionally.
Whether it's a RAG pipeline, analytics migration, or AI workflow — let's talk.