State-space models: from linear recurrences to S4 and Mamba

Before this lesson: lesson 07 (convolution), lesson 08 (linear recurrences, spectral radius, gates) and lesson 09 (the $O(L^2)$ cost of attention and the KV cache). Math: the matrix exponential and linear ODEs in linear-algebra lesson 16, circulant matrices and FFT convolution in linear-algebra lesson 19, and eigenvalues in linear-algebra lesson 05. Time: about 80 minutes, plus the notebook. This is a heavy lesson; the folded boxes can be skipped on a first reading.

---

Why this matters #

Attention costs $O(L^2)$ in the sequence length $L$, and its KV cache grows with every generated token (lesson 09). Many signals are very long: one hour of a seismogram at 100 Hz has 360,000 samples, and attention over it would compare $1.3 \times 10^{11}$ pairs per layer. An SSM trains in parallel, like a CNN or a transformer, and generates one step at a time with constant memory, like an RNN.

The model $x' = Ax + Bu$, $y = Cx$ is the linear time-invariant system of control theory and of the Kalman filter. S4 (Gu, Goel, Ré 2022) turned it into a deep-learning layer. It was the first model to do better than chance on Path-X, the hardest task of the Long Range Arena benchmark (sequences of 16,384 steps), where earlier models stayed at 50%. Mamba (Gu and Dao 2023) made the parameters depend on the input. It reached transformer quality on language with about 5 times higher generation throughput. Because the model lives in continuous time, $\Delta$ is a physical time step: S4 kept about 96% accuracy on speech sampled at half the training rate, without retraining, just by doubling $\Delta$. Interviews ask: "how does Mamba differ from a transformer and from an LSTM?", "why can S4 use an FFT and Mamba not?", "what is a parallel scan?"

1. Predict first #

2. The picture #

The simplest SSM: a leaky integrator #

Take one number: $x'(t) = a\,x(t) + b\,u(t)$ with $a < 0$. The input $u$ pushes the state; the term $a\,x$ pulls it back to 0. Starting from rest, the solution is

This is an average of the past input with exponentially decaying weights. The memory time is $1/|a|$. Engineers call it a leaky integrator or a low-pass filter.

Sample it. Choose a step $\Delta$ and keep the input constant during each step. Then the state after each step follows

This is the scalar linear RNN of lesson 08, with $w = \bar a$. Unrolled, $x_k = \sum_j \bar a^{\,j}\, \bar b\, u_{k-j}$: a convolution of the input with the exponential kernel $\bar b\,\bar a^{\,j}$.

Three views of one map #

The same model can be written in three ways:

• an ODE $x' = ax + bu$: the continuous-time parameters, and the meaning of $\Delta$;
• a recurrence $x_k = \bar a x_{k-1} + \bar b u_k$: one step at a time, cheap generation with a fixed-size state;
• a convolution $y_k = \sum_j K_j u_{k-j}$: the whole sequence at once, parallel training with an FFT.

Interactive figure: loading…

A richer kernel from complex eigenvalues #

One real number gives only a decaying exponential. Take a state $x \in \mathbb{C}^N$ and a diagonal $A$ with eigenvalues $\lambda_n$. Each mode is a scalar SSM, and the kernel adds them: $K_j = \sum_n C_n \bar B_n\, e^{j\Delta\lambda_n}$. A complex $\lambda_n = \alpha + i\omega$ gives $e^{j\Delta\lambda_n} = e^{j\Delta\alpha}\big(\cos(j\Delta\omega) + i \sin(j\Delta\omega)\big)$: a damped oscillation. So the kernel is a sum of $N$ damped oscillations. With $N$ in the tens, this is a long, structured filter with only $O(N)$ parameters. For a real output, modes come in conjugate pairs $\lambda, \bar\lambda$.

Stability for free. If $\operatorname{Re}\lambda < 0$, then $|e^{\Delta\lambda}| = e^{\Delta\operatorname{Re}\lambda} < 1$ for every $\Delta > 0$. The left half of the complex plane maps into the unit disk, so the recurrence never explodes. This is the spectral-radius problem of lesson 08, solved by the choice of parameters.

Interactive figure: loading…

A deep SSM applies one SSM per channel (a depthwise convolution with a kernel as long as the sequence), then mixes the channels with linear layers and nonlinearities. The SSM is linear in its input; the network is not.

3. Definitions and main results #

The continuous model and two discretizations #

Per channel, with input $u(t) \in \mathbb{R}$, state $x(t) \in \mathbb{R}^N$ and output $y(t) \in \mathbb{R}$:

with $A \in \mathbb{R}^{N\times N}$, $B \in \mathbb{R}^{N\times 1}$, $C \in \mathbb{R}^{1\times N}$. The term $Du$ is a skip connection; we drop it in the analysis. The discrete model, in the S4 convention, is

Theorem (zero-order hold, ZOH). If $u(t) = u_k$ is constant on the $k$-th step of length $\Delta$, then the exact solution of the ODE satisfies the recurrence with

Proof. (The matrix exponential is in linear-algebra lesson 16.)

1. By the product rule, $\frac{d}{dt}\big(e^{-At}x(t)\big) = e^{-At}\big(x'(t) - Ax(t)\big) = e^{-At}Bu(t)$.
2. Integrate over one step $[t_0, t_0 + \Delta]$ and multiply by $e^{A(t_0+\Delta)}$: $x(t_0 + \Delta) = e^{A\Delta}x(t_0) + \int_{t_0}^{t_0+\Delta} e^{A(t_0+\Delta-s)}Bu(s)\, ds$.
3. With $u(s) = u_k$ and the substitution $\tau = t_0 + \Delta - s$, the integral is $\big(\int_0^\Delta e^{A\tau}d\tau\big)Bu_k$.
4. If $A$ is invertible, $\frac{d}{d\tau}\big(A^{-1}e^{A\tau}\big) = e^{A\tau}$, so $\int_0^\Delta e^{A\tau}d\tau = A^{-1}(e^{\Delta A} - I)$. $\blacksquare$

For a scalar $a \ne 0$: $\bar a = e^{a\Delta}$ and $\bar b = \frac{e^{a\Delta} - 1}{a}\, b$.

When $A$ is not invertible

The integral exists for every $A$; only the closed form needs $A^{-1}$. For the integrator $A = 0$ it gives $\bar B = \Delta B$. In general, use the series $\int_0^\Delta e^{A\tau}d\tau = \sum_{n\ge0} A^n\Delta^{n+1}/(n+1)!$. Or note that, while $u$ is constant, the pair $(x, u)$ solves $\frac{d}{dt}(x, u) = M(x, u)$ with $M = \begin{pmatrix} A & B \\ 0 & 0\end{pmatrix}$, so $e^{\Delta M} = \begin{pmatrix} \bar A & \bar B \\ 0 & 1\end{pmatrix}$: one matrix exponential gives both.

Bilinear rule (Tustin; the choice of the original S4). Apply the trapezoidal rule to the state: $x_k - x_{k-1} = \frac{\Delta}{2}A(x_k + x_{k-1}) + \Delta B u_k$. Solve for $x_k$:

For a scalar with $z = a\Delta$: $\frac{1 + z/2}{1 - z/2} = e^z + \frac{z^3}{12} + O(z^4)$. So the bilinear $\bar a$ matches $e^{a\Delta}$ to second order. It also maps the left half-plane into the unit disk, because $|1 + z/2| < |1 - z/2|$ exactly when $\operatorname{Re} z < 0$.

The two limits of $\Delta$ (the key to Mamba). For a stable scalar, $a < 0$, with the ZOH:

• $\Delta \to 0$: $\bar a \to 1$ and $\bar b \to 0$. The state is kept and the input is ignored.
• $\Delta \to \infty$: $\bar a \to 0$ and $\bar b \to -b/a$. The old state is forgotten; the state jumps to the steady state $-bu_k/a$ of the current input.

Theorem: the recurrence is a convolution #

With $x_{-1} = 0$ and fixed $\bar A, \bar B, C$:

Proof.

1. Claim: $x_k = \sum_{j=0}^{k} \bar A^{\,j}\bar B\, u_{k-j}$. For $k = 0$: $x_0 = \bar A\cdot 0 + \bar B u_0$. ✓
2. Step: $x_k = \bar A x_{k-1} + \bar B u_k = \sum_{j=0}^{k-1}\bar A^{\,j+1}\bar B u_{k-1-j} + \bar A^{\,0}\bar B u_k$. Rename $j + 1 \to j$ in the sum: it becomes $\sum_{j=1}^{k}\bar A^{\,j}\bar B u_{k-j}$, and the last term is the $j = 0$ term. So $x_k = \sum_{j=0}^{k}\bar A^{\,j}\bar B u_{k-j}$.
3. Multiply by $C$: $y_k = \sum_j (C\bar A^{\,j}\bar B)\, u_{k-j}$. $\blacksquare$

The kernel depends only on the lag $j$, because $\bar A$, $\bar B$, $C$ are the same at every step. This property is called linear time invariance (LTI). The kernel $K$ is the impulse response: the output for the input $u = (1, 0, 0, \dots)$.

Computing it with the FFT. Direct summation costs $O(L^2)$. The discrete Fourier transform (DFT) turns a convolution into a product, but a length-$L$ DFT gives a circular convolution: the computed $y_0$ would also receive $K_1 u_{L-1} + K_2 u_{L-2} + \dots$, future inputs wrapped around to the start. The causal convolution is multiplication by a lower-triangular Toeplitz matrix; the length-$L$ DFT product is multiplication by a circulant matrix (linear-algebra lesson 19). The fix: pad $K$ and $u$ with zeros to length at least $2L - 1$, multiply the DFTs, invert, and keep the first $L$ outputs. The wrapped part then meets only zeros. Cost: $O(L\log L)$ per channel.

Interactive figure: loading…

Two modes of the same layer. For training, the whole sequence is known: compute $K$ once and convolve in parallel with the FFT. For generation, run the recurrence: $O(N)$ work per step for a diagonal $A$, and only the state $x_k$ is stored.

Stability by parametrization. If every eigenvalue $\lambda$ of $A$ has $\operatorname{Re}\lambda < 0$, then every eigenvalue $e^{\Delta\lambda}$ of $\bar A = e^{\Delta A}$ has $|e^{\Delta\lambda}| = e^{\Delta\operatorname{Re}\lambda} < 1$, for every $\Delta > 0$. So $\rho(\bar A) < 1$ and $\bar A^{\,j} \to 0$: the kernel decays. Models enforce this by storing the real part as $\operatorname{Re}\lambda = -\exp(\theta)$ with a free parameter $\theta$, so training can never make it positive.

S4: why $A$ needs structure #

Initialization matters. A random $A$ works badly: the state forgets quickly and keeps no useful summary. The S4 paper reports that changing $A$ from random to the HiPPO matrix took an earlier SSM on sequential MNIST from 60% to 98% accuracy. The HiPPO-LegS matrix is

HiPPO derives $A$ so that the state holds the coefficients of an approximation of the input history by Legendre polynomials: the state is an online, compressed summary of the whole past. (The time-invariant form used in S4 weights the past with exponentially warped Legendre polynomials; Gu et al. 2022, "How to Train Your HiPPO".)

The cost problem. Computing $K = (C\bar B, C\bar A\bar B, \dots)$ naively needs $L$ products with an $N\times N$ matrix: $O(N^2 L)$ work. A diagonal $A$ would make it cheap: then $K_j = \sum_n C_n \bar B_n \bar\lambda_n^{\,j}$, $O(NL)$ work. A change of basis $x = V\tilde x$ gives the same map, $(A, B, C) \to (V^{-1}AV, V^{-1}B, CV)$, so why not diagonalize HiPPO? Because its eigenvector matrix is extremely ill-conditioned (S4, Lemma 3.2): with the natural scaling, its entries grow like $2^{4N/3}$. In floating point the diagonalization is useless:

$N$	eigenvalues of HiPPO-LegS	condition number of the eigenvector matrix $V$
4	$-1, -2, -3, -4$	$9 \times 10^{1}$
8	$-1, \dots, -8$	$8 \times 10^{4}$
16	$-1, \dots, -16$	$8 \times 10^{10}$
24	$-1, \dots, -24$	$9 \times 10^{16}$ (the limit of float64)

The matrix is lower triangular, so its eigenvalues are the diagonal entries $-1, -2, \dots, -N$. The trouble is only in the eigenvectors, which are nearly parallel (linear-algebra lesson 05).

S4's answer: normal plus low rank

HiPPO-LegS can be written as $A = V\Lambda V^* - PQ^\top$ with $V$ unitary, $\Lambda$ diagonal and $P, Q$ of rank 1 (normal plus low rank, NPLR). A unitary change of basis is perfectly conditioned. S4 then computes the DFT of the kernel from the resolvent $(I - \bar A\zeta)^{-1}$ at the roots of unity $\zeta$. It removes the low-rank term with the Woodbury identity, reduces the diagonal part to a Cauchy kernel, and inverts the DFT with an FFT: $\tilde O(N + L)$ work.

Diagonal SSMs. DSS (Gupta et al. 2022) and S4D (Gu et al. 2022) simply drop the low-rank term and use a diagonal $A$. S4D proves that the normal part of HiPPO-LegS alone gives the same kernel as $N \to \infty$; the kernel computation is then two lines of code. Common initializations: S4D-Lin, $\lambda_n = -\frac12 + i\pi n$ (damped Fourier modes), and S4D-Real, $\lambda_n = -(n+1)$ (Mamba's default). The real part is stored as $-\exp(\cdot)$, so it stays negative.

Mamba: selection makes the model time-varying #

The limit of LTI. One fixed kernel weights each input by its distance in time, never by its content. Gu and Dao show the failure on selective copying: copy the coloured tokens, which sit at random positions between noise tokens. An S4 layer without gating reaches 18.3%; its selective version (S6) reaches 97.0%.

The selective SSM (S6). Per channel and per state index, the state follows $x_t = \bar A_t x_{t-1} + \bar B_t u_t$, but now the parameters depend on the current input. With $d$ channels and state size $N$ per channel:

• $B_t = \operatorname{Linear}_N(u_t)$ and $C_t = \operatorname{Linear}_N(u_t)$, shared by all channels;
• $\Delta_t = \operatorname{softplus}\big(\text{bias} + \operatorname{Linear}(u_t)\big) \in \mathbb{R}^d$, one step size per channel (a low-rank projection in practice); softplus is $\log(1 + e^z)$, a smooth positive function;
• $A \in \mathbb{R}^{d \times N}$ stays a parameter (diagonal for each channel), but $\bar A_t = \exp(\Delta_t A)$ now depends on $t$;
• per channel: $x_t = \bar A_t \odot x_{t-1} + \bar B_t u_t$, $\ y_t = C_t x_t + D u_t$.

The paper uses the ZOH. The reference code uses the ZOH for $A$ and Euler, $\bar B_t = \Delta_t B_t$, for $B$ ("$A$ is the more important term").

Theorem (Mamba, Theorem 1: selection generalizes gating). Take $N = 1$, $A = -1$, $B = 1$ and $\Delta_t = \operatorname{softplus}(z_t)$ with $z_t = \operatorname{Linear}(u_t)$. Then the ZOH recurrence is

where $s$ is the logistic sigmoid.

Proof.

1. $\bar A_t = e^{-\Delta_t} = e^{-\log(1 + e^{z_t})} = \frac{1}{1 + e^{z_t}} = 1 - s(z_t)$.
2. $\bar B_t = (\Delta_t A)^{-1}(e^{\Delta_t A} - 1)\,\Delta_t B = \frac{e^{-\Delta_t} - 1}{-1} = 1 - e^{-\Delta_t} = s(z_t)$. $\blacksquare$

This is the gate of the GRU in lesson 08, derived from discretization. A small $\Delta_t$ ignores the token and keeps the memory; a large $\Delta_t$ resets the state to the token. Resetting also stops separate documents, packed into one training sequence, from leaking into each other.

Interactive figure: loading…

Why the convolution view is lost. Unroll the recurrence (per channel and state index):

The weight of $u_s$ depends on $t$, on $s$ and on all inputs in between, not only on the lag $t - s$. There is no single kernel, so there is no FFT.

Theorem: the associative scan #

The recurrence $x_t = a_t x_{t-1} + b_t$ (elementwise; $a_t = \bar A_t$, $b_t = \bar B_t u_t$) is linear in $x$. Read $e_t = (a_t, b_t)$ as the affine map $x \mapsto a_t x + b_t$, and define "first $e$, then $e'$":

Then $e_0 \bullet e_1 \bullet \dots \bullet e_t$ maps $x_{-1}$ to $x_t$. With $x_{-1} = 0$, the state $x_t$ is its second component.

Check with numbers. $e_0 = (0.9, 1)$ and $e_1 = (0.5, 0)$. The loop gives $x_0 = 1$ and $x_1 = 0.5$. The operator gives $e_0 \bullet e_1 = (0.5 \cdot 0.9,\ 0.5 \cdot 1 + 0) = (0.45, 0.5)$, and its second component is $x_1 = 0.5$. ✓

Theorem. $\bullet$ is associative: $(e \bullet e') \bullet e'' = e \bullet (e' \bullet e'')$.

Proof.

1. Write $(a, b)$ as the matrix $M(a, b) = \begin{pmatrix} a & b \\ 0 & 1\end{pmatrix}$. It maps $(x, 1)$ to $(ax + b, 1)$.
2. "First $e$, then $e'$" is the product $M(e')M(e)$: $\begin{pmatrix} a' & b' \\ 0 & 1\end{pmatrix}\begin{pmatrix} a & b \\ 0 & 1\end{pmatrix} = \begin{pmatrix} a'a & a'b + b' \\ 0 & 1\end{pmatrix} = M(e \bullet e')$.
3. Matrix multiplication is associative, so $\bullet$ is associative. It is not commutative: order matters. $\blacksquare$

Parallel prefix. Associativity lets us regroup the long product into a tree, and the nodes of one level of the tree can be computed at the same time.

• Hillis–Steele: for $s = 1, 2, 4, \dots$, replace every $e_t$ with $t \ge s$ by $e_{t-s}\bullet e_t$, all $t$ at once. Depth $\lceil\log_2 L\rceil$ rounds, work $O(L\log L)$ combines.
• Blelloch: an up-sweep (a tree of totals) and a down-sweep (prefixes passed down the tree). Depth about $2\log_2 L$ rounds, work $O(L)$ combines.

Interactive figure: loading…

A scan needs each step to stay in a family of maps that is closed under composition. Affine maps are closed: composing two gives another one, stored as one pair $(a, b)$. The map $h \mapsto \tanh(Wh + Ux_t)$ is not closed: composing two tanh steps gives a new kind of function that no fixed-size pair can store. So a tanh RNN stays sequential. A diagonal $\bar A_t$ keeps the combine elementwise; a dense one would need $N\times N$ matrix products in every combine.

The Mamba block and the hardware-aware scan

The block. Project the input to two branches of width $Ed$ ($E = 2$). Branch 1: a short causal depthwise convolution (kernel 4), SiLU, the selective SSM. Branch 2: SiLU, used as a gate. Multiply the branches and project back to $d$. Most parameters sit in the projections, $3Ed^2$ per block, so two Mamba blocks match the $12d^2$ of one transformer layer (attention plus MLP).

The kernel. The expanded state has shape (batch, $L$, $d$, $N$), $N$ times larger than the input. Mamba's fused GPU kernel loads $\Delta$, $A$, $B$, $C$ into fast on-chip memory (SRAM), discretizes, scans and multiplies by $C$ there, and writes only $y$ back to the main GPU memory (HBM). This saves a factor of about $N$ in memory traffic; the paper reports 20–40 times faster than a plain scan. The backward pass recomputes the states instead of storing them. Long sequences are scanned in chunks, and the state is carried from chunk to chunk.

Costs #

Per layer, for length $L$, width $d$ and state size $N$ (the linear projections, $O(Ld^2)$, are common to all):

	Training	Per generated token	Stored for generation
Attention	$O(L^2 d)$	$O(Ld)$	KV cache, $2Ld$
S4 (LTI)	$O(dL\log L)$ with the FFT	$O(dN)$	state, $dN$
Mamba	$O(LdN)$ work, $O(\log L)$ depth with a parallel scan	$O(dN)$	state, $dN$ (plus a short convolution buffer)

The price of the fixed state: it must compress the whole past, so exact recall from a long context is weaker than with attention. Hybrid models keep a few attention layers; Jamba, for example, uses one attention layer for every seven Mamba layers.

Answers to section 1. (1) The SSM stores $dN = 16{,}384$ numbers per layer, whatever $L$ is. The KV cache stores $2Ld \approx 2 \times 10^9$, about 125,000 times more. (2) In general, no. A fixed kernel can only weight inputs by their distance, and here the distance is random. This is the selective-copying failure that motivated Mamba. (3) No. Tanh steps are not closed under composition, so there is no fixed-size combine operator; a tanh RNN must run step by step.

4. Worked example #

(a) Discretize $x' = -2x + u$, $y = x$ (so $a = -2$, $b = c = 1$) with $\Delta = 0.25$, so $a\Delta = -0.5$.

• ZOH: $\bar a = e^{-0.5} = 0.60653$ and $\bar b = \frac{e^{-0.5} - 1}{-2} = \frac{-0.39347}{-2} = 0.19673$.
• Bilinear: $\bar a = \frac{1 + (-0.5)/2}{1 - (-0.5)/2} = \frac{0.75}{1.25} = 0.6$ and $\bar b = \frac{0.25}{1.25} = 0.2$. Close to the ZOH.
• Euler for $b$ (Mamba's code): $\bar b = \Delta b = 0.25$, which is 27% larger than the ZOH value at this step.

(b) Kernel (ZOH). $K_j = c\,\bar a^{\,j}\,\bar b$: $K_0 = 0.19673$, $K_1 = 0.60653 \cdot 0.19673 = 0.11933$, $K_2 = 0.60653 \cdot 0.11933 = 0.07237$.

(c) Recurrence = convolution for the input $u = (1, 0, 2)$:

• recurrence: $x_0 = 0.19673 \cdot 1 = 0.19673$; $x_1 = 0.60653 \cdot 0.19673 + 0 = 0.11933$; $x_2 = 0.60653 \cdot 0.11933 + 0.19673 \cdot 2 = 0.07237 + 0.39347 = 0.46584$;
• convolution: $y_2 = K_0u_2 + K_1u_1 + K_2u_0 = 0.39347 + 0 + 0.07237 = 0.46584$. ✓

(d) Selection by hand (the Theorem 1 system). Logits $z = (2, -3, 0)$ and inputs $u = (1, 5, 2)$.

• $\Delta = \operatorname{softplus}(z) = (2.127, 0.0486, 0.693)$ and $g = s(z) = (0.881, 0.047, 0.5)$.
• $x_1 = (1 - 0.881)\cdot 0 + 0.881 \cdot 1 = 0.881$.
• $x_2 = 0.953 \cdot 0.881 + 0.047 \cdot 5 = 0.839 + 0.237 = 1.076$.
• $x_3 = 0.5 \cdot 1.076 + 0.5 \cdot 2 = 1.538$.
• The large input $u_2 = 5$ got a tiny $\Delta$ and was mostly ignored.
• The weights in $x_3$: $u_1$ gets $g_1(1 - g_2)(1 - g_3) = 0.420$, $u_2$ gets $g_2(1 - g_3) = 0.024$, $u_3$ gets $g_3 = 0.5$. An LTI model would give $u_1$ a weight that depends only on its lag, 2.

(e) Counting a scan for $L = 1024$: Hillis–Steele needs 10 rounds and $\sum_{k=0}^{9}(1024 - 2^k) = 9217$ combines; Blelloch about $2L$ combines in about 20 rounds; the loop 1023 combines in 1023 rounds.

5. In code #

SciPy does the discretization, the recurrence and the convolution, so you can check the worked example and the theorem. The notebook builds each piece from scratch.

import numpy as np
from scipy.signal import cont2discrete, dlsim, fftconvolve

# (a) x' = -2x + u, y = x, step 0.25
A, B, C, D = np.array([[-2.0]]), np.array([[1.0]]), np.array([[1.0]]), np.array([[0.0]])
for method in ["zoh", "bilinear"]:
    Ab, Bb, Cd, Dd, _ = cont2discrete((A, B, C, D), 0.25, method=method)
    print(method, Ab.item(), Bb.item(), Cd.item(), Dd.item())
# zoh 0.6065306597126334 0.1967346701436833 1.0 0.0
# bilinear 0.6 0.2 0.8 0.1     <- SciPy also changes C and D for 'bilinear'; S4 keeps C

# recurrence = convolution for a random stable system with N = 4
rng = np.random.default_rng(0)
N, L = 4, 256
A = -np.diag(np.arange(1.0, N + 1)) + 0.5 * rng.standard_normal((N, N))
B, C = rng.standard_normal((N, 1)), rng.standard_normal((1, N))
print(np.linalg.eigvals(A).real.max() < 0)                       # True: stable
u = rng.standard_normal(L)
impulse = np.zeros(L); impulse[0] = 1.0
for method in ["zoh", "bilinear"]:
    Ab, Bb, *_ = cont2discrete((A, B, C, D), 0.1, method=method)
    # S4: x_k = Ab x_{k-1} + Bb u_k, y_k = C x_k. dlsim uses s_{k+1} = Ab s_k + Bb u_k,
    # y_k = C' s_k + D' u_k, so take s_k = x_{k-1}, C' = C Ab and D' = C Bb.
    sys = (Ab, Bb, C @ Ab, C @ Bb, 0.1)
    y_rnn = dlsim(sys, u)[1].ravel()                  # recurrent mode
    K = dlsim(sys, impulse)[1].ravel()                # kernel = impulse response
    y_cnn = fftconvolve(u, K)[:L]                     # convolutional mode (padded FFT)
    print(method, np.abs(np.linalg.eigvals(Ab)).max() < 1, f"{np.abs(y_rnn - y_cnn).max():.0e}")
# zoh True 5e-17
# bilinear True 7e-17

Why the FFT needs padding, and why S4 does not just diagonalize HiPPO:

from scipy.linalg import circulant, toeplitz

L = 6
K = 0.6 ** np.arange(L)
T = toeplitz(K, np.zeros(L))                  # causal: y = T u uses u_0..u_k only
Cm = circulant(K)                             # a length-L DFT product multiplies by this
print(np.allclose(np.tril(Cm), T), np.abs(np.triu(Cm, 1)).max())   # True 0.6: future inputs leak in
print(np.allclose(circulant(np.r_[K, np.zeros(L)])[:L, :L], T))     # True: padding fixes it

def hippo_legs(N):
    n = np.arange(N)
    A = -np.sqrt((2 * n[:, None] + 1) * (2 * n[None, :] + 1))
    return np.tril(A, -1) - np.diag(n + 1.0)

for N in [4, 8, 16, 24]:
    lam, V = np.linalg.eig(hippo_legs(N))
    print(N, lam.real.max(), lam.real.min(), f"cond(V) = {np.linalg.cond(V):.0e}")
# 4 -1.0 -4.0 cond(V) = 9e+01
# 8 -1.0 -8.0 cond(V) = 8e+04
# 16 -1.0 -16.0 cond(V) = 8e+10
# 24 -1.0 -24.0 cond(V) = 9e+16

Mamba's Theorem 1 and the two limits of $\Delta$:

z = np.linspace(-6, 6, 13)
delta = np.logaddexp(0, z)                    # softplus(z) = log(1 + e^z)
A_bar = np.exp(-delta)                        # ZOH with A = -1
B_zoh = 1 - np.exp(-delta)                    # ZOH for B (B = 1)
g = 1 / (1 + np.exp(-z))
print(np.allclose(A_bar, 1 - g), np.allclose(B_zoh, g))     # True True
for i in [0, 6, 12]:                          # z = -6, 0, 6
    print(f"z={z[i]:+.0f} Delta={delta[i]:.4f} A_bar={A_bar[i]:.4f} B_zoh={B_zoh[i]:.4f}")
# z=-6 Delta=0.0025 A_bar=0.9975 B_zoh=0.0025   keep the state, ignore the token
# z=+0 Delta=0.6931 A_bar=0.5000 B_zoh=0.5000
# z=+6 Delta=6.0025 A_bar=0.0025 B_zoh=0.9975   reset the state to this token

6. Where it shows up #

• Long sequences. S4 and S4D for long-range benchmarks, raw audio (16,000-step speech; SaShiMi for audio generation) and medical time series. H3 and Hyena are convolutional relatives.
• Language and biology. Mamba for language, DNA and audio; in the paper, Mamba-3B matches transformers of twice its size. Mamba-2 (Dao and Gu 2024) restricts $\bar A_t$ to a scalar times the identity. This links the selective SSM to masked linear attention and lets it run on matrix multiplications. Hybrid language models such as Jamba mix Mamba, attention and mixture-of-experts layers.
• ML for science. Long, densely sampled records: seismograms, EEG and ECG, climate and weather station series, genomes. The continuous-time view fits physical data: $\Delta$ can follow the sampling rate, so the same weights can serve sensors with different rates. $x' = Ax + Bu$ is also the linear Gaussian state-space model of data assimilation (the Kalman filter), and a learned linear evolution in a lifted state is close to Koopman-operator models. A deep SSM is a natural backbone for long waveforms, or the temporal part of a spatio-temporal emulator.
• Practice. The official kernels are in the mamba_ssm package. mamba-minimal is a short PyTorch reference with a sequential loop, useful for reading.

7. Common mistakes #

• No zero padding in the FFT. A length-$L$ FFT gives a circular convolution: future inputs leak into early outputs. Pad to at least $2L - 1$.
• Inverting $A$ for the ZOH. $A^{-1}(e^{\Delta A} - I)B$ fails for a singular $A$; the integral (series or block exponential) always works.
• $\bar B = \Delta B$ as "the" discretization. This is Euler, accurate only for $\Delta\|A\| \ll 1$. The ZOH and bilinear $\bar B$ depend on $A$.
• Comparing with SciPy's bilinear $C$ and $D$. cont2discrete(..., method="bilinear") rescales $C$ and adds a $D$. S4's $\bar A$ and $\bar B$ match SciPy; its $C$ does not change.
• "Mamba makes $A$ input-dependent." $A$ is a parameter; $\bar A_t$ depends on the input only through $\Delta_t$.
• "An SSM trains sequentially because it is an RNN." LTI models train as convolutions, selective models with a parallel scan.
• "Any RNN can use a parallel scan." Only if each step lives in a family closed under composition, such as affine maps. Not a tanh RNN.
• Diagonalizing HiPPO in floating point. The eigenvectors are exponentially ill-conditioned. Use NPLR, or a diagonal initialization such as S4D-Lin or S4D-Real.
• Letting $\operatorname{Re}\lambda$ become positive. The kernel then grows exponentially. Parametrize $\operatorname{Re}\lambda = -\exp(\theta)$.

8. Check yourself #

1. The integrator $x' = u$ ($A = 0$, $B = 1$). Give the ZOH $\bar A$ and $\bar B$.
   Answer

   $\bar A = e^0 = 1$ and $\bar B = \int_0^\Delta 1\, d\tau = \Delta$. The closed form $A^{-1}(e^{\Delta A} - I)B$ is undefined, but its limit as $A \to 0$ is also $\Delta$.
2. $L = 8$, and you multiply length-8 FFTs of $K$ and $u$ without padding. What does the computed $y_0$ contain?
   Answer

   $K_0u_0 + K_1u_7 + K_2u_6 + \dots + K_7u_1$: the circular sum adds seven terms with future inputs. The correct $y_0$ is $K_0u_0$.
3. An S4D-Lin mode has $\lambda = -\frac12 + 3\pi i$ and $\Delta = 0.1$. What is $|e^{\Delta\lambda}|$, after how many steps has the mode decayed by half, and what is its period in steps?
   Answer

   $|e^{\Delta\lambda}| = e^{-0.05} = 0.951$. Half-life: $\ln 2/0.05 \approx 13.9$ steps. Period: $2\pi/(0.1 \cdot 3\pi) \approx 6.7$ steps.
4. In the Theorem 1 system, which $\Delta_t$ gives the gate $g_t = 0.9$?
   Answer

   $g = 1 - e^{-\Delta}$, so $\Delta = -\ln 0.1 = \ln 10 \approx 2.303$. It comes from $z = \ln 9$, because $\operatorname{softplus}(\ln 9) = \ln(1 + 9) = \ln 10$.
5. S4 is trained on speech at 16 kHz. How do you run it at 8 kHz without retraining?
   Answer

   Double $\Delta$. The samples are now twice as far apart in time, $u_k = u(k\Delta)$, and the continuous parameters $A$, $B$, $C$ stay the same.
6. Why can Mamba not compute its kernel once and use the FFT, as S4 does?
   Answer

   $\bar A_t$, $\bar B_t$, $C_t$ change with $t$, so the weight of $u_s$ in $y_t$ is $C_t(\prod_{r=s+1}^t \bar A_r)\bar B_s$. It depends on $(t, s)$ and on the inputs between them, not on $t - s$ alone. There is no shared kernel.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. Write the continuous SSM and its discrete recurrence (S4 convention).

   Show answer

2. Give the ZOH discretization of $(A, B)$ with step $\Delta$.

   Show answer

3. Compute the ZOH $\bar a$ and $\bar b$ for $x' = -2x + u$ with $\Delta = 0.25$.

   Show answer

4. What is the convolution kernel of an LTI SSM?

   Show answer

5. Why must you zero-pad before computing the SSM convolution with an FFT?

   Show answer

6. Why is an SSM with $\operatorname{Re}\lambda(A) < 0$ stable for every step size?

   Show answer

7. Why does S4 not simply diagonalize the HiPPO matrix?

   Show answer

8. What does "selection" mean in Mamba?

   Show answer

9. Give the scan operator for $x_t = a_tx_{t-1} + b_t$ and say why it allows a parallel scan.

   Show answer

10. What does an SSM store per layer to generate the next token, compared with attention?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Explain Mamba in two minutes." — A linear recurrence with a diagonal state per channel whose $\Delta$, $B$, $C$ come from the input (selection = learned gating), computed by a fused parallel scan; linear-time training and constant-memory generation.
• "S4, Mamba and attention: what are the costs?" — The table in section 3: $O(L^2d)$ against $O(dL\log L)$ against $O(LdN)$ for training; a KV cache of $2Ld$ against a state of $dN$ for generation.
• "Why does S4 need HiPPO and structure?" — A random $A$ forgets; HiPPO compresses the history; the NPLR or diagonal forms make the kernel cheap and stable to compute.
• "What is a parallel scan?" — The affine-map operator $(a,b)\bullet(a',b') = (a'a, a'b+b')$ is associative, so the prefixes form a tree of depth $O(\log L)$.
• "When would you still use attention?" — For exact retrieval and copying from a long context (in-context recall); hence hybrid models.

10. Practice #

The notebook for this lesson runs in Colab (NumPy and SciPy only). It has five exercises, each with a check cell and a folded solution:

1. Write the ZOH and bilinear discretizations and check them against scipy.signal.cont2discrete.
2. Compute a random stable diagonal SSM as a recurrence and as an FFT convolution, and show that they agree up to rounding error (below $10^{-12}$).
3. Write the scan operator, a sequential scan and a Hillis–Steele scan, and show that they give the same states.
4. Write a Mamba-style selective scan with $N = 4$ and check it against a naive loop and against the gate form of Theorem 1.
5. Build the HiPPO-LegS matrix and measure the condition number of its eigenvectors as $N$ grows.

11. Watch and read #

EPFL AI Center x CLAIRE - Structured State Space Models for Deep Sequence Modeling - Prof. Albert GuEPFL AI Center · Albert Gu · 61 min

The author of S4 and Mamba explains the three views (continuous, recurrent, convolutional), why structure is needed, and selection.

Mamba and S4 Explained: Architecture, Parallel Scan, Kernel Fusion, Recurrent, Convolution, MathUmar Jamil · 74 min

A step-by-step derivation of discretization, the convolution view, selection and the parallel scan.

Mamba: Linear-Time Sequence Modeling with Selective State Spaces (Paper Explained)Yannic Kilcher · 41 min

A reading of the Mamba paper: why selection is needed, and how the hardware-aware scan works.

• Gu, Goel, Ré, Efficiently Modeling Long Sequences with Structured State Spaces (S4): §2 (the three views, HiPPO, the bilinear rule, the kernel), §3 (why not diagonalize, NPLR, complexity), §4 (results, including the sampling-rate change).
• Rush and Karamcheti, The Annotated S4: the discrete SSM, the convolution, HiPPO and the diagonal case, in code.
• Gu, Gupta, Goel, Ré, On the Parameterization and Initialization of Diagonal State Space Models (S4D): S4D-Lin and S4D-Real.
• Gu and Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces: §3.1–3.3 (selection, the scan, the fused kernel), §3.5 (Theorem 1).
• Dao and Gu, Transformers are SSMs (Mamba-2): the link between selective SSMs and linear attention.
• Guy Blelloch, Prefix Sums and Their Applications: the work-efficient scan.
• mamba-minimal: a short, readable PyTorch Mamba (ZOH for $A$, Euler for $B$ in selective_scan).