On the Parameterization and Initialization of Diagonal State Space Models

• URL: http://arxiv.org/abs/2206.11893v1
• Date: Thu, 23 Jun 2022 17:58:39 GMT
• Title: On the Parameterization and Initialization of Diagonal State Space Models
• Authors: Albert Gu, Ankit Gupta, Karan Goel, Christopher R\'e
• Abstract summary: We show how to parameterize and initialize diagonal state space models. We show that the diagonal restriction of S4's matrix surprisingly recovers the same kernel in the limit of infinite state dimension.
• Score: 35.68370606343843
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• Provable Benefits of Complex Parameterizations for Structured State Space Models [51.90574950170374]
  Structured state space models (SSMs) are linear dynamical systems adhering to a specified structure. In contrast to typical neural network modules, whose parameterizations are real, SSMs often use complex parameterizations. This paper takes a step towards explaining the benefits of complex parameterizations for SSMs by establishing formal gaps between real and complex diagonal SSMs.
  arXiv  Detail & Related papers  (2024-10-17T22:35:50Z)
• Robustifying State-space Models for Long Sequences via Approximate Diagonalization [47.321212977509454]
  State-space models (SSMs) have emerged as a framework for learning long-range sequence tasks. diagonalizing the HiPPO framework is itself an ill-posed problem. We introduce a generic, backward-stable "perturb-then-diagonalize" (PTD) methodology.
  arXiv  Detail & Related papers  (2023-10-02T23:36:13Z)
• Deep Latent State Space Models for Time-Series Generation [68.45746489575032]
  We propose LS4, a generative model for sequences with latent variables evolving according to a state space ODE. Inspired by recent deep state space models (S4), we achieve speedups by leveraging a convolutional representation of LS4. We show that LS4 significantly outperforms previous continuous-time generative models in terms of marginal distribution, classification, and prediction scores on real-world datasets.
  arXiv  Detail & Related papers  (2022-12-24T15:17:42Z)
• Liquid Structural State-Space Models [106.74783377913433]
  Liquid-S4 achieves an average performance of 87.32% on the Long-Range Arena benchmark. On the full raw Speech Command recognition, dataset Liquid-S4 achieves 96.78% accuracy with a 30% reduction in parameter counts compared to S4.
  arXiv  Detail & Related papers  (2022-09-26T18:37:13Z)
• Long Range Language Modeling via Gated State Spaces [67.64091993846269]
  We focus on autoregressive sequence modeling over English books, Github source code and ArXiv mathematics articles. We propose a new layer named Gated State Space (GSS) and show that it trains significantly faster than the diagonal version of S4.
  arXiv  Detail & Related papers  (2022-06-27T01:50:18Z)
• How to Train Your HiPPO: State Space Models with Generalized Orthogonal Basis Projections [22.421814045703147]
  Linear time-invariant state space models (SSM) have been shown to be very promising in machine learning. We introduce a theoretically rich class of SSMs that also lets us derive more intuitive S4 variants for other bases. These insights improve S4's performance to 86% on the Long Range Arena benchmark, with 96% on the most difficult Path-X task.
  arXiv  Detail & Related papers  (2022-06-24T02:24:41Z)
• Diagonal State Spaces are as Effective as Structured State Spaces [3.8276199743296906]
  We show that our $textitDiagonal 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. 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.
  arXiv  Detail & Related papers  (2022-03-27T16:30:33Z)
• Efficiently Modeling Long Sequences with Structured State Spaces [15.456254157293836]
  We propose a new sequence model based on a new parameterization for the fundamental state space model. S4 achieves strong empirical results across a diverse range of established benchmarks, including (i) 91% accuracy on sequential CIFAR-10 with no data augmentation or auxiliary losses, on par with a larger 2-D ResNet.
  arXiv  Detail & Related papers  (2021-10-31T03:32:18Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.