36 · S4 — Structured State Space Sequence Model

一、核心背景

1.1 什么是 State Space Model

State Space Models (SSMs) originate from control theory. They model a system's dynamics through a latent state that evolves over time:

Where A is the state transition matrix (N x N), B is the input projection (N x 1), and C is the output projection (1 x N). The state h(t) is a vector that compresses the entire history of the input signal.

1.2 Discretization for Deep Learning

To use SSMs in a discrete deep learning setting, we discretize using the Zero-Order Hold (ZOH) method with step size Δt:

This gives us the recurrent form used in neural networks:

1.3 Convolution View

Since the SSM is linear and time-invariant, the recurrence can be unrolled into a convolution, enabling parallel training:

This dual representation is key: recurrent for fast inference, convolutional for parallel training.

二、核心数学框架

2.1 The Long-Range Dependency Problem

Previous SSMs struggled with long-range dependencies because the state matrix A was either diagonal (exponential decay, no memory retention) or dense (prohibitively expensive O(N²) computation). The Copy task was the litmus test: can the model remember a token from position 1 and reproduce it at position L?

Traditional SSMs either failed entirely (diagonal A loses signal exponentially) or required state dimension N proportional to sequence length L (defeating their efficiency advantage).

2.2 HiPPO Theory

The key insight of S4: certain structured matrices have ideal properties for capturing long-range dependencies. The HiPPO (High-Order Polynomial Projection Operators) framework provides a theoretical foundation for designing state matrices that compress history optimally.

The HiPPO matrix (also called the Legendre matrix) is defined as:

This matrix encodes a sliding window of the input history using orthogonal polynomial basis functions. The state vector h(t) stores coefficients of the projection of the input signal's history onto Legendre polynomials. As new input arrives, the coefficients update to maintain the projection over a sliding window.

2.3 Normal Plus Low-Rank (NPLR) Decomposition

The HiPPO matrix is special: it can be decomposed into a normal matrix (unitarily diagonalizable) plus a low-rank term:

This decomposition is the computational key that makes S4 tractable:

• The normal part VΛV* captures the structured memory dynamics
• The low-rank part PQ^T captures interactions that would otherwise require a full dense matrix
• The Woodbury identity enables O(N) instead of O(N²) computation of the convolution kernel

三、训练与推理

3.1 Training Mode (Convolution)

For training, S4 uses the convolution view: the entire kernel K is computed in one shot using the SSM's structure, and then the output is computed via FFT convolution (O(L log L)):

# S4 training: compute kernel once, convolve
# N = state dimension, L = sequence length
K = compute_ssm_kernel(A_bar, B_bar, C, L)  # O(N log N + L)
y = FFT_convolve(x, K)                       # O(L log L)

3.2 Inference Mode (Recurrence)

At inference time, the SSM unrolls as a recurrence, requiring O(1) per step regardless of past history:

# S4 inference: constant-time per step
h = zeros(N)           # state initialization
for t in range(T):
    y_t = C @ h        # output from current state
    h = A_bar @ h + B_bar * x_t  # state update

四、实验结果

4.1 Long Range Arena (LRA)

S4 was the first SSM to achieve competitive results on all LRA tasks:

五、局限性

5.1 Time-Invariance: The Fundamental Gap

S4 is a Linear Time-Invariant (LTI) system. Its parameters A, B, C, and Δ are learned during training but fixed at inference. This means:

• The model cannot selectively focus on relevant versus irrelevant inputs
• Every token is processed with the same dynamics — no "attention" mechanism
• There is no content-based reasoning: the model cannot "decide to remember this token but forget that one"
• Performance degrades on tasks requiring precise retrieval from specific positions

This is the fundamental weakness that Mamba would later address.

5.2 Other Limitations

• No Content-Based Addressing: Unlike attention, which retrieves information based on query-key similarity, S4 treats all positions equally
• Fixed State Size: The recurrence has a fixed state dimension N, placing a hard cap on memory capacity
• Numerical Sensitivity: The matrix exponential in discretization can be numerically unstable for large state dimensions

六、上下游关联

上一篇：HiPPO (Gu et al., 2020) — introduced the polynomial projection theory for memory. S4 is the practical application of HiPPO theory to deep learning.

下一篇：Mamba (Gu & Dao, 2023) — S4 + input-dependent selectivity (B, C, Δ become functions of x). Inherits HiPPO initialization, NPLR decomposition, and the convolutional-recurrent duality from S4.

引用/采用：S4 spawned the entire structured SSM family: S4D (diagonal S4), DSS, Mega (simplified gated SSM), and ultimately Mamba and Mamba-2. It is also a key building block in state-space vision models.

七、个人思考

S4 is a pivotal paper not because it achieved the best numbers (Mamba later surpassed it), but because it proved the SSM hypothesis: that a well-designed recurrent model can match Transformers on long-range tasks. Before S4, the community largely believed that attention was necessary for long-range reasoning. S4 showed that the bottleneck was not the architecture but the parameterization of the state transition.

The NPLR decomposition is the most technically beautiful part: it's a textbook example of using mathematical structure (normal matrix theory + Woodbury identity) to turn an O(N²) algorithm into O(N). This type of "structure-aware algorithm design" is what separates groundbreaking ML papers from incremental ones.

The lesson for future researchers: don't blame the architecture class, fix the parameterization.