[論文レビュー] 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

ひとこと要約

この論文は対角SSM（対角状態空間モデル）をどのようにパラメータ化し初期化するかを分析し、単純な対角SSM（S4D）がS4と同等の性能を発揮できることを示す。画像、音声、医療時系列タスクおよび Long Range Arena で85%の強力な結果を得ており、理論的洞察と対角SSMと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のバリアントを経験的に評価する。
• DPLRアプローチに匹敵する対角SSMの簡便なカーネル計算を提供する。

提案手法

• 対角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）。
• 対角SSMのカーネル計算は Vandermonde 的積を用いて簡単で、数行のコードで実装可能。
• 離散化の選択とBの訓練は、Aの初期化と比較して影響が小さく、S4Dの単純さを支持する。
• InvおよびLin初期化を用いるS4Dの派生形は解釈可能な基底（減衰したフーリエ様）を提供し、Long Range Arenaでほぼ最先端性能を達成（85%）。
• DSSと比較して、S4Dはソフトマックス正規化を回避し、より単純なカーネル計算と安定した性能をもたらす。