S4: Efficiently Modeling Long Sequences with Structured State Spaces

A linear state space layer that keeps 64 numbers of state per channel and solved a 16,384-step task every earlier model failed.

The layer is a linear differential equation from control theory. The paper's contribution is an algorithm that builds a convolution kernel as long as the input sequence without ever raising the state matrix to a power.

Path-X hands a model 16,384 pixels one at a time and asks whether two circles in the picture are joined by a path. Every model benchmarked on it before S4 scored 50%, the rate of guessing. S4 scored 96.35%.

Path-X is the hardest task in the Long Range Arena benchmark (Tay et al., 2020). The input is a 128×128 black-and-white image with two small circles and a tangle of dashed curves, and the label says whether one curve connects the circles. A convolutional network handed the 2-D image would find this easy. The benchmark does not hand over the image: it flattens it row by row into a sequence of 16,384 pixel values and feeds the model one value per position, so any 2-D structure has to be recovered from the sequence.

Flattening makes the task long. Two pixels directly above and below each other in the picture sit 128 positions apart in the sequence, and the two circles can be thousands of positions apart. Drag the teal circle in the figure and the strip below the picture shows the gap in tokens; one step down in the picture costs a full row.

The paper by Gu, Goel and Ré is mostly a numerical linear algebra paper. It takes the state space model, which control engineers have used since the 1960s, observes that nobody could compute it fast enough to train as a deep learning layer, and spends most of its length on an algorithm that does. That algorithm, and the initialization it was built to make affordable, are the starting point of the state space family that includes Mamba and Mamba-2.

The page follows the paper's order: what a state space model computes, how a step size turns it into a recurrence, which state matrix makes it remember long histories, why that matrix defeats the standard speed-up, the decomposition and algorithm that get around it, and then the layer, its cost, and the results.

Why sequence length is the hard part

To relate a token to one 16,000 positions earlier, the two standard architectures each have a problem. Attention scores every position against every other, so any token can read any earlier token directly, at any distance. The cost is $L^2$ scores per head per layer. At $L = 16{,}384$ that is 268,435,456 scores, and a matrix of that size to store. FlashAttention later removed the storage by never materializing the matrix; the 268 million multiply-adds remain. The efficient-Transformer variants in Long Range Arena approximate the quadratic term with sparsity or low-rank structure. All of them scored 50.00 on Path-X, some because they ran out of memory or compute and some because they trained and did not learn the task.

A recurrent network carries a fixed-size hidden state forward and does a constant amount of work per step, so memory does not grow with length. For an input at step 1 to affect the output at step 16,000, its signal has to pass through 16,000 applications of the recurrent update, and the gradient back to it is a product of 16,000 Jacobians. A product that long either shrinks toward zero or blows up unless every factor stays very close to norm one, which is the vanishing and exploding gradient problem.

S4 is a recurrence too, but a linear one whose state matrix is chosen so that the decay over many steps is slow and controlled, and whose linearity lets the whole recurrence be computed as one convolution during training.

How a state space model works

The state space model is two equations:

$u(t)$ is a scalar input signal that changes over continuous time. $x(t)$ is a vector of $N$ numbers, the state. $y(t)$ is a scalar output. The first equation gives the rate of change of each state coordinate: a linear combination of the current state through the $N \times N$ matrix $\bm{A}$, plus the current input scaled by the $N \times 1$ vector $\bm{B}$. The second reads the output off the state with the $1 \times N$ vector $\bm{C}$, plus a direct copy of the input scaled by the number $\bm{D}$.

In control theory a state is a sufficient summary of the past: given $x(t)$ and the inputs after time $t$, the earlier inputs are never needed again. The $N$ numbers are therefore neither a window over recent inputs nor a cache of them. They are a fixed-size compression of the entire history, and the model can only use what survives the compression. S4 sets $N = 64$ for the Long Range Arena experiments.

$\bm{D}u$ bypasses the state, so the paper sets $\bm{D} = 0$ in the exposition and treats it as a skip connection. The released code adds it back as a trained per-channel parameter.

Two properties of (1) drive the rest of the paper. It is linear: the response to the sum of two inputs is the sum of the two responses, because nothing in the equations is a gate, a $\tanh$, or any other nonlinearity. And it is time-invariant: $\bm{A}, \bm{B}, \bm{C}$ do not depend on $t$ or on the input, so a spike at step 10 produces the same response as a spike at step 5,000, shifted by 4,990 steps. A system with both properties is called linear time-invariant, or LTI.

An LTI system computes a convolution. Write the input as a sum of single spikes, one per time step, each scaled by the input value there. By linearity the output is the sum of the responses to the individual spikes. By time-invariance each of those responses is one fixed shape, shifted to the spike's position and scaled by its value. Adding up shifted, scaled copies of one shape is convolution with that shape, which signal processing calls the impulse response.

To compute it, turn the ODE into a recurrence over discrete steps (the next section derives $\bm{\overline{A}}, \bm{\overline{B}}, \bm{\overline{C}}$) and unroll it from a zero initial state:

The bars mark discretized quantities. Each entry $\bm{\overline{C}}\bm{\overline{A}}^{i}\bm{\overline{B}}$ is a $1 \times N$ row times an $N \times N$ matrix power times an $N \times 1$ column, so it is a single number, and $\bm{\overline{K}}$ is a filter of $L$ numbers, as long as the sequence. Feed in a 1 followed by zeros and the output is $\bm{\overline{K}}$ itself.

No approximation happens between (3) and (4). With the initial state at zero, which the paper assumes and the code does, the recurrence and the convolution compute the same function. S4 uses both: during training the whole sequence is available, so it computes all $L$ outputs at once with the convolution; during generation the next input does not exist yet, so it runs the recurrence at a fixed cost per token. The figure runs a two-state system both ways. Switch between the recurrent and convolutional views and compare the output rows digit by digit.

The smallest version can be computed by hand, with the HiPPO matrix introduced two sections below. Take $N = 2$, step size $\Delta = 1$, and a readout that takes the first state coordinate:

The discretization formula of the next section gives

Since $\bm{C} = [1\ 0]$, the output is the first coordinate, whose update is $x_k = \tfrac13 x_{k-1} + \tfrac23 u_k$, so the kernel is the geometric sequence $\bm{\overline{K}}_i = \tfrac23(\tfrac13)^i$:

Feed it the pulse $u = (1,1,1,1,0,0,0,0)$. Unrolling the recurrence, convolving directly, or convolving by FFT all give

While the pulse is on, the output climbs toward 1; after it ends, the output falls by a factor of 3 per step. This is a leaky running average of recent input. At $N = 64$ and $L = 16{,}384$, computing $\bm{\overline{K}}$ this way would mean 16,383 multiplications by a 64×64 matrix, which is the problem the rest of the paper solves.

Discretization and the step size Δ

Equation (1) is written in continuous time, and a sequence model receives $u_0, u_1, u_2, \dots$. The bridge is one number, the step size $\Delta$, which says how much continuous time passes between consecutive tokens: $u_k = u(k\Delta)$. The data does not determine $\Delta$, so S4 learns it, one value per channel.

The simplest discretization is forward Euler, which replaces the derivative with a finite difference:

S4 uses the bilinear transform (Tustin, 1947) instead, which averages the derivative at both ends of the step, as the trapezoidal rule does:

This display carries the same number as the recurrence above because the paper prints them together as one equation group. The factor $\bm{I}+\tfrac{\Delta}{2}\bm{A}$ is a forward half-step, and the inverse $(\bm{I}-\tfrac{\Delta}{2}\bm{A})^{-1}$ is a backward half-step.

Why not Euler? Both discretizations are rational functions of $\bm{A}$, so each eigenvalue $\lambda$ of $\bm{A}$ maps to an eigenvalue of $\bm{\overline{A}}$ by the same scalar formula:

A continuous system is stable when every eigenvalue has negative real part. A discrete recurrence is stable when every eigenvalue of $\bm{\overline{A}}$ has magnitude below 1; otherwise that mode is multiplied by more than 1 at every step and grows without bound. The bilinear map sends the whole left half-plane into the open unit disk, for every $\Delta > 0$. Euler sends $\lambda = -n$ to $1 - \Delta n$, which leaves the disk once $\Delta > 2/n$.

The HiPPO matrix of the next section has eigenvalues exactly $-1, -2, \dots, -N$. At $N = 64$ and $\Delta = 0.1$, the top of S4's initialization range, Euler sends $\lambda = -64$ to $1 - 6.4 = -5.4$; that mode is multiplied by 5.4 per token, about $10^{732}$ over 1,000 tokens. The bilinear transform sends it to $(1-3.2)/(1+3.2) = -0.524$. At the small end, $\lambda = -1$ with $\Delta = 0.1$ goes to 0.9048. In the figure, the shaded band is the stable region. Sweep $\Delta$ and switch $N$: the amber Euler points fall out of the band once $\Delta$ passes $2/N$, while the teal bilinear points stay inside at every setting.

Besides stability, $\Delta$ sets how far back the layer looks. The slowest mode of the HiPPO matrix has $\lambda = -1$, and the bilinear map sends it to approximately $e^{-\Delta}$, so it decays by a factor of $e$ over about $1/\Delta$ tokens. The released code initializes one $\Delta$ per channel, drawn log-uniformly from $[0.001,\ 0.1]$, which spreads those horizons from about 10 to about 1,000 tokens across the channels of a single layer.

The continuous-time view also explains a Speech Commands result. The dataset is recorded at 16 kHz. An S4 model trained at that rate scores 96.30% on test audio resampled to 8 kHz, with no retraining, after its $\Delta$ is doubled. Halving the sample rate doubles the real time between samples, and doubling $\Delta$ makes the discretization use that spacing, while the learned $\bm{A}, \bm{B}, \bm{C}$ describe the same continuous system as before. The rescale is manual: someone has to know the new sample rate and change $\Delta$.

The HiPPO matrix

Everything so far holds for any $\bm{A}$, and most choices of $\bm{A}$ train badly. The LSSL, the paper's immediate predecessor, reports 62% on permuted MNIST with a random, suitably normalized $\bm{A}$, against 98.76% with the matrix below.

That matrix comes from HiPPO (Gu et al., 2020), which asked a different question: with a budget of $N$ numbers, what is the best summary of a signal's history, and can it be updated online as new input arrives? HiPPO's answer is a running least-squares polynomial approximation. Keep the $N$ coefficients of the best polynomial approximation of the history so far, in an orthogonal polynomial basis. HiPPO proves that those coefficients obey a linear ODE, so they can be updated by integrating the ODE instead of refitting. For the scaled Legendre measure (LegS) the matrix is:

The square roots normalize the Legendre polynomials. The lower-triangular shape means coefficient $n$ is driven by the lower-degree coefficients $0, \dots, n-1$. The diagonal entry $-(n+1)$ is the decay rate of coefficient $n$, so higher-degree coefficients, which carry finer detail, decay faster. Because the matrix is triangular its eigenvalues are its diagonal entries, $-1, -2, \dots, -N$, the eigenvalues used in Figure 3.

S4 changes one thing about how HiPPO uses this matrix, and the change alters what the state remembers. HiPPO's LegS system is time-varying. Written with the $\bm{A}$ and $\bm{B}$ of (2), it is

where the $1/t$ stretches the approximation window as time passes, so the coefficients always summarize the whole history from 0 to $t$ with uniform weight. S4 keeps the matrix and drops the $1/t$, which gives back equation (1) with a constant $\bm{A}$: the system becomes time-invariant, which is what the convolution requires. Run equation (2) as a fixed linear system and its impulse response has a closed form,

where $P_n$ is the degree-$n$ Legendre polynomial. Its argument is $e^{-t}$, not $t$: the time-invariant system fits polynomials to the past in a warped time variable in which recent history is stretched out and receives most of the resolution, and distant history is squeezed toward zero. The memory fades smoothly with age instead of weighting all of history equally. The S4 paper does not discuss the change; the closed form was worked out in the follow-up paper How to Train Your HiPPO.

The figure draws the input history on that warped axis, with "now" at the right edge and older input compressed toward the left. Raise $N$ from 1 and watch which part of the history the reconstruction resolves first.

The matrix also never amplifies the state, even briefly. Let $\bm{P}_n = (2n+1)^{1/2}$, the same vector as $\bm{B}$. The symmetric part of (2) is

and a symmetric part bounded by $-\tfrac12\bm{I}$ gives $\lVert e^{t\bm{A}}\rVert_2 \le e^{-t/2}$ for every $t \ge 0$. The norm of the state falls monotonically from any starting point. Many non-normal matrices first amplify the state by a large factor before it decays, which wrecks a recurrence in floating point; this one cannot. The vector $\bm{P}$ returns two sections below as the low-rank part of the decomposition.

Why diagonalizing the HiPPO matrix fails

Building $\bm{\overline{K}}$ from (5) directly takes $L$ successive multiplications by an $N \times N$ matrix: $O(N^2L)$ operations, and $O(NL)$ memory to keep the intermediate states for backpropagation. The LSSL used $N = 256$, and at that size it needed orders of magnitude more memory than an RNN or CNN of similar size, which kept it from being run at the sizes it needed.

The textbook speed-up is diagonalization. If $\bm{A} = \bm{V}\bm{\Lambda}\bm{V}^{-1}$ with $\bm{\Lambda}$ diagonal, then $\bm{A}^i = \bm{V}\bm{\Lambda}^i\bm{V}^{-1}$, and a power of a diagonal matrix is $N$ scalar powers. The change of basis costs nothing in the model. The paper's Lemma 3.1 states that

compute the same map from $u$ to $y$, because in $\bm{C}\bm{V}(\bm{V}^{-1}\bm{A}\bm{V})^i\bm{V}^{-1}\bm{B}$ every $\bm{V}\bm{V}^{-1}$ pair cancels.

Equation (2) is diagonalizable. Its eigenvalues $-1, \dots, -N$ are real and distinct, distinct eigenvalues guarantee a full set of independent eigenvectors, and Lemma 3.2 even gives the eigenvector matrix in closed form, $\bm{V}_{ij} = \binom{i+j}{i-j}$.

In floating point it fails anyway. The eigenvectors are independent but nearly parallel, so $\bm{V}$ is badly ill-conditioned: writing an ordinary vector in that basis takes huge coordinates of opposite sign that nearly cancel, and float64 keeps about 16 significant digits, so the cancellation destroys the result. Lemma 3.2 bounds the entries of $\bm{V}$ by $2^{4N/3}$, about $4.9 \times 10^{25}$ at $N = 64$. The condition number of $\bm{V}$, which we computed in 260-digit arithmetic because float64 cannot measure it, is about $7.7\times10^4$ at $N=8$, $8.3\times10^{10}$ at $N=16$, $1.2\times10^{23}$ at $N=32$, and $3.2\times10^{47}$ at $N=64$. A condition number of $3.2\times10^{47}$ means about 158 bits of precision are needed to diagonalize at S4's size, and float64 has 53.

Nothing overflows, the matrix is not defective, and the eigenvalues stay exact at every $N$; only the computed eigenvectors lose relative precision. The figure's meter shows it directly: it takes the left matrix apart into its exactly known eigenvalues and eigenvectors, multiplies them back together in float64, and reports the relative error. Move $N$ up and watch the error pass 100% at about $N = 26$.

The matrices that diagonalize safely form a known class. A matrix is normal when it commutes with its conjugate transpose, $\bm{A}\bm{A}^{*} = \bm{A}^{*}\bm{A}$. The spectral theorem says a matrix can be diagonalized by a unitary $\bm{V}$ (a rotation, with condition number 1, so nothing is lost to rounding) exactly when it is normal. Symmetric and skew-symmetric matrices are normal; equation (2) is not, and the figure's "not normal" readout measures by how much.

Normal plus low-rank

Requiring $\bm{A}$ to be normal would rule out the HiPPO matrix. The paper asks for less: a normal matrix plus a low-rank correction. It calls this form NPLR, normal plus low-rank. After the unitary change of basis that diagonalizes the normal part it becomes DPLR, diagonal plus low-rank, which is the form S4 stores:

For the HiPPO matrix the correction has rank one, and the decomposition is exact. With $\bm{P}_n = (2n+1)^{1/2}$, add $\tfrac12\bm{P}\bm{P}^{\top}$ to equation (2) and check the three kinds of entry. Above the diagonal, where $\bm{A}$ is zero, the sum is half the outer-product entry, $\tfrac12\bm{P}_n\bm{P}_k$. Below the diagonal it is

which is the mirror entry with the opposite sign. On the diagonal it is $-(n+1) + \tfrac12(2n+1) = -\tfrac12$ in every row. So

A skew-symmetric matrix has purely imaginary eigenvalues and is normal, and adding a multiple of the identity keeps it normal and shifts every eigenvalue by that multiple. Every eigenvalue of the corrected matrix is therefore $-\tfrac12 + i\omega$ for some real $\omega$. At $N = 8$ they are $-0.5 \pm 19.857i$, $-0.5 \pm 5.354i$, $-0.5 \pm 1.958i$ and $-0.5 \pm 0.427i$. All modes of the normal part decay at the same rate and differ only in how fast they oscillate. In the DPLR basis, $\bm{\Lambda}$ holds these eigenvalues on its diagonal, and the original eigenvalues $-1, \dots, -N$ are recovered only once the rank-one term is subtracted back.

The trainable parameters of one S4 channel are $\bm{\Lambda}, \bm{P}, \bm{Q}, \bm{B}, \bm{C}$, five complex vectors of length $N$, which the paper counts as $5N$ parameters. The released code differs in ways that matter when reading it: $\bm{Q}$ is tied to $\bm{P}$ (the $\bm{\Lambda} - \bm{P}\bm{P}^{*}$ form from the SaShiMi paper), only one member of each complex conjugate pair is stored, and $\Delta$ and $\bm{D}$ are trained as well, outside the $5N$. The kernel $\bm{\overline{K}}$ is never a parameter; it is recomputed from these vectors on every forward pass. At Path-X scale, with $H = 256$ channels, roughly 66,000 stored numbers generate a $256 \times 16{,}384$ filter bank, 4,194,304 taps.

How S4 computes the kernel

Storing $\bm{A} = \bm{\Lambda} - \bm{P}\bm{Q}^{*}$ does not by itself make (5) cheap, because powers of a diagonal-plus-rank-one matrix do not stay diagonal plus rank one. Expand the square:

The part that is not $\bm{\Lambda}^2$ has columns in the span of $\bm{P}$ and $\bm{\Lambda}\bm{P}$, so it has rank two. In general the off-diagonal part of the $k$-th power lives in the span of $\bm{P}, \bm{\Lambda}\bm{P}, \dots, \bm{\Lambda}^{k-1}\bm{P}$, so its rank grows by one per multiplication until it reaches $N$ and the matrix is effectively dense again. The paper puts it as "powering up this sum is still slow".

Inverses behave the opposite way. The Woodbury identity says that a rank-$p$ change to a matrix changes its inverse by rank $p$, and the only new inverse needed is $p \times p$:

At $p = 1$ the inner inverse is one scalar division. The figure computes both in float64 on an 8×8 example, $\bm{M} = \bm{\Lambda} - \bm{P}\bm{Q}^{*}$ with a real scalar $s$: singular values of the off-diagonal part of $\bm{M}^k$ on the left, and of the off-diagonal part of $(s\bm{I} - \bm{M})^{-1}$ on the right. Step $k$ up and one more singular value rises above the rounding floor each time; move $s$ and the right panel keeps a single bar.

So S4 rewrites the kernel in terms of an inverse, using a generating function: put the $L$ kernel values into the coefficients of a polynomial in $z$ and evaluate the polynomial instead of computing the coefficients:

The right side is the finite geometric series, $\textstyle\sum_{i<L} r^i$ equal to $(1-r^L)/(1-r)$, with the matrix $\bm{\overline{A}}z$ in place of $r$, valid wherever the inverse exists. The $L$ powers of $\bm{\overline{A}}$ have become one inverse and one fixed power $\bm{\overline{A}}^L$. Two choices make the right side cheap to evaluate.

First, evaluate at the $L$-th roots of unity, $z_k = e^{-2\pi i k/L}$. At those points the values of $\hat{\mathcal{K}}_L$ are exactly the discrete Fourier transform of $\bm{\overline{K}}$, so one inverse FFT recovers the $L$ kernel values with nothing lost. Also at those points $z^L = 1$, so the factor $\bm{I}-\bm{\overline{A}}^{L}z^{L}$ becomes the constant $\bm{I}-\bm{\overline{A}}^{L}$ and can be absorbed into the output vector once:

S4 learns $\bm{\tilde{C}}$ directly and never computes $\bm{\overline{A}}^L$.

Second, expand the remaining inverse with Woodbury. With $\bm{A} = \bm{\Lambda} - \bm{P}\bm{Q}^{*}$, the bilinear $\bm{\overline{A}}$, and the substitution

the generating function at each root of unity reduces to

All four bracketed quantities have the form $\bm{v}^{*}\bm{R}\bm{w}$, and since $\bm{\Lambda}$ is diagonal, each one is a sum of $N$ divisions:

Evaluating such a sum at $L$ points is multiplication by a Cauchy matrix, the matrix with entries $1/(s_i - \lambda_j)$, and numerical analysis has fast, stable algorithms for it. That is Theorem 3: the kernel reduces to four Cauchy multiplies, $\tilde{O}(N+L)$ operations and $O(N+L)$ memory, against $O(N^2L)$ and $O(NL)$ for the direct computation.

The divisions are also numerically safe. For $z = e^{-i\theta}$ on the unit circle, the node is $s = \tfrac{2i}{\Delta}\tan(\theta/2)$, purely imaginary, and every $\lambda_j$ has real part $-\tfrac12$. So $\lvert s - \lambda_j\rvert \ge \tfrac12$ and no denominator comes near zero.

The whole algorithm, in NumPy:

# Algorithm 1: the length-L kernel from N diagonal poles plus a rank-one term.
# Lam, P, B, Ct are complex vectors of length N. Ct is the learned C-tilde.

z = np.exp(-2j * np.pi * np.arange(L) / L)     # L-th roots of unity (minus sign)
s = (2 / dt) * (1 - z) / (1 + z)               # bilinear image: purely imaginary

def cauchy(v):                                 # sum_j v_j / (s - Lam_j), every s
    return (v[None, :] / (s[:, None] - Lam[None, :])).sum(-1)

k00 = cauchy(Ct.conj() * B)                    # the four Cauchy sums
k01 = cauchy(Ct.conj() * P)
k10 = cauchy(P.conj() * B)                     # released code ties Q = P
k11 = cauchy(P.conj() * P)

Khat = 2 / (1 + z) * (k00 - k01 / (1 + k11) * k10)   # Woodbury: one divide
K = np.fft.ifft(Khat).real                     # length-L kernel, time domain

Algorithm 1 prints the roots of unity as $e^{+2\pi i k/L}$; paired with its own inverse FFT that returns the kernel reversed in time, a filter that looks at the future. Lemma C.2 and the released code use $e^{-2\pi i k/L}$, as above. And when $L$ is even, one node is $z = -1$, where $2/(1+z)$ divides by zero. The singularity is removable: $s$ goes to infinity there, the resolvent shrinks at the same rate, and the limit is $\bm{\tilde{C}}^{*}(\bm{I}+\bm{\overline{A}})^{-1}\bm{\overline{B}}$, so that node keeps a finite value.

The $\tilde{O}(N+L)$ of Theorem 3 assumes a fast-multipole Cauchy algorithm. Appendix A says those algorithms were not implemented for GPUs, and the release computes the Cauchy sums directly in $O(NL)$ time, with $O(N+L)$ memory because the $N \times L$ matrix is never stored. At $N = 64$ that is a modest cost, and it is paid once per batch, since the kernel does not depend on the input.

The figure runs the listing above in the browser, on the HiPPO initialization at $N = 64$. The left panel is the resulting kernel; the right panel shows the 64 diagonal entries of $\bm{\Lambda}$ after the bilinear map, the poles of the four Cauchy sums. Drag $\Delta$ across the initialization range and past both ends. At small $\Delta$ the kernel stretches over more than a thousand taps and the poles bunch near $z = 1$; at large $\Delta$ they spread around the circle and the kernel becomes a short, jagged burst. The poles stay strictly inside the circle at every $\Delta$.

The recurrent view needs the same structure. Appendix C.2 writes the bilinear $\bm{\overline{A}}$ as a product of two factors: the forward half-step $\bm{I} + \tfrac{\Delta}{2}\bm{A}$, which is diagonal plus rank one, and the backward half-step $(\bm{I} - \tfrac{\Delta}{2}\bm{A})^{-1}$, which Woodbury also turns into diagonal plus rank one. Multiplying a vector by each factor costs $O(N)$, so one generation step costs $O(N)$ (Theorem 2). The readout during generation has to use $\bm{\overline{C}}$, not the learned $\bm{\tilde{C}}$, which contains the $(\bm{I}-\bm{\overline{A}}^L)$ factor meant only for the convolution. Recover it once by solving the definition of $\bm{\tilde{C}}$ for $\bm{\overline{C}}$, an $N \times N$ linear system. The factor is far from the identity: at $\Delta = 0.001$, $L = 1024$, $N = 64$, the norm of $\bm{\overline{A}}^L$ is about 0.60.

# generation: the same parameters, run as a recurrence
AbarL = matrix_power(Abar, L)          # once, before generating
C = solve((I - AbarL).conj().T, Ct)   # undo C-tilde = (I - Abar^L)^* C
x = zeros(N, complex)                 # the whole memory of the past: N numbers
for u_k in stream:
    x = Abar @ x + Bbar * u_k         # O(N): Abar is a product of two DPLR factors
    y_k = (C.conj() @ x).real + D * u_k

The S4 layer and its cost

One SSM maps a scalar sequence to a scalar sequence, and a network layer has to transform $H$ features per position. S4 runs $H$ independent SSMs, one per channel, each with its own parameters and its own $\Delta$; mixes the channels with a position-wise linear layer; and applies a nonlinearity. This is the structure of a depthwise-separable convolution, one filter per channel followed by a $1 \times 1$ convolution across channels, and the code builds the mixer as a kernel-size-1 Conv1d.

S4's per-channel filters differ from a CNN's in three ways. They are as long as the sequence. They are generated from about $5N$ numbers per channel instead of stored tap by tap, so the parameter count does not grow with $L$, and the same parameters can produce the filter at a different length or a different $\Delta$. And each filter has an equivalent recurrence, which a CNN filter does not.

The state update has no gate. In $x_k = \bm{\overline{A}}x_{k-1} + \bm{\overline{B}}u_k$ nothing depends on the input, so one fixed kernel can represent the layer. All of a deep S4 network's nonlinearity sits between SSM layers. Later S4 blocks add a GLU after the channel mixer, still outside the SSM. Mamba removes exactly this property: it makes the SSM's parameters depend on the input.

The convolution runs through FFTs, and the FFT computes a circular convolution, in which the end of the output wraps around and adds onto the beginning. Take $K = (1, 2)$ and $u = (3, 4)$. The linear convolution is $(3, 10, 8)$. The length-2 circular convolution is $(11, 10)$: the 8 wrapped around onto the 3. Zero-padding both signals to at least $2L - 1$ before the transform and keeping the first $L$ outputs gives the linear result; the code pads to $2L$. Without the padding, each early output picks up contributions from inputs near the end of the sequence, an error as large as the signal itself. Here is one layer on a Path-X batch, with shapes:

# one S4 layer on a Path-X batch: B=32 sequences, H=256 channels, L=16384
u   = x.transpose(1, 2)               # (32, 256, 16384)  one scalar signal per channel
K   = kernel(Lam, P, B, Ct, dt, L)    # (256, 16384)      one length-L filter per channel
U_f = rfft(u, n=2*L)                  # zero-pad to 2L: FFT convolution is circular
K_f = rfft(K, n=2*L)
y   = irfft(U_f * K_f, n=2*L)[..., :L] + D[:, None] * u   # (32, 256, 16384)
y   = gelu(y)
out = W_mix @ y                       # (256 x 256) position-wise mix, a 1x1 conv

The gradient flows through every line of that listing and through Algorithm 1 into $\bm{\Lambda}, \bm{P}, \bm{B}, \bm{\tilde{C}}$ and $\Delta$, plus $\bm{D}$ and the mixer weights.

On cost: attention grows with the square of the length. The naive SSM was already linear in $L$, so linearity was not S4's contribution. S4 removed the $N^2$ factor on the work and the $NL$ of stored states. Slide the length in the figure and the three curves separate; switch to memory to see the naive SSM's problem.

Two qualifications on that chart. Training compute for S4 is $O(L \log L)$ because of the FFT; only memory and the per-token generation cost are linear and constant. And the paper's complexity table (Table 1) leaves out the channel mixer, which costs $\Theta(BLH^2)$ and is the larger of S4's two terms at realistic widths. At $B = 32$, $H = 256$, $L = 16{,}384$:

Against the LSSL at layer width 512, S4 takes 29.6× less time per training step and 392× less memory (the introduction rounds these to 30× and 400×). Against efficient Transformers, parameter-matched at length 4,096 on one A100, S4 runs 5.19× as fast as a vanilla Transformer on 9.1% of its memory, in the same range as the Performer (3.79×, 8.6%) and the Linear Transformer (5.35×, 6.7%).

Results

On Long Range Arena, S4 has the best score on all six tasks and averages 86.09%; every baseline averages under 60%, the best being Luna-256 at 59.37%. On Path-X the baselines are at 50.00, the score of a coin flip, and S4 is at 96.35. Pick a task in the figure to see the full column.

The paper is inconsistent about the 86.09. Section 4.2 still says 80.48%, the average from the first arXiv version; Table 4 reports 86.09; and the appendix table prints both rows. Appendix D.5 explains that the table was refreshed with results from the S4D and How to Train Your HiPPO follow-ups, for the same model with minor hyperparameter changes, and asks readers to read it before citing. The Path-X change behind 88.10 to 96.35 is the $\Delta$ range, lowered from $[0.001, 0.1]$ to $[0.0001, 0.01]$. By the horizon estimate above, the slowest mode at $\Delta = 0.0001$ keeps $e^{-1.64} \approx 0.19$ of its value across all 16,384 tokens, while at $\Delta = 0.001$ it keeps about $7.7 \times 10^{-8}$. That calculation is our reading of why the change helps; the paper reports only the change and the score.

The other experiments test generality across data types:

• Raw speech (Speech Commands, 10 classes, 16,000-sample waveforms with no feature extraction): 98.32%, higher than every baseline that was given MFCC features, which are 100× shorter. A specialized audio CNN, the paper's improved WaveGAN discriminator, gets 96.25% with 26.3M parameters, almost 90× S4's 0.3M.
• Sequential CIFAR-10 (one pixel at a time, no 2-D structure): 91.13%, competitive with a 2-D ResNet-18 that has more parameters (11.0M against S4's 7.9M).
• CIFAR-10 density estimation: 2.85 bits per dimension with no 2-D bias, equal to PixelSNAIL, which uses 2-D convolutions and attention.
• WikiText-103 language modeling: 20.95 test perplexity at 249M parameters, against 20.51 for a 247M Transformer. S4 does not win here, and the paper describes the result as closing the gap. It is the best attention-free result by more than 2 perplexity (TaLK Convolution is next at 23.3). The prose says "within 0.8 ppl", which refers to an earlier run at 21.28; the retrained model's gap is 0.44.

The recurrent view makes generation fast. A Transformer producing token $k$ attends over all $k$ earlier tokens, so each token costs more than the last; S4 updates an $N$-number state per channel at the same cost for every token. On WikiText-103 that gives 48,000 tokens per second against 800, a 60× speedup. On CIFAR the speedup ranges from 65.1× for the base model down to 10.5× for the large 2.85-bpd model. These are throughputs at the largest batch that fits on one A100, not single-sequence latency.

The ablation, S4D and Mamba

Section 4.4 asks which part of S4 produces the results, on sequential CIFAR-10 with models of at most 100K parameters. Random matrices in S4's NPLR form perform poorly. Training $\bm{A}$ improves every initialization, and with $\bm{A}$ trained every initialization reaches perfect training accuracy, yet their validation accuracies differ by more than 15 points. The models share a parameter count and an optimizer, so the starting matrix alone moves validation accuracy by that much. With 0.1 dropout added, the same trends hold, and full S4 with the HiPPO initialization reaches 84.27% test accuracy with 100K parameters. The paper concludes that the gains come mainly from the HiPPO initialization, not from the NPLR parameterization; NPLR makes that initialization affordable to compute.

Follow-up work simplified the parameterization within months. DSS (Gupta et al., 2022) and then S4D (Gu et al., 2022) dropped the low-rank term and kept only a diagonal $\bm{\Lambda}$ initialized from the HiPPO spectrum, and landed within about a point of S4 on the LRA average. A diagonal state matrix turns the kernel into a Vandermonde product and removes the Woodbury correction, the Cauchy sums, and most of Appendix C. The simplification costs something on the longest task: S4D-LegS averages 84.89 against S4's 86.09 and trails by 4.4 points on Path-X (91.95 against 96.35), and S4D-Lin, whose eigenvalues have linearly spaced imaginary parts, fails Path-X outright.

The larger change went the other way. Everything in this paper depends on time-invariance: the same $\bm{A}, \bm{B}, \bm{C}$ at every position, so one kernel represents the layer and one FFT computes it. A filter that applies the same weights at every position cannot down-weight one token or retain another based on its content, which is a plausible reason S4 trails Transformers on language. Mamba makes $\bm{B}$, $\bm{C}$ and $\Delta$ functions of the input. With parameters that change at every token there is no single kernel to FFT, so Mamba replaces the convolution with a hardware-aware parallel scan, and it discretizes $\bm{\overline{A}}$ with the exact exponential $e^{\Delta\bm{A}}$ instead of the bilinear transform. Mamba-2 restricts the state matrix to a scalar times the identity so the scan can be computed with matrix multiplications, and Mamba-3 brings back a trapezoid-rule discretization of the input term.