A state space model is the oldest idea in signal processing wearing a new hat: carry a small hidden state forward through time, update it linearly at every step, and read the output off that state. What makes SSMs interesting for sequence modelling is a duality — the same parameters can be evaluated as a recurrence (fixed memory per token, perfect for generation) or as a single long convolution (O(L log L) parallel training). That is what lets S4 and Mamba train like a Transformer and run like an RNN. This article builds the model from the continuous equations, derives both views and the discretization connecting them, works a numeric example, and shows why the choice of the matrix A — not the architecture around it — is what made SSMs finally work.
The model: a linear system with a hidden state
Start in continuous time. A single-input, single-output linear time-invariant (LTI) system maps u(t) to y(t) through a hidden state x(t):
x′(t) = A x(t) + B u(t)
y(t) = C x(t) + D u(t)
A: [N, N] B: [N, 1] C: [1, N] D: scalar
x: [N] u, y: scalarsN is the state dimension — typically small, 8 to 64 — and it is completely independent of the sequence length. A says how the state evolves on its own, B how new input enters it, C how it is read out, and D is just a skip connection. In a deep model this scalar system is replicated per channel: with d channels you run d independent copies in parallel, so the layer’s state is [d, N]. Everything is linear — nonlinearity lives between layers, not inside the recurrence — and that linearity is exactly what buys the convolutional form later.
From continuous time to discrete steps
Tokens arrive at discrete positions, so the system must be discretized with a step size Δ. Solving the ODE exactly over one interval, with u held constant across it (zero-order hold, ZOH), gives:
ZOH: Ā = exp(ΔA)
B̄ = A^-1 (exp(ΔA) − I) B
Bilinear: Ā = (I − Δ/2 · A)^-1 (I + Δ/2 · A)
B̄ = (I − Δ/2 · A)^-1 ΔBThe bilinear (Tustin) rule is a rational approximation of the matrix exponential, cheaper than ZOH and used by the original S4; ZOH is standard in Mamba. Δ is a learned parameter, stored in log space so it stays positive, and it sets the timescale: a small Δ means Ā ≈ I, so the state barely moves per token and the model remembers far back; a large Δ turns the state over quickly and focuses on recent input.
The recurrent view: constant memory per token
Once discretized, the model is a plain linear RNN:
x_k = Ā x_{k-1} + B̄ u_k
y_k = C x_k + D u_kEach step is one matrix–vector product: O(N^2) for a dense Ā, or O(N) when Ā is diagonal, as in S4D and Mamba. Crucially, cost and memory are independent of how many tokens came before: no cache grows, and the entire past is compressed into the [d, N] state. That is the property attention cannot offer, and the whole reason SSMs are attractive for long-context or on-device generation. The price is sequentiality — x_k depends on x_{k-1}, so training a length-L sequence would take L serial steps. The convolutional view removes exactly that obstacle.
The convolutional view: one kernel, one FFT
Because the recurrence is linear and time-invariant, you can unroll it in closed form. With x_{-1} = 0:
x_0 = B̄ u_0 → y_0 = C B̄ u_0
x_1 = Ā B̄ u_0 + B̄ u_1 → y_1 = C Ā B̄ u_0 + C B̄ u_1
y_k = Σ_{j=0..k} (C Ā^j B̄) u_{k-j}
K̄ = (C B̄, C Ā B̄, C Ā² B̄, …, C Ā^{L-1} B̄) → y = K̄ * uSo the whole layer is a single causal convolution with a kernel as long as the sequence. Once K̄ is known, evaluating y is one FFT convolution: O(L log L), fully parallel. That is the SSM duality in a line — train as a convolution, generate as a recurrence, identical parameters, identical outputs. The only hard part is producing K̄ cheaply, since writing out L powers of Ā naively costs O(N²L).
A worked one-dimensional example
Take N = 1 so everything is a scalar: A = −0.5, B = 1, C = 1, Δ = 0.5. ZOH gives:
Ā = exp(ΔA) = exp(−0.25) ≈ 0.7788
B̄ = (exp(−0.25) − 1) / (−0.5) ≈ 0.4424
K̄_k = C Ā^k B̄ = 0.4424 × 0.7788^k
K̄ = [0.4424, 0.3446, 0.2684, 0.2090, 0.1628, …]Check it against the recurrence with an impulse u = [1, 0, 0, …]: x_0 = 0.4424, x_1 = 0.7788 × 0.4424 = 0.3446, x_2 = 0.2684 — the same numbers, as promised. The kernel is a decaying exponential whose memory horizon is readable from the pole: the response falls to 1/e after k = 1/0.25 = 4 steps. To remember a thousand steps you would need Ā ≈ 0.999. That observation is the entire problem with SSMs, and the next section.
Why a randomly initialized A forgets
The kernel is built from powers of Ā, so the eigenvalues of Ā decide everything. If any |λ| > 1 the state explodes; if |λ| < 1 the contribution of a token decays geometrically as |λ|^k. A random A scatters eigenvalues well inside the unit disc, so effective memory is a handful of steps — and the gradient flowing back k steps is scaled by that same |λ|^k, the classic vanishing-gradient failure of RNNs restated in spectral terms.
Nor can you fix it by pushing every eigenvalue to 1 − ε: a state whose modes all decay at nearly the same slow rate is a blurry running average, not a memory. What you want is a matrix whose modes span many timescales at once. That is a design problem, not a training problem.
HiPPO: choosing A so the state is an optimal summary
HiPPO (High-order Polynomial Projection Operators) answers it directly. Given N numbers, what is the best online summary of everything seen so far? Project the input history onto the first N Legendre polynomials and keep the coefficients. Differentiating that projection yields an ODE of exactly the SSM form, with a specific matrix — HiPPO-LegS:
A_nk = −√((2n+1)(2k+1)) if n > k
A_nk = −(n + 1) if n = k
A_nk = 0 if n < kSo A is lower-triangular with a negative, scale-graded diagonal: each state coordinate is tied to a different decay rate, giving a spectrum of timescales in one matrix. Initializing with HiPPO instead of randomly is the single change that took a linear RNN from unusable to state-of-the-art on Long Range Arena. The architecture was never the bottleneck; the initialization was.
S4: making the kernel actually computable
HiPPO’s matrix is dense, and computing K̄ = (C Ā^k B̄)_{k<L} from a dense Ā costs O(N²L) — far too slow. S4’s contribution is structural: write A as diagonal plus low-rank (DPLR), A = Λ − PQ*, a form the HiPPO matrix admits. You then never build the kernel step by step; instead you evaluate its truncated generating function K̂(z) = Σ_k K̄_k z^k at the L roots of unity, where the DPLR form collapses the sum into a Cauchy matrix–vector product, and take one inverse FFT to recover K̄. Total cost drops to Õ(N + L).
S4D then showed most of that machinery is optional: a purely diagonal A with a HiPPO-inspired initialization matches S4 on most tasks, makes Ā^k an elementwise power, and is what modern implementations — Mamba included — actually use.
Cost accounting against attention
The comparison is cleanest as a table, with L the sequence length, d the model width and N the state dimension (N << L, typically 16):
| Self-attention | SSM | |
|---|---|---|
| Training compute | O(L²d) | O(Ld log L) — FFT conv |
| Sequential depth | O(1) | O(1) in conv form |
| Per-token generation | O(Ld) | O(Nd) |
| State carried | KV cache, grows as O(Ld) | fixed O(Nd) |
The last row is the one that matters. A 24-layer, d = 1024 Transformer in fp16 stores 2 × 24 × 1024 × 2 = 96 KB of KV per token — 780 MB at 8k tokens, still growing. An SSM of comparable width carries a few megabytes, constant forever. Attention buys exact random-access recall; the SSM buys a bounded, lossy summary.
Selectivity: the one thing LTI cannot do
Time-invariance is what gives the convolution, and also the model’s ceiling. Because Ā, B̄, C are the same at every position, the kernel is fixed: the layer applies an identical filter regardless of what the token says, so it cannot remember this token and discard that one. On content-based tasks like selective copying or induction heads — things attention does trivially — a vanilla LTI SSM fails.
Mamba’s fix is to make Δ, B and C functions of the input (B_t = Linear(u_t), and so on). The transition now depends on content — a large Δ_t resets the state, a small one preserves it — which is exactly the gating an RNN needs. The cost is that the system is no longer time-invariant, so there is no single convolution kernel; training uses a hardware-aware parallel scan instead (O(L) work, O(log L) depth).
Why this matters on CPU and for small models
On a CPU the binding constraint is memory bandwidth, not FLOPs. An attention layer must stream the entire KV cache from RAM for every token generated, so per-token cost grows with context. An SSM step touches a fixed [d, N] state that fits in L2 cache, so token 8000 costs exactly what token 8 cost.
That flat profile is the argument for SSMs in small on-device models: predictable latency, a fixed memory budget you can size in advance, no cache eviction policy to design. With a diagonal Ā the recurrent step is a few elementwise multiply-adds per channel — trivially vectorizable. The caveat is numerical: repeated multiplication by Ā compounds low-precision rounding, so the state is usually kept at higher precision than the weights.
Pitfalls
The errors that actually bite. Skipping discretization: learning Ā directly, without Δ and the exponential map, throws away the timescale parameterization and trains worse. Unconstrained Δ: it must stay positive — parameterize it as exp(·) or softplus(·). Ignoring stability: the real parts of A’s eigenvalues must be negative, so enforce it explicitly (for diagonal A, store −exp(a_real)) or the state diverges.
Forgetting causal padding: the FFT convolution is circular, so pad kernel and input to 2L and discard the wrapped tail, or position 0 silently sees the end of the sequence. And the two paths must share identical (Ā, B̄, C, Δ) — a discretization mismatch yields a model that trains well and generates nonsense.
x′ = Ax + Bu, y = Cx — discretized by a learned step size Δ into Ā = exp(ΔA). Its power comes from a duality: the same parameters run as a recurrence with fixed O(Nd) state for generation, or as a single length-L convolution with kernel K̄ = (C Ā^k B̄) for O(L log L) parallel training. What made the idea work was not the architecture but the initialization: a random A forgets in a few steps, while HiPPO supplies a matrix whose modes span many timescales, and S4/S4D make its kernel computable. The trade against attention is exact random-access recall versus a bounded lossy summary — which is why per-token cost stays flat at any context length, and why Mamba had to sacrifice time-invariance to buy back content-based selection.