State Space Models（SSM）とMamba

Transformerの代替アーキテクチャとして注目されるState Space Models。S4による長距離依存の学習、Mambaの選択的状態空間による言語モデリング性能。線形計算量とRNNライクな推論効率。

最終更新：2025年11月

1. 背景と動機

1.1 Transformerの限界

  • 計算量：Self-Attentionは$O(n^2)$で系列長に対して二乗
  • メモリ：KVキャッシュが推論時に線形増加
  • 推論：各トークン生成で全履歴を参照

長文脈（100K+トークン）では深刻なボトルネック。

1.2 RNNの再評価

RNNの利点：

  • 固定サイズの隠れ状態
  • 線形計算量$O(n)$
  • 推論時のKVキャッシュ不要

課題は長距離依存の学習困難さ。これをSSMで解決できないか？

1.3 SSMへの期待

制御理論・信号処理のState Space Modelを深層学習に応用。

  • RNNの効率性
  • Transformerに匹敵する長距離モデリング
  • 畳み込みによる並列訓練

2. SSMの基礎

2.1 連続時間SSM

制御理論の線形時不変（LTI）システム：

$$\dot{h}(t) = Ah(t) + Bx(t)$$ $$y(t) = Ch(t) + Dx(t)$$
  • $h(t) \in \mathbb{R}^N$：隠れ状態
  • $x(t) \in \mathbb{R}$：入力
  • $y(t) \in \mathbb{R}$：出力
  • $A \in \mathbb{R}^{N \times N}$, $B \in \mathbb{R}^{N \times 1}$, $C \in \mathbb{R}^{1 \times N}$：パラメータ

2.2 離散化

デジタル信号処理のために離散化。Zero-Order Hold (ZOH)：

$$\bar{A} = \exp(\Delta A)$$ $$\bar{B} = (\Delta A)^{-1}(\exp(\Delta A) - I) \cdot \Delta B$$

$\Delta$：サンプリング間隔（学習可能）

離散時間更新式：

$$h_k = \bar{A}h_{k-1} + \bar{B}x_k$$ $$y_k = Ch_k$$

2.3 畳み込みとしての解釈

LTIシステムは畳み込みと等価：

$$y = x * \bar{K}$$ $$\bar{K} = (C\bar{B}, C\bar{A}\bar{B}, C\bar{A}^2\bar{B}, ...)$$

畳み込みカーネル$\bar{K}$を事前計算可能。FFTで$O(n \log n)$で計算。

2.4 二つのモード

  • 訓練時：畳み込みモード（並列処理）
  • 推論時：再帰モード（逐次処理、固定メモリ）

この柔軟性がSSMの強み。

3. S4：Structured State Space

3.1 HiPPO理論

Gu et al. (2020)。長期記憶のための状態空間初期化。

HiPPO行列：入力履歴の直交多項式近似を維持。

特に HiPPO-LegS：ルジャンドル多項式による圧縮。

3.2 S4の構造

Gu et al. (2022)。HiPPOを実用化。

課題：$N \times N$行列$A$の計算が$O(N^2)$

解決：対角+低ランク構造（DPLR）

$$A = \Lambda - PQ^*$$

$\Lambda$：対角行列、$P, Q$：低ランク行列

これにより$O(N)$で計算可能。

3.3 S4の性能

Long Range Arena (LRA)：長距離依存ベンチマーク

  • Path-X（16Kトークン）：S4が解けた初のモデル
  • 画像、テキスト、数学など多様なタスクで優位

しかし言語モデリングではTransformerに劣る。

3.4 S4の変種

  • S4D：対角SSM（さらに単純化）
  • S5：並列スキャン実装
  • H3：SSM + Attention のハイブリッド

主要論文：

  • Gu et al. (2020) "HiPPO: Recurrent Memory with Optimal Polynomial Projections", NeurIPS
  • Gu et al. (2022) "Efficiently Modeling Long Sequences with Structured State Spaces (S4)", ICLR

4. Mamba：Selective SSM

4.1 S4の限界

S4はLTI（線形時不変）：パラメータ$(A, B, C)$が入力に依存しない。

問題：コンテンツ依存の選択ができない。

例：「重要なトークンに注目」「無関係な情報を忘却」が困難。

4.2 Selective SSM

Gu & Dao (2023)。パラメータを入力依存に。

$$B_t = f_B(x_t), \quad C_t = f_C(x_t), \quad \Delta_t = f_\Delta(x_t)$$

$f_B, f_C, f_\Delta$：線形射影

直感：

  • $\Delta$が大きい→入力を強く記憶
  • $\Delta$が小さい→以前の状態を維持
  • 入力依存で選択的に記憶・忘却

4.3 ハードウェア効率的実装

入力依存でLTI性が失われる→FFT畳み込みが使えない。

解決：Parallel Scan（並列プレフィックス和）+ カーネル融合

  • SRAMに状態を保持
  • HBMとの転送を最小化
  • FlashAttentionと同様のIO-aware実装

4.4 Mambaアーキテクチャ

SSMブロックの構成：

  1. 線形射影で次元拡大
  2. 1D畳み込み
  3. Selective SSM
  4. 乗算ゲート
  5. 線形射影で次元削減

TransformerのAttention + FFNをSSMブロックに置換。

4.5 性能

言語モデリング：

  • Transformerに匹敵するperplexity
  • 特に長文脈で優位
  • 推論速度：最大5倍高速

その他のタスク：

  • DNA配列モデリング
  • 音声認識
  • 時系列予測

4.6 Mamba-2

Dao & Gu (2024)。さらなる効率化。

  • 状態空間双対性（SSD）：SSMと線形Attentionの統一理論
  • より効率的な実装
  • スケーリング性の改善

主要論文：

  • Gu & Dao (2023) "Mamba: Linear-Time Sequence Modeling with Selective State Spaces", arXiv
  • Dao & Gu (2024) "Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality", ICML

5. Transformerとの比較

5.1 計算量とメモリ

特性 Transformer Mamba
訓練計算量 $O(n^2 d)$ $O(n d N)$
推論（1トークン） $O(n d)$（KVキャッシュ参照） $O(d N)$（固定）
KVキャッシュ $O(n d)$（系列長に比例） $O(d N)$（固定）
長距離依存 直接参照 状態に圧縮

$n$：系列長、$d$：隠れ次元、$N$：状態次元

5.2 能力の比較

Transformerが優位：

  • In-Context Learning
  • 長距離のコピー・検索
  • 精密な位置参照

Mambaが優位：

  • 超長文脈処理
  • 推論効率
  • 連続的なストリーム処理

5.3 スケーリングの検証

Mambaは現時点で〜3Bパラメータで主に検証。

100B+でのTransformerとの比較は限定的。

2024-2025年に大規模検証が進行中。

6. ハイブリッドアーキテクチャ

6.1 Attention + SSMの組み合わせ

両者の長所を活かす設計。

6.2 Jamba（AI21 Labs, 2024）

  • Transformer層とMamba層を交互に配置
  • MoEも組み込み
  • 256Kトークン文脈
  • 52B総パラメータ、12B活性化

主要論文：

  • Lieber et al. (2024) "Jamba: A Hybrid Transformer-Mamba Language Model"

6.3 その他のハイブリッド

  • Zamba：類似のハイブリッド設計
  • Griffin（DeepMind）：ゲート付きリニアRNN
  • RWKV：RNNライクな効率性 + Transformer性能

6.4 設計上の考慮

  • Attention層の配置頻度
  • どの機能をどちらに任せるか
  • 訓練安定性への影響

7. 応用と展望

7.1 現在の主な応用

  • 言語モデリング：効率的なLLM
  • DNA配列：ゲノム解析
  • 音声・オーディオ：長時間音声処理
  • 時系列：センサーデータ、金融
  • ビジョン：Vision Mamba（ViM）

7.2 未解決の課題

  • 大規模（100B+）でのスケーリング検証
  • In-Context Learningの改善
  • マルチモーダル対応
  • 最適なハイブリッド比率の探索

7.3 今後の方向性

  • SSM + Attentionの最適な組み合わせ
  • ハードウェア最適化（専用アクセラレータ）
  • 新しい状態空間構造の探索
  • 理論的理解の深化

8. 参考文献

SSM基礎

  • Gu et al. (2020) "HiPPO", NeurIPS
  • Gu et al. (2022) "S4", ICLR

Mamba

  • Gu & Dao (2023) "Mamba", arXiv
  • Dao & Gu (2024) "Mamba-2 / SSD", ICML

ハイブリッド

  • Lieber et al. (2024) "Jamba"
  • De et al. (2024) "Griffin: Mixing Gated Linear Recurrences with Local Attention", arXiv

解説リソース

  • Sasha Rush "The Annotated S4"
  • Albert Gu "Lecture: State Space Models"

関連ページ