On the parameterization and initialization of diagonal state space models
对角状态空间模型的参数化与初始化
状态空间模型(SSM)最近被证明是一种非常有效的深度学习层,作为序列模型(如RNN、CNN或Transformers)的一种有前途的替代方案。显示这一潜力的第一个版本是S4模型,它通过使用称为HiPPO矩阵的规定状态矩阵,对涉及长距离依赖性的任务特别有效。虽然这有一个可解释的数学机制来建模长依赖关系,但它引入了一种难以实现的自定义表示和算法。另一方面,S4的一个最新变体称为DSS,表明当使用基于近似S4矩阵的特定初始化时,将状态矩阵限制为完全对角仍然可以保持原始模型的性能。这项工作试图系统地理解如何参数化和初始化这种对角状态空间模型。虽然从经典结果可以看出,几乎所有SSM都具有等效的对角形式,但我们表明初始化对性能至关重要。我们通过证明S4矩阵的对角限制令人惊讶地在无限状态维度的极限下恢复了相同的内核,从而从数学上解释了DSS的工作原理。我们还系统地描述了在参数化和计算对角SSM时的各种设计选择,并进行了一项消除这些选择影响的受控经验研究。我们的最终模型S4D是S4的简单对角版本,其内核计算只需要2行代码,在几乎所有设置中的性能都与S4相当,在图像、音频和医疗时间序列领域具有最先进的结果,在Long Range Arena基准上的平均值为85%。
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.
网友评论