State-Space Models and S4

Attention scales quadratically with sequence length — double the sequence and the compute needed roughly quadruples — which makes it expensive on sequences of tens of thousands of steps, like tick-level order flow over a full session. Recurrent networks scale linearly, which is better, but their gradients through Backpropagation Through Time decay over long spans, so they struggle to actually use information from thousands of steps back even when the compute is cheap enough to try. S4 (Structured State Space Sequence model) asks whether there's a way to process very long sequences that is both linear in cost like an RNN and free of the gradient decay that cripples RNNs over long spans. Its answer borrows an idea from control engineering that predates deep learning by half a century.

The analogy: a thermostat with memory of the whole day

A building's temperature at any moment isn't a random number — it's the outcome of a continuous underlying process: heat leaking in and out, the heating system responding, all evolving smoothly over time according to a small number of physical rules. A control engineer models this with a state: a compact summary (say, "current temperature and its rate of change") that evolves according to fixed equations and that, together with external inputs (heater on/off), fully determines the future. Crucially, you don't need every past reading in memory to predict tomorrow — the current state already summarizes everything relevant about the past that matters for what happens next. A state-space model treats a data sequence the same way: a hidden state evolves smoothly according to a small, fixed set of linear equations, and the output at each moment is read off that state. S4's contribution was finding a way to parameterize and initialize those equations so the state can actually carry information across thousands of steps without decaying — and to compute the whole thing efficiently.

The maths

A continuous-time linear state-space model describes a hidden state $x(t)$ evolving under an input $u(t)$:

$$x'(t) = A x(t) + B u(t), \qquad y(t) = C x(t) + D u(t)$$

In words: the state's rate of change is a fixed linear combination of its current value (via matrix $A$) plus the current input (via $B$); the output $y(t)$ is a fixed linear readout of the state (via $C$), plus a direct pass-through of the input (via $D$, often set to zero). $A$, $B$, $C$, $D$ are the model's learned parameters — a small, fixed set, regardless of how long the sequence is.

For a discrete sequence (which is what any real dataset is), this gets discretized with a step size $\Delta$ into a recurrence:

$$x_k = \bar{A} x_{k-1} + \bar{B} u_k, \qquad y_k = C x_k$$

In words: this looks exactly like an RNN's update rule — new state from old state plus new input — but with one crucial difference: $\bar A$ and $\bar B$ come from a fixed, structured continuous-time system rather than being an arbitrary learned matrix passed through a nonlinearity. That structure is what S4 exploits.

Because the recurrence is linear (no $\tanh$ or sigmoid squashing it at each step, unlike an RNN or LSTM), it can be unrolled into an explicit formula for $y_k$ in terms of all past inputs:

$$y_k = \sum_{j=0}^{k} C \bar{A}^{\,k-j} \bar{B}\, u_j$$

In words: the output at step $k$ is a weighted sum of every past input $u_j$, where the weight on an input from $k-j$ steps ago is $C\bar A^{k-j}\bar B$ — a single fixed formula, not a chain of $k-j$ separate multiplications each passed through a saturating nonlinearity. This is exactly a convolution, with a kernel of length $k$ built directly from $A$, $B$, $C$. S4's key engineering insight is choosing a special structured form for $A$ (derived from a mathematical object called the HiPPO matrix, designed specifically to compress a long history into a fixed-size state with minimal information loss) that makes this convolution kernel computable efficiently — in roughly $O(L \log L)$ for a sequence of length $L$, rather than the $O(L^2)$ that plain attention costs.

Worked example 1: unrolling a tiny state-space recurrence by hand

Take a 1-dimensional toy system with $\bar A = 0.8$, $\bar B = 1.0$, $C = 1.0$, and inputs $u_1 = 2, u_2 = 1, u_3 = 3$, starting from $x_0 = 0$.

$$x_1 = 0.8(0) + 1.0(2) = 2.0, \qquad y_1 = 1.0(2.0) = 2.0$$ $$x_2 = 0.8(2.0) + 1.0(1) = 1.6 + 1 = 2.6, \qquad y_2 = 2.6$$ $$x_3 = 0.8(2.6) + 1.0(3) = 2.08 + 3 = 5.08, \qquad y_3 = 5.08$$

Now verify the convolution formula gives the same $y_3$ directly, without stepping through $x_1, x_2$:

$$y_3 = C\bar A^{2}\bar B\, u_1 + C\bar A^{1}\bar B\, u_2 + C\bar A^{0}\bar B\, u_3$$ $$= (0.64)(2) + (0.8)(1) + (1)(3) = 1.28 + 0.8 + 3 = 5.08$$

Both routes give $5.08$. The recurrence and the convolution are two ways of computing the identical answer — the recurrence is cheap for generating one step at a time (inference), the convolution form is cheap for processing a whole sequence in parallel (training), and S4's engineering is precisely about making that convolution kernel fast to compute for very long $L$ even though $A$ is not this simple a scalar in practice.

Worked example 2: why the decay rate doesn't compound the way an RNN's does

In worked example 1, the weight on an input from $d$ steps back is $\bar A^{d} = 0.8^{d}$ — decaying, but not catastrophically:

$$0.8^{20} \approx 0.0115$$

Compare this to the earlier RNN vanishing-gradient calculation (in Backpropagation Through Time) where a chain of $\tanh$-saturated Jacobians around 0.6 gave $0.6^{20} \approx 0.0000366$ — roughly 300 times smaller. The gap exists because S4's $\bar A$ comes from a structured, carefully-designed continuous system (via the HiPPO initialization) chosen specifically so its eigenvalues support slow decay across long spans, rather than from an arbitrary trained matrix squashed through a saturating nonlinearity at every step. This is a deliberate design choice, not a coincidence — the entire point of the HiPPO-based initialization is to make long-range memory the default behavior of the untrained model, something gradient descent then only has to preserve rather than discover from scratch.

Picture the kernel weight $\bar A^d$ as a function of distance $d$ back in time — it settles toward zero but far more gently than a saturating-nonlinearity chain, similar to how a running average's envelope narrows gradually rather than collapsing sharply; that gentle decay is what lets an S4 layer keep a usable signal from thousands of steps back.

What this means in practice

State-space layers are a strong fit for very long, densely-sampled sequences where attention's quadratic cost is genuinely prohibitive — full-session tick data, long macro time series, or any pipeline where "look back thousands of steps efficiently" is a hard requirement rather than a nice-to-have. Because the underlying system is linear, an S4 layer alone can't represent input-dependent, content-based routing the way attention or a gated RNN can — it treats every sequence with the same fixed dynamics regardless of what's actually in it, which is precisely the limitation Mamba and Selective State Spaces was built to address.