Diagonal State Spaces are as Effective as Structured State Spaces

By Ankit Gupta, Jonathan Berant. Neural Information Processing Systems, 2022

508 citations recorded by the source corpus.

Modeling long range dependencies in sequential data is a fundamental step towards attaining human-level performance in many modalities such as text, vision, audio and video. While attention-based models are a popular and effective choice in modeling short-range interactions, their performance on tasks requiring long range reasoning has been largely inadequate. In an exciting result, Gu et al. (ICLR 2022) proposed the $\textit{Structured State Space}$ (S4) architecture delivering large gains over state-of-the-art models on several long-range tasks across various modalities. The core proposition of S4 is the parameterization of state matrices via a diagonal plus low rank structure, allowing efficient computation. In this work, we show that one can match the performance of S4 even without the low rank correction and thus assuming the state matrices to be diagonal. Our $\textit{Diagonal State Space}$ (DSS) model matches the performance of S4 on Long Range Arena tasks, speech classification on Speech Commands dataset, while being conceptually simpler and straightforward to implement.

Subfields: Computational Theory and Mathematics, Computer Networks and Communications

View Paper: https://www.semanticscholar.org/paper/71e15a9a52dcafca57bff5f310b95e2c7d0cfc87

How researchers see this paper

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

In Computational Theory and Mathematics

In Computer Networks and Communications

Citations in the corpus

On the Parameterization and Initialization of Diagonal State Space Models (2022, methodology): “We first present a simplified view on the expressivity of diagonal SSMs mentioned by [11].”
Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model (2024, unknown): “…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…”
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.”
Simplified State Space Layers for Sequence Modeling (2022, methodology): “Further, the initialization of the timescales can be important [20, 22], and limiting to sampling a single initial parameter may lead to a poor initialization.”
RS-Mamba for Large Remote Sensing Image Dense Prediction (2024, background): “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].”
LocalMamba: Visual State Space Model with Windowed Selective Scan (2024, background): “Despite initial challenges in training, owing to their computational and memory intensity, recent advancements [9–11, 16, 43] have significantly ameliorated these issues, positioning deep SSMs as formidable competitors against CNNs and Transformers.”
Long-term Forecasting with TiDE: Time-series Dense Encoder (2023, methodology): “, 2019b) and applied parameter efficient SSMs (Gu et al.; Gupta et al., 2022) to modeling long range dependencies in sequences.”
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].”
Long Range Language Modeling via Gated State Spaces (2022, methodology): “Note that similar to [Gu et al., 2022a, Gupta et al., 2022], we used a constant learning rate of 1e-3 and set weight decay rate to 0.0 for state space parameters part of GSS Layer .”
Gated Linear Attention Transformers with Hardware-Efficient Training (2023, unknown): “…(Wu et al., 2019), Long Convolu-tion (Fu et al., 2023c; Qin et al., 2023a; Poli et al., 2023; Massaroli et al., 2023; Li et al., 2023d; Romero et al., 2021), and State Space Models (Gu et al., 2021a; Gupta & Berant, 2022; Gu et al., 2021b; Hasani et al., 2022; Smith et al., 2023; Ma et al., 2023).”
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].”
Hyena Hierarchy: Towards Larger Convolutional Language Models (2023, background): “• H3: Implicit parametrization using state-space models (SSMs), in particular S4 (Gu et al., 2021; Gupta et al., 2022).”
Hungry Hungry Hippos: Towards Language Modeling with State Space Models (2022, unknown): “To construct the filter f from A,B,C efficiently, A is often constrained to be diagonal (Gupta et al., 2022; Gu et al., 2022b), or diagonal plus low-rank (Gu et al.”
Avey-B (2026, unknown): “…(e.g., via using linear attention (Katharopoulos et al., 2020; Choromanski et al., 2021; Peng et al., 2025) and RNN-inspired architectures (Gu et al., 2021; Gupta et al., 2022; Fu et al., 2022; Gu & Dao, 2023)), but little of it has been adapted to the bidirectional, encoder-only paradigm.”