State Space Models Math

A sequence as a dynamical system

Attention treats a sequence as a set and lets every token look at every other token — powerful, but O(N^2) in length. State space models take the opposite stance, inherited from signal processing: a sequence is a signal in time, and you summarize everything seen so far in a small hidden state that you update as each new input arrives. Nothing attends to anything; information flows forward through the state.

Concretely, an SSM maps a 1-D input signal u(t) to a 1-D output y(t) through an N-dimensional latent state x(t). A deep SSM layer just runs many such single-input single-output systems in parallel — one per feature channel — and mixes channels between layers, exactly as a Transformer stacks attention and MLP blocks. The whole game is choosing the four small matrices that govern one channel’s system, and then computing it fast. Because the state is a fixed size no matter how long the sequence, the memory cost of generation does not grow with context — the structural advantage that motivates the entire family.

The continuous linear system

The foundational object is a linear, time-invariant (LTI) ordinary differential equation. For state x(t) ∈ R^N, scalar input u(t), and scalar output y(t):

x′(t) = A x(t) + B u(t)     (state equation)
y(t)  = C x(t) + D u(t)     (output equation)

with shapes A: [N,N], B: [N,1], C: [1,N], D: [1,1]. Read it physically: A is the state-transition matrix — how the state evolves on its own, its eigenvalues deciding what decays, oscillates, or persists; B injects the current input into the state; C reads the state out into the output; and D is a direct input→output skip term (often dropped, or folded into a residual connection, so D = 0). Everything ‘remembered’ lives in x(t), and the linearity of the map is exactly what will later let us collapse the whole system into a convolution. This is one channel; a real layer learns many.

Discretization: from calculus to a recurrence

Neural sequences are discrete tokens, not a continuous signal, so we sample the ODE at a step size Δ (the time between tokens, itself a learnable parameter). Discretization converts the continuous (A, B) into discrete matrices A_d, B_d that advance the state exactly one step. The standard choice, zero-order hold (ZOH), assumes the input is constant across each interval and integrates the ODE in closed form:

A_d = exp(Δ A)
B_d = (Δ A)^-1 (exp(Δ A) − I) · Δ B
    = A^-1 (exp(Δ A) − I) B

An alternative is the bilinear (Tustin) transform, a rational approximation to the matrix exponential: A_d = (I − Δ/2 A)^-1 (I + Δ/2 A) and B_d = (I − Δ/2 A)^-1 Δ B. Both share the same intuition — a small Δ means fine time resolution and a longer effective memory, a large Δ coarsens it. C and D pass through unchanged. Crucially, because the system is time-invariant, A_d, B_d are computed once and reused at every step.

The recurrent view: O(1)-state streaming

With the discrete matrices in hand, the continuous ODE becomes an ordinary linear recurrence over the token index k = 0, 1, 2, …:

x_k = A_d x_{k-1} + B_d u_k
y_k = C x_k  (+ D u_k)

This is how you run an SSM at inference time, and it is the source of its efficiency pitch. Each step multiplies the fixed N×N matrix by the state and adds the new input: O(N) work per token (or O(N^2) for a dense A, but S4 makes A diagonal, so it is O(N)), and a state of just N numbers that never grows. Contrast the Transformer’s KV cache, which grows with every token generated: attention’s decode memory is O(N) and rising, while an SSM’s is O(N_state) and constant. For long generations, on CPUs and small devices especially, that bounded, cache-free state is the whole appeal — the model streams tokens at a steady cost per step no matter how far into the sequence it is.

The convolutional view: one big kernel

The recurrence is inherently sequential — step k needs step k−1 — which is fine for generation but painfully slow to train, since you cannot parallelize across a length-L sequence. Here the linearity pays off. Unroll the recurrence from x_{-1} = 0:

x_k = Σ_{j=0}^{k} A_d^{k-j} B_d u_j
y_k = Σ_{j=0}^{k} (C A_d^{k-j} B_d) u_j

The bracketed scalars depend only on the gap k−j, not on k itself — the hallmark of an LTI system. So the output is a convolution of the input with a single fixed kernel:

K_d = (C B_d, C A_d B_d, C A_d^2 B_d, …, C A_d^{L-1} B_d)
y   = K_d ∗ u

Once K_d is known, computing y for the whole sequence is one convolution, done in O(L log L) with an FFT — fully parallel across positions. Training uses this mode; inference uses the recurrence.

The duality: two faces of one model

The central insight of structured SSMs is that these two computations are the same function with the same parameters, chosen per phase:

You train the model as a giant parallel convolution, then flip a switch and run it as a constant-memory recurrence for generation — no re-derivation, just two algebraically equivalent ways to evaluate one linear system. This is the property attention cannot match: attention parallelizes beautifully in training but has no cheap constant-state recurrent form for decode. The one hard requirement for the convolutional view to exist is time-invariance — the kernel entries must depend only on the gap k−j. The moment A_d, B_d, C vary with the input, the tidy fixed kernel disappears, which is precisely the design boundary between S4 and Mamba.

HiPPO: initializing A for long-range memory

A linear recurrence can remember the distant past only if A is chosen well. Initialize A randomly and the state either explodes or, far more often, forgets: eigenvalues inside the unit circle make old inputs decay geometrically, so information from thousands of steps ago vanishes long before it matters. Early SSM experiments were mediocre for exactly this reason.

The fix that unlocked the family is HiPPO (High-order Polynomial Projection Operators). HiPPO derives a specific structured A (and B) so that the hidden state maintains an optimal compression of the entire input history — concretely, the coefficients of the best polynomial (Legendre) approximation to everything seen so far. Rather than a random matrix, A becomes a principled lower-triangular operator whose entries encode how to update that running polynomial summary as each new sample arrives. Initializing the SSM from the HiPPO matrix is what let these models capture dependencies over tens of thousands of steps and win long-range benchmarks. The exact entries depend on the measure and sign convention, but the load-bearing idea is simple: A is designed for memory, not learned from scratch.

S4: making the structured kernel computable

HiPPO fixes memory but reintroduces a cost problem. Building the convolution kernel K_d naively means forming powers A_d^0, A_d^1, …, A_d^{L-1} of a dense N×N matrix — roughly O(N^2 L) work and memory, which is prohibitive for long sequences. S4 (Structured State Spaces, Gu et al. 2021) is the breakthrough that makes the structured kernel cheap.

S4’s move is to parameterize A as diagonal-plus-low-rank (DPLR). Under that structure, the kernel can be computed through its generating function — a Cauchy-kernel / rational evaluation — in near-linear O(N + L) time instead of O(N^2 L), sidestepping the repeated matrix powers entirely. A later simplification, S4D, showed a purely diagonal A (a diagonal approximation of HiPPO) works nearly as well and is far simpler to implement, which is why most modern SSMs use diagonal state matrices. S4 is thus the foundational structured SSM: HiPPO for memory, DPLR/diagonal structure for speed, and the recurrent/convolutional duality for deployment.

Complexity and a worked intuition

Put numbers to it. Take a state size N = 64 and a sequence of length L = 16{,}000 for one channel. Attention over this sequence forms an L×L score matrix: ~2.5×10^8 interactions — the quadratic wall. The SSM never builds anything of size L^2.

Training (convolution): build the length-L kernel in O(N + L) with S4 structure, then one FFT convolution in O(L log L) — here log2(16000) ≈ 14, so on the order of L × 14 operations, dramatically below L^2.

Inference (recurrence): per token, one diagonal state update of size N = 64 plus a readout — a few hundred flops and a 64-number state that never grows. Generating token 16,000 costs exactly what generating token 10 cost, whereas a Transformer at that point is streaming a KV cache of 16,000 entries per layer per head. That flat, bounded per-step cost is why the architecture is attractive for long-context, CPU-bound inference.

The LTI ceiling — and where Mamba begins

Everything above rests on one assumption: the system is time-invariant. A_d, B_d, C and Δ are fixed — the same at every position, independent of what the input actually says. That is exactly what buys the convolutional view and the FFT training speedup, but it is also a real ceiling. An LTI system applies the identical dynamics to every token, so it cannot selectively decide, based on content, to remember this word and ignore that one. It compresses history by a fixed rule, not by relevance.

This is the precise gap that selective SSMs (Mamba) close, and the boundary between this article and its companion. Mamba makes B, C, and Δ functions of the input — input-dependent, and therefore time-varying. That restores content-based gating (the state can focus or forget on demand) but breaks time-invariance, so the convolution kernel no longer exists; Mamba recovers parallelism instead with a hardware-aware parallel scan. The S4-style LTI SSM in this article is the foundation; selection is the modification that made SSMs competitive with attention on language.