[Paper Review] On the Parameterization and Initialization of Diagonal State Space Models

Albert Gu, Ankit Gupta+2 more|arXiv (Cornell University)|Jun 23, 2022

EEG and Brain-Computer Interfaces73 citations

TL;DR

The paper analyzes how to parameterize and initialize diagonal state space models (SSMs), showing that a simple diagonal SSM (S4D) can match S4 in performance, with strong results on image, audio, and medical time-series tasks and 85% on Long Range Arena. It provides theoretical insights and an empirical study comparing diagonal versus DPLR SSMs.

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.

Motivation & Objective

• Understand why a diagonal restriction can preserve S4 performance via initialization and kernel structure.
• Systematically categorize parameterization and computation choices for diagonal SSMs.
• Show that a diagonal SSM can replicate S4 dynamics in the infinite-state limit and provide practical initialization schemes.
• Empirically evaluate diagonal SSM variants across image, audio, and medical time-series tasks.
• Provide a simple, implementable kernel computation for diagonal SSMs that rivals DPLR approaches.

Proposed method

• Define and analyze diagonal SSMs (A diagonal, B, C) and their kernel K(t)=C e^{tA} B.
• Show that kernel computation reduces to a Vandermonde-based matrix operation enabling O(N+L) or O(NL) time depending on implementation.
• Compare parameterization choices (discretization, B and C handling, eigenvalue constraints) between S4, DSS, and S4D.
• Prove that the diagonal approximation of the HiPPO-based A matrix recovers the same kernel as S4 in the limit N→∞ (Theorem 3).
• Propose S4D variants (S4D-Inv, S4D-Lin) with simple A initializations and analyze their empirical performance.
• Conduct ablations on initialization, discretization, and training of B to isolate core effects on performance.

Experimental results

Research questions

• RQ1 Can a diagonal SSM match the performance of the original S4 when parameterized and initialized appropriately?
• RQ2 What are the essential design choices in parameterizing and computing diagonal SSMs, and how do they affect performance?
• RQ3 How does the diagonal HiPPO-based initialization relate to S4 dynamics in the large-state limit?
• RQ4 Do simple, implementable diagonal SSMs (S4D) achieve competitive results across image, audio, and time-series domains?
• RQ5 What is the impact of training versus freezing components like B in diagonal SSM parameterizations?

Key findings

• S4D, a diagonal SSM, can perform comparably to S4 across most settings and achieves strong results on image, audio, and medical time-series benchmarks.
• The diagonal approximation of HiPPO-based matrices recovers the same kernel as S4 in the infinite-state limit (Theorem 3).
• Kernel computation for diagonal SSMs via Vandermonde-like products is simple and can be implemented in a few lines of code.
• Ablations show that discretization choices and training B have limited impact relative to initialization of A, supporting the simplicity of the S4D approach.
• S4D variants with Inv and Lin initializations offer interpretable bases (damped Fourier-like) and achieve near-state-of-the-art performance on Long Range Arena (85%).
• Compared to DSS, S4D avoids softmax normalization, yielding simpler kernel computation and stable performance.