RNNs & Sequence Models2022advanced10 min read

Efficiently Modeling Long Sequences with Structured State Spaces

النمذجة الكفوءة للتسلسلات الطويلة عبر فضاءات الحالة المُهيكلة

Gu, A. · Goel, K. · Ré, C. — ICLR

The problem

requires capturing dependencies between elements that may be thousands of steps apart. Transformers solve this with , but attention costs grow quadratically with sequence length — making them impractical for sequences of 10,000+ steps. RNNs process sequences step by step with constant memory, but suffer from vanishing gradients that destroy long-range memory. By 2021, no model could efficiently handle truly long-range dependencies: the Long Range Arena benchmark, designed to test sequences of 1,000 to 16,000 steps, defeated every existing architecture on its hardest tasks.

The contribution

S4 introduces a new parameterization of state space models based on the framework. The key insight is that the HiPPO matrix — which is mathematically designed to remember long histories — can be decomposed into (DPLR) form. This decomposition reduces the SSM computation to evaluating a , which can be done in O(N log N) time. The result is a model that runs as a during (parallelizable like a ) and as a recurrence during (efficient like an ). S4 achieved 91% on sequential CIFAR-10, solved the Path-X task (length 16,384) that defeated all prior models, and set state-of-the-art on every Long Range Arena task, averaging 80.48% vs less than 60% for all baselines.

The impact

S4 revived state space models as a competitive alternative to Transformers for sequence modeling. It demonstrated that classical control theory ideas — continuous-time dynamical systems, , and structured matrices — could outperform attention on long sequences. S4 directly inspired a wave of follow-up work: S4D simplified the parameterization to purely diagonal, S5 introduced parallel scans, and ultimately Mamba added input-dependent selectivity to create the first SSM competitive with Transformers on language modeling. The SSM paradigm S4 launched now powers models across audio, genomics, video, and time series.

Imagine two ways to remember a long lecture. The way: you record every word and, whenever you need an answer, you search the entire recording — effective, but impossibly slow for a 10-hour lecture. The RNN way: you take notes on a single sticky note, overwriting as you go — fast, but by hour ten you have forgotten hour one.

S4 is a structured notebook: a set of pages mathematically designed so that each new sentence updates every page with just the right amount of old information preserved. The notebook never overflows, never forgets the opening, and — here is the magic — can be read both page-by-page (like an RNN at inference) or all at once (like a CNN during training).

The problem: why long sequences defeat existing models

Sequence modeling sits at the heart of machine learning: language, audio, time series, and even images can be treated as sequences. The central challenge is capturing long-range dependencies — patterns where an event at step 1 influences the outcome at step 10,000.

Transformers handle this by letting every position attend to every other position. This is powerful but costs $O(L^2)$ in computation and memory, where $L$ is the sequence length. For a sequence of length 16,000, that means 256 million pairwise comparisons — per layer.

RNNs (including LSTMs) process one step at a time, compressing the entire history into a fixed-size . This is $O(L)$ in computation, but the problem means that gradients shrink exponentially through time, making it nearly impossible to learn dependencies beyond a few hundred steps.

What the field needed was a model that combines the RNN's linear scaling with the Transformer's long-range memory. That model is S4.

Open in Lab

Compare how Transformers (quadratic), RNNs (linear but forgetful), and S4 (linear with long memory) scale with sequence length.

The demo wakes as you arrive…

State space models: the mathematical foundation

A (SSM) is borrowed from control theory and signal processing. It describes a system with a hidden state $x(t)$ that evolves over continuous time. The system takes an input signal $u(t)$ , updates its hidden state, and produces an output $y(t)$ .

Think of it as a water tank system: the input $u(t)$ is water flowing in, the state $x(t)$ is the water level in multiple connected tanks, and the output $y(t)$ is a measurement of the system. Matrix $A$ controls how the tanks are connected — how water flows between them. Matrix $B$ controls where the input enters, and matrix $C$ controls what we measure.

The continuous-time SSM is defined by two equations. The first is the state equation describing how the hidden state evolves:

x'(t) = A\,x(t) + B\,u(t)

State equation — how the hidden state evolves — The state

x(t) \in \mathbb{R}^N

captures the system's memory. Matrix

A \in \mathbb{R}^{N \times N}

controls internal dynamics — how the state interacts with itself. Matrix

B \in \mathbb{R}^{N \times 1}

maps the scalar input into the state space. The derivative

x'(t)

means this update happens continuously.

The second equation is the output equation — what we can observe from the hidden state:

y(t) = C\,x(t) + D\,u(t)

Output equation — projecting the state to a prediction — Matrix

C \in \mathbb{R}^{1 \times N}

reads out from the state. The skip connection

D\,u(t)

lets the input pass directly to the output. In practice

D

is often treated as a simple skip connection and the core learning happens through

A

B

, and

C

Open in Lab

Watch how input flows through the state space: water enters (input), distributes across tanks (state), and produces a reading (output). Adjust matrix A to see how connectivity changes the system's memory.

The demo wakes as you arrive…

Discretization: from continuous time to discrete steps

The continuous SSM is elegant but cannot directly process discrete data like text tokens or audio samples. We need to discretize it — convert the continuous equations into a step-by-step recurrence that operates on a discrete sequence $(u_0, u_1, u_2, \ldots)$ .

S4 uses the bilinear (Tustin) method with a learnable step size $\Delta$ to transform the continuous matrices $(A, B)$ into discrete matrices $(\bar{A}, \bar{B})$ :

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

Bilinear discretization — converting continuous dynamics to discrete steps — The step size

\Delta

controls the resolution: small

\Delta

captures fast-changing signals, large

\Delta

sees slow trends. Crucially,

\Delta

is *learnable* — each layer discovers the timescale that best serves its task.

After discretization, the SSM becomes a simple linear recurrence — exactly like an RNN, but with structured matrices:

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

Discrete SSM recurrence — the RNN view — At each step

k

, the state

x_k

is updated by multiplying the previous state by

\bar{A}

and adding the new input scaled by

\bar{B}

. This is the *recurrent view* — efficient for autoregressive generation (one new token at a time).

Open in Lab

See how unrolling the recurrence produces a convolution kernel. Toggle between the RNN view (step-by-step) and the CNN view (parallel convolution).

The demo wakes as you arrive…

HiPPO: the matrix that remembers everything

The matrix $A$ is the heart of the SSM — it determines what the model remembers and forgets. A random matrix leads to rapid forgetting. The breakthrough came from the HiPPO (High-order Polynomial Projection Operator) framework, introduced in prior work by the same authors.

The intuition: at every moment $t$ , the state $x(t)$ should store the best polynomial approximation of the entire input history $u(s)$ for $s \leq t$ . Imagine fitting a polynomial curve through all the data you have seen so far, and updating it continuously as new data arrives. The HiPPO matrix $A$ is the exact matrix that achieves this optimal online approximation.

Concretely, the HiPPO-LegS matrix (Legendre measure, scaled) is:

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

HiPPO-LegS matrix — optimal polynomial memory — Each row

n

of the state corresponds to the

n

-th Legendre polynomial coefficient of the running input history. Low-order coefficients capture broad trends; high-order coefficients capture fine details. Together,

N

coefficients approximate the full history with provably optimal accuracy.

Open in Lab

Watch how HiPPO continuously updates its polynomial approximation of the input history. Each coefficient tracks a different frequency of the signal.

The demo wakes as you arrive…

The computational challenge: making HiPPO fast

HiPPO gives us the perfect memory matrix, but computing the convolution kernel $\bar{K}$ naively requires materializing all powers $\bar{A}^k$ for $k = 0, \ldots, L-1$ . With a dense $N \times N$ matrix and a sequence of length $L$ , this costs $O(N^2 L)$ — far too slow for practical use with large state sizes.

The key insight of S4 is a structural decomposition. The authors proved that all HiPPO matrices can be written in Normal Plus Low-Rank (NPLR) form:

$A = V \Lambda V^* - PQ^\top$

where $V$ is unitary, $\Lambda$ is diagonal, and $P, Q$ have rank $r = 1$ or $2$ . By conjugating with $V$ , this becomes Diagonal Plus Low-Rank (DPLR):

$\tilde{A} = \Lambda - \tilde{P}\tilde{Q}^*$

This structure is the breakthrough: it reduces the kernel computation to evaluating a Cauchy kernel — a well-studied problem with fast $O(N \log N)$ algorithms via the FFT.

\hat{\bar{K}}_\Omega(z) = \frac{2\Delta}{1+z} \left[ \tilde{C}^\top \left( \frac{2}{\Delta} \frac{1-z}{1+z} - \tilde{A} \right)^{-1} \tilde{B} \right]

S4 kernel in frequency domain — the Cauchy kernel connection — In the z-transform domain, the SSM kernel takes this closed form. The inverse

(\omega I - \tilde{A})^{-1}

with DPLR structure reduces to a Cauchy-like dot product that can be evaluated at all

L

frequencies simultaneously using the FFT.

Open in Lab

See how the HiPPO matrix decomposes into diagonal + low-rank components, enabling the fast Cauchy kernel computation.

The demo wakes as you arrive…

The S4 deep architecture

A single SSM layer maps one input channel to one output channel. To build a practical deep model, S4 stacks multiple layers in a structure similar to a Transformer block:

Each S4 block consists of: (1) a linear SSM layer that processes the sequence, (2) a nonlinear (typically GELU), (3) a position-wise feedforward network for mixing features, (4) residual connections and (either pre-norm or post-norm).

Multiple input channels are handled by running independent SSM copies in parallel — one per channel — similar to . A linear mixing layer then combines information across channels.

The full model stacks 4–6 of these blocks. Despite its simplicity, this architecture achieves strong results because the SSM layer — with its HiPPO initialization and learned step size — provides something no other layer can: provably optimal long-range memory with efficient computation.

Open in Lab

Explore the S4 block architecture: click each component to see how SSM, activation, FFN, and residual connections work together.

The demo wakes as you arrive…

Results: solving the unsolvable

S4's results were dramatic. On the Long Range Arena (LRA) benchmark — six tasks designed to stress-test long-range reasoning over sequences of 1K–16K steps — S4 outperformed all 11 Transformer variants and every prior model on every single task, averaging 80.48% vs less than 60% for all baselines.

Most striking was Path-X (sequence length 16,384): a task requiring the model to determine whether two points on a 128×128 image are connected by a path. Every prior model — including all efficient Transformers — scored at chance level (50%). S4 achieved 94.20%.

On sequential CIFAR-10 (classifying images processed one pixel at a time, 1,024 steps), S4 reached 91.13% accuracy — matching a 2D ResNet despite seeing the image as a flat 1D sequence.

For speech commands (raw audio classification, 16,000 samples), S4 reached 98.32% on the full 35-class task, outperforming specialized audio architectures.

Open in Lab

Compare S4 against Transformers and RNNs across Long Range Arena tasks. Hover over each task to see the performance gap.

The demo wakes as you arrive…

S4 in code: the core computation

Simplified S4 kernel computationpython

Simplified to show the idea — not the real implementation.

import torch
import torch.nn as nn
import torch.fft as fft

class S4Kernel(nn.Module):
    """Compute the S4 convolution kernel using DPLR structure."""
    def __init__(self, N: int, L: int):
        super().__init__()
        # Diagonal part of A (complex, learned)
        self.Lambda = nn.Parameter(torch.randn(N, dtype=torch.cfloat))
        # Low-rank correction
        self.P = nn.Parameter(torch.randn(N, 1, dtype=torch.cfloat))
        self.Q = nn.Parameter(torch.randn(N, 1, dtype=torch.cfloat))
        self.B = nn.Parameter(torch.randn(N, 1, dtype=torch.cfloat))
        self.C = nn.Parameter(torch.randn(1, N, dtype=torch.cfloat))
        # Learnable step size
        self.log_dt = nn.Parameter(torch.randn(1))
        self.L = L

    def forward(self):
        dt = self.log_dt.exp()  # Step size
        # Evaluate kernel at L frequencies via Cauchy kernel
        omega = torch.arange(self.L, dtype=torch.cfloat)
        # Simplified: Cauchy dot product at each frequency
        # Full version uses Woodbury for low-rank correction
        z = 2 * (1 - omega) / (1 + omega)  # bilinear transform
        vals = (z.unsqueeze(-1) - self.Lambda.unsqueeze(0))
        K_hat = (self.C @ (1.0 / vals * self.B)).squeeze(-1)
        # Back to time domain
        K = fft.ifft(K_hat, n=self.L).real
        return K  # shape: (L,)

The SSM revolution: from HiPPO to Mamba

2020
HiPPO — the memory foundation
Gu et al. introduced the HiPPO framework showing that specific matrices can provably approximate the history of a continuous signal using polynomial projections. This solved the theoretical memory problem for SSMs.
2021
LSSL — combining RNNs, CNNs, and SSMs
The Linear State Space Layer showed that SSMs can be viewed as both recurrent and convolutional models, establishing the dual computational view that S4 would exploit.
2022
S4 — the structured breakthrough (this paper)
S4 solved the computational bottleneck with the DPLR decomposition, achieving state-of-the-art on Long Range Arena and solving Path-X for the first time. Published at ICLR 2022 as an Oral presentation.
2022
S4D and DSS — diagonal simplifications
Gupta et al. and Gu et al. showed that purely diagonal parameterizations can match S4's performance, dramatically simplifying implementation while retaining the HiPPO-inspired initialization.
2023
Mamba — selective state spaces
Gu and Dao introduced input-dependent selection to SSMs, making the matrices B and C functions of the input. This broke the linear time invariance assumption but made SSMs competitive with Transformers on language modeling for the first time.

S4 proved that the gap between classical dynamical systems theory and modern deep learning was not a gap at all — it was a bridge waiting to be built. By showing that structured matrices from control theory could outperform attention on the hardest sequence tasks, S4 opened the door to an entirely new family of sequence models that now rivals the Transformer paradigm.

CitationGu, Goel, Ré. Efficiently Modeling Long Sequences with Structured State Spaces. ICLR, 2022.

Terms in this paper

State Space Modelنموذج فضاء الحالة
HiPPOHiPPO
Discretizationالتمييز
Convolutionالالتفاف الرقمي
Recurrent Neural Network (RNN)الشبكة العصبية التكرارية
Long-Term Dependencyالاعتمادية البعيدة المدى
Cauchy Kernelنواة كوشي
Diagonal Plus Low-Rankقُطرية زائد منخفضة الرتبة
Linear Time Invarianceثبات الزمن الخطّي
Kalman Filterمرشِّح كالمن
Parallel Scanالمسح المتوازي
Convolutional Neural Network (CNN)الشبكة العصبية الالتفافية
Vanishing Gradientاضمحلال متجهات الميل
Sequence Modelingنمذجة التتابعات
Activation Functionدالة التنشيط