S4: Efficiently Modeling Long Sequences with Structured State Spaces

A long memory needs more than a long loop

A recurrent model can carry a state through sixteen thousand inputs without constructing a sixteen-thousand-by-sixteen-thousand attention matrix. That does not mean its state will retain useful information from the beginning. A convolution can process an entire sequence in parallel, but an arbitrary long filter needs a way to learn and construct its coefficients. S4 connects these two problems: design a compact dynamical memory, then compute its long convolution kernel efficiently.

The starting point is a linear state-space model. The state summarizes the past; the readout chooses what to expose. During full-sequence training the same linear operator can be evaluated as a convolution. During streaming generation it can be evaluated recurrently. The central contribution is the structured parameterization and kernel algorithm that make moving between these views practical, rather than the recurrence/convolution identity alone. See the paper, §§2–3.

Albert Gu, Karan Goel and Christopher Ré first posted the work in 2021; it appeared at ICLR 2022. This review follows arXiv v3, August 2022, whose appendix distinguishes original and updated results. The source walkthrough uses the official repository pinned to e757cef. That later implementation contains several SSM variants, so a current repository default should not be assumed to be the original S4 experiment.

HiPPO gives the state something useful to remember

Suppose a signal's entire history must fit into a fixed number of coordinates. Keeping only the newest samples discards the distant past. Keeping one running average preserves a broad trend but loses shape. HiPPO instead motivates representing history through coefficients in an orthogonal polynomial basis: low-order coordinates summarize broad structure, while higher-order coordinates represent finer variation. The number of coordinates controls the resolution of the memory, not the number of timesteps it can read.

Change the number of modes below, then switch from a smooth signal to a step. The reconstructed curve illustrates what a fixed-dimensional polynomial summary can and cannot preserve.

This figure projects one fixed window using approximate numerical integration. Actual HiPPO defines an online projection under a chosen measure; in the scaled Legendre case the history interval expands over time. Those details matter: a static plot of polynomial fitting is intuition, not a numerical reproduction of that online theorem. S4 uses the resulting structured matrix as an initialization and learns the state-space parameters; it does not guarantee that every learned state remains an exact polynomial coefficient. The provenance is HiPPO, NeurIPS 2020, and S4's use of that prior is described in §2.2.

The important matrix here is lower triangular. Its diagonal is A_{nn}=-(n+1), and below the diagonal it contains A_{nk}=-\sqrt{(2n+1)(2k+1)}. These couplings distinguish it from a bag of independent decays. The problem is that the same matrix is awkward to compute with efficiently. Before resolving that issue, it helps to see precisely what it computes.

Continuous dynamics become a sampled sequence model

For a single input and output channel, the continuous model is

The state has dimension N; A moves information within that state, B injects the input, and C reads it out. The direct term Du(t) is a separate skip connection, not stored memory. S4 learns parameters of this system inside a neural sequence layer.

A discrete input sequence needs a step size. S4 uses the bilinear transform to obtain \bar A and \bar B, rather than simply treating the continuous matrix as the recurrent transition. Increase the step size below and compare bilinear discretization with explicit Euler on the same decaying system.

For a stable scalar rate, the bilinear map places the discrete pole inside the unit circle for every positive step size. A very large step can still yield a poor approximation to the underlying trajectory. In particular, its discrete pole can become negative, producing alternating samples even when the continuous signal decays monotonically. Numerical stability, approximation fidelity and useful learned memory are three different properties. The formulas and discretization convention follow S4 equation (3).

The learnable step size sets a timescale. A smaller step makes one sequence index represent less continuous time and generally stretches the filter in index space. This also gives a principled way to change the discretization for a changed sampling rate, provided the signal and task support that interpretation. It does not guarantee that every learned model will extrapolate to an arbitrary resolution or sequence length.

Recurrence and convolution are two evaluations of one operator

Using the paper's indexing convention, the initial state is x_{-1}=0, and each input updates the state before its output is read:

Unroll three steps. The first output contains C\bar B u_0. The second contains C\bar A\bar B u_0+C\bar B u_1. Each earlier input is multiplied by a power of the same fixed transition. Define the impulse response once and the whole sequence becomes a causal convolution:

Inspect an output index below. The stems show the individual signed terms that sum to the recurrent readout; the output curves lie on top of one another.

The equivalence is exact for this linear time-invariant subsystem. It assumes fixed parameters and the stated initial condition. Nonzero initial state adds its own homogeneous response. The direct skip adds Du_t. An input-dependent transition or readout generally destroys the single fixed convolution kernel, which helps explain why later selective state-space models use a different execution strategy. This derivation follows S4 equations (4–5).

Full-sequence training can apply the causal convolution with FFTs. The input and kernel must be zero-padded sufficiently to prevent circular wraparound, then the causal prefix is retained. Streaming generation instead carries the state and performs one update per new input. Neither representation makes an autoregressive dependency between generated tokens disappear; recurrence simply avoids recomputing the entire past operator for each token.

There is still a missing piece. A fast input convolution helps only after the long kernel exists. Computing its taps by repeatedly multiplying a dense state matrix costs work proportional to LN^2. S4's main algorithm addresses that construction cost.

Structure makes the useful memory computationally accessible

A tempting solution is to diagonalize the HiPPO matrix: powers of a diagonal matrix are cheap. But an algebraically valid eigenbasis can be numerically disastrous. S4 shows that direct diagonalization of this particular triangular matrix has badly scaled eigenvectors. Negative eigenvalues alone do not solve the conditioning problem; the matrix is non-normal. See §3.1 and Lemma 3.2.

The useful observation is that adding a low-rank correction produces a normal matrix. For HiPPO-LegS, take p_n=\sqrt{n+1/2}. Then M=A+pp^\top satisfies M+M^\top=-I: it is a skew-symmetric matrix shifted by minus one-half of the identity. Such a matrix is normal and has a unitary eigenbasis over the complex numbers. Inspect each term below; the correction is a dense matrix but has rank one.

If M=V\Lambda V^*, changing state coordinates by the unitary matrix yields the original operator in diagonal-plus-low-rank form:

The readout and input vectors must change with the state basis. Diagonalizing only A while leaving them untouched changes the model. Dropping the low-rank term also changes the model. “Normal plus low rank” describes the decomposition before the basis change; “diagonal plus low rank” describes it afterward. S4's Theorem 1 establishes these representations for the relevant HiPPO families, with rank one or two depending on the family.

The structure is valuable for recurrence as well: the bilinear update can be applied through diagonal and low-rank operations without materializing a dense transition, giving linear work in the state dimension for fixed rank. The dense solver used in this review deliberately sacrifices that efficiency to make the algebra easy to verify.

From a resolvent to a finite kernel

Diagonal-plus-low-rank matrices are convenient to invert using Woodbury, but not convenient to raise to many powers. S4 therefore changes the question. Instead of computing every kernel coefficient directly, evaluate its generating function at frequency nodes, then recover the coefficients with an inverse FFT. The resolvent replaces repeated powers with an inverse.

For rank one, the needed expression reduces to four sums of terms with denominators s-\lambda_n, plus a small correction. These are Cauchy matrix-vector products. The spectral parameter comes from the bilinear change of variable, not from an independently chosen Fourier filter. The diagonal part supplies simple denominators; Woodbury restores the low-rank coupling. See Algorithm 1 and Appendix C.3.

One detail is essential: a frequency sample of the infinite impulse response is not automatically a frequency sample of the desired finite kernel. At the length's roots of unity, its missing tail is absorbed into \widetilde C=C(I-\bar A^L) in the row-vector convention used here. Remove that correction below, select the inverse transform, and observe the aliased tail.

This demonstration actually evaluates the Cauchy/Woodbury expression and compares its inverse DFT with a separately constructed impulse response. It uses a small authored real DPLR system, not a trained HiPPO system, and dense preprocessing rather than fast kernel construction. This separation lets the example verify the identity without pretending to reproduce the production numerical backend.

S4's near-linear kernel-construction theorem relies on fast Cauchy multiplication, with logarithmic factors suppressed. A naive all-pairs implementation still performs work proportional to NL and can materialize a large intermediate. The paper's Appendix A distinguishes its implementation from asymptotically faster multipole algorithms; the official source dispatches among CUDA, KeOps and naive Cauchy backends. A complexity theorem and a measured GPU latency are different evidence.

A linear memory lives inside a nonlinear network

A single SSM is linear in its input. That makes the convolution identity possible, but linearity alone cannot supply the nonlinear feature transformations needed for rich tasks. S4 places state-space processing inside a deep neural architecture with channel mixing, nonlinearities and residual connections. The per-channel temporal operator handles long-range mixing; pointwise transformations combine features across channels. The exact block and readout depend on the task. See §4 and the experimental configuration appendix.

The training setup also matters. The official repository's optimizer guidance calls out special treatment of state-space parameters, including lower learning rates in some settings and zero weight decay for those registered parameters. An experiment that swaps the block into an unrelated training recipe is not a controlled test of the paper's architecture. Likewise, the current repository's small S4D examples use a later diagonal variant and should not be relabeled as a reproduction of the original DPLR S4.

What the long-range results establish

Long Range Arena tests several kinds of dependency: symbolic expressions, text, retrieval and flattened visual tasks. Path-X serializes a 128-by-128 image into 16,384 positions and asks whether marked points are connected by a path. It is a striking setting for a global memory model because flattening places related spatial positions far apart in sequence order. It is also a specific synthetic connectivity task, not a proxy for every kind of reasoning over a sixteen-thousand-token document.

Compare the two explicitly labeled S4 rows below. All six task accuracies are retained so the mean has a visible denominator.

In v3 Table 10, the original S4 row averages 80.48% across six tasks and reports 88.10% on Path-X. The updated row averages 86.09% and reports 96.35% on Path-X. V3's main Table 4 uses the updated row, while some opening claims retain the earlier rounded Path-X result. Appendix D.5 attributes the updated S4-LegS runs to follow-up work and describes minor hyperparameter changes, rather than presenting them as the identical original training runs. Comparing an updated headline against an original ablation without checking the revision would obscure what was actually tested.

The appendix describes task-specific hyperparameter searches and best configurations. Its efficiency comparison uses a parameter-tied S4 setup at a matched depth and hidden size; that setup should not be silently equated with every best-accuracy configuration. It also reports single-GPU experiments with hardware varying by task. Baseline results in the LRA table come from earlier benchmark literature rather than a uniform local rerun. These facts limit causal conclusions about architecture, tuning budget and hardware. See Appendix D.1–D.2 and Table 11.

The clean empirical takeaway is that a learned structured state-space network can preserve and use information across long serialized inputs on these benchmarks. The evidence does not establish that every small linear memory can retrieve arbitrary past content, that sequence length alone predicts task difficulty, or that S4 uniformly outperforms attention under matched modern training conditions.

Follow the mathematics into executable code

The five rows connect the original definitions to the corresponding implementation choices. They use a consistent row-vector readout convention and state clearly where the reference trades speed for inspectability. The full Python file supplies common linear algebra helpers and deterministic checks; individual regions are excerpts, not independent programs.

Each Python panel is generated from a named region of examples/s4.py. Run the complete file, which supplies shared linear algebra helpers and checks. This reference uses only the Python standard library; it provides neither training nor a fast Cauchy implementation.

Download complete Python referenceReference notes and limitations

1. Initialize a structured memory operator

Original paper

Enlarge figure

Math and tensor shapes

A_{nk}=-\begin{cases}\sqrt{(2n+1)(2k+1)}&n>k\\n+1&n=k\\0&n<k\end{cases}

N → A [N,N]; indices begin at zero

For N = 2 the rows are [−1, 0] and [−√3, −2]. B[n] = √(2n+1) is the LegS input vector used in the equivalence fixture. The learned readout C chooses which memory combinations contribute to the output.

Pinned official implementation

Python implementation

python

def hippo(size):
    """Construct the real HiPPO-LegS matrix of S4 equation (2)."""
    if not isinstance(size, int) or not 1 <= size <= 32:
        raise ValueError("State size must be 1 through 32")
    return [
        [
            0.0 if i < j else -(i + 1.0) if i == j
            else -math.sqrt((2 * i + 1) * (2 * j + 1))
            for j in range(size)
        ]
        for i in range(size)
    ]

2. Discretize with a linear solve

Original paper

Enlarge figure

Math and tensor shapes

\bar A=(I-\Delta A/2)^{-1}(I+\Delta A/2),\quad \bar B=(I-\Delta A/2)^{-1}\Delta B

A [N,N], B [N], Δ scalar → Ā [N,N], B̄ [N]

One augmented right-hand side solves for both outputs. This is a dense reference, not the O(N) structured recurrent update. For A = −2, B = 3 and Δ = 0.3, Ā = 0.5384615 and B̄ = 0.6923077. C and D remain unchanged.

Pinned official implementation

Python implementation

python

def discretize(matrix, b, delta):
    """Bilinear discretization; solve instead of forming an inverse."""
    if not math.isfinite(delta) or delta <= 0:
        raise ValueError("Positive finite step required")
    size = len(matrix)
    minus = [[float(i == j) - delta * matrix[i][j] / 2
              for j in range(size)] for i in range(size)]
    plus = [[float(i == j) + delta * matrix[i][j] / 2
             for j in range(size)] + [delta * b[i]] for i in range(size)]
    result = solve(minus, plus)
    return [row[:size] for row in result], [row[-1] for row in result]

3. Check the indexing before changing execution mode

Original paper

Enlarge figure

Math and tensor shapes

K_j=C\bar A^j\bar B,\quad y_t=\sum_{j=0}^{t}K_{t-j}u_j+Du_t

u [L], state [N], K [L] → y [L]

The recurrence accepts a nonzero initial state and an optional direct skip. The convolution panel assumes both are zero. With a nonzero initial state add CĀ^(t+1)x[−1]; with D nonzero add Du[t]. A finite chunk of a stream must carry its previous state instead of silently resetting it.

Pinned official implementation

Python implementation

python

def recurrence(matrix, b, c, inputs, initial=None, skip=0.0):
    """Update before reading; the initial state is x[-1]."""
    state = [0.0] * len(b) if initial is None else list(initial)
    output = []
    for value in inputs:
        state = [x + injection * value
                 for x, injection in zip(matvec(matrix, state), b)]
        output.append(dot(c, state) + skip * value)
    return output


def impulse_kernel(matrix, b, c, length):
    """Return K[j] = C Abar**j Bbar, starting with K[0] = C Bbar."""
    state = list(b)
    kernel = []
    for _ in range(length):
        kernel.append(dot(c, state))
        state = matvec(matrix, state)
    return kernel


def convolve(kernel, inputs):
    """Direct causal convolution; no circular wraparound."""
    return [sum(kernel[t - j] * inputs[j] for j in range(t + 1))
            for t in range(len(inputs))]

4. Keep the low-rank correction inside the resolvent

Original paper

Enlarge figure

Math and tensor shapes

R=(sI-\Lambda)^{-1},\quad C(sI-A)^{-1}B=CRB-CRP(1+Q^*RP)^{-1}Q^*RB

rank-one fixture: Λ,P,Q,B,C [N], s complex → scalar complex

The first helper verifies A + ppᵀ is normal for LegS. The second helper uses a separate real DPLR fixture to isolate Woodbury algebra. Each pair is a Cauchy sum over 1/(s−λ[n]). These are complementary checks, not an implementation of eigendecomposition. Complex production parameters require conjugation and conjugate-pair bookkeeping.

Pinned official implementation

Python implementation

python

def corrected_hippo(size):
    """A + pp^T = -I/2 + skew-symmetric, hence normal."""
    matrix = hippo(size)
    p = [math.sqrt(i + 0.5) for i in range(size)]
    normal = [[matrix[i][j] + p[i] * p[j]
               for j in range(size)] for i in range(size)]
    return normal, p


def resolvent(lambdas, p, q, b, c, s):
    """Return C(sI-A)^-1 B for the real fixture A=diag(lambda)-pq^T.

    Four Cauchy sums and one rank-one Woodbury correction. General complex
    DPLR uses q conjugates; here q is real, so conjugation is identity.
    """
    r = [1 / (s - value) for value in lambdas]

    def pair(left, right):
        return sum(x * y * pole for x, y, pole in zip(left, right, r))

    return pair(c, b) - pair(c, p) * pair(q, b) / (1 + pair(q, p))

5. Correct the finite horizon before the inverse transform

Original paper

Enlarge figure

Math and tensor shapes

\widetilde C=C(I-\bar A^L),\quad s=\frac{2}{\Delta}\frac{1-z}{1+z},\quad \widehat K_L(z)=\frac{2}{1+z}\widetilde C(sI-A)^{-1}B

L frequency nodes → spectrum [L] → inverse DFT → K [L]

This uses a row-vector C convention throughout. The paper switches to a conjugate-transposed column convention in Appendix C. The reference computes C̃ densely and evaluates a naive DFT so the identity is inspectable. Production code handles length-dependent C setup, conjugate symmetry and optimized kernels. Recovering K does not itself apply the input convolution.

Pinned official implementation

Python implementation

python

def kernel_spectrum(length=17, delta=0.3, truncate=True):
    """Verify Algorithm 1 on an authored real DPLR system.

    Odd length avoids z=-1, where the bilinear change of variable needs its
    removable-singularity limit. Dense preprocessing is intentionally slow.
    """
    if not isinstance(length, int) or length < 3 or length % 2 == 0:
        raise ValueError("Odd length >= 3 required")
    lambdas, p = [-0.5, -2.0, -4.0], [0.4, 0.3, 0.2]
    b, c = [1.0, 0.5, 0.25], [1.0, -0.6, 0.2]
    matrix = [[(value if i == j else 0) - p[i] * p[j]
               for j in range(3)] for i, value in enumerate(lambdas)]
    abar, bbar = discretize(matrix, b, delta)
    c_power = list(c)
    for _ in range(length):
        c_power = matvec(transpose(abar), c_power)
    c_tilde = [x - (tail if truncate else 0) for x, tail in zip(c, c_power)]
    spectrum = []
    for k in range(length):
        z = cmath.exp(-2j * math.pi * k / length)
        s = 2 * (1 - z) / (delta * (1 + z))
        spectrum.append(2 / (1 + z) * resolvent(
            lambdas, p, p, b, c_tilde, s))
    expected = impulse_kernel(abar, bbar, c, length)
    recovered = [value.real for value in dft(spectrum, inverse=True)]
    error = max(abs(x - y) for x, y in zip(expected, recovered))
    return expected, spectrum, recovered, error

Original paper excerpts: Albert Gu, Karan Goel, Christopher Ré. Paper v3. Cropped/rasterized from the original PDF; no scientific content redrawn.

Run from the repository root: python3 examples/s4.py. Numeric tests: node --test test/s4-model.test.mjs. These commands check authored arithmetic; reported LRA results come from the cited paper.

Where the approach helps, and where it can lose information

S4 is attractive when a long signal benefits from a stable, shared temporal filter and efficient recurrent evaluation: broad temporal structure can be accumulated without retaining every past activation at inference. Its continuous parameterization also separates the underlying dynamics from a particular sample spacing. Those are useful capabilities even when a Transformer remains more accurate for a particular task.

The same design imposes a constraint. Within one linear SSM, the temporal operator is fixed after conditioning on its learned parameters; it does not choose a new transition based on the current token. A finite state must compress the history, so content that was not retained cannot later be recovered merely because a query becomes important. Attention can instead compare a current query with stored past representations, at the cost of keeping and accessing them. This is a mechanism-level distinction, not proof of a universal accuracy ordering.

Nor does a structured initialization guarantee benign optimization forever. Learned rates, readouts, time scales and numerical precision all affect the actual response. Very long horizons can stress discretization, cancellation and kernel generation. Any deployment claim should therefore specify sequence length, state size, batch, precision, backend and whether timing includes kernel setup, convolution, or recurrent decoding. The figures here report arithmetic and cited accuracies; they report no measured latency or checkpoint inference.

S4's lasting contribution is to show that a useful memory prior and an efficient computational representation can be designed together. Its recurrence, convolution and continuous-time views are valuable because the structure makes the conversions tractable. Later work asks how much of the structure is necessary and how to make the memory content-dependent.

Related reading

• HiPPO: Recurrent Memory with Optimal Polynomial Projections develops the memory-projection framework behind the initialization.
• Linear State-Space Layers establishes the continuous, recurrent and convolutional perspective that S4 makes more practical.
• S4D: On the Parameterization and Initialization of Diagonal State Space Models investigates a simpler diagonal construction. Its results should be attributed to the later method.
• Convolution operation and dilated convolutions provide the filtering perspective; self-attention gives a contrasting content-addressed mechanism.
• VMamba follows the later selective-state-space branch into images, where scan order and spatial alignment introduce additional design questions.