Skip to content

S4

Introduced in Efficiently Modeling Long Sequences with Structured State Spaces (Gu et al., 2022).

The Structured State Space sequence model (S4) makes the State Space Model practical for long sequences. The LSSL with a HiPPO transition matrix models long-range dependencies well, but computing its recurrence or convolution kernel with a dense \(A\in\mathbb{R}^{N\times N}\) is expensive in time and memory and numerically unstable. S4 solves this through a structured parameterization of \(A\).

Structured Transition Matrix

The HiPPO matrices are normal plus low-rank (NPLR). They can be conjugated into a diagonal plus low-rank (DPLR) form

\[ A=\Lambda-PQ^{*} \]

with diagonal \(\Lambda\) and low-rank \(P,Q\). With this structure, the discretized recurrence and the convolution kernel can be computed efficiently: the low-rank correction is handled with the Woodbury identity, and the diagonal part reduces to a Cauchy kernel. This brings the cost from dense matrix powers down to roughly linear in state size and sequence length.

S4 was the first model to solve all tasks of the Long Range Arena benchmark, including Path-X with sequences of length 16k, and performed strongly on raw-audio speech classification.

S4D

Introduced in On the Parameterization and Initialization of Diagonal State Space Models (Gu et al., 2022).

S4D drops the low-rank term and uses a purely diagonal complex \(A=\mathrm{diag}(a_{1},\ldots,a_{N})\), so the recurrence becomes elementwise. With a suitable initialization, it matches S4 on most tasks while being much simpler. Common initializations:

  • S4D-LegS: the diagonal part of the HiPPO-LegS matrix.
  • S4D-Lin: \(a_{n}=-\frac{1}{2}+i\pi n\), damped oscillators with linearly spaced frequencies.
  • S4D-Real: \(a_{n}=-(n+1)\), purely real decay.

Each \(a_{n}\) is a damped oscillator: \(\mathrm{Re}(a_{n})<0\) sets the memory length, \(\mathrm{Im}(a_{n})\) the frequency. Mamba builds on the diagonal S4D formulation (S4D-Real for real, S4D-Lin for complex variants). The parameterization used in the oscillatory SSM project is worked out in Theory.

Applications

  • SaShiMi uses S4 blocks in an hourglass architecture for raw-waveform audio modeling and is the basis for several LTI-SSM speech enhancement models (see SSMs for Speech Enhancement).