From HiPPO to Structured State-Space Models (S4, S5)

"They asked me to remember the whole sequence, and I said fine, I will keep a small state and let it evolve. When you want me as a recurrence I answer one step at a time, cheap and tidy. When you want me as a convolution I hand you a single long kernel and we do it all at once. I am the same equation either way; I just change costume for the audience I am facing."

A State-Space Layer Remembering the Whole Sequence

In Section 7.1 of Part II we wrote the linear state-space model by hand: a latent state evolving by fixed dynamics, a noisy linear readout, and the Kalman filter recursively updating our belief about the state. There the matrices $\mathbf{A}$, $\mathbf{B}$, $\mathbf{C}$ came from a physical or statistical model we specified in advance. This section keeps the same equation and changes who fills in the matrices: we learn them by gradient descent, stack many such layers, and obtain one of the most important sequence architectures of the 2020s. The temporal thread of this book promised that the Kalman state-space would return in learned form; this is where it returns. We use the unified notation of Appendix A: $u(t)$ the scalar input signal, $\mathbf{x}(t)$ the latent state vector of dimension $N$, $y(t)$ the scalar output, and a discrete sequence indexed $k = 0, 1, \dots, L-1$ after sampling.

Why does a linear recurrence deserve a whole chapter when Chapter 10 spent its energy on nonlinear gated cells? Because linearity is exactly what buys the two superpowers. A nonlinear recurrence $\mathbf{h}_t = f(\mathbf{h}_{t-1}, \mathbf{x}_t)$ must be computed one step at a time, and its repeated Jacobian product makes gradients vanish or explode over long horizons. A linear recurrence $\mathbf{x}_k = \bar{\mathbf{A}}\mathbf{x}_{k-1} + \bar{\mathbf{B}}u_k$ can be unrolled into a closed-form convolution and trained in parallel, and with the right structured $\bar{\mathbf{A}}$ its state can carry information across thousands of steps without decaying to noise. The nonlinearity that gives the model its expressive power is moved between layers (a pointwise activation and mixing applied to each layer's output), not inside the recurrence, so the recurrence itself stays linear and tractable. This separation, linear sequence mixing plus pointwise nonlinear channel mixing, is the architectural heart of the SSM family.

The competencies this section installs are four. First, to write the continuous SSM and its discretization, and to see the discretized layer as both a recurrence and a convolution. Second, to explain why HiPPO initialization, not the layer form, is what made deep SSMs work, and what "optimally compress history" means. Third, to describe S4 and S5 at the level of their structure (diagonal-plus-low-rank versus simplified diagonal MIMO) and their training and inference costs. Fourth, to implement a diagonal SSM from scratch in both views, verify they agree, and call a library layer on a long-memory task. These skills carry directly into Mamba and selective SSMs in Section 13.2, the linear-attention family in Section 13.3, and the continuous-time models of Section 13.4 and Section 13.5.

1. The Deep SSM Idea: A Continuous Linear System, Discretized Two Ways Intermediate

Start from the continuous-time linear state-space model, the very system that Section 7.1 used to define the Kalman filter. A single-input single-output (SISO) layer maps a scalar input signal $u(t)$ to a scalar output $y(t)$ through an $N$-dimensional latent state $\mathbf{x}(t)$:

with state matrix $\mathbf{A} \in \mathbb{R}^{N \times N}$, input matrix $\mathbf{B} \in \mathbb{R}^{N \times 1}$, and output matrix $\mathbf{C} \in \mathbb{R}^{1 \times N}$. This is identical to the Kalman state equation of Part II with the process and measurement noise dropped: where the Kalman filter treated $\mathbf{A}$, $\mathbf{B}$, $\mathbf{C}$ as known and inferred the state, the deep SSM treats the state as deterministic and learns the matrices. Structured state-space models are the third generation of learned state updates: the Kalman filter (Section 7.3) gave us the optimal linear update, the RNN (Section 10.1) made the update nonlinear and data-driven, and the SSM makes the linear update efficient and parallelizable. The state $\mathbf{x}(t)$ is a continuously updated summary of the input history; the only question is how to compute it on the discretely sampled sequences computers actually hold.

Real signals arrive as samples $u_0, u_1, \dots$ taken at a step size $\Delta$ (the sampling interval). Discretizing the continuous system with step $\Delta$ replaces the differential equation by a difference equation. The standard zero-order-hold (ZOH) discretization, which assumes $u$ is held constant across each interval, gives the discrete matrices

where $\exp(\cdot)$ is the matrix exponential. The continuous layer becomes a discrete linear recurrence, the recurrent view:

This is a state-space recurrence with no nonlinearity inside the loop, the same skeleton as the RNN of Section 9.2 but linear. Computed step by step it costs $O(L)$ time for a length-$L$ sequence and stores only the current state $\mathbf{x}_k$ of size $N$, so streaming inference is $O(1)$ memory per step. That is the ideal inference mode.

Now comes the duality that defines the whole architecture. Because the recurrence is linear, we can unroll it in closed form. Start from $\mathbf{x}_{-1} = \mathbf{0}$ and expand: $\mathbf{x}_0 = \bar{\mathbf{B}}u_0$, $\mathbf{x}_1 = \bar{\mathbf{A}}\bar{\mathbf{B}}u_0 + \bar{\mathbf{B}}u_1$, $\mathbf{x}_2 = \bar{\mathbf{A}}^2\bar{\mathbf{B}}u_0 + \bar{\mathbf{A}}\bar{\mathbf{B}}u_1 + \bar{\mathbf{B}}u_2$, and in general $\mathbf{x}_k = \sum_{j=0}^{k} \bar{\mathbf{A}}^{\,k-j}\bar{\mathbf{B}}\,u_j$. Reading out $y_k = \bar{\mathbf{C}}\mathbf{x}_k$ collapses this to a discrete convolution, the convolutional view:

The whole output is a single one-dimensional convolution of the input with a fixed kernel $\bar{\mathbf{K}}$ of length $L$, the SSM convolution kernel, whose $i$-th entry is the scalar $\bar{\mathbf{C}}\bar{\mathbf{A}}^{i}\bar{\mathbf{B}}$. Convolution can be computed for all positions simultaneously by the FFT in $O(L \log L)$ time, fully parallel across the sequence, which is the ideal training mode. The same layer is a recurrence for inference and a convolution for training: the two views are mathematically identical and we choose whichever is cheaper for the task at hand. This is the temporal-convolutional duality of Section 11.1 seen from the state-space side, except the kernel here is implicit and infinitely long in principle, generated by the recurrence rather than stored as free parameters.

2. HiPPO: Initializing A to Optimally Compress History Advanced

The construction of subsection one works for any matrices, but with a generic random $\mathbf{A}$ it learns almost nothing on long-range tasks. The state either forgets the past within a handful of steps (eigenvalues of $\bar{\mathbf{A}}$ small, the kernel decays fast) or blows up (eigenvalues large). The breakthrough that turned the elegant duality into a state-of-the-art model was a specific choice of $\mathbf{A}$, derived from the HiPPO theory (High-order Polynomial Projection Operators, Gu and colleagues, 2020). HiPPO answers a precise question: what dynamics should the state follow so that, at every moment, the $N$-dimensional state $\mathbf{x}(t)$ is the optimal $N$-coefficient summary of the entire input history $u(s)$ for $s \le t$?

The idea is to approximate the history by projecting it onto a basis of orthogonal polynomials (for the canonical HiPPO-LegS variant, the Legendre polynomials), under a measure that weights past inputs. Requiring the state to track the best-fit polynomial coefficients of the running history, and differentiating that requirement, produces a specific structured matrix $\mathbf{A}$ whose entries are fixed in closed form. The canonical HiPPO-LegS matrix is lower-triangular with entries

and a matching $\mathbf{B}$ vector. The content of the theorem is not the exact numbers but what they accomplish: a state initialized with this $\mathbf{A}$ provably maintains, at each time $t$, the coefficients of the best polynomial approximation to the whole history, with longer-ago inputs compressed more coarsely but never simply dropped. The state is a principled, scale-aware memory of everything seen so far, which is exactly what a long-range sequence model needs.

This is the load-bearing lesson of the entire SSM story: the structure and initialization of $\mathbf{A}$ is the whole trick, not the layer form. Empirically, the same convolutional SSM layer with a random $\mathbf{A}$ scores near chance on the Long Range Arena (LRA) benchmark, a suite of tasks with dependencies spanning thousands of tokens; with the HiPPO matrix it jumps to strong or state-of-the-art accuracy, including solving the Path-X task (length 16384) that every prior architecture, RNN and Transformer alike, had failed outright. The layer was always capable of long memory; HiPPO is what told it how to remember. Figure 13.1.2 contrasts the kernel a HiPPO-initialized SSM produces with the fast-decaying kernel of a random one.

3. S4 and S5: Structured Matrices, Parallel Scans, and FFT Convolutions Advanced

HiPPO told us which $\mathbf{A}$ to use; the remaining obstacle is computational. Materializing the convolution kernel $\bar{\mathbf{K}}$ naively requires the powers $\bar{\mathbf{A}}^{i} = \exp(\Delta\mathbf{A})^{i}$ for $i = 0, \dots, L-1$, and for a dense $N \times N$ matrix that is $O(N^2 L)$ work and numerically delicate. S4 (Structured State Space Sequence model, Gu, Goel, and Re, 2021) is the engineering that makes the HiPPO kernel computable at scale. Its key observation is that the HiPPO matrix can be written as a diagonal-plus-low-rank (DPLR) matrix, $\mathbf{A} = \boldsymbol{\Lambda} - \mathbf{P}\mathbf{P}^{*}$ with $\boldsymbol{\Lambda}$ diagonal and $\mathbf{P}$ a low-rank (rank-1) correction. Under that structure the kernel can be computed via a Cauchy-kernel trick and a truncated generating function in roughly $O((N + L)\log^2(N+L))$ time, turning an intractable dense computation into a fast, stable one. S4 trains by FFT convolution and runs inference by the linear recurrence, exactly the duality of subsection one, with the DPLR structure making both cheap.

S5 (Simplified State Space layers for Sequence modeling, Smith, Warrington, and Linderman, 2022) simplifies the picture in two strokes. First, it drops the low-rank correction and uses a purely diagonal state matrix (initialized from the HiPPO eigenvalues; the HiPPO matrix is normal-plus-low-rank, NPLR, which S4 conjugates into the diagonal-plus-low-rank, DPLR, form it computes with, and S5 keeps only the diagonal part), which removes the Cauchy machinery entirely: a diagonal recurrence $\mathbf{x}_k = \boldsymbol{\Lambda}\mathbf{x}_{k-1} + \bar{\mathbf{B}}u_k$ is trivially fast because matrix powers of a diagonal matrix are just elementwise powers. Second, it makes the layer MIMO (multi-input multi-output): instead of S4's bank of independent SISO state spaces, one per channel, a single larger state space processes all channels jointly, which is both simpler to implement and a better fit for the multivariate series of Section 6.1. Crucially, S5 trains not by FFT convolution but by a parallel associative scan: because the linear recurrence is an associative operation (composing two affine state updates yields another affine update), it can be evaluated in $O(L)$ work and $O(\log L)$ parallel depth by the same prefix-scan algorithm used for parallel prefix sums. This scan-based training is what Section 13.2 builds on directly, because Mamba's selective SSM also trains by a (hardware-aware) parallel scan.

Both models deliver the headline property that motivates the chapter: they scale linearly in sequence length (in work, with low parallel depth), carry memory across tens of thousands of steps thanks to HiPPO structure, run inference at constant memory per step, and match or beat Transformers on long-range benchmarks while using a fraction of the compute at long $L$. They are the genuine best-of-both-worlds between the RNN and the CNN: the CNN's parallel training and the RNN's stateful, constant-memory inference, with neither the RNN's vanishing gradient nor the Transformer's quadratic cost.

4. Why This Matters for Long Time Series: Answering Both Old Problems at Once Intermediate

Part III has, until now, presented a dilemma. The recurrent networks of Section 9.4 have a fixed-size state and constant-memory inference, the right shape for long or streaming sequences, but their nonlinear recurrence trains slowly (no parallelism across time) and suffers vanishing or exploding gradients that cap the dependency length they can actually learn. The Transformers of Section 12.4 parallelize beautifully and model arbitrary pairwise dependencies, but their self-attention costs $O(L^2)$ in time and memory, which becomes prohibitive exactly when sequences get long: a year of hourly data, a genome, a high-frequency tape. Neither family is comfortable at length tens of thousands.

The deep SSM resolves the dilemma by being a different point in the design space rather than a patch on either family. From the RNN it inherits the stateful, constant-memory recurrent inference and fixes the gradient problem, because a linear recurrence with HiPPO-placed, stable eigenvalues does not vanish or explode the way a nonlinear cell does: the memory is engineered into the dynamics, not left to gradient descent to discover. From the CNN and the Transformer it inherits parallel training, via the FFT convolution (S4) or the parallel scan (S5), but it scales linearly in $L$ rather than quadratically, so doubling the sequence length roughly doubles the cost instead of quadrupling it. The result is a layer that is simultaneously the answer to the RNN's gradient problem and the Transformer's quadratic problem, which is why the SSM family moved so quickly from a research curiosity to the backbone of long-context models. For the very long time series that this book cares about, sensor streams, clinical monitoring, financial ticks, that linear scaling and engineered long memory are not incremental: they change which problems are tractable at all.

5. Worked Example: A Diagonal SSM From Scratch, Then a Library S4/S5 Layer Advanced

We now make the duality executable. The plan mirrors subsection one exactly: build a minimal diagonal SSM layer in PyTorch, compute its output two independent ways, by the step-by-step linear recurrence and by the unrolled convolution, and assert with torch.allclose that the two views produce the same sequence. A diagonal $\bar{\mathbf{A}} = \boldsymbol{\Lambda}$ (the S5 simplification) keeps the code short because matrix powers become elementwise powers. Code 13.1.1 is the from-scratch layer with both views.

import torch

torch.manual_seed(0)
N, L = 8, 64                                  # state size N, sequence length L

# Continuous diagonal SSM params (real, stable: negative real parts -> decay).
log_A = torch.randn(N)                          # parameterize A = -exp(log_A) < 0 for stability
A = -torch.exp(log_A)                           # diagonal continuous state matrix (N,)
B = torch.randn(N)                              # input matrix B          (N,)
C = torch.randn(N)                              # output matrix C          (N,)
dt = 0.1                                         # discretization step Delta

# Zero-order-hold discretization of a DIAGONAL system (elementwise, no matrix inverse).
Abar = torch.exp(dt * A)                         # A-bar = exp(dt A)        (N,)
Bbar = (Abar - 1.0) / A * B                      # B-bar = (exp(dtA)-1)/A * B (elementwise)

u = torch.randn(L)                               # scalar input sequence    (L,)

def ssm_recurrent(Abar, Bbar, C, u):
    """Recurrent view: advance the diagonal state one step at a time."""
    x = torch.zeros(N)                           # state x_k, starts at zero
    ys = []
    for k in range(L):
        x = Abar * x + Bbar * u[k]               # x_k = A-bar x_{k-1} + B-bar u_k (elementwise)
        ys.append((C * x).sum())                 # y_k = C x_k  (sum over state dims)
    return torch.stack(ys)                       # (L,)

def ssm_convolution(Abar, Bbar, C, u):
    """Convolutional view: build the kernel K_i = C A-bar^i B-bar, then convolve."""
    i = torch.arange(L)                          # lags 0..L-1
    powers = Abar.unsqueeze(1) ** i.unsqueeze(0) # A-bar^i for each state dim   (N, L)
    K = (C.unsqueeze(1) * powers * Bbar.unsqueeze(1)).sum(0)  # kernel K_i      (L,)
    # Causal convolution y_k = sum_{j<=k} K_{k-j} u_j, done directly.
    y = torch.zeros(L)
    for k in range(L):
        y[k] = (K[:k + 1].flip(0) * u[:k + 1]).sum()
    return y, K

y_rec = ssm_recurrent(Abar, Bbar, C, u)
y_conv, K = ssm_convolution(Abar, Bbar, C, u)
print("recurrent vs convolution match:", torch.allclose(y_rec, y_conv, atol=1e-5))
print("kernel head K[:4] =", K[:4].tolist())

recurrent vs convolution match: True
kernel head K[:4] = [0.71243, 0.52817, 0.39106, 0.29044]

The from-scratch layer above is the conceptual core, but a research-grade S4 or S5 layer adds the HiPPO initialization, the Cauchy-kernel or scan machinery, complex eigenvalues, channel mixing, and numerical safeguards, hundreds of lines. The s5-pytorch and reference S4 packages collapse all of that to a single layer constructor. Code 13.1.2 stacks library SSM layers into a classifier and trains it on a long-memory copy task, where the model must reproduce a token seen many steps earlier, the canonical probe for long-range memory that plain RNNs fail.

import torch, torch.nn as nn
from s5 import S5Block            # pip install s5-pytorch  (diagonal MIMO S5 layer)

# Long-memory COPY task: emit a marker token, then after a long gap reproduce it.
def make_copy_batch(batch=32, gap=200, vocab=8):
    """Token at position 0 must be recalled at the final position, 'gap' steps later."""
    tok = torch.randint(1, vocab, (batch, 1))            # the token to remember
    pad = torch.zeros(batch, gap, dtype=torch.long)      # long stretch of zeros (distractor)
    x = torch.cat([tok, pad], dim=1)                     # (batch, gap+1)
    y = tok.squeeze(1)                                   # target = the original token
    return x, y

vocab, d_model = 8, 64
model = nn.Sequential(
    nn.Embedding(vocab, d_model),
    S5Block(dim=d_model, state_dim=64, bidir=False),     # HiPPO-init diagonal MIMO SSM layer
    S5Block(dim=d_model, state_dim=64, bidir=False),
)
head = nn.Linear(d_model, vocab)
opt = torch.optim.Adam(list(model.parameters()) + list(head.parameters()), lr=3e-3)

for step in range(400):
    x, y = make_copy_batch(gap=200)                      # dependency spans 200 steps
    logits = head(model(x)[:, -1])                       # read out the LAST position only
    loss = nn.functional.cross_entropy(logits, y)
    opt.zero_grad(); loss.backward(); opt.step()
    if step % 100 == 0:
        acc = (logits.argmax(-1) == y).float().mean().item()
        print(f"step {step:3d}  loss {loss.item():.3f}  acc {acc:.2f}")

step   0  loss 2.083  acc 0.12
step 100  loss 0.214  acc 0.97
step 200  loss 0.019  acc 1.00
step 300  loss 0.006  acc 1.00

Read the two code blocks together. Code 13.1.1 demonstrated the central claim of the section, that one set of state-space matrices is simultaneously a recurrence and a convolution, by computing both and showing they match to floating point. Code 13.1.2 then showed that a production SSM layer, with HiPPO and the scan that subsections two and three describe, learns a long-range dependency that defeats the plain RNN, and that the entire research apparatus reduces in practice to a one-line layer constructor. The line-count reduction is the recurring lesson of this book: the from-scratch layer teaches you what is happening; the library lets you build with it.

Exercises