Özer Cem KELAHMETKELAHMET
Tools · Architectures

S4

S4 is a structured state-space sequence model that handles very long inputs efficiently. Learn how it works, why it matters and how it led to Mamba.

Overview

S4 is the structured state-space sequence model. It reads a long sequence through a hidden state that follows a linear rule, and it is built so that the same layer can be computed in two ways: as a recurrence for fast generation and as a convolution for fast training. A careful choice of structure in its main matrix lets it remember information over thousands of steps.

When it appeared, S4 solved long-range tasks that Transformers and recurrent networks had failed on, and it did so at near-linear cost in sequence length. It became the base for a whole family of later models, including Mamba. If you want to understand why modern state-space layers look the way they do, S4 is the place to start. It also has a close link to continuous-time models such as the Neural ODE.

What it is

One layer, two views

A state-space layer takes an input signal, keeps a hidden state and produces an output signal. The rule for the state is linear. In continuous time it reads: the rate of change of the state equals a matrix A times the state plus a matrix B times the input. The output is a matrix C times the state.

S4 turns this into a neural network layer with these properties:

  • A, B and C are learned.
  • A has a special structure that lets it hold long memory.
  • The layer runs as a recurrence at inference time.
  • The layer runs as a long convolution at training time.

An everyday picture

Imagine a bathtub with a tap and a drain. The water level is the state. The tap adds water according to the input. The drain removes water at a rate set by A. The reading on the side of the tub is the output. If the drain is slow, the tub remembers what was poured in long ago. If it is fast, only recent pours matter.

S4 uses many small tubs side by side, with drains tuned so that together they remember the recent past in fine detail and the distant past in coarse detail.

What "structured" means

A general matrix A is expensive to use and hard to keep stable. S4 forces A into a special form: a normal matrix plus a low-rank correction. The normal part is easy to work with, because it can be diagonalised. The low-rank part is small. Together they keep the model both stable and cheap.

The role of HiPPO

The starting point for A comes from a framework called HiPPO, which stands for high-order polynomial projection operators. In plain words, it gives a specific matrix A for which the hidden state is the best polynomial summary of the input history so far. Starting from this matrix, instead of random values, is what lets S4 remember over long spans.

What it is not

  • It is not attention. It has no pairwise comparison between positions.
  • It is not a standard recurrent network. Its recurrence is linear, so it can also be written as a convolution.
  • It is not input-dependent. The same dynamics apply at every step. This is the main difference from Mamba.

Where it came from

Control theory roots

State-space models have been a standard tool in control and signal processing for decades. A system with a hidden state, linear dynamics and an output is the classic description of many physical systems.

The long-memory problem

Recurrent networks such as LSTM and GRU struggle to remember over very long ranges. Gradients vanish or explode over many steps. Transformers can look far back but their cost grows with the square of length. By the end of the 2010s, long-range sequence modelling was an open problem.

HiPPO

In 2020, a paper introduced HiPPO, a way to build a memory that provably approximates the whole history of a signal using orthogonal polynomials. It was first used inside recurrent cells and gave large gains on tasks with very long dependencies.

The S4 paper

In 2021, the paper "Efficiently Modeling Long Sequences with Structured State Spaces" introduced S4. It was presented at a major conference in 2022. Its main contributions were:

  • A new parameterisation of A as normal plus low-rank, which made the computation efficient and stable.
  • A fast algorithm for the convolution kernel.
  • Strong results on the Long Range Arena benchmark, including a task that no earlier model had solved.
  • Good results on raw audio and on sequential image tasks.

Simplifications and successors

After S4, many follow-ups appeared:

  • Diagonal versions (such as S4D and DSS) showed that a plain diagonal matrix, with a good initialisation, works almost as well and is far simpler.
  • S5 used a multi-input, multi-output state-space layer with a parallel scan.
  • H3 and related gated designs moved state-space layers into language modelling.
  • Mamba added selection, so the dynamics depend on the input. See Mamba.

How it works

Step 1: start from the continuous system

The continuous system is:

  • dx/dt = A x + B u, where u is the input and x is the state.
  • y = C x (plus an optional direct term D u).

The state x has a modest size, say 64 dimensions per channel. Each input channel gets its own copy of the system with its own parameters.

Step 2: discretise with a step size

Real data are samples, so the system must be discretised. A step size, called delta, converts the continuous matrices into discrete ones. A common choice is the bilinear transform, which keeps the system stable. A learned delta lets the model choose its time scale. A small delta means a long memory in samples, and a large delta means a short one.

After discretisation the layer is a recurrence:

  • x[k] = Abar x[k-1] + Bbar u[k]
  • y[k] = C x[k]

Step 3: see that it is a convolution

Unroll the recurrence. The output at step k is a weighted sum of all earlier inputs, with weights C Abar^j Bbar for lag j. These weights form a kernel. The whole layer is therefore a convolution of the input with a long kernel.

An analogy: ringing a bell. Each strike produces a decaying tone. The sound you hear now is the sum of all earlier tones, each faded by its age. The fade pattern is the kernel.

Step 4: compute the kernel fast

A long kernel could be built by powering the matrix Abar again and again, which is slow. S4 uses the special structure of A to compute the kernel quickly:

  • Move to a frequency-domain description using a generating function.
  • Use the normal plus low-rank form to reduce the problem to a Cauchy matrix product, with a correction by the Woodbury identity.
  • Evaluate it at roots of unity.
  • Use a fast Fourier transform to get the kernel.

The result is a kernel of length L computed in about L log L time with modest memory. You do not need to follow every step to use the layer. The key idea is that the structure of A makes this possible.

Step 5: train with the convolution, run with the recurrence

Mode Form used Why
Training Convolution via FFT Fully parallel over the sequence
Autoregressive generation Recurrence Constant time and memory per step

The two views give identical outputs, up to numerical error. This duality is one of the main selling points.

Step 6: wrap in a deep network

An S4 layer is used as a block, usually with:

  • A linear mixing layer across channels after the convolution.
  • A nonlinearity such as GELU, often with a gate.
  • Residual connections and normalisation.
  • Dropout.

Blocks are stacked, from a few up to dozens, and a head reads the result.

Step 7: choose the initialisation

The initial value of A matters a great deal. Starting from a HiPPO matrix gives long memory from the first step of training. Random starts often fail on hard long-range tasks. In diagonal versions, the same effect is achieved by choosing eigenvalues that follow a pattern based on HiPPO.

Key concepts and terms

Term Plain meaning
State-space model A hidden state updated by a linear rule, with an output read from it
State The memory vector carried across steps
A, B, C, D State, input, output and skip matrices
Discretisation Converting the continuous system into step-by-step form
Delta The step size, which sets the time scale of memory
HiPPO A method to choose A so the state summarises history well
Normal plus low-rank (NPLR) The structure S4 gives to A
Convolution kernel The weights that map past inputs to the current output
FFT Fast Fourier transform, used for fast long convolutions
Cauchy kernel A matrix type that S4 reduces its computation to
Long Range Arena A benchmark suite for long-sequence models
Path-X A very hard long-range task of the benchmark
Linear time-invariant (LTI) Dynamics that do not change over time
Diagonal SSM A simpler variant using a diagonal A
Recurrent mode Running the layer one step at a time

Time invariance and its consequence

Because the dynamics are the same at every step, S4 is a linear time-invariant system. This is why a convolution works. It is also why S4 cannot decide what to ignore based on content. That limit motivated the selective design of Mamba.

The family at a glance

Model Main change from S4 Why it matters
DSS and S4D Diagonal state matrix Much simpler code with similar accuracy
S5 Multi-input multi-output state, parallel scan Fewer components, works with irregular steps
H3 Gated design with shift and diagonal parts First step toward language modelling
Hyena Implicit long convolutions with gating Attention-free long-context layer
Mamba Input-dependent dynamics Selection and content-aware memory

Frequency view in plain words

Every kernel can be seen as a filter. Some filters let slow changes through and block fast ones. Others do the opposite. By learning the eigenvalues of A, an S4 layer learns a bank of filters. The imaginary part of an eigenvalue sets a pitch, and the real part sets how fast the tone fades. A layer with many eigenvalues is like a set of tuning forks of different pitches and decay times. The sum of their ringing is the kernel.

Why long memory is hard

  • A plain recurrent network multiplies its state by a matrix at every step. If that matrix shrinks the state, old information vanishes. If it grows the state, training blows up.
  • S4 picks A so that its eigenvalues sit just inside the stable region, with a range of decay rates. Some channels fade fast and others fade very slowly.
  • The training signal can then reach far back without exploding.

A worked example

A numeric walk-through of the recurrence

Take a one-dimensional state with discrete values Abar = 0.9, Bbar = 1 and C = 1. The input sequence is 1, 0, 0, 0, 0.

Step Input State after update Output
1 1 0.9 * 0 + 1 = 1.000 1.000
2 0 0.9 * 1.000 = 0.900 0.900
3 0 0.9 * 0.900 = 0.810 0.810
4 0 0.9 * 0.810 = 0.729 0.729
5 0 0.9 * 0.729 = 0.656 0.656

The kernel is 1, 0.9, 0.81, 0.729, 0.656, and so on. A single strike at step 1 leaves a decaying trace. Convolving the input with this kernel gives the same outputs as the recurrence. Both views agree.

Checking the equivalence in code

python
import numpy as np

Abar, Bbar, C = 0.9, 1.0, 1.0
u = np.array([1.0, 0.0, 0.0, 0.0, 0.0])
L = len(u)

# recurrence
x, y_rec = 0.0, []
for k in range(L):
    x = Abar * x + Bbar * u[k]
    y_rec.append(C * x)

# convolution with kernel K[j] = C * Abar**j * Bbar
K = np.array([C * (Abar ** j) * Bbar for j in range(L)])
y_conv = np.convolve(u, K)[:L]

print(np.allclose(y_rec, y_conv))   # True

A small S4D-style layer in PyTorch

The diagonal variant is simple enough to write in a few lines. This version builds the kernel directly and applies it with an FFT.

python
import torch
import torch.nn as nn

class TinyS4D(nn.Module):
    def __init__(self, d_model=32, d_state=16):
        super().__init__()
        # S4D-Lin style initial eigenvalues: -1/2 + i*pi*n
        n = torch.arange(d_state)
        A = -0.5 + 1j * torch.pi * n
        self.A = nn.Parameter(torch.view_as_real(A.repeat(d_model, 1)))
        self.B = nn.Parameter(torch.ones(d_model, d_state))
        self.C = nn.Parameter(torch.randn(d_model, d_state, dtype=torch.cfloat) * 0.1)
        self.log_dt = nn.Parameter(torch.empty(d_model).uniform_(-4.6, -2.3))
        self.D = nn.Parameter(torch.randn(d_model))

    def kernel(self, L):
        dt = torch.exp(self.log_dt).unsqueeze(-1)                  # (d, 1)
        A = torch.view_as_complex(self.A)                           # (d, n)
        dA = dt * A
        Cbar = self.C * (torch.exp(dA) - 1) / A                     # (d, n)
        t = torch.arange(L).view(1, 1, L)
        K = (Cbar.unsqueeze(-1) * self.B.unsqueeze(-1) * torch.exp(dA.unsqueeze(-1) * t)).sum(1)
        return K.real                                               # (d, L)

    def forward(self, u):                                           # u: (batch, d, L)
        L = u.shape[-1]
        K = self.kernel(L)
        n = 2 * L
        y = torch.fft.irfft(torch.fft.rfft(u, n=n) * torch.fft.rfft(K, n=n), n=n)[..., :L]
        return y + u * self.D.unsqueeze(-1)

layer = TinyS4D()
print(layer(torch.randn(4, 32, 1024)).shape)   # torch.Size([4, 32, 1024])

This sketch is for learning. Production code adds careful initialisation, a stable real-valued parameterisation, and an efficient custom kernel.

What this example shows

  • The learned delta sets how long each channel remembers.
  • Imaginary parts in A create oscillations, so the kernel can encode periodic patterns.
  • A single FFT turns a length 1024 sequence into a few fast operations.

A second example: two time scales

Suppose two channels with Abar = 0.5 and Abar = 0.99, both given a single strike of size 1.

Steps after the strike Channel with 0.5 Channel with 0.99
1 0.500 0.990
5 0.031 0.951
10 0.001 0.904
50 near 0 0.605
100 near 0 0.366

The first channel is a short-term memory. The second is a long-term memory. A layer with a spread of such values can model patterns at many scales. This is what log-spaced step sizes aim for.

Reading the result for a table

Imagine a table of daily sensor readings. A channel with a fast decay tracks the last few days. A channel with a slow decay tracks the seasonal level. A classifier or regressor on top of the states can then use both.

Strengths and limitations

Strengths

  • Handles sequences of tens of thousands of steps in a single layer.
  • Near-linear training cost, through the FFT convolution.
  • Constant-cost generation, through the recurrent form.
  • Strong on long-range benchmarks and on raw signals such as audio.
  • Can be applied at different sampling rates, since it is built from a continuous system.
  • Few parameters per channel.
  • Mathematically clean, with clear stability properties.

Limitations

  • Linear time invariance: it cannot choose to ignore or emphasise content, so it does worse on language than attention models.
  • The original NPLR implementation is intricate. Many users prefer the diagonal variants.
  • Sensitive to initialisation. A poor starting A can break long-range learning.
  • Needs complex number support, which some frameworks handle poorly.
  • Weak at exact recall of a specific earlier item.
  • Reading order is fixed unless a bidirectional design is used.

Why selection matters

Think of a tape recorder with a fixed playback filter. It treats every sound the same way. A language model needs a filter that responds to meaning. S4 has the first kind of filter. Mamba added the second kind.

How it compares with alternatives

Model Cost in length Parallel training Content-dependent memory Strong on
S4 Near linear Yes, via FFT No Long-range signals
Mamba Linear Yes, via scan Yes Language, long context
Transformers Quadratic Yes Yes, by attention General tasks
LSTM Linear No Yes, by gates Short to medium sequences
GRU Linear No Yes, by gates Short to medium sequences
Neural ODE Solver dependent Limited Nonlinear dynamics Irregular continuous time
CNN (long kernel) L log L with FFT Yes No Local patterns
PatchTST Moderate Yes Yes Forecasting

S4 and convolutional networks

An S4 layer is a convolution with a very long, implicitly defined kernel. A normal convolutional layer stores each kernel weight. S4 stores a handful of parameters per channel and generates the whole kernel from them. This is why it can cover a huge receptive field with so few parameters.

S4 and Neural ODEs

A Neural ODE learns a nonlinear rate of change and needs a numerical solver. S4 uses a linear rate of change and solves it in closed form with its step size. The cost is less flexibility. The gain is speed and exact parallelism.

Role in tabular foundation models

How S4 connects to tables

A tabular foundation model learns from raw tables with rows and columns, mixed numeric and categorical cells, missing values and often no column names. These models are mostly attention-based, such as TabPFN and TabICL. S4 is not their standard backbone. Its relevance is as a way to read very long sequences cheaply, and as the ancestor of selective state-space layers that are now being tried on table data.

Row and column structure

Reading a table as a sequence can be done in two ways. You can scan along the columns of one row, or down the rows of the table. The second is long. S4 can summarise thousands of rows in a single pass. The catch is order: rows in a table have no natural order, and S4 is order-sensitive. Pre-training with random row shuffles, and averaging predictions over several shuffles, helps remove order artifacts.

Time-indexed tables

Many real tables are time series in disguise: sensor readings, transactions and logs with a time stamp column. For these, S4's origins in signal processing are a strength. It handles long histories at a fixed cost and can adapt to a different sampling rate by rescaling delta. See also Chronos and PatchTST for forecasting models on time series.

Mixed types and missing values

Cells are encoded into vectors before any sequence layer sees them. Numbers, categories and missing flags all become vectors of the same size. S4 then sees only vectors. For missing cells in a time-indexed table, the continuous-time origin of S4 helps: the step size can be set to match the gap between observations, so the state decays across the gap in a principled way. Plain S4 uses a fixed grid, so irregular gaps need a time-aware variant.

No column names

Without names, a model must infer each column's role from values alone. A state-space layer over a column's values acts as a learned summariser: it compresses a long list of values into a short vector that captures scale, rhythm and drift. That vector can feed a classifier that guesses the kind of column.

Pre-training and in-context learning

Foundation models for tables are pre-trained on many synthetic and real tables. They then do in-context learning: given labelled example rows, they predict for new rows without updating weights. S4's linear time-invariant memory is a weaker fit for this than attention, because the model cannot choose which example to retrieve. Selective layers are closer. Still, S4 shows that a fixed-size state can hold a useful summary of a very long context, and that insight carries to the newer designs. See Mamba.

Continual learning

A recurrent form lets a model absorb new rows by updating a state. This is cheap and suits growing tables. Forgetting in the weights still needs separate methods. See Continued pre-training.

Imputation and anomaly detection

For imputation, a long-memory layer can look across many past rows to predict a missing cell. For anomaly detection, the model predicts the next value, and a large error flags a surprise. The long memory lets it model slow seasonal patterns that short-window methods miss. See Anomaly detection, Imputation networks and Masked cell modeling.

Sector and semantic classification

Classifying what domain a table or column belongs to means reading many values and summarising them. A state-space encoder pools a long column in near-linear time. A head such as the ones described in Semantic sector heads can then use the summary. For wider context on how models are built and measured, see /model/architecture, /model/research and /model/benchmarks.

An honest summary of the role

S4 is best seen as an important building block and teaching case. It explains why state-space layers can remember so much and how to compute them quickly. In current tabular work, its descendants are more relevant than S4 itself.

Practical guidance and best practices

Pick the variant

  • Start with a diagonal variant such as S4D. It is simple and works well.
  • Use the full NPLR form only if you need its exact behaviour or are reproducing a result.
  • Use a selective model if the task needs content-dependent memory.

Set the step size

  • Initialise delta on a log scale across channels, so some channels remember short ranges and others long ones.
  • Rescale delta when the sampling rate of the data changes.
  • Let delta be learned but with a smaller learning rate than the main weights.

Initialise A carefully

  • Use a HiPPO-based or S4D-Lin style initialisation.
  • Keep the real part of the eigenvalues negative for stability.
  • Do not apply weight decay to A and delta.

Architecture tips

  • Use a gated or GLU-style output.
  • Add normalisation and residual paths around each layer.
  • Use dropout across the channel dimension.
  • Use bidirectional layers for tasks where the full sequence is available at once.

Data handling for tables

  • Embed cells first, then scan.
  • Shuffle row order each batch during pre-training.
  • Provide a missing indicator channel.
  • Pass time gaps as an extra feature if the table is time-indexed.

Testing

  • Compare the recurrent and convolutional modes on a short input.
  • Test longer lengths than seen in training.
  • Add a simple baseline, such as a convolution or a recurrent model.

A short decision table

Situation Suggested choice
Long raw signals such as audio or sensors S4 or S4D
Language or content-aware filtering Mamba or attention
Short sequences Plain convolution or small attention
Irregular time stamps Time-aware state-space variant or Neural ODE
Static table, small Tree ensemble or tabular foundation model

Common mistakes

  • Using a random A and expecting long memory.
  • Applying weight decay to the dynamics parameters.
  • Using one time scale for all channels.
  • Forgetting that S4 is time invariant, then expecting content-based filtering.
  • Passing data at a different sampling rate without adjusting the step size.
  • Ignoring numerical error between the two computation modes.
  • Feeding unordered table rows in a fixed order.
  • Choosing S4 for short sequences where a simple convolution would do.
  • Implementing the full NPLR form from scratch when a diagonal form is enough.
  • Skipping a fair comparison with an attention baseline of the same size.

Frequently asked questions

What is S4 in simple terms?

S4 is a neural network layer based on a hidden state that follows a linear rule. It reads a sequence, keeps a compact memory and produces an output. Its special structure lets it remember over thousands of steps and be computed quickly in training.

What does the S4 name stand for?

The name comes from the four S letters in structured state space sequence model. It refers to a sequence model that uses a state-space layer with a structured matrix. The structure is what makes the layer both stable and fast.

Why can S4 be computed as a convolution?

Because its dynamics are linear and do not change over time, the output at any step is a fixed weighted sum of earlier inputs. Those weights form a kernel, so the whole layer equals a convolution. A fast Fourier transform then computes it in parallel for the full sequence.

What is HiPPO?

HiPPO is a method for choosing the main state matrix so that the hidden state is the best polynomial approximation of the input history. Starting S4 from such a matrix gives it long memory from the beginning of training. Later variants use simpler matrices that follow the same pattern.

How is S4 different from Mamba?

S4 uses the same dynamics at every step, while Mamba computes its step size and read and write maps from the current input. That makes Mamba able to ignore or emphasise content, which helps with language. S4 relies on a convolution for speed, and Mamba uses a hardware-aware scan.

Is S4 better than a Transformer?

On long-range signal tasks, S4 matched or beat Transformers at much lower cost. On language and other tasks that need content-dependent recall, attention models have been stronger. The right choice depends on sequence length and on what the task requires.

Can S4 be used on tabular data?

It can be used to scan along rows or along time-indexed records, and it is helpful for long histories. It is not the usual backbone for tabular foundation models, which mostly use attention. The ideas from S4 matter more than the layer itself in that area.

What is S4D?

S4D is a simplified version of S4 that uses a diagonal state matrix. It gives close to the same accuracy with much simpler code and faster computation. Many practical systems use it instead of the full structured form.

Does S4 work with irregular time stamps?

The plain layer assumes a regular grid. Because it comes from a continuous system, it can be adapted by making the step size follow the actual time gaps, but this needs care. Neural ODE and controlled-equation models handle irregular time more directly.

What is the Long Range Arena?

It is a benchmark suite of tasks with long inputs, such as long text, image pixels read as sequences and a path-finding task. S4 was the first model to solve its hardest task, which helped drive interest in state-space models.

Glossary

  • Bilinear transform: a way to discretise a continuous system while keeping it stable.
  • Convolution kernel: the weights that map past inputs to the current output.
  • Delta: the step size that sets the memory time scale.
  • Diagonal SSM: a state-space layer with a diagonal state matrix.
  • FFT: fast Fourier transform, used for fast convolution.
  • HiPPO: a framework for building a memory that summarises history.
  • Linear time-invariant: a system whose rule is linear and fixed over time.
  • Long Range Arena: a benchmark suite for long-sequence models.
  • NPLR: normal plus low-rank, the structure S4 gives to its state matrix.
  • Recurrent mode: running the layer one step at a time.
  • Selection: letting parameters depend on the input, as in Mamba.
  • State: the hidden memory vector.
  • State-space model: a model built on a hidden state with a linear update.
  • Woodbury identity: a matrix identity used to handle a low-rank correction.

Summary

S4 is a structured state-space sequence model. It uses a carefully shaped state matrix, often started from HiPPO, to remember over very long ranges. It trains as a fast convolution and runs as a cheap recurrence, giving near-linear cost in length.

Its main limit is that the same dynamics apply to every step, so it cannot choose content to keep, which led to selective successors such as Mamba. For tabular foundation models, S4 is mostly a conceptual ancestor: it shows how a fixed-size state can summarise a long context, and it informs designs that scan long row lists or time-indexed records. Related reading: Neural ODE, Transformers, /model/architecture and /model/benchmarks.

It belongs to Architectures in my stack, under AI and machine learning. Others in the same group:

Type to search.
    ↑↓ navigate↵ open