Structured State Space Models

Structured State Space Models (SSMs) are sequence models built from a linear recurrence with a learned state. Models like S4 and Mamba use them as an alternative to attention that scales linearly with sequence length while still capturing long range dependencies.

The State Space Model

The state space model is defined by a differential equation where $u(t)$ is a 1D input signal that is mapped to an N-dimensional latent state $x(t)$ before being projected to a 1D output signal $y(t)$: $x^{\prime }(t)=Ax(t)+Bu(t)$ $y(t)=Cx(t)+Du(t)$ A, B, C, D are matrices learned via gradient descent.

Discretization and the Recurrent View

To run on sampled sequences we discretize with step size $\Delta$, giving matrices $\bar{A},\bar{B}$ (for example via the zero order hold rule $\bar{A}=\exp (\Delta A)$, $\bar{B}=(\Delta A{)}^{-1}(\exp (\Delta A)-I)\Delta B$). The model becomes a linear recurrence:

$$x_k=\bar{A}{x}_{k-1}+\bar{B}{u}_k,y_k=C{x}_k$$

This runs in $O(L)$ time and $O(1)$ memory per step, ideal for autoregressive inference, exactly like an RNN but linear.

The Convolutional View

Because the recurrence is linear and time invariant, unrolling it gives a convolution with a fixed kernel $\bar{K}$:

$$y=\bar{K}\ast u,\bar{K}=(C\bar{B},C\bar{A}\bar{B},C{\bar{A}}^2\bar{B},\dots ,C{\bar{A}}^{L-1}\bar{B})$$

During training the whole sequence is known, so we compute $\bar{K}$ once and convolve (via FFT) in $O(L\log L)$, fully parallel over the sequence. The same weights switch to the recurrent form at inference. This duality (parallel conv for training, cheap recurrence for generation) is the core appeal of SSMs.

HiPPO and S4 Intuition

A randomly initialized $A$ forgets the past quickly. HiPPO theory derives a specific structured $A$ whose state optimally compresses the input history onto a basis of orthogonal polynomials, so the state is a running memory of everything seen so far. S4 makes this practical by parameterizing $A$ as a diagonal plus low rank matrix, which allows fast kernel computation and gives strong results on long range benchmarks.

Mamba: Selective SSM

Selectivity

Classic SSMs are time invariant: $\bar{A},\bar{B},C$ are the same at every step, so the model cannot choose what to remember based on content. Mamba makes $B$, $C$, and $\Delta$ functions of the input $u_k$. This input dependent gating lets the model focus on or ignore tokens, closing much of the quality gap with attention.

The price is that a content dependent recurrence is no longer a fixed convolution, so Mamba uses a hardware aware parallel scan instead of FFT convolution.

Scan Code Sketch

The recurrence is an associative scan over the sequence. A simple sequential reference (a real kernel uses a parallel scan):

import torch
 
def ssm_scan(A_bar, B_bar, C, u):
    # A_bar: (L, N, N)   B_bar: (L, N)   C: (L, N)   u: (L,)
    L, N = C.shape
    x = torch.zeros(N)
    ys = []
    for k in range(L):
        x = A_bar[k] @ x + B_bar[k] * u[k]   # x_k = A x_{k-1} + B u_k
        ys.append(torch.dot(C[k], x))        # y_k = C x_k
    return torch.stack(ys)

For a diagonal $A$ the state update is elementwise and the whole scan parallelizes across $L$ with a Blelloch style prefix scan, which is how Mamba stays fast on GPUs.

SSMs vs Attention

Aspect	SSM (Mamba)	Transformers attention
Sequence scaling	$O(L)$	$O(L^2)$
Inference state	fixed size $x$	KV cache grows with $L$
Content routing	via selectivity	native (every pair)
Recall of exact tokens	weaker	strong

Trade-off

SSMs compress history into a fixed size state, so they excel at long, smooth signals but can struggle with tasks needing exact copy or precise lookup of distant tokens, where attention’s explicit pairwise access wins. Hybrid architectures interleave SSM and attention layers to get both.