[论文解读] 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 Interfaces被引用 73

一句话总结

本文分析如何对对角线状态空间模型（SSMs）进行参数化和初始化，结果表明一个简单的对角线 SSM（S4D）在性能上可以达到 S4 的水平，在图像、音频和医疗时间序列任务上取得出色结果，在 Long Range Arena 上达到 85%。同时给出理论洞见并比较对角线与 DPLR SSM 的经验研究。

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.

研究动机与目标

• 理解为何对角线约束通过初始化和核结构能够维持 S4 的性能。
• 系统性地对对角线 SSM 的参数化与计算选择进行分类。
• 证明在无限状态极限下，对角线 SSM 能复制 S4 动力学，并提供实用的初始化方案。
• 在图像、音频和医疗时间序列任务上进行对角线 SSM 变体的经验评估。
• 提供一种简单、可实现的对角线 SSM 核计算，能够媲美 DPLR 方法。

提出的方法

• 定义并分析对角线 SSM（A 为对角矩阵，B，C）及其核 K(t)=C e^{tA} B。
• 证明核计算简化为基于 Vandermonde 的矩阵运算，使时间复杂度在实现不同下达到 O(N+L) 或 O(NL)。
• 比较 S4、DSS 与 S4D 在参数化选择（离散化、B 和 C 的处理、特征值约束）方面的差异。
• 证明基于 HiPPO 的 A 矩阵的对角近似在极限 N→∞ 时恢复与 S4 相同的核（定理 3）。
• 提出 S4D 的变体（S4D-Inv、S4D-Lin），采用简单的 A 初始化并分析它们的经验性能。
• 对初始化、离散化以及对 B 的训练进行消融实验，以分离对性能的核心影响。

实验结果

研究问题

• RQ1 在参数化和初始化恰当时，对角线 SSM 能否匹配原始 S4 的性能？
• RQ2 在对角线 SSM 的参数化与计算中有哪些关键设计选择，它们如何影响性能？
• RQ3 基于 HiPPO 的对角初始化在大状态极限下与 S4 动力学有何关系？
• RQ4 简单、可实现的对角线 SSM（S4D）是否在图像、音频和时间序列领域获得竞争性结果？
• RQ5 在对角线 SSM 参数化中训练与冻结如 B 的组件有何影响？

主要发现

• S4D，作为对角线 SSM，可以在大多数设置下与 S4 相媲美，并在图像、音频和医疗时间序列基准上取得突出结果。
• 基于 HiPPO 的矩阵的对角近似在无限状态极限下恢复与 S4 相同的核（定理 3）。
• 通过 Vandermonde 风格乘积对角线 SSM 的核计算简单，可以用几行代码实现。
• 消融研究表明离散化选择和对 B 的训练相较于对 A 的初始化影响有限，支持 S4D 方法的简洁性。
• 带有 Inv 和 Lin 初始化的 S4D 变体提供可解释的基底（阻尼傅里叶样式），在 Long Range Arena 上达到近乎最先进的性能（85%）。
• 与 DSS 相比，S4D 避免了 softmax 归一化，核计算更简单，性能稳定。