[論文レビュー] Diagonal State Spaces are as Effective as Structured State Spaces

Ankit Gupta, Gu, Albert+1 more|arXiv (Cornell University)|Mar 27, 2022

Topic Modeling被引用数 64

ひとこと要約

本論文は、状態空間の完全対角パラメータ化（DSS）が長距離タスクにおいてS4モデルの性能に匹敵しつつ、実装と解析がより単純であることを示す。

ABSTRACT

Modeling long range dependencies in sequential data is a fundamental step towards attaining human-level performance in many modalities such as text, vision, audio and video. While attention-based models are a popular and effective choice in modeling short-range interactions, their performance on tasks requiring long range reasoning has been largely inadequate. In an exciting result, Gu et al. (ICLR 2022) proposed the $ extit{Structured State Space}$ (S4) architecture delivering large gains over state-of-the-art models on several long-range tasks across various modalities. The core proposition of S4 is the parameterization of state matrices via a diagonal plus low rank structure, allowing efficient computation. In this work, we show that one can match the performance of S4 even without the low rank correction and thus assuming the state matrices to be diagonal. Our $ extit{Diagonal State Space}$ (DSS) model matches the performance of S4 on Long Range Arena tasks, speech classification on Speech Commands dataset, while being conceptually simpler and straightforward to implement.

研究の動機と目的

• モダリティ横断の長距離依存性モデリングを動機付ける（テキスト、画像、音声）。
• DLPRベースのS4アーキテクチャの性能と同等の対角状態空間パラメータ化が実現できるかを評価する。
• 長距離ベンチマークで競争力のある精度を持つ、S4のより単純で透明性の高い代替案を提供する。

提案手法

• 対角状態行列を強制し、シーケンス処理のカーネルを導出することで、対角State Space (DSS)モデルを導入する。
• 訓練を安定化させるための2つのDSSバリアント：DSS exp（対角Aを実部が負の場合に使用し、HiPPOに由来する固有構造を持つ）とDSS softmax（行ごとに指数関数を正規化して成長を抑える）。
• 離散化されたカーネルと入力列の高速畳み込み（FFTベース）により、状態空間カーネルを効率的に計算する。
• 初期化スキーム（Skew-Hippo）と学習を安定化させ、優れた性能を達成するための訓練の詳細を提供する。
• 自己回帰的デコード時に、閉形式の再帰を持つ対角Aで状態を取得する方法を説明する。

実験結果

リサーチクエスチョン

• RQ1 完全に対角な状態空間パラメータ化（DSS）は、長距離シーケンスタスクでDLPR S4モデルと同等の性能を達成できるか？
• RQ2 状態空間パラメータ化を簡素化（対角のみ）しても表現力を保ちつつ、単純さと解釈性を向上させるか？
• RQ3 初期化とカーネル長が、テキスト・画像・音声の各モダリティにまたがる長距離タスクの性能にどう影響するか？
• RQ4 長いシーケンスにおけるDSSとS4および注意機構ベースのトランスフォーマーとの実用的なトレードオフ（計算効率、安定性）は何か？

主な発見

モデル	ListOps	Text	Retrieval	Image	Pathfinder	Path-X	平均
Transformer	36.37	64.27	57.46	42.44	71.40	✗	53.66
Reformer	37.27	56.10	53.40	38.07	68.50	✗	50.56
BigBird	36.05	64.02	59.29	40.83	74.87	✗	54.17
Linear Trans.	16.13	65.90	53.09	42.34	75.30	✗	50.46
Performer	18.01	65.40	53.82	42.77	77.05	✗	51.18
FNet	35.33	65.11	59.61	38.67	77.80	✗	54.42
Nyströmformer	37.15	65.52	79.56	41.58	70.94	✗	57.46
Luna-256	37.25	64.57	79.29	47.38	77.72	✗	59.37
H-Transformer-1D	49.53	78.69	63.99	46.05	68.78	✗	61.41
S4 (as in [ GGR22 ] )	58.35	76.02	87.09	87.26	86.05	88.10	80.48
S4 ( our run )	57.6	75.4	87.6	86.5	86.2	88.0	80.21
DSS softmax (ours)	60.6	84.8	87.8	85.7	84.6	87.8	81.88
DSS exp (ours)	59.7	84.6	87.6	84.9	84.7	85.6	81.18
DSS exp-no-scale (ours)	59.3	82.4	86.0	81.2	81.3	✗	65.03

• DSS softmaxは平均LRA精度81.88を達成し、S4（80.21）にほぼ匹敵し、Transformerベースラインを約20ポイント上回る。
• DSS expも競争力のある成績（平均81.18）を示し、訓練がより単純になり得る一方、DSS exp-no-scaleは長距離タスクで通常劣る。
• Speech Commandsの生の音声では、DSSバリアントが98.2%（DSS exp）および97.7–98.2%（他）で、S4（98.1–98.32%）と同程度であり、ドメイン横断の有効性を示す。
• アブレーション実験では、対角固有値のランダム初期化が性能を低下させ、カーネル長の切り詰めが精度を低下させることが示され、長距離モデリングと適切な初期化の重要性を強調している。
• カーネルの可視化は、Path-Xのような一部タスクが長距離相互作用に依存する一方、ImageやSCはより局所的であることを示しており、カーネル長の感度のタスクごとの差を説明している。