On the Parameterization and Initialization of Diagonal State Space Models

By Albert Gu, Ankit Gupta, Karan Goel, Christopher R'e. Neural Information Processing Systems, 2022

629 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Cognitive Neuroscience, Electrical and Electronic Engineering

View Paper: https://www.semanticscholar.org/paper/ca444821352a4bd91884413d8070446e2960715a

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Cognitive Neuroscience

In Electrical and Electronic Engineering

Citations in the corpus

Simplified State Space Layers for Sequence Modeling (2022, methodology): “[19] Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré.”
Efficiently Modeling Long Sequences with Structured State Spaces (2021, neutral): “The results reported in Table 4 have been updated to results from the papers [19, 20].”
VMamba: Visual State Space Model (2024, background): “Since then, many flavors of structured state space models sprang up with different structures like complex-diagonal structure [19, 16], multiple-input multiple output supporting [43], decomposition of diagonal plus low-rank operations [21], selection mechanism [14].”
RS-Mamba for Large Remote Sensing Image Dense Prediction (2024, methodology): “Notable among these are models incorporating complex-diagonal structures to improve temporal modeling capabilities [32], [33], as well as those supporting multiple-input multiple-output configurations to increase model flexibility [30].”
Hungry Hungry Hippos: Towards Language Modeling with State Space Models (2022, methodology): “The diagonal SSM constrains A to be diagonal and initializes it from the diagonal version of HiPPO (S4D [26]).”
PointMamba: A Simple State Space Model for Point Cloud Analysis (2024, methodology): “To alleviate computational burden, HTTYH [19], DSS [22], and S4D [16] propose employing a diagonal matrix within S4, maintaining performance without excessive computational costs.”
Resurrecting Recurrent Neural Networks for Long Sequences (2023, methodology): “…and inference), while exponentiating dense matrices is computationally expensive, and (b) while nonlinear recurrences must be computed sequentially, unrolling a linear recurrence can be parallelized using associative scans resulting in faster training (Gupta et al., 2022a; Smith et al., 2022).”
Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model (2024, neutral): “…SSM-based methods, such as the linear state-space layers (LSSL) [22], structured state space sequence model (S4) [21], diagonal state space (DSS) [24], and S4D [23], are proposed to process sequence data across a wide range of tasks and modalities, particularly on modeling long-range dependencies.”
Long Range Language Modeling via Gated State Spaces (2022, background): “Further, alternative models based on linear dynamical systems have shown great promise [Gu et al., 2022a, Gupta et al., 2022, Gu et al., 2022b].”
U-KAN Makes Strong Backbone for Medical Image Segmentation and Generation (2024, background): “Recently, Mamba [18] has achieved a groundbreaking milestone with its linear-time inference and efficient training process by integrating selection mechanism and hardware-aware algorithms into previous works [19, 21, 57].”