[論文レビュー] Simplified State Space Layers for Sequence Modeling

Jimmy T. H. Smith, Andrew Warrington+1 more|arXiv (Cornell University)|Aug 9, 2022

Parallel Computing and Optimization Techniques被引用数 77

ひとこと要約

この論文は S5 レイヤを導入します。多入力多出力の状態空間レイヤで、S4 と効率性を同等に保ちつつ、完全に再帰的な時系列ドメインアプローチを用い、長距離シーケンスタスクで最先端の結果を達成します。

ABSTRACT

Models using structured state space sequence (S4) layers have achieved state-of-the-art performance on long-range sequence modeling tasks. An S4 layer combines linear state space models (SSMs), the HiPPO framework, and deep learning to achieve high performance. We build on the design of the S4 layer and introduce a new state space layer, the S5 layer. Whereas an S4 layer uses many independent single-input, single-output SSMs, the S5 layer uses one multi-input, multi-output SSM. We establish a connection between S5 and S4, and use this to develop the initialization and parameterization used by the S5 model. The result is a state space layer that can leverage efficient and widely implemented parallel scans, allowing S5 to match the computational efficiency of S4, while also achieving state-of-the-art performance on several long-range sequence modeling tasks. S5 averages 87.4% on the long range arena benchmark, and 98.5% on the most difficult Path-X task.

研究の動機と目的

• より簡潔で再帰的な時系列ドメインレイヤによる、長距離シーケンスモデリングの効率化を動機づけ、実現する。
• S4 の bank of SISO SSMs を単一の MIMO SSM に置換し、効率とスケーラビリティを改善する。
• 長距離依存性を捕捉するために diagonalized HiPPO に触発された初期化を活用する。
• 標準的な長距離ベンチマークで競争力のある、または優れた性能を示すとともに、不規則なサンプリングと可変観測間隔をサポートする。

提案手法

• S4 bank of independent SISO SSMs を、潜在サイズ P の diagonalized MIMO SSM 一つに置換する。
• 状態行列 A を対角化して時系列ドメインでの効率的な並列スキャンを可能にする。
• 対角化された HiPPO に触発された初期化（HiPPO-N）を行い、指針となる仮定の下で S5 のダイナミクスを S4 と関連づける。
• 連続時間 SSM を零次ホールドで離散化し、学習可能な時間スケール Δ を用いる。
• 非線形活性化と混合ステップを用いて SSM の前活性化からレイヤ出力を生成する。
• 潜在サイズ P = O(H) の場合、S5 は実行時間とメモリの複雑さで S4 に匹敵する。

実験結果

リサーチクエスチョン

• RQ1 対角化ダイナミクスを持つ単一の MIMO SSM は、S4 の銀行の SISO SSMs と同等の機能容量を実現できるか？
• RQ2 HiPPO に触発された初期化（対角化 HiPPO-N）は、MIMO 設定において S5 の効果的な開始点を提供するか？
• RQ3 S5 レイヤは、シーケンス長に対して線形の複雑さを保ちつつ、不規則なサンプリングと時間変化パラメータを扱えるか？
• RQ4 S4 および他のベースラインと比較して、長距離シーケンスベンチマークで S5 の性能はどうか？
• RQ5 S5 設計のコア前提（例：同一の A、同一 Δ、P ≈ H）を緩和した場合の実用的影響は？

主な発見

モデル	ListOps	テキスト取得	画像	Pathfinder	Path-X	Avg.
S5	62.15	89.31	91.40	88.00	95.33	98.58	87.46

• S5 は長距離シーケンスタスクで線形のシーケンス長に対して実行時間を要し、効率性は S4 と同等で最先端の性能を達成する。
• Long Range Arena ベンチマークで、S5 は平均スコア 87.4%、難易度が高い Path-X タスクで 98.5% を達成。
• S5 は音声分類でいくつかのベースラインを上回り、複数のタスクで S4 のバリアントと同等またはそれを上回る。
• 対角化 HiPPO-N 初期化は S4 で用いられる HiPPO-LegS 初期化と比較して競争力のある性能を提供し、安定した対角化可能ダイナミクスを可能にする。
• S5 は時間ドメインのパラレルスキャンを通じて不規則サンプリングと時刻変化の SSM をサポートする、という点で、畳み込みベースの S4 カーネルとは異なる。