Simplified State Space Model (S5)

Updated 19 February 2026

Simplified State Space Model (S5) is a sequence modeling technique that replaces ensembles of SISO SSMs with a unified MIMO architecture to enhance efficiency and accuracy.
It leverages a diagonalizable MIMO core and parallel scan methods to achieve linear computational complexity even on long-sequence tasks.
S5 supports robust quantization and scalable hardware implementations, delivering state-of-the-art results on benchmarks like Long Range Arena and sMNIST.

The Simplified State Space Model (S5) is a sequence modeling architecture that replaces ensembles of single-input single-output (SISO) structured state-space sequence (S4) models with a single multi-input multi-output (MIMO) state-space model. S5 directly addresses computational and modeling inefficiencies present in earlier SSM layers by leveraging a diagonalizable MIMO core, enabling both competitive accuracy and significant implementation simplification. S5 achieves linear complexity for long-sequence tasks and matches or exceeds S4-style architectures on established benchmarks, while opening new avenues for efficient hardware deployment and robust system identification (Smith et al., 2022, Yu et al., 2023, Abreu et al., 2024, Mattes et al., 2023).

1. Mathematical Formulation

S5 is formulated as a continuous-time, linear, time-invariant state-space system: $\frac{dx(t)}{dt} = A\,x(t) + B\,u(t), \quad y(t) = C\,x(t) + D\,u(t)$ where $x(t)\in\mathbb{R}^P$ is the latent state, $u(t)\in\mathbb{R}^H$ is the input, $y(t)\in\mathbb{R}^H$ is the output, $A\in\mathbb{R}^{P\times P}$ , $B\in\mathbb{R}^{P\times H}$ , $C\in\mathbb{R}^{H\times P}$ , $D\in\mathbb{R}^{H\times H}$ are learnable parameters (Smith et al., 2022, Abreu et al., 2024).

Discretization of the system (usually via zero-order hold for timestep vector $\Delta$ ) yields

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

with diagonalization $A=V\Lambda V^{-1}$ , resulting in

$\overline{\Lambda} = \exp(\Lambda \Delta),\quad \overline{B} = \Lambda^{-1}(\overline{\Lambda} - I)\,\widetilde B,\quad \overline{C} = \widetilde C,\quad \overline{D} = D$

Parameters are typically maintained in this ‘eigen’ basis for computational efficiency.

2. Relationship to S4, HiPPO, and Diagonalization

S5 generalizes the S4 approach by transitioning from an $H$ -fold bank of SISO SSMs to a singular, potentially lower-rank MIMO SSM. In S4, each input channel is processed independently with a distinct SISO system, followed by a final mixing step. S5 instead allows direct parameter sharing and latent interaction within the MIMO SSM, collapsing the S4 structure while maintaining expressivity (Smith et al., 2022).

HiPPO initialization yields strong long-contextual biasing but central S4/S4D matrices (e.g., HiPPO-LegS) are not stably diagonalizable. S5 circumvents this by diagonalizing the normal (unitarily diagonalizable) component, referred to as HiPPO-N, while optionally incorporating block-diagonal or hierarchical schemes to relax assumptions and enable multi-scale modeling (Smith et al., 2022, Yu et al., 2023). However, only weak convergence to the HiPPO target is achieved in direct diagonalizations, motivating further advances discussed below.

3. Computational Complexity and Implementation

S5’s diagonal recurrence permits use of associative parallel scan primitives, avoiding heavy reliance on FFTs or Cauchy kernel tricks found in S4. Offline (full-sequence) operation executes SSMs in $O(PHL)$ work and $O(\log L)$ parallel time, matching S4’s asymptotic complexity (assuming latent size $P=O(H)$ ). Online, autoregressive modes have per-step cost $O(PH)$ (Smith et al., 2022, Mattes et al., 2023).

Crucially, implementation reduces to diagonal matrix multiplications and parallel prefix operations, removing the need for FFTs. This not only streamlines code paths (e.g., JAX or CUDA implementation), but also allows for precise hardware mappings and aggressive quantization (Abreu et al., 2024).

4. Extensions: Robustness and Initialization via PTD

Direct diagonalization of non-normal HiPPO matrices is numerically ill-posed, with the potential for catastrophic frequency response spikes and non-robust extrapolation (Yu et al., 2023). The “perturb-then-diagonalize” (PTD) methodology introduces a controlled perturbation $E$ such that $A+E$ is diagonalizable with well-conditioned eigenvectors and the overall system closely approximates the original HiPPO transfer function in operator norm: $|H_{PTD}(i\omega) - H_{DPLR}(i\omega)| \leq (2 \ln n + 4) \epsilon + O(\sqrt{\ln n}\,\epsilon^2)$ where $\epsilon$ parametrizes the spectral perturbation (Yu et al., 2023).

Empirical results show S5-PTD maintains performance under adversarial frequency shift and outperforms standard S5, S4D on noisy and long-range tasks, averaging 87.6% accuracy on the Long Range Arena benchmark (Yu et al., 2023).

5. Quantization for Efficient Deployment

The Q-S5 effort systematically evaluates quantization-aware training (QAT) and post-training quantization (PTQ) for S5, targeting deployment on resource-constrained hardware (Abreu et al., 2024). Uniform symmetric quantization is applied to all weights and activations, with key findings:

The recurrent diagonal $\overline{A}$ must be held at $\geq 8$ bits to avoid collapse across most tasks.
Activations and non-recurrent weights tolerate lower precision; components such as $B, C, D$ and most non-SSM weights can be quantized to 2–4 bits with marginal impact (except on critical temporal/dynamical benchmarks).
PTQ is effective on language-based tasks (≤1pp accuracy drop), but QAT is required for strong performance on signal or structure tasks.
End-to-end quantized S5 models achieve <1% drop on sMNIST and LRA tasks (with high-precision $\overline{A}$ ).

Task	Full-prec.	PTQ W8A8	QAT W8A8	QAT W4A8 ( $\overline{A}$ @8)	QAT W2A8 ( $\overline{A}$ @8)
sMNIST	99.65	96.27	99.54	99.63	99.56
ListOps	62.15	26.65	39.05	36.80	36.80
Text	89.31	88.49	57.39	50.72	52.21
Retrieval	91.40	89.87	86.26	82.78	72.15
sCIFAR	88.00	44.83	86.95	87.20	85.57
Pathfinder	95.33	50.90	50.81	95.06	94.34

Performance on structural or dynamical tasks is highly sensitive to the quantization of the recurrent core; QAT enables near-parity with full-precision S5 when this constraint is observed (Abreu et al., 2024).

6. Empirical Performance and Applications

On Long Range Arena (LRA), S5 achieved the highest linear-complexity average accuracy (87.46%) and outperformed S4, S4D, and other baseline architectures on the hardest Path-X task (98.58%). S5 also attained state-of-the-art or competitive results on speech classification (96.52% 16 kHz Speech Commands, 94.53% zero-shot 8 kHz transfer) and pixel-level sequence tasks (99.65% on sMNIST) (Smith et al., 2022).

Hierarchical imagination models, such as Hieros, utilize S5 as a world model, exploiting its capacity for both parallel sequence encoding and O(n²) fast iterative rollout. In this context, S5 outperforms GRU-based and Transformer-based world models on Atari 100k, with strong exploration and sample efficiency (Mattes et al., 2023).

S5 further enables efficient interpolation on irregularly-sampled dynamical systems, with speedups up to 86× against convolutional recurrent units and improved mean squared error (Smith et al., 2022).

7. Discussion, Limitations, and Future Directions

The main advantages of S5 are:

A “single MIMO SSM” design collapses S4’s $H$ -fold ensemble into one compact layer.
State updates are efficiently implementable via diagonal recurrences and parallel scan methods.
S5 natively supports time-varying SSMs, irregular sampling, and continuous-time inputs.
Empirical parity (or better) with S4 on established sequence tasks, with reduced implementation and memory overhead (Smith et al., 2022).

Potential limitations include dependence on complex-valued parameters for eigen decomposition, the necessity of judicious parameter initialization (especially latent dimension $P$ selection), and limited expressivity if $P\ll H$ . Direct diagonalization–based initializations (without PTD or similar regularizations) are fragile in adversarial settings or on frequency-shifted data (Yu et al., 2023).

Open research directions involve exploring adaptive and block-diagonal HiPPO initializations, connections to classical filtering, and mesh extensions to higher-dimensional structured convolutions (Smith et al., 2022). Hardware-optimized variants and robust, low-precision S5 layers remain active areas for further development (Abreu et al., 2024).

References: (Smith et al., 2022, Yu et al., 2023, Mattes et al., 2023, Abreu et al., 2024)