Mamba: Transformerに挑む状態空間モデル

MambaはTransformerに挑む状態空間モデル(State Space Model)

• AIは現在世界を支配しており、その中心にはTransformerがある
• Mambaは状態空間モデル(SSM)という代替的なモデルクラスに属する
• MambaはTransformerと似た性能とスケーラビリティを持ちながら、長いシーケンス長でも実行可能
• 注目すべき点は、Mambaが「アテンションメカニズム」の「二乗ボトルネック」を取り除き、長いコンテキストを可能にしていること
• MambaはTransformerより最大5倍高速に実行される

Transformerの問題点 - アテンションだけでは十分でないかもしれない

• Transformerではすべてのトークンが以前のトークンを参照できるため、コンテキストが大きくなるほどモデルは遅くなる
• このようなKVキャッシュの保存には、空間計算量としても O(n) が必要
• 既存のTransformerのボトルネックを緩和する技術はあるが、根本的な解決には別のアプローチが必要

基礎モデルのバックボーン

• 優れたMLアーキテクチャのバックボーンにおける2つの重要な構成要素
  • トークン間の 通信(Communication)
  • トークン内の 計算(Computation)
• TransformerブロックはAttentionとMLPsで構成されている
• Mambaは通信のために制御理論に着想を得たSSMを使い、計算のためにはMLPスタイルの射影を維持する

Mambaの動機 - Temple Runに戻る

• 状態(state) とは、システムの将来の振る舞いを決定するのに必要な変数を意味する
• 状態は過去について知るべきすべてを圧縮したものであり、マルコフ決定過程へと変換される

離散化 - 量子化された世界で生きる

• 連続時間の微分方程式を離散時間の差分方程式へ変換する過程を 離散化(discretisation) と呼ぶ
• Mambaはゼロ次ホールド(ZOH)離散化を使用する

SSM行列を理解する

• A、B、C、D行列は、それぞれ状態遷移、新しい入力を状態にマッピング、状態をSSM出力にマッピング、新しい入力を出力へ通過させる役割を持つ

効率性 vs 効果性: Attention is Focus, Selectivity is Prioritisation（注意は集中、選択性は優先順位付け）

• Transformerは非常に 効果的 だが 効率的 ではない
• Mambaアーキテクチャは、効率性と効果性のパレートフロンティアを押し広げるソリューションを提供する

選択メカニズム

• 選択性(Selectivity) は、各トークンをその必要性に応じて状態へ変換できるようにする
• MambaはA、B、C行列をxの関数にすることで、静的ではなくコンテキスト依存にしている

選択性の問題点

• 選択メカニズムを適用すると、非選択的なSSMより計算が遅くなる可能性がある
• ハードウェア最適化によって、Mambaは同程度のサイズのTransformerより高速に実行できる

機械学習と政治経済学 - 状態の大きさはどれくらいあるべきか

• シーケンスモデルにおける効率性と効果性のトレードオフは、状態をどれだけうまく圧縮できるかによって特徴づけられる
• 状態表現が重要であり、状態を選択的かつ動的に圧縮することが鍵となる

Transformer対Mambaの情報フロー

• Transformerは訓練データとコンテキストデータを通じて学習する
• Mambaでは訓練データとコンテキストデータが圧縮・フィルタリングされてアクセス可能になる

新しいプロンプティング・パラダイムとしての状態の置き換え

• Mambaのようなモデルを使えば、専門データから生成された状態ライブラリを共有できる
• 状態によって、バックプロップなしでも推論時に無限のコンテキスト学習を適用できる

Mambaとメカニズム解釈可能性

• Mambaの解釈可能性は、トークン間で情報がどう移動するかを理解することに焦点を当てる

MambaとSSMが次にやること

• Mambaのようなモデルは、非常に長いコンテキストと長期記憶が必要なシナリオで優れた性能を発揮する可能性がある

エージェントとAI安全性

• 言語モデルは本質的に安全だが、長期的なシーケンス推論の可能性は、エージェントベースのAI安全性の重要性を改めて浮上させる

TransformerとMambaの最良の協業

• Mambaの長いコンテキストと、Transformerの短いシーケンスに対する高解像度を組み合わせることには価値がある

GN⁺の見解

• MambaはTransformerのボトルネックを解決し、長いシーケンス処理に有効な代替案を提示する
• この技術は特に、医療、遺伝学、自然言語処理など、データの長いシーケンスが重要な分野で有用になりうる
• Mambaの選択メカニズムが実際に有効かどうかについて、さらなる研究が必要
• Mambaの選択性は、Transformerが提供する高い精度と効率性のあいだでバランスを見つける助けになる可能性がある