# On the Parameterization and Initialization of Diagonal State Space Models

arXiv: https://arxiv.org/abs/2206.11893 · HTML: https://mlsanity.com/paper/2206.11893/

## Key facts

- Cited by 424 papers in the map (data through 14 Oct 2026).
- Builds on 17 of its 40 references that are in the map.
- First posted on arXiv 23 Jun 2022.

## Builds on (most cited first, 17 of 17)

- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [Language Modeling with Gated Convolutional Networks](https://mlsanity.com/paper/1612.08083/) (2016)
- [Long Range Arena: A Benchmark for Efficient Transformers](https://mlsanity.com/paper/2011.04006/) (2020)
- [HiPPO: Recurrent Memory with Optimal Polynomial Projections](https://mlsanity.com/paper/2008.07669/) (2020)
- [Neural Controlled Differential Equations for Irregular Time Series](https://mlsanity.com/paper/2005.08926/) (2020)
- [Diagonal State Spaces are as Effective as Structured State Spaces](https://mlsanity.com/paper/2203.14343/) (2022)
- [It's Raw! Audio Generation with State-Space Models](https://mlsanity.com/paper/2202.09729/) (2022)
- [Neural Rough Differential Equations for Long Time Series](https://mlsanity.com/paper/2009.08295/) (2020)
- [CKConv: Continuous Kernel Convolution For Sequential Data](https://mlsanity.com/paper/2102.02611/) (2021)
- [Lipschitz Recurrent Neural Networks](https://mlsanity.com/paper/2006.12070/) (2020)
- [FlexConv: Continuous Kernel Convolutions with Differentiable Kernel Sizes](https://mlsanity.com/paper/2110.08059/) (2021)
- [Learning Longer-term Dependencies in RNNs with Auxiliary Losses](https://mlsanity.com/paper/1803.00144/) (2018)
- [UnICORNN: A recurrent model for learning very long time dependencies](https://mlsanity.com/paper/2103.05487/) (2021)
- [Trellis Networks for Sequence Modeling](https://mlsanity.com/paper/1810.06682/) (2018)

## Built on by (most cited first, 30 of 424)

- [Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model](https://mlsanity.com/paper/2401.09417/) (2024)
- [VMamba: Visual State Space Model](https://mlsanity.com/paper/2401.10166/) (2024)
- [Simplified State Space Layers for Sequence Modeling](https://mlsanity.com/paper/2208.04933/) (2022)
- [Hungry Hungry Hippos: Towards Language Modeling with State Space Models](https://mlsanity.com/paper/2212.14052/) (2022)
- [Resurrecting Recurrent Neural Networks for Long Sequences](https://mlsanity.com/paper/2303.06349/) (2023)
- [Long Range Language Modeling via Gated State Spaces](https://mlsanity.com/paper/2206.13947/) (2022)
- [Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models](https://mlsanity.com/paper/2402.19427/) (2024)
- [LocalMamba: Visual State Space Model with Windowed Selective Scan](https://mlsanity.com/paper/2403.09338/) (2024)
- [A Survey on Efficient Inference for Large Language Models](https://mlsanity.com/paper/2404.14294/) (2024)
- [Mega: Moving Average Equipped Gated Attention](https://mlsanity.com/paper/2209.10655/) (2022)
- [U-KAN Makes Strong Backbone for Medical Image Segmentation and Generation](https://mlsanity.com/paper/2406.02918/) (2024)
- [Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling](https://mlsanity.com/paper/2406.07522/) (2024)
- [PointMamba: A Simple State Space Model for Point Cloud Analysis](https://mlsanity.com/paper/2402.10739/) (2024)
- [Protein Design with Guided Discrete Diffusion](https://mlsanity.com/paper/2305.20009/) (2023)
- [EfficientVMamba: Atrous Selective Scan for Light Weight Visual Mamba](https://mlsanity.com/paper/2403.09977/) (2024)
- [Liquid Structural State-Space Models](https://mlsanity.com/paper/2209.12951/) (2022)
- [A Survey on Large Language Model Acceleration based on KV Cache Management](https://mlsanity.com/paper/2412.19442/) (2024)
- [HGRN2: Gated Linear RNNs with State Expansion](https://mlsanity.com/paper/2404.07904/) (2024)
- [How to Train Your HiPPO: State Space Models with Generalized Orthogonal Basis Projections](https://mlsanity.com/paper/2206.12037/) (2022)
- [The Mamba in the Llama: Distilling and Accelerating Hybrid Models](https://mlsanity.com/paper/2408.15237/) (2024)
- [Caduceus: Bi-Directional Equivariant Long-Range DNA Sequence Modeling](https://mlsanity.com/paper/2403.03234/) (2024)
- [Motion Mamba: Efficient and Long Sequence Motion Generation](https://mlsanity.com/paper/2403.07487/) (2024)
- [Video Mamba Suite: State Space Model as a Versatile Alternative for Video Understanding](https://mlsanity.com/paper/2403.09626/) (2024)
- [Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks](https://mlsanity.com/paper/2402.04248/) (2024)
- [Hierarchically Gated Recurrent Neural Network for Sequence Modeling](https://mlsanity.com/paper/2311.04823/) (2023)
- [VL-Mamba: Exploring State Space Models for Multimodal Learning](https://mlsanity.com/paper/2403.13600/) (2024)
- [In-Context Language Learning: Architectures and Algorithms](https://mlsanity.com/paper/2401.12973/) (2024)
- [Mamba-3: Improved Sequence Modeling using State Space Principles](https://mlsanity.com/paper/2603.15569/) (2026)
- [What Makes Convolutional Models Great on Long Sequence Modeling?](https://mlsanity.com/paper/2210.09298/) (2022)
- [ZigMa: A DiT-style Zigzag Mamba Diffusion Model](https://mlsanity.com/paper/2403.13802/) (2024)

## Abstract (arXiv)

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.