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.