Source-linked AI summary
On the Parameterization and Initialization of Diagonal State Space Models
Albert Gu, Ankit Gupta, Karan Goel, Christopher Ré
TL;DR
S4D addresses the difficulty of implementing S4 by systematically studying diagonal SSM parameterizations and initialization. It explains why DSS works theoretically and yields a simple diagonal model that performs comparably to S4 in almost all settings, including 85% on Long Range Arena.
Problem
S4’s DPLR representation and algorithm are difficult to understand and implement, while diagonal SSM performance depends critically on initialization.
Method
The paper systematically categorizes diagonal SSM design choices, analyzes DSS’s initialization, and introduces S4D by combining S4-style computation and parameterization with DSS’s initialization.
Results
85% on the Long Range Arena benchmark, with S4D competitive with S4 in almost all settings and near state-of-the-art results across image, audio, and medical time-series benchmarks.
Takeaways & Limitations
Diagonal SSMs can retain strong S4-like performance while using a much simpler, theoretically motivated formulation.
Takeaways & Limitations
Expressive equivalence does not guarantee optimization performance, and different initializations with the same spectrum can perform very differently.
Abstract
from arXiv · showhide
State space models (SSM) have recently been shown to be very effective as a deep learning layer as a promising alternative to sequence models such as RNNs, CNNs, or Transformers. The first version to show this potential was the S4 model, which is particularly effective on tasks involving long-range dependencies by using a prescribed state matrix called the HiPPO matrix. While this has an interpretable mathematical mechanism for modeling long dependencies, it introduces a custom representation and algorithm that can be difficult to implement. On the other hand, a recent variant of S4 called DSS showed that restricting the state matrix to be fully diagonal can still preserve the performance of the original model when using a specific initialization based on approximating S4's matrix. This work seeks to systematically understand how to parameterize and initialize such diagonal state space models. While it follows from classical results that almost all SSMs have an equivalent diagonal form, we show that the initialization is critical for performance. We explain why DSS works mathematically, by showing that the diagonal restriction of S4's matrix surprisingly recovers the same kernel in the limit of infinite state dimension. We also systematically describe various design choices in parameterizing and computing diagonal SSMs, and perform a controlled empirical study ablating the effects of these choices. Our final model S4D is a simple diagonal version of S4 whose kernel computation requires just 2 lines of code and performs comparably to S4 in almost all settings, with state-of-the-art results for image, audio, and medical time-series domains, and averaging 85\% on the Long Range Arena benchmark.
1 Introduction
S4D investigates whether diagonal state space models can retain S4’s effectiveness while avoiding its difficult DPLR representation and algorithm. It combines systematic parameterization choices, a theoretically motivated initialization, and controlled experiments to produce a simpler competitive model.
- S4 achieves strong long-sequence modeling, but its DPLR parameterization and computation require sophisticated linear-algebra techniques that are difficult to understand and implement.
- Diagonal SSMs are algebraically expressive, but their numerical performance depends critically on initialization.Random or alternative diagonal parameterizations can perform much worse empirically.
- S4D combines S4’s computation and parameterization with DSS’s initialization to create a simple, theoretically principled, empirically effective diagonal SSM.
- The paper systematically categorizes diagonal SSM parameterization and computation choices, allowing S4 and DSS variants to be compared within one framework.
- The diagonal restriction of S4’s HiPPO matrix produces the same dynamics as S4 when the state size goes to infinity.
- 85% on the Long Range Arena benchmark, while the best S4D methods remain competitive with S4 across almost all settings.
2 Background
The background introduces SSMs as linear time-invariant sequence transformations represented by differential equations or convolutions. S4 makes them effective through HiPPO-based structure and DPLR computation, while DSS demonstrates that a carefully initialized diagonal restriction can also perform well.
- SSMs map input signals to output signals and can be represented either as linear ordinary differential equations or continuous convolutions.
- A convolution kernel can be viewed as a linear combination of basis kernels controlled by A, B, and C; for diagonal SSMs, each basis kernel is e^(tA_n)B_n.
- S4 uses the HiPPO-LegS matrix so its basis kernels correspond to normalized Legendre polynomials and support long-range modeling through an interpretable decomposition.
- S4 decomposes the HiPPO matrix into a normal plus rank-1 matrix and computes its DPLR convolution kernel with a specialized efficient algorithm.
- DSS showed empirically that removing the low-rank portion of the HiPPO-LegS DPLR form yields a diagonal matrix performing comparably to S4.
- The paper studies DSS components systematically and explains the theoretical interpretation of its diagonal state matrix.
3 Parameterizing Diagonal State Spaces
Diagonal SSMs separate state-space parameterization, discretization, and kernel computation into modular design choices. The resulting S4D formulation uses efficient Vandermonde multiplication, stability-oriented parameterization, and flexible choices that match S4’s computational complexity.
- Discretization: There is little empirical difference between the discretization methods, although bilinear discretization smooths the DSS kernel to match the S4 kernel.
- Kernel computation: Diagonal SSM kernels can be computed through Vandermonde matrix-vector multiplication rather than S4’s specialized DPLR procedure.
- Kernel computation: Theoretical diagonal-kernel computation requires O~(N + L) operations and O(N + L) space, matching the complexity of DPLR SSMs.
- Parameterization: Parameterizing the real part of A to be negative enforces stability, while unconstrained real parts may become unstable outside fixed-length sequence settings.
- Parameterization: S4D trains B separately from C, unlike DSS’s direct product parameterization; training B provides a minor but consistent performance improvement.
- S4D: S4D matches S4’s parameterization in its base version, while its diagonal state matrix corresponds to masking out S4’s DPLR low-rank component.
- S4D: S4D’s design choices consolidate multiple structured-SSM variants, and diagonal SSMs remain highly competitive with full DPLR versions.
- Discretization: S4D separates discretization from kernel computation, allowing any discretization, whereas DSS requires ZOH for its softmax-based computation.
4 Initialization of Diagonal State Matrices
Diagonal SSMs are algebraically expressive, but effective performance depends critically on initialization and numerical representation. The paper develops and analyzes S4D initializations whose structured frequencies and simple computation support diagonal models comparable to S4.
- Expressivity and limitations: Almost all SSMs are equivalent to a complex diagonal SSM, but this algebraic expressivity does not guarantee strong performance after optimization.Randomly initialized dense real or diagonal complex matrices can perform poorly despite being fully expressive.
- Expressivity and limitations: Two initializations with the same spectrum can have very different performance, showing that numerical representation matters beyond eigenvalues.This limitation motivated S4’s low-rank correction rather than pure diagonalization.
- S4D-LegS: As N →∞, diagonalizing the normal-HiPPO component recovers the same SSM basis as the HiPPO-LegS matrix with its low-rank correction.The diagonal matrix is unitarily equivalent to the normal-HiPPO matrix, preserving stability and timescale.
- Alternative initializations: S4D-Inv approximates S4D-LegS using an inverse-law scaling for imaginary eigenvalue components, while S4D-Lin uses Fourier-like frequencies.For S4D-Inv, the n-th positive imaginary component scales as Θ(n^-1); S4D-Lin produces damped Fourier basis functions.
- Design principles: The real part of A controls decay, whereas its imaginary part controls oscillation frequency and should be spread out rather than randomly initialized.A_n = −1/2 gives a constant timescale with envelope e^-t/2; S4D-Inv and S4D-Lin provide structured frequency choices.
5 Experiments
Experiments show that carefully initialized diagonal SSMs are highly competitive with DPLR S4 across diverse tasks, while initialization choices matter substantially more than most parameterization choices.
- Full Comparisons of S4D and S4 Methods: S4D is competitive with S4 on almost all settings and achieves near-state-of-the-art results across image, audio, and medical time-series benchmarks.The study evaluates image pixels, biosignal time series, and audio waveforms across sequence lengths from 1K to 16K.
- S4D Initialization Ablations: Simple changes to the S4D-LegS eigenvalue scaling law perform worse, whereas a Fourier-frequency linear law also performs well.The inverse scaling λ_n ∝ n^-1 approximates original S4 dynamics; changing its constant or using a quadratic law reduces performance.
- S4D Initialization Ablations: Changing the imaginary-part scale by factors of 0.01 or 100.0 substantially reduces S4D performance in all settings.These changes largely preserve the basis-function shape, indicating that the precise frequency scaling remains important empirically.
- S4D Initialization Ablations: Randomizing imaginary parts causes minor degradation, randomizing real parts causes minor but consistent degradation, and randomizing both degrades performance further.Random imaginary initialization replaces uniformly spaced frequency nodes with uniformly random values.
- Parameterization, Computation, Discretization: Most parameterization choices have little overall effect for fixed initialization, while training B gives a minor improvement and unconstrained real parts can cause instability.Full DPLR versions are often slightly better than diagonal versions during training, but this is not attributed to their larger parameter count.
6 Conclusion
S4D-style state space models offer a simpler way to work with S4-style sequence models while retaining broad applicability. The paper presents S4D as a practical and theoretically promising generic sequence model.
- 6 Conclusion: S4D provides a simpler implementation framework for S4-style models across sequential data domains.The work covers core structures and kernel computation while reducing the implementation burden for practitioners.
- 6 Conclusion: S4D is positioned as a strong generic sequence model for a variety of domains.
- 6 Conclusion: The paper presents S4D as opening theoretical directions for state space models.
A.1 Proofs
The proofs explain why diagonalizing the HiPPO-LegS state matrix can recover S4 dynamics in the infinite-state limit. They also establish that this behavior depends on HiPPO-LegS’s special structure and does not generally extend to other low-rank perturbations.
- A.1 Proofs: Theorem 3 analyzes why approximating the HiPPO-LegS DPLR matrix with its diagonal component can preserve the original dynamics.The result is specific to the special structure of this matrix rather than a general property of DPLR systems.
- A.1 Proofs: The HiPPO-LegS system admits a convolutional interpretation whose kernels are linked to shifted, scaled Legendre polynomials.These kernels form a complete orthonormal basis under an exponentially weighted measure.
- A.1 Proofs: The state represents projections of the input onto a complete orthonormal basis, enabling reconstruction from the current state in the infinite-state limit.
- A.1 Proofs: Removing the low-rank correction from HiPPO-LegS recovers the same SSM kernel as state size approaches infinity.The result explains mathematically why the diagonal DSS initialization can approximate S4-LegS.
- A.1 Proofs: Random rank-1 perturbations produce substantially different kernels that diverge as perturbation magnitude increases, unlike the HiPPO-LegS correction.The analogous low-rank removal for S4-FouT also dramatically changes its kernels.
B Experiment Details
The experiments compare S4D and S4 under controlled parameterization, initialization, and training settings across ablation datasets and Long Range Arena. They document both implementation choices and dataset-specific protocol variations.
- B Experiment Details: The ablation protocol uses a four-layer architecture with hidden dimension H = 128 and approximately 100K trainable parameters.The A and B parameters are tied across the H SSM copies.
- B Experiment Details: S4 and S4D variants are evaluated with closely matched parameter counts, differing mainly through the DPLR low-rank tensor.The extra S4 parameters arise from P in A = Λ − P P^T.
- B Experiment Details: Results are averaged over 2 or 3 seeds, and reported values are generally best validation accuracy except for large models in Table 4.
- B Experiment Details: Speech Commands results use the full 35-class dataset, whereas earlier S4-related work commonly used a smaller 10-class subset or preprocessing features.The paper therefore notes that those results are not directly comparable.
- B Experiment Details: Long Range Arena uses bidirectional S4-style layers, GLU feedforward mixing, and cosine learning-rate scheduling rather than the original unidirectional and plateau-scheduler setup.The bidirectional layer concatenates two convolution kernels back-to-back.
- B Experiment Details: Dataset-specific settings vary, including larger state size N = 256 for BIDMC and deeper or narrower models for ListOps.ListOps uses 8 layers and H = 128, while BIDMC increases state size to capture higher-frequency features.