[論文レビュー] Theoretical Foundations of Deep Selective State-Space Models

Nicola Muca Cirone, Antonio Orvieto+3 more|arXiv (Cornell University)|Feb 29, 2024

Complex Systems and Decision Making被引用数 4

ひとこと要約

本論文は、粗いパス理論を用いて、深層選択的状態空間モデル（SSM）の理論的基盤を確立し、Mamba や GLA における入力制御線形再帰が、入力を系列のシグネチャに投影することにより、時間スケールにわたる非線形相互作用を明示的に捉えることを証明している。主な結果として、ランダムに初期化された広い選択的 SSM は、追加の MLP がなくても完全表現可能であることが示され、対角形バージョンは深さを用いて高次統計を効率的に計算する。

ABSTRACT

Structured state-space models (SSMs) such as S4, stemming from the seminal work of Gu et al., are gaining popularity as effective approaches for modeling sequential data. Deep SSMs demonstrate outstanding performance across a diverse set of domains, at a reduced training and inference cost compared to attention-based transformers. Recent developments show that if the linear recurrence powering SSMs allows for multiplicative interactions between inputs and hidden states (e.g. GateLoop, Mamba, GLA), then the resulting architecture can surpass in both in accuracy and efficiency attention-powered foundation models trained on text, at scales of billion parameters. In this paper, we give theoretical grounding to this recent finding using tools from Rough Path Theory: we show that when random linear recurrences are equipped with simple input-controlled transitions (selectivity mechanism), then the hidden state is provably a low-dimensional projection of a powerful mathematical object called the signature of the input -- capturing non-linear interactions between tokens at distinct timescales. Our theory not only motivates the success of modern selective state-space models such as Mamba but also provides a solid framework to understand the expressive power of future SSM variants.

研究の動機と目的

• 現代の選択的 SSM を解析するための理論的枠組みを、粗いパス理論のツールを用いて提供すること。
• SSM における入力制御線形再帰が、入力系列の連続関数を近似する上で完全表現可能であることを証明すること。
• 対角形 SSM（例：Mamba）が深さを用いて高次統計を効率的に計算し、密行列型と同等の表現力を達成できることを示すこと。
• 制御された微分方程式（CDE）に基づく共通の数学的定式化により、S4、Mamba、GLA および関連モデルの分析を統一すること。
• 将来的な SSM の変種の理解と設計に一般化可能な理論的基盤を提供すること。

提案手法

• 入力系列によって駆動される制御された微分方程式（CDE）として SSM をモデル化し、隠れ状態が入力依存の線形遷移に従って進化することを仮定する。
• 粗いパス理論を適用して、隠れ状態を入力のシグネチャへの低次元射影として形式化し、非線形な時間的相互作用を捉える。
• 状態遷移行列 $ W_{s,t} $ を、入力駆動ベクトル場の反復積分の和として定義し、S4 や Mamba の再帰を一般化する。
• ワロンスキー行列と経路積分表現を用いて、加法的入力項を伴う CDE の明示的解を導出する。
• ランダムに初期化された密行列型入力制御再帰が、隠れ状態に十分な入力統計を捉えることで完全表現可能であることを証明する。
• 対角 SSM ブロックの連結が高次グローバル統計を計算することを示し、深さの極限において密行列型と同等の表現力を達成する。

Figure 1 : Example of the dynamics of the S4 and Mamba CDEs derived in Sec. 3.1 . The continuous-time approximation of these algorithms can be written as a Linear CDE. In S4, $\omega^{\text{X}}_{t}=t$ , while $\xi^{\text{X}}_{t}=\int_{0}^{t}X_{s}ds$ . In Mamba, $\omega^{\text{X}}_{t}=\int_{0}^{t}\si

実験結果

リサーチクエスチョン

• RQ1 Mamba のような現代的選択的 SSM の表現力は、どのように理論的に裏付けられるか？
• RQ2 選択的 SSM が長期間にわたる系列において非線形相互作用をモデル化できる背後にある数学的構造は何か？
• RQ3 追加の非線形部品（例：MLP）がなくても、入力制御線形再帰は完全表現可能か？
• RQ4 対角形 SSM と密行列型の両者は、高次入力統計をどのように捉えているか？
• RQ5 制御された微分方程式と粗いパス理論に基づく統一的枠組みを用いて、多様な SSM アーキテクチャを分析・比較できるか？

主な発見

• 選択的 SSM における入力制御線形再帰は、明示的に入力のシグネチャへの低次元射影として隠れ状態を表現し、非線形な時間的相互作用を捉える。
• ランダムに初期化された広い密行列型入力制御 SSM は完全表現可能である：追加の MLP がなくても、入力系列から出力への任意の連続関数を近似可能である。
• Mamba のような対角形 SSM は、深さを用いて高次統計を効率的に計算し、深さの極限において密行列型と同等の表現力を達成する。
• 選択的 SSM の隠れ状態は、数学的に制御された微分方程式（CDE）を解くことと同等であり、経路積分表現とワロンスキー行列を用いて明示的解が導出される。
• 本フレームワークにより、S4、Mamba、GLA および関連モデルが共通の理論的基盤に統一され、将来的な SSM 変種の厳密な比較・分析が可能になる。
• 理論により、選択的 SSM が標準的な S4 やアテンション機構よりも長距離モデリングで優れる理由が説明される：構造的で入力依存の遷移により、より豊かな入力統計を捉えているためである。

Figure 3 : Comparison of Linear CDE, Mamba and S5 on anti-symmetric signature prediction tasks. As suggested by the theory Mamba struggles to immediately capture these high order statistics, but performance improves with chaining . On the other hand S5 does not learn such non-linear features even af