LTI SSMs: S4, S4D, S5 as discretized continuous systems

8.1 State-space to discrete: the S4 setup

Chapter 7 closed on a recurrence and called it the recurrent view of an S4-style layer. Make the layer explicit. S4 fixes a single-input, single-output (SISO) continuous system

with $\statemat \in \R^{N\times N}$ the HiPPO-LegS matrix of §7.3, $\inputmat \in \R^{N\times 1}$, $\outputmat \in \R^{1\times N}$, and $D$ a scalar feedthrough. To run it on a sampled sequence $u_0, u_1, \dots$ we discretize at step $\stepsize$. The S4 paper itself uses Chapter 4’s bilinear transform (§4.4), $\discA = (I - \tfrac{\stepsize}{2}\statemat)^{-1}(I + \tfrac{\stepsize}{2}\statemat)$ Gu et al. (2022) ; this book standardizes on the zero-order hold (§4.3) — the rule S5 Smith et al. (2023) and the Mamba line Gu & Dao (2024) later adopted, and one whose difference from bilinear the S4D ablations measure as negligible Gu et al. (2022) :

computed through the augmented matrix-exponential trick of §4.3 (the companion’s discretize_zoh stacks $\big[\begin{smallmatrix}\statemat\stepsize & \inputmat\stepsize\\ 0 & 0\end{smallmatrix}\big]$ and exponentiates, sidestepping $\statemat^{-1}$).

What makes this an architecture rather than a fixed filter is which parts train. S4 learns $\outputmat$, $D$, and $\log\stepsize$ (the timescale, kept positive by the log), while $\statemat, \inputmat$ are initialized from HiPPO and often frozen or updated slowly. Stability is inherited for free: by the §7.7 spectrum every eigenvalue of $\statemat$ has negative real part, so Chapter 4’s ZOH proposition gives $\abs{\lambda(\discA)} = e^{\Re(\lambda(\statemat))\stepsize} < 1$ for every $\stepsize > 0$. The slowest mode (eigenvalue $-1$) decays per step by $e^{-\stepsize}$ — for the companion’s default $\stepsize = 0.1$ that is $0.905$, a long but strictly contracting memory.

This is one way to compute the layer. The rest of the chapter is, in large part, other ways to compute the same thing — and the surprising fact that they are equal.

8.2 The structured state matrix, revisited

Before the second view, recall why the HiPPO $\statemat$ is special enough to deserve fast algorithms. It is dense and lower-triangular (§7.3), so a black-box $e^{\statemat\stepsize}$ and a black-box kernel would cost more than the structure warrants. The structure S4 exploits is that the HiPPO-LegS matrix is normal-plus-low-rank (NPLR): up to an orthogonal change of basis it is a normal matrix plus a rank-one correction, and after diagonalizing the normal part it becomes diagonal-plus-low-rank (DPLR),

with $\Lambda$ the (complex) diagonal of eigenvalues and $P Q^*$ a rank-one term Gu et al. (2022) . That the HiPPO-LegS matrix admits this diagonal-plus-rank-one form is not evident from the §7.3 closed form — it is a construction of the S4 paper, which we take as given rather than re-derive here. This is the same non-normality that §7.5 flagged as a conditioning hazard, now read constructively: the eigenvector basis is oblique (hence ill-conditioned), but the correction that makes $\statemat$ non-normal is only rank one, and a rank-one correction is exactly what the Woodbury identity handles cheaply.

The DPLR form is the hinge for the rest of the chapter. Section 8.4 uses the full $\Lambda - PQ^*$ to compute S4’s kernel quickly; §8.5 throws away the $PQ^*$ term entirely and keeps only $\Lambda$, which is S4D.

8.3 The convolution–recurrence duality

The recurrence of §8.1 is sequential: $\statevec_{k+1}$ needs $\statevec_k$. Unrolling it from zero initial state exposes a second, fully parallel description. With $\statevec_0 = 0$,

Define the SSM convolution kernel $K \in \R^{L}$ by $K_i = \outputmat\, \discA^{\,i}\, \discB$. Then the output is a causal convolution plus the feedthrough:

The proof (Exercise 8.4, §8.9) is the unrolling above made rigorous by induction. The two forms are not approximations of each other — they are the same numbers. The companion confirms it to machine precision: the recurrent lax.scan output and the FFT convolution agree to a measured residual $\sim 10^{-15}$ in float64, which test_conv_recurrence_duality pins below $10^{-12}$.

The duality is why S4 is practical: train as a convolution, deploy as a recurrence. Training sees whole sequences at once, so the parallel $O(L\log L)$ FFT wins; autoregressive inference produces one token at a time, so the $O(L)$ scan with its $O(N)$ state wins. One set of parameters, two runtimes.

8.4 The Cauchy kernel: structure buys speed

The duality leaves one cost open: building the kernel $K$. Naively, $K_i = \outputmat\,\discA^{\,i}\,\discB$ iterates the state vector $\discA^{\,i}\discB$ — one $N\times N$ mat-vec per tap, $O(N^2 L)$ in total. For the teaching regime $N \le 64$ that is exactly what the companion does, and it is plenty fast. S4’s published contribution is to compute the same kernel in near-linear time, $\tilde O(N + L)$, by using the DPLR structure of §8.2 Gu et al. (2022) .

The route is a generating function. Collect the taps into a polynomial and evaluate it on the unit circle: the truncated generating function is

a geometric series in the matrix $\discA z$. Recovering $K$ from $\hat K$ evaluated at the $L$-th roots of unity is one inverse FFT. The remaining cost is the resolvent $(I - \discA z)^{-1}$, and here DPLR pays off: with $\discA$ built from $\Lambda - PQ^*$ the Woodbury identity turns the resolvent into a Cauchy matrix–vector product — entries $1/(\zeta_k - \lambda_n)$ — plus a rank-one correction. The whole kernel is then a Cauchy/Woodbury solve in $N$ plus the inverse FFT in $L$ (the §8.4 figure isolates the $N$-scaling at fixed $L$); the dense $\discA^{\,i}$ never appears.

We present this kernel but do not implement it: the companions use the transparent naive kernel throughout (§8.10). The point of §8.4 is the shape of the cost, not the engineering — and the shape is what motivates S4D.

8.5 S4D: the diagonal restriction

The Cauchy machinery is the price of keeping the rank-one term $PQ^*$. S4D asks what happens if we simply drop it Gu et al. (2022) — keep the diagonal $\Lambda$, discard the low-rank correction. The state matrix is now genuinely diagonal, $\statemat = \Lambda = \diag(\lambda_1, \dots, \lambda_N)$, and everything decouples mode by mode. The kernel collapses from a Cauchy computation to a Vandermonde sum:

with $\bar B_n$ the per-mode diagonal ZOH input (this book’s convention; the S4D paper itself treats the integration rule as interchangeable — its base version keeps S4’s bilinear — and measures no noticeable difference between the two). In practice the $N$ real states are carried as $N/2$ complex conjugate pairs and the real kernel is reconstructed as $K_l = 2\,\Re\!\big(\sum_n \outputmat_n \bar B_n \bar\Lambda_n^{\,l}\big)$.

The derivation and the stability claim are Exercise 8.5 (§8.9). What makes S4D more than “S4 with worse expressivity” is how it secures the stability condition. The S4D-Lin initialization places the modes at $\lambda_n = -\tfrac12 + i\pi n$, and — crucially — stores them in the parameterization

so that $\Re(\lambda_n) = -e^{(\log A_{\text{real}})_n} < 0$ for any real value the parameters take. Stability is not hoped for; it is structural.

This is the constructive answer to the §7.5 conditioning hazard: where S4 must diagonalize an ill-conditioned non-normal matrix and hope the basis behaves, S4D never forms that basis at all — it parameterizes the eigenvalues directly and pins their real parts negative.

Diagonal state space models are not a downgrade in practice: with the right initialization they match full S4 on long-range benchmarks while being markedly simpler to implement Gupta et al. (2022) . The companion plots the S4D-Lin spectrum against the full HiPPO spectrum to make the restriction visible.

8.6 S5: the parallel associative scan

S4D made the state diagonal but still computed it as a convolution. S5 keeps the diagonal state and abandons the convolution for a parallel scan Smith et al. (2023) . Two changes. First, S5 is MIMO: one shared diagonal state of size $P$ drives $H$ input/output channels through $\inputmat \in \C^{P\times H}$ and $\outputmat \in \C^{H\times P}$, instead of $H$ independent SISO systems. Second, it computes the recurrence

(diagonal $\discA$, so $\odot$ is elementwise) with an associative scan. The trick is that a linear recurrence is an associative operation on $(\text{transition}, \text{input})$ pairs:

which composes “apply the first affine update, then the second.” Scanning $\bullet$ over the pairs $(\discA, \discB u_k)$ yields every state $\statevec_k$ at once.

The proof is Exercise 8.6 (§8.9). The companion pins the equivalence to machine precision — the associative scan and the sequential scan agree to about $3\times10^{-15}$ (pinned below $10^{-12}$) — and contrasts the two on the algorithmic cost that actually differs: critical-path depth.

S5 deliberately stops at one diagonal SSM run in parallel. Its MIMO channel mixing and, above all, the move to input-dependent transitions are Chapter 9’s subject — and that move is what makes the associative scan indispensable rather than merely convenient.

8.7 What’s next

The thread of this chapter is that S4, S4D, and S5 are three computational treatments of one linear time-invariant system: a convolution and a recurrence that the duality proves equal (§8.3), a diagonal restriction whose stability is free (§8.5), and a parallel scan that reorders the same arithmetic into $O(\log L)$ depth (§8.6). Time-invariance is what every one of these relies on — a fixed $\discA$ is what lets a single kernel exist and a single transition seed the scan.

Chapter 9 breaks time-invariance on purpose: Mamba makes $\statemat, \inputmat, \outputmat$ functions of the input Gu & Dao (2024) (and the structured state-space duality Dao & Gu (2024) later reframes the result), which destroys the convolutional view and leaves the selective scan — the §8.6 primitive, now load-bearing — as the only way to compute the layer. Chapter 10 then replaces ZOH with the exponential-trapezoidal scheme of Chapter 4 when the selective dynamics turn stiff, in Mamba-3 Lahoti et al. (2026) ; that chapter is where the symplectic-integrator pilot finds its empirical anchor. Everything downstream inherits the discretized-LTI scaffold built here.

8.8 Exercises

Six problems mixing computation and theory. Short/numerical (8.1–8.3) have inline collapsible solutions; the theory exercises (8.4–8.6) have full worked solutions in §8.9.

Exercise 8.1 (computation)

Take a diagonal toy system $\statemat = \diag(-1, -2)$, $\inputmat = (1, 1)^\top$, $\outputmat = (1, 1)$, $D = 0$, and step $\stepsize = 0.1$. Compute $\discA$ and $\discB$ by the diagonal ZOH, then the first four kernel taps $K_0, \dots, K_3$. Confirm the taps decay.

Exercise 8.2 (computation + code)

The S4D kernel sums one geometric sequence per mode. (a) Show that $\abs{\bar\Lambda_n^{\,l}} = e^{\Re(\lambda_n)\stepsize\, l}$, so each mode decays at its own rate. (b) Show that a complex mode $\lambda$ with weight $w = \outputmat\bar B$ and its conjugate $\bar\lambda$ with weight $\bar w$ together contribute $w\,\bar\Lambda^{\,l} + \bar w\,\overline{\bar\Lambda}^{\,l} = 2\Re(w\,\bar\Lambda^{\,l})$, explaining the real $2\Re$ form. (c) Run companions/ch08/jax/s4d_kernel.py and confirm the Vandermonde kernel matches the recurrence oracle.

Exercise 8.3 (code)

Run companions/ch08/jax/s4_duality.py and confirm the recurrent and convolutional outputs agree (the companion reports a residual below $10^{-15}$). Then explain why the agreement would break if the FFT used length $L$ instead of $2L$.

Exercise 8.4 (theory) — solution in §8.9

Prove the convolution–recurrence duality (Proposition, §8.3): for zero initial state, the SISO recurrence equals $y = K * u + D u$ with $K_i = \outputmat\,\discA^{\,i}\,\discB$.

Exercise 8.5 (theory) — solution in §8.9

Derive the S4D Vandermonde kernel from the diagonal recurrence, prove the discrete system is stable iff $\Re(\lambda_n) < 0$ for all $n$, and show that the S4D-Lin parameterization $\lambda_n = -e^{(\log A_{\text{real}})_n} + i(A_{\text{imag}})_n$ enforces stability for every parameter value.

Exercise 8.6 (theory) — solution in §8.9

Prove that the scan operator $(a_1,b_1)\bullet(a_2,b_2) = (a_2 a_1,\, a_2 b_1 + b_2)$ is associative, that its inclusive prefix scan reproduces the sequential states $\statevec_k$, and that a balanced reduction achieves parallel depth $O(\log L)$.

8.9 Full solutions to theory exercises

Solution to Exercise 8.4

Start from $\statevec_0 = 0$ and $\statevec_{k+1} = \discA\statevec_k + \discB u_k$. Claim: $\statevec_{k+1} = \sum_{j=0}^{k} \discA^{\,k-j}\discB\, u_j$. Induct on $k$. Base ($k=0$): $\statevec_1 = \discA\cdot 0 + \discB u_0 = \discA^{\,0}\discB\, u_0$. Step: assuming $\statevec_k = \sum_{j=0}^{k-1}\discA^{\,k-1-j}\discB u_j$,

Now the output, using $y_k = \outputmat\statevec_{k+1} + D u_k$:

with $K_i = \outputmat\discA^{\,i}\discB$. The sum $\sum_{j=0}^k K_{k-j} u_j$ is precisely the causal convolution $(K * u)_k$, so $y = K * u + D u$. The lower-triangular Toeplitz structure of the map $u \mapsto y - Du$ is the convolution. $\blacksquare$

Solution to Exercise 8.5

For diagonal $\statemat = \diag(\lambda_n)$ the matrix exponential is diagonal, $\discA = \diag(e^{\lambda_n\stepsize})$, and the ZOH input $\discB = \statemat^{-1}(e^{\statemat\stepsize} - I)\inputmat$ acts mode-by-mode: $\bar B_n = (e^{\lambda_n\stepsize} - 1)/\lambda_n \cdot B_n$. The kernel inherits the diagonality,

a sum of geometric sequences — the Vandermonde form (the real $2\Re$ version follows by pairing conjugates, Exercise 8.2). Stability: the kernel and the state stay bounded iff each geometric ratio satisfies $\abs{\bar\Lambda_n} < 1$. Since $\abs{\bar\Lambda_n} = \abs{e^{\lambda_n\stepsize}} = e^{\Re(\lambda_n)\stepsize}$ and $\stepsize > 0$, this holds iff $\Re(\lambda_n) < 0$ for every $n$. Finally, under S4D-Lin the real part is $\Re(\lambda_n) = -e^{(\log A_{\text{real}})_n}$, and a negative exponential is strictly negative for every real exponent — so $\Re(\lambda_n) < 0$ holds for all parameter values, with no constraint to enforce during training. Stability is built into the parameterization rather than maintained by it. $\blacksquare$

Solution to Exercise 8.6

Associativity. Write the action of a pair $(a, b)$ on a state $h$ as the affine map $h \mapsto a h + b$. Composition of affine maps is associative because function composition is. Concretely, with $\bullet$ elementwise,

and the two right-hand sides are identical. Prefix scan reproduces the states. Let $e_k = (\discA, \discB u_k)$ and $s_k = e_0 \bullet \cdots \bullet e_k$ be the inclusive scan. By induction the second component of $s_k$ equals $\statevec_k$: for $k = 0$ it is $\discB u_0 = \statevec_0$ (with $\statevec_{-1}=0$); and $s_k = s_{k-1} \bullet e_k$ has second component $\discA\cdot(\text{2nd comp of } s_{k-1}) + \discB u_k = \discA\statevec_{k-1} + \discB u_k = \statevec_k$. Depth. Because $\bullet$ is associative, computing all prefixes is a prefix-sum problem, which a balanced binary reduction tree evaluates in $O(\log L)$ sequential rounds with $O(L)$ total work Blelloch (1990) — versus the $O(L)$ rounds of the naive recurrence. $\blacksquare$

8.10 Companion code

Two language tracks for Chapter 8 — JAX (the reference, with the parallel scan) and PyTorch (the nn.Module operators). Julia is omitted here: the S4D reference and the associative-scan primitive both live in the JAX/PyTorch world, and a stdlib-Julia rebuild would add a third spelling without a third lesson. The JAX and torch companions are pinned bit-for-bit against each other, the cross-framework lesson Chapter 7 introduced.

JAX (companions/ch08/jax/):

• s4_core.py — the harness-free S4 numerical core ported from the Week 4 source: make_hippo_legs, discretize_zoh / discretize_bilinear, the naive $O(N^2 L)$ kernel, and the recurrent and convolutional views. The §8.3 duality lives here.
• s4_duality.py — emits kernel-duality.png (§8.3): kernel taps plus the two views overlaid.
• s4_kernel_cost.py — times the naive kernel across $N$; emits kernel-vs-n.png (§8.4).
• s4d_kernel.py — the S4D-Lin diagonal init and Vandermonde kernel (§8.5), complex128; emits s4d-spectrum.png.
• s5_scan.py — the S5 associative scan via jax.lax.associative_scan, with the sequential reference; emits s5-scan-depth.png (§8.6).
• tests/ — test_s4.py pins the duality (recurrent $\equiv$ convolutional), the kernel oracle, ZOH stability, and the FFT-pad guard; test_s4d_s5.py pins S4D by-construction stability, the Vandermonde-vs-recurrence cross-check, and S5 scan-equivalence.

PyTorch (companions/ch08/torch/):

• s4d_kernel.py — the S4D kernel as an nn.Module (fixed-init buffers), the Vandermonde computation ported from the reference S4D, in complex128.
• s5_sequential.py — the diagonal MIMO S5 as a sequential scan (PyTorch has no native parallel scan); the states match the JAX associative scan bit-for-bit.
• tests/test_s4d_torch.py — cross-framework parity: the torch S4D kernel and S5 states equal their JAX counterparts to within $10^{-8}$.

To run from the repo root:

# JAX (uses the uv .venv with jax, numpy, matplotlib)
PYTHONPATH=. python companions/ch08/jax/s4_duality.py
PYTHONPATH=. python companions/ch08/jax/s4_kernel_cost.py
PYTHONPATH=. python companions/ch08/jax/s4d_kernel.py
PYTHONPATH=. python companions/ch08/jax/s5_scan.py

# PyTorch (needs the .venv [torch] extra)
PYTHONPATH=. python companions/ch08/torch/s4d_kernel.py
PYTHONPATH=. python companions/ch08/torch/s5_sequential.py

All figures emit to public/figures/ch08/.