A state-space model is a sequence layer that keeps a fixed-size summary of everything it has read, called its state, and updates that state with each new token. The update is linear in the state and the same at every step, so the same layer can run one token at a time like a recurrent network or over a whole sequence at once like a convolution. This section builds such a layer from a continuous-time system, the kind of equation engineers use for circuits and springs, and then chooses the system so that the state is a faithful record of the distant past. The result, called S4, was the first model to solve a task whose answer depends on inputs 16,000 steps apart, and the state-space layers inside today's hybrid language models descend from it.

Here is the task. The Long Range Arena benchmark (Tay et al. 2021) collects tasks built to break efficient sequence models. The hardest, Path-X, shows a model a 128 by 128 black-and-white image one pixel at a time, 16,384 steps in all, and asks whether two marked dots are joined by a dashed path. The Transformer and every faster variant of it that the benchmark tested stayed at 50%, the accuracy of a coin flip. The answer depends on pixels thousands of steps apart, and nothing carried the signal that far. A year later S4, the structured state-space sequence model, reached 88% (Gu, Goel & Ré 2021). The same model classified raw one-second recordings of ten spoken words, 16,000 numbers per clip with no hand-made audio features, at 98.3% accuracy.

Two earlier sections set up what such a layer must look like. A recurrence is linear in its state when each step multiplies the old state by a number and adds a weighted input, $h_t = ah_{t-1} + ux_t$, with a decay $a$ and an input weight $u$. §21.1 showed that such a recurrence unrolls into a lower-triangular matrix, so it trains in parallel, and that its memory is set by the decay. An input's weight halves every $t_{1/2} = \ln 2/(-\ln a)$ steps, its half-life. §21.2 showed what happens with no decay at all. A state that only adds is crowded out by its own past.

Neither section said how to choose the decays. A half-life of 16,000 steps needs a decay within about $4\times10^{-5}$ of 1. Noticing what happened three steps ago needs fast decays too. And a few fading sums of the inputs are a crude record of the past. We would rather have a state from which the whole history could be redrawn.

This section gets all of that from the continuous-time linear system of physics and control theory. We solve it exactly and turn it into a recurrence over tokens in three ways, two of which stay stable at any step size. We show that the same layer is a long convolution, which the fast Fourier transform computes in $O(T\log T)$ time. We derive a state matrix, called HiPPO, that makes the state the best polynomial summary of the past. We see why S4 needed a structured form of it, and why its successors got away with a plain diagonal matrix. The section ends with the one thing this whole family cannot do, which §21.4 fixes.

A system that integrates its input

Fill a leaky bucket from a tap. The water level rises with the inflow and drains in proportion to how full the bucket is. Write $h(t)$ for the level at time $t$, $x(t)$ for the inflow, and $\dot h(t)$ for the rate at which the level changes. Then $\dot h(t) = -h(t) + x(t)$. Start empty and pour at a constant rate of 1. The level is $h(t) = 1 - e^{-t}$. It is 0.63 after one unit of time and 0.95 after three, and it approaches 1, the level at which drain balances inflow.

The level at any moment is a weighted record of all past inflow, with recent inflow weighted most. A resistor charging a capacitor, a mass on a damped spring, and a thermometer obey the same kind of equation. Each one integrates its input and forgets at a rate set by its own constants. Each one is also linear. Twice the input gives twice the response, and the response to a sum of inputs is the sum of the responses.

Now give the system $N$ internal quantities that feed into each other, like $N$ connected tanks. The state becomes a vector, and the rates become a matrix.

In plain words, the state changes at a rate that is a fixed linear function of the state and the current input, and the output is a fixed linear readout of the state, plus a direct path from the input.

A few notation rules hold for the whole chapter. $\mathbf{A}$ is always a state matrix, never a matrix of attention weights. $N$ is the state size, never a parameter count. Bold $\mathbf{h}(t)$ with parentheses is the continuous state, and $\mathbf{h}_t$ with a subscript will be its value at token $t$. SSM is short for state-space model.

Solving it exactly

For a single number $a$, the equation $\dot h = ah$ is solved by $h(t) = e^{at}h(0)$. A square matrix has an exponential too, defined by the same power series, and it solves the matrix equation in the same way.

For a diagonal matrix the third property is all there is. With $\mathbf{A} = \operatorname{diag}(-1, -2)$, $e^{\mathbf{A}t} = \operatorname{diag}(e^{-t}, e^{-2t})$, two leaky buckets that drain at different rates.

In plain words, the state now is what is left of the old state, plus every past input pushed through the system's response $e^{\mathbf{A}(t-s)}\mathbf{B}$ for the time $t - s$ since it arrived. Check it on the bucket with constant inflow, where $A = -1$, $B = 1$, and $h(0) = 0$. The integral is $\int_0^t e^{-(t-s)}\,ds = 1 - e^{-t}$, the curve above.

In the basis of eigenvectors, each coordinate of the state is its own leaky bucket with rate $\lambda_n$. We call such a coordinate a mode. An eigenvalue does not have to be real. A real matrix can have complex eigenvalues, in conjugate pairs $\lambda = \sigma \pm \mathrm{i}\omega$. Then $e^{\lambda t} = e^{\sigma t}(\cos\omega t + \mathrm{i}\sin\omega t)$ is a rotation at frequency $\omega$ whose radius shrinks at rate $|\sigma|$. It is the complex rotation of rotary position embeddings (§5.5), now with a decay.

Real eigenvalues give a memory that fades smoothly. Complex ones give a memory that oscillates as it fades, which lets a few state numbers represent wiggly patterns in the past. This is the eigenvalue analysis of a linear RNN from §3.3, carried into continuous time. The powers $\lambda^k$ of a discrete step have become exponentials $e^{\lambda t}$, and the condition $|\lambda| < 1$ has become $\operatorname{Re}\lambda < 0$. In practice the memory of a state-space layer is set by the eigenvalues of $\mathbf{A}$, and an interviewer who asks how long such a layer remembers wants to hear about their real parts.

From continuous time to tokens

A sequence of tokens is not a function of time. To use the equation we pretend that it is a signal sampled every $\Delta$ units of time. Token $x_t$ is held constant over the interval $((t-1)\Delta, t\Delta]$ that ends at step $t$, and the state after step $t$ is $\mathbf{h}_t = \mathbf{h}(t\Delta)$. So input $t$ affects output $t$, the update-then-read convention of this chapter. The step size $\Delta > 0$ is a learned parameter, and we will see that it is the most important one.

Turning one step of the equation into an update rule is called discretizing it. There are three standard ways.

Zero-order hold (ZOH). Hold the input constant over the step and solve the equation exactly. Apply the solution formula from $t_0 = (t-1)\Delta$ to $t\Delta$ with $x(s) = x_t$,

where the second form substitutes $\tau = t\Delta - s$. By property 1, $\frac{d}{d\tau}\big(\mathbf{A}^{-1}e^{\mathbf{A}\tau}\big) = e^{\mathbf{A}\tau}$ when $\mathbf{A}$ is invertible, so the integral is $\mathbf{A}^{-1}(e^{\Delta\mathbf{A}} - \mathbf{I})$. Nothing is approximated. The ZOH update is exact for an input that really is held constant over each step.

Euler. Replace the derivative by a forward difference, $\mathbf{h}_t - \mathbf{h}_{t-1} \approx \Delta(\mathbf{A}\mathbf{h}_{t-1} + \mathbf{B}x_t)$. This is the first thing anyone would try, and we will see that it is the worst.

Bilinear, also called the trapezoid rule. Average the state term over the two ends of the step, $\mathbf{h}_t - \mathbf{h}_{t-1} \approx \frac{\Delta}{2}\mathbf{A}(\mathbf{h}_t + \mathbf{h}_{t-1}) + \Delta\mathbf{B}x_t$, and solve for $\mathbf{h}_t$. That gives $(\mathbf{I} - \frac{\Delta}{2}\mathbf{A})\mathbf{h}_t = (\mathbf{I} + \frac{\Delta}{2}\mathbf{A})\mathbf{h}_{t-1} + \Delta\mathbf{B}x_t$. This is the rule S4 used.

In plain words, the discrete decay of each mode depends only on $\Delta$ times its eigenvalue. A discretized SSM is the linear recurrence of §21.1, with a decay $\bar{\mathbf{A}}$ and an input weight $\bar{\mathbf{B}}$. They are no longer free numbers. They are functions of a continuous system and a step size.

The cell repeats the example for all three rules.

Which rules are stable at every step size

The step size is learned, so it moves during training. A recurrence run for 16,000 steps turns any discrete eigenvalue of size above 1 into overflow. So we need a stable continuous system to stay stable after discretization, whatever $\Delta$ is. In the complex plane, the stable continuous eigenvalues fill the left half-plane ($\operatorname{Re}\lambda < 0$). The stable discrete ones fill the unit disk, the complex numbers of size less than 1.

In plain words, ZOH and bilinear map the whole left half-plane into the unit disk, so you can train $\Delta$ freely and never leave the stable region. Euler scales the plane by $\Delta$ and then shifts it right by 1. Only the disk of radius 1 around $-1$ lands inside the unit disk, so a large step pushes fast modes out.

Stability is one property and accuracy is another. Compare each rule with the exact factor $e^z = 1 + z + z^2/2 + z^3/6 + \cdots$, where $z = \Delta\lambda$.

• Euler's $1 + z$ is already wrong at the $z^2$ term. Its error per step is $O(\Delta^2)$, and over the $t/\Delta$ steps needed to reach a fixed time $t$ the errors add up to $O(\Delta)$. A rule like this is called first order.
• Bilinear agrees with $e^z$ up to the $z^2$ term and is first wrong at $z^3$, so its error over a fixed time is $O(\Delta^2)$. It is second order.
• ZOH is exact for the state.

The next cell measures all three on a decaying rotation with no input, and then looks at a fast mode.

Halving $\Delta$ divides the bilinear error by 4, the signature of a second-order rule. It divides the Euler error by 3.4, then 2.5, then 2.2, approaching the factor 2 of a first-order rule. The ZOH error is at rounding level.

The fast mode shows something the orders do not. At $\Delta|\lambda| = 5$, ZOH decays it almost completely ($e^{-5} = 0.0067$). Euler explodes, with $\bar a = -4$. Bilinear keeps it stable but sends it to $\bar a = -0.43$, a mode that flips sign every step. As $\Delta|\lambda| \to \infty$ the bilinear $\bar a$ tends to $-1$, so fast modes ring instead of dying. That is one reason the later diagonal models and Mamba prefer ZOH.

The last column names the models of this chapter, each introduced where it first matters. S4, its predecessor LSSL, and the diagonal models DSS, S4D, and S5 appear later in this section. Mamba (§21.4) keeps the exact $\bar{\mathbf{A}} = e^{\Delta\mathbf{A}}$ but uses $\bar{\mathbf{B}} = \Delta\mathbf{B}$, the first term of $(e^{z} - 1)/z = 1 + z/2 + \cdots$. Mamba-3 (Lahoti et al. 2026) later named that rule exponential-Euler. It replaced it with the exponential-trapezoidal rule of the last row, which blends the current and previous inputs and which §21.6 describes.

In practice this theorem is why every modern SSM learns a continuous eigenvalue and a step size, and never the discrete decay directly. An interviewer will ask you to derive ZOH, and then to explain why Euler is unsafe at a large step.

The step size is a timescale

Under ZOH, a continuous mode with rate $\operatorname{Re}\lambda < 0$ becomes a decay of size $|\bar\lambda| = e^{\Delta\operatorname{Re}\lambda}$ per step. Put that into the half-life formula $t_{1/2} = \ln 2/(-\ln|\bar\lambda|)$ of §21.1, and it becomes

In plain words, $\Delta$ converts the system's own clock into token steps. A small $\Delta$ means that many tokens pass per unit of the system's time, so the state remembers for many tokens. A large $\Delta$ means that each token is a big jump in time, so the state forgets fast and tracks the latest input.

The step size gives these models a property that attention and discrete RNNs lack. The continuous system does not depend on $\Delta$. If the data are recorded at a different rate, changing $\Delta$ in proportion gives the same continuous model. S4 was trained on speech sampled 16,000 times a second. On the same speech sampled at half that rate, with its $\Delta$ doubled and no retraining, it still scored 96.3% (Gu, Goel & Ré 2021).

The widget shows all of this on one eigenvalue. The left panel is the continuous plane, where you drag $\lambda$ (its conjugate moves with it) inside the stable left half. The right panel is the discrete plane with the unit circle, and it marks where each rule sends $\lambda$ for the current $\Delta$. A faint trail shows where each rule sends $\lambda$ as $\Delta$ grows, and Sweep Δ runs the points along it. Below, the chosen rule's mode is started at 1 and drawn solid (its real part) against the exact continuous mode, dashed. Up to a constant factor this is the mode's impulse response, what it does after a single unit input. Watch it decay, ring, or blow up.

import numpy as np

def discretize(lam, B, delta, method):
    raise NotImplementedError

The same layer is a convolution

A recurrence with fixed parameters weighs each past input by how long ago it arrived and by nothing else. That is the definition of a convolution, and it gives a second way to compute the layer. Start with the bucket.

The same is true in general. Unroll the discrete recurrence from $\mathbf{h}_0 = \mathbf{0}$. The first state is $\bar{\mathbf{B}}x_1$, the second $\bar{\mathbf{A}}\bar{\mathbf{B}}x_1 + \bar{\mathbf{B}}x_2$, and in general $\mathbf{h}_t = \sum_{s=1}^{t}\bar{\mathbf{A}}^{t-s}\bar{\mathbf{B}}x_s$. That is the unrolled form of §21.1 with a constant decay. Multiply by $\mathbf{C}$ and write $j = t - s$ for the lag.

Proof. The unrolled state follows by induction. It holds at $t = 1$. If it holds at $t - 1$, then $\mathbf{h}_t = \bar{\mathbf{A}}\sum_{s=1}^{t-1}\bar{\mathbf{A}}^{t-1-s}\bar{\mathbf{B}}x_s + \bar{\mathbf{B}}x_t = \sum_{s=1}^{t}\bar{\mathbf{A}}^{t-s}\bar{\mathbf{B}}x_s$. Multiplying by $\mathbf{C}$ and substituting $j = t - s$ gives the sum. For a diagonal $\bar{\mathbf{A}}$, $\mathbf{C}\bar{\mathbf{A}}^j\bar{\mathbf{B}} = \sum_n C_n\bar\lambda_n^j\bar B_n$. $\blacksquare$

In plain words, the whole layer is one fixed filter slid along the sequence. The weight $\bar K_j$ says how much a token $j$ steps ago contributes to the output now, and it is the same at every position. The kernel is also the layer's impulse response. Feed it a single 1 followed by zeros, and the outputs are $\bar K_0, \bar K_1, \dots$ (plus $D$ at the first step).

So a time-invariant SSM has two faces. As a recurrence it decodes one token at a time with $O(N)$ work and a state of $N$ numbers, like the RNNs of §3.3. As a convolution it trains on a whole sequence at once, like the convolutional networks of §0.8, but with a filter as long as the sequence. A CNN learns each filter weight separately, and a filter of 16,000 free weights would be far too many to learn well. Here the filter is generated by $O(N)$ parameters, however long $T$ is. For a diagonal $\mathbf{A}$ they are the $N$ eigenvalues, $\mathbf{B}$, $\mathbf{C}$, $\Delta$, and $D$. Fast parallel training and cheap decoding in one layer is the reason this line of work existed.

Conjugate pairs make the kernel real and halve its cost. A real $\mathbf{A}$ has complex eigenvalues in conjugate pairs. After diagonalizing, its $\mathbf{B}$ and $\mathbf{C}$ come in matching conjugate pairs too. So the kernel's terms pair up as $q + q^* = 2\operatorname{Re}q$, where $q^*$ is the complex conjugate (we avoid the overbar, which here means discretized), and the kernel is real. Implementations therefore store only one mode of each pair, $N/2$ complex modes, and take twice the real part of their sum (Gu et al. 2022, S4D).

Convolution by FFT, without wraparound

Computed directly, the convolution of two sequences of length $T$ costs $O(T^2)$, no better than attention. The Fourier transform turns a convolution into an entrywise product, and a fast algorithm computes the transform in $O(T\log T)$. There is one catch. The transform's convolution wraps around the end of the sequence, and a causal model must not.

The definitions come first. For a sequence $\mathbf{f}$ of length $L$, indexed from 0, the discrete Fourier transform (DFT) is $\hat f_m = \sum_{n=0}^{L-1}f_ne^{-2\pi\mathrm{i}mn/L}$. It is invertible, $f_n = \frac1L\sum_{m=0}^{L-1}\hat f_me^{2\pi\mathrm{i}mn/L}$. The circular convolution of $\mathbf{f}$ and $\mathbf{g}$ is $(\mathbf{f}\circledast\mathbf{g})_n = \sum_{j=0}^{L-1}f_jg_{(n-j)\bmod L}$, in which an index that runs off the start wraps around to the end. The convolution theorem says that the DFT of a circular convolution is the entrywise product $\hat f_m\hat g_m$. So a convolution costs two forward transforms, a product, and one inverse transform.

Computed from its definition, a DFT costs $O(L^2)$. The fast Fourier transform (FFT) computes a DFT or its inverse in $O(L\log L)$ operations, and §35.2 derives the algorithm.

Now the wraparound. Take $T = 2$, a kernel $(K_0, K_1)$, and inputs $(x_0, x_1)$. The causal output is $y_0 = K_0x_0$. A circular convolution of length 2 gives $y_0 = K_0x_0 + K_1x_1$ instead, because the index $0 - 1$ wraps around to 1. The future input $x_1$ has leaked into the first output. The cure is to pad. Append zeros so both sequences have length 4. Now the wrapped index lands on a zero, and $y_0 = K_0x_0 + K_1\cdot 0$ is right.

In general, zero-pad both sequences to a length $L \ge 2T - 1$. Every true lag lies in $[-(T-1), T-1]$, every wrapped index lands on a zero, and the first $T$ entries of the circular convolution equal the causal one. Three FFTs of length $L$ cost $O(T\log T)$.

The next cell builds a diagonal SSM with 8 conjugate pairs ($N = 16$) and runs it as a recurrence over 4,096 tokens. It compares the result with the padded FFT convolution and with an unpadded one.

The recurrence and the padded FFT agree to rounding. The kernel's imaginary part is at rounding level, because the modes come in conjugate pairs. Without padding, the first outputs are off by 0.1 to 0.3, as large as the outputs themselves, because they have absorbed the tail of the sequence. This bug never crashes. In a causal model it leaks the future into the past, which is why it deserves its own test.

In practice, an SSM layer trains in convolution mode and decodes in recurrent mode, with the same parameters. An interviewer will ask you to show that the two give the same outputs, and what goes wrong without padding.

import numpy as np

def ssm_kernel(lam, B, C, delta, T):
    raise NotImplementedError

def causal_fft_conv(x, K):
    raise NotImplementedError

What should the state remember

We now have a layer that is both a recurrence and a convolution, and stable at any step size. What it is good for depends entirely on the state matrix $\mathbf{A}$, and a random one is nearly useless. The linear state-space layer, LSSL for short (Gu et al. 2021), was the direct predecessor of S4. The S4 paper reports what happened when the LSSL read handwritten digits one pixel at a time (Gu, Goel & Ré 2021). With a random state matrix it was 60% accurate. With the matrix derived below it was 98% accurate. The LSSL paper's own figure for random matrices is 62%, on a version of the task with the pixels shuffled. Where does the good matrix come from?

Gu, Dao, Ermon, Rudra & Ré 2020 asked a cleaner question. Suppose that at every moment $t$ we want $N$ numbers from which we could redraw the entire history $x(s)$, $0 \le s \le t$, as well as possible. And we want to update those numbers online, as time moves forward, without storing the history. Which $N$ numbers should we keep, and how should they change?

A natural answer is the best polynomial fit. Fit a polynomial of degree less than $N$ to the history on $[0, t]$ by least squares, and keep its coefficients. As $t$ grows, the interval stretches, the fit is redone, and the coefficients change. The surprise is that they change by a linear differential equation. So the coefficients are the state of an SSM, and the equation tells us its matrix. The framework is called HiPPO, for high-order polynomial projection operators.

To make the fit easy, write the polynomial in a basis of polynomials that are orthonormal on the interval. Each has unit size, and any two are perpendicular in the sense of an inner product, an average of their product. Then each coefficient of the fit can be computed on its own.

The projection theorem of §0.10 says that the least-squares fit by a combination of orthonormal vectors has coefficients equal to the inner products with those vectors. That section works in $\R^n$ with the dot product. Its proof uses only the inner product (the residual is orthogonal to the subspace, then Pythagoras), so it holds word for word for functions on $[0, t]$ with $\langle f, g\rangle_t$. With the basis $g_0, \dots, g_{N-1}$, requiring the residual $x - \sum_kc_kg_k$ to be orthogonal to each $g_n$ gives $c_n = \langle x, g_n\rangle_t$. So the state we want is

and the history is redrawn as $x(s) \approx \sum_{n} h_n(t)\,g_n(s, t)$.

Start with $N = 1$, where you can do it in your head. Then $g_0 = 1$, and $h_0(t) = \frac1t\int_0^t x(s)\,ds$ is the running mean of the input. By the product rule,

That is a one-dimensional SSM whose rate $1/t$ slows down as the history grows, so that every past moment keeps an equal share of the mean. The general case is the same calculation with more bookkeeping.

In plain words, the coefficients of the best polynomial fit to everything seen so far obey a linear equation whose matrix is lower triangular. Coefficient $n$ is driven by the new input with weight $\sqrt{2n+1}$ and pulled by the lower-order coefficients. The whole system runs on a clock that slows like $1/t$, so the newest moment and the oldest count equally in the fit. The name LegS is short for Legendre polynomials on a scaled window, the window $[0, t]$ that stretches with time.

This is the matrix that took the LSSL from 60% to 98%. A random matrix makes the state a few arbitrary fading sums. This one makes it a summary from which the past can be redrawn. In practice, HiPPO is the reason S4 is initialized the way it is, and an interviewer will ask you to say in one sentence what the state of a HiPPO system holds.

The HiPPO matrix up close

Write $\mathbf{A} = -\mathbf{A}_{\text{LegS}}$ for the matrix with the minus sign, the form S4 uses as its state matrix. It is lower triangular, so its eigenvalues are its diagonal entries, $-1, -2, \dots, -N$. They are distinct and in the left half-plane, so $\mathbf{A}$ is stable, and it is diagonalizable, $\mathbf{A} = \mathbf{V}\boldsymbol\Lambda\mathbf{V}^{-1}$. In the eigenvector basis the system would be $N$ independent leaky buckets, and the kernel would be the cheap Vandermonde product of the convolution theorem. Why not diagonalize it once and be done?

Because $\mathbf{V}$ is catastrophically ill-conditioned. Gu, Goel & Ré 2021, in their Lemma 3.2, show that its entries grow like $2^{4N/3}$. Recall from §0.10 that the condition number $\kappa$ of a matrix is the ratio of its largest stretch to its smallest, and that working with the matrix costs about $\log_{10}\kappa$ decimal digits of accuracy. In the diagonal basis, the parameters $\mathbf{V}^{-1}\mathbf{B}$ and $\mathbf{C}\mathbf{V}$ would be enormous numbers that cancel, and float32 has only about 7 digits to cancel with.

The condition number passes $10^{10}$ at $N = 16$, and by $N = 24$ it passes $10^{16}$, beyond what float64 can resolve (so that last value is itself only a rough estimate). S4 used $N = 64$. But the cell also shows the way out. Add a rank-one matrix to $\mathbf{A}$, and the result is a multiple of the identity plus a skew-symmetric matrix, one with $\mathbf{S}^\top = -\mathbf{S}$. That matrix can be diagonalized by a unitary $\mathbf{V}$, a complex matrix whose columns are orthonormal, with condition number exactly 1 even at $N = 64$.

In plain words, the HiPPO matrix is a well-behaved system of rotations and decays, every mode decaying at rate $\frac12$ and rotating at its own frequency, plus a rank-one correction that ruins the conditioning. Keep the two pieces apart and nothing is ill-conditioned. That structure is what the word structured in S4's name refers to.

import numpy as np

def hippo_legs(N):
    raise NotImplementedError

def normal_plus_low_rank(N):
    raise NotImplementedError

Running HiPPO online, exactly

This subsection and the next are the deepest part of the section. The first proves that a recurrence over samples really does keep the best polynomial fit of a history it never stores. The second shows what S4's time-invariant version of the system remembers instead. On a first read you can skip ahead to the diagonal state spaces.

The factor $1/t$ looks like trouble, because it makes the system change with time. But $\frac{1}{t}\,dt = d(\ln t)$. On the clock $\tau = \ln t$ the same equation reads $\frac{d\mathbf{h}}{d\tau} = -\mathbf{A}_{\text{LegS}}\mathbf{h} + \mathbf{p}\,x$, a time-invariant SSM with state matrix $-\mathbf{A}_{\text{LegS}}$. So the discretizations above apply, with a step that is no longer constant. Count samples from $k = 0$, and hold sample $x_k$ on $[k, k+1)$. The step from $t = k$ to $t = k+1$ then has length $\ln\frac{k+1}{k}$ on the new clock, and ZOH gives the exact update.

One more fact makes the update tidy. Column 0 of $\mathbf{A}_{\text{LegS}}$ is $(1, \sqrt3, \sqrt5, \dots) = \mathbf{p}$, so $\mathbf{A}_{\text{LegS}}\mathbf{e}_0 = \mathbf{p}$, where $\mathbf{e}_0 = (1, 0, \dots, 0)$. So the ZOH input term $\mathbf{A}_{\text{LegS}}^{-1}(\mathbf{I} - \bar{\mathbf{A}}_k)\mathbf{p}$ collapses to $(\mathbf{I} - \bar{\mathbf{A}}_k)\mathbf{e}_0$.

In plain words, each step pulls the state toward "the history is the constant $x_k$" by the factor $\bar{\mathbf{A}}_k$, which gets closer to the identity as $k$ grows. The newest sample moves the state less and less, because it is a smaller and smaller share of the whole history.

The HiPPO paper prints a forward-Euler version, $\mathbf{h}_{k+1} = (\mathbf{I} - \frac1k\mathbf{A}_{\text{LegS}})\mathbf{h}_k + \frac1k\mathbf{p}\,x_k$, for illustration. Its experiments use the bilinear rule, which its ablation finds accurate enough. Both are approximations. Euler's matrix also has diagonal entries $1 - (n+1)/k$ of size above 1 while $k$ is less than half of $N$.

The next two cells check the Legendre identity used in the proof of the HiPPO equation. Then they run the exact update on a signal of 500 samples and compare its state with the projection computed offline from the whole history.

The online state equals the offline projection to rounding at every $N$. So the recurrence really does keep the best polynomial fit of a history it no longer has. The signal contains a cosine with seven cycles, and no fit with fewer than about 25 polynomials can follow it. That is why the error stays large up to $N = 16$ and then falls to $2\times10^{-4}$ at $N = 32$.

The Euler update stalls near 0.16 however large $N$ is. Its error per step grows with the degree, like $((n+1)/k)^2$ for coefficient $n$. So the high-degree coefficients that carry the fine detail are exactly the ones it computes worst.

The widget runs the same memory on several signals. The faint curve is the whole signal, the dark curve the part seen up to time $t$, and the colored curve what the $N$ numbers of the state can redraw of it. The bars underneath are the state itself. Pick a signal and a state size $N$, then run or scrub $t$.

import numpy as np
from scipy.linalg import expm
from numpy.polynomial import legendre as leg

def legs_online(x, N):
    raise NotImplementedError

def reconstruct(h, t, s):

The time-invariant version, and what it remembers

There is a catch. S4 does not use the time-varying LegS system. It puts the same matrix in a time-invariant SSM, $\dot{\mathbf{h}} = \mathbf{A}\mathbf{h} + \mathbf{B}x$ with a learned $\Delta$, so that the convolution view applies. Dropping the $1/t$ throws away the optimality argument above. For a while the HiPPO matrix in S4 had no explanation, only results. Gu, Johnson, Timalsina, Rudra & Ré 2022 supplied one, and the log clock lets us prove it.

In plain words, the time-invariant system compresses the past onto Legendre polynomials in the variable $e^{-r}$, with a weight $e^{-r}$ that fades exponentially. Most of the interval $(0, 1]$ of $w$ is spent on the last couple of time units, and everything older is squeezed toward $w = 0$. So the memory is sharp for the recent past and blurry for the distant past. With a step $\Delta$, one time unit is $1/\Delta$ tokens, so $\Delta$ again sets how far back "recent" reaches. The scaled LegS memory of the earlier widget is the other choice, equally blurry everywhere but never forgetting. The cell checks the formula for the impulse response at three lags.

The widget below draws this memory over the scaled one, both with 16 numbers about the same signal. Before you drag $\Delta$, predict which part of the history each one redraws well.

The HiPPO paper also proves that in the scaled memory, the gradient of the state at time $t$ with respect to an input at any earlier time is of order $1/t$ (Gu et al. 2020, Proposition 5). It falls polynomially with time, not exponentially as in a classical RNN. That is the theoretical counterpart of the 16,000-step memory we wanted at the start.

To run this time-invariant layer as a convolution, S4 needs the kernel of a matrix that is not diagonal. It gets it from the normal-plus-low-rank split proved above, with an algorithm that the next box describes. You do not need its details for the rest of the chapter, because the diagonal models that followed made it unnecessary.

Diagonal state spaces

S4's machinery is heavy. A few months after it, Gupta, Gu & Berant 2022 asked whether the low-rank term was needed at all. Their diagonal state space model (DSS) kept only a diagonal state matrix and matched S4 on Long Range Arena and on speech. Gu, Gupta, Goel & Ré 2022 then explained why, and made the model simple. Their diagonal S4, called S4D, rests on a proof about the limit $N \to \infty$. Recall that S4's matrix is a normal part, which a unitary change of basis turns into the diagonal matrix $\boldsymbol\Lambda$, minus a rank-one term. Keeping only $\boldsymbol\Lambda$ and halving the input vector gives the same kernel as the full HiPPO matrix in that limit. Their experiments showed that what matters is how the eigenvalues are initialized, not the low-rank structure.

So a state-space layer can be a set of independent complex leaky buckets, each a decaying rotation, whose kernel is the Vandermonde product of the convolution theorem. The definition below fixes the two standard ways to set one up.

The S4D-Lin initialization repays a second look, because it says how big $N$ and $\Delta$ must be for a given job.

The loss falls from 0.84 at the start, about the 0.80 of always predicting zero, to 0.013 after 300 steps. The trained layer gives the same outputs whether it runs as an FFT convolution (the training mode) or one token at a time (the decoding mode), up to float32 rounding.

Two more simplifications, S5 and the LRU

The simplified state space layer, S5 (Smith, Warrington & Linderman 2022), changed a different choice. An S4 layer is $d$ independent single-input single-output SSMs, a block-diagonal system with $dN$ state numbers. S5 uses one multi-input multi-output (MIMO) SSM instead, with $\mathbf{B} \in \mathbb{C}^{N\times d}$ and $\mathbf{C} \in \mathbb{C}^{d\times N}$. Its state matrix is diagonal, initialized from the diagonalized normal part of HiPPO. S5 computes the recurrence with a parallel scan, an algorithm that combines the steps of a linear recurrence in a tree of depth $O(\log T)$, instead of a convolution. A scan does not need a fixed kernel, so S5 can change $\bar{\mathbf{A}}$ at every step, for example for data sampled at irregular times. That is a first step toward the input-dependent dynamics of §21.4, and §21.5 develops the scan.

The LRU of Orvieto et al. 2023 went furthest. Starting from a plain RNN, the authors removed the nonlinearity from the recurrence and diagonalized it over the complex numbers. They found that three choices closed the remaining gap to S4, and all three are in the definition box.

• The exponential parameterization keeps $|\lambda| = e^{-e^{\nu}} < 1$ for any real $\nu$, so training can never leave the stable region.
• The ring initialization samples $\lambda$ uniformly by area over the ring $r_{\min} \le |\lambda| \le r_{\max}$, close to the unit circle for long memory. Their best result on sequential CIFAR-10 (images classified from their pixels read one at a time) used $r_{\min} = 0.9$ and $r_{\max} = 0.999$.
• The normalization $\gamma$ stops the state from blowing up as $|\lambda| \to 1$, for the reason in the next box.

Proof. Unrolling, $h_t = \gamma\sum_{j=0}^{t-1}\lambda^jx_{t-j}$. Then $\mathbb{E}|h_t|^2 = \mathbb{E}[h_th_t^*] = \gamma^2\sum_{i,j}\lambda^i(\lambda^*)^j\,\mathbb{E}[x_{t-i}x_{t-j}] = \gamma^2\sum_{j=0}^{t-1}|\lambda|^{2j}$, since only the terms with $i = j$ survive. The geometric sum gives the formula. $\blacksquare$

In plain words, a mode with $|\lambda| = 0.999$ sums about a thousand past inputs. Without normalization its state is about 22 times larger than that of a mode with $|\lambda| = 0$, since $\sqrt{1/(1 - 0.999^2)} \approx 22.4$. The factor $\gamma$ puts every mode on the same scale, whatever its memory, so the readout $\mathbf{C}$ does not have to undo a scale that varies by orders of magnitude.

import numpy as np

def lru_params(nu, theta):
    raise NotImplementedError

def ring_init(N, r_min, r_max, rng, max_phase=2 * np.pi):
    raise NotImplementedError

What S4 achieved, and where time invariance breaks

The S4 row gives the average of the original paper and the one in its third arXiv version. The LRU paper prints no average, so its entry is our mean of the six scores in its Table 3. DSS is from its Table 1, and the other rows are from each paper's abstract.

On language the story was different. A language model is scored by its perplexity, the exponential of its average next-token loss, so lower is better. On WikiText-103, a standard benchmark of Wikipedia articles, S4 came within 0.8 perplexity of a Transformer, and within 0.4 after retraining in the paper's final version. That was a strong result for a model with no attention, but the gap would not close. Fu, Dao, Saab, Thomas, Rudra & Ré 2022 found the reason with synthetic tasks. SSMs struggled to recall a token seen earlier and to compare tokens across the sequence.

Their layer, called H3 (Hungry Hungry Hippos), added two things. A shift SSM has a kernel that only delays its input by a few steps. Multiplicative interactions, products of SSM outputs with projections of the input, make what passes through depend on content. H3 came within 0.4 perplexity of a Transformer on OpenWebText, a corpus of web pages. A 125M-parameter hybrid that kept two attention layers beat the Transformer by 1.0. Hyena (Poli et al. 2023) replaced the SSM by long convolutions whose filters are generated by a small network, interleaved with gates computed from the data. It matched Transformer quality at a context of 2,048 tokens with 20% less training compute.

Both fixes point at the same root cause, which the convolution theorem makes plain. The weight with which token $s$ reaches output $t$ is $\bar K_{t-s}$, a number fixed by the parameters and the lag. It cannot depend on what token $s$ is. Take associative recall. The input is a list of key-value pairs followed by a query key, and the answer is the value stored next to that key. The position of the answer depends on the input, but a time-invariant layer weighs each position by its distance alone. And since the layer is linear in $\mathbf{x}$, it cannot multiply "this key matches the query" by "this is the value".

Nonlinear layers stacked between SSMs can approximate this product only with much depth or width. H3 and Hyena patched the gap with gates around a time-invariant core. Mamba goes to the root and lets $\Delta$, $\mathbf{B}$, and $\mathbf{C}$ depend on the current token. That costs the convolution view, and it is the subject of §21.4. An interviewer will ask exactly this, why a time-invariant SSM cannot do associative recall and what it takes to fix it.

Read more

• Gu, Goel & Ré (2021). Efficiently Modeling Long Sequences with Structured State Spaces. The S4 paper, with the NPLR parameterization, the Cauchy-kernel algorithm, Lemma 3.2 on the ill-conditioned eigenvectors, and the Path-X result.
• Gu, Dao, Ermon, Rudra & Ré (2020). HiPPO: Recurrent Memory with Optimal Polynomial Projections. Online function approximation as an ODE, the LegS matrix, its discretizations, and the gradient bounds.
• Gu, Johnson, Timalsina, Rudra & Ré (2022). How to Train Your HiPPO: State Space Models with Generalized Orthogonal Basis Projections. What the time-invariant HiPPO system computes, and how to initialize $\Delta$.
• Gu, Gupta, Goel & Ré (2022). On the Parameterization and Initialization of Diagonal State Space Models. S4D, the diagonal limit of S4, the S4D-Lin and S4D-Inv initializations, and a careful ablation of every design choice.
• Gupta, Gu & Berant (2022). Diagonal State Spaces are as Effective as Structured State Spaces. DSS, the first evidence that the low-rank term can be dropped.
• Smith, Warrington & Linderman (2022). Simplified State Space Layers for Sequence Modeling. S5, one MIMO SSM computed by a parallel scan.
• Orvieto et al. (2023). Resurrecting Recurrent Neural Networks for Long Sequences. The LRU, the stable exponential parameterization, ring initialization, and the normalization derived above.
• Fu, Dao, Saab, Thomas, Rudra & Ré (2022). Hungry Hungry Hippos: Towards Language Modeling with State Space Models. Why SSMs lagged on language (recall and comparison) and the gated H3 layer.
• Poli et al. (2023). Hyena Hierarchy: Towards Larger Convolutional Language Models. Long implicit convolutions with data-controlled gates, the other answer to the recall gap.
• Tay et al. (2021). Long Range Arena: A Benchmark for Efficient Transformers. The six tasks, Path-X among them, on which every Transformer variant failed.
• Rush & Karamcheti. The Annotated S4. S4 implemented line by line, including the DPLR kernel this section only states.