What Is a State Space Model?

The idea is older than deep learning

State Space Models come from classical control theory — Kalman, Bellman, the whole 1960s LQR/LQG lineage. The core idea predates neural networks by half a century: at every timestep, compress the entire history into a fixed-size hidden state vector, and update that state based on the current input.

The continuous-time form is two coupled differential equations: ḣ(t) = A·h(t) + B·x(t) for the state evolution, y(t) = C·h(t) for the output. Discretize and you get the recurrence h_t = Ā·h_{t-1} + B̄·x_t, y_t = C·h_t. That's it. A is the state transition matrix, B is the input projection, C is the output projection. The state h has a fixed dimension — call it N — that does not grow with sequence length.

Why this is interesting for sequence modeling

The fixed state dimension is the headline feature. At inference time, you maintain one h vector of size N regardless of whether you've processed 100 tokens or 100 million. That gives you O(1) memory per step. Compute is O(n) over the sequence because each step does one matrix-vector multiply against A and B. Compare that to attention's O(n²) compute and O(n) KV-cache and the appeal is obvious — if you can match Transformer quality.

The catch is in that "if". Compressing arbitrarily long history into a fixed-size vector is, by definition, lossy. The mathematical question SSM research has been answering for the last five years is: which choices of A, B, C, and which discretization scheme, lose the least useful information? S4, S4D, H3, Mamba, Mamba-2, Mamba-3 are all answers to that question — each one less lossy than the last for the workloads that matter.

Code

A naive SSM forward pass — the math, in 8 lines of PyTorch·python

import torch

def ssm_naive(x, A, B, C):
    # x: (batch, seq_len, d_model)
    # A, B, C: (d_state, d_state), (d_state, d_model), (d_model, d_state)
    B_, T, _ = x.shape
    h = torch.zeros(B_, A.size(0), device=x.device)
    ys = []
    for t in range(T):
        h = h @ A.T + x[:, t] @ B.T   # state update
        ys.append(h @ C.T)             # output projection
    return torch.stack(ys, dim=1)

Exercise

Implement the naive SSM forward pass above and run it on a random (1, 1024, 16) input with d_state=64. Verify it doesn't blow up (state matrix A initialized to 0.99 * eye works as a stable starting point). Then time it. You'll find this naive version is hopelessly slow on a GPU because the for-loop is sequential. The next lesson will explain how the SSM community recovered training-time parallelism without giving up the recurrent inference form.

The idea is older than deep learning

Why this is interesting for sequence modeling

Code

External links

Exercise

Progress

Comments 0