Mambaモデルとは

Mambaモデルは、おそらくディープラーニング・アーキテクチャーとして初めて、Transformerモデルがその名声を得た本来のタスク、すなわち言語モデリングにおいて、その有効性に匹敵する存在となります。特筆すべきは、Mambaアーキテクチャーが主要なLLMベンチマーク評価において、同規模のTransformerと同等の性能を示しつつ、レイテンシーやメモリー要件の面で大幅に効率的であることを実証している点です。

Mambaアーキテクチャーは、Tri Dao氏とAlbert Gu氏によって2023年の論文「Mamba: Linear-Time Sequence Modeling with Selective State Spaces」で初めて提唱されました。1年後、彼らは元のMamba論文に続き、SSMとTransformerの関連を探るとともに、Mambaアーキテクチャーを改良し大幅に高速化したバージョンを提示し、それをMamba-2と名付けました。

オリジナルのMamba論文発表から2年が経過した現在も、LLMにおける主流はTransformerであり続けていますが、このアーキテクチャーはオープンソース・モデルにおいて採用される例が増えています。Mistral AI社のCodestral Mambaのように、純粋なMambaモデルもあります。さらに、AI2I社のJambaシリーズやIBM® Granite 4.0を含む多くのモデルは、アテンション（Transformer）レイヤーとSSM（Mamba）レイヤーの両方を組み込んだハイブリッド・モデルです。性能面でのメリットに加え、Mambaベースのモデルの普及は、比較的安価なハードウェア上でも円滑に動作することから、AIへのアクセスの民主化を実現することが期待されています。

SSMはもともと、入力に基づいて電気信号、気象パターン、移動物体の軌道といった連続的な系列の次の状態を予測するために設計されました。概念的にも数学的にも、SSMは2017年にTransformerが登場する以前に自然言語処理（NLP）を支配していたリカレント・ニューラル・ネットワーク（RNN）と関連しており、さらに畳み込みニューラル・ネットワーク（CNN）や隠れマルコフモデル（HMM）などの他の機械学習アルゴリズムとも関係しています。

名前が示すとおり、SSMは状態空間をモデル化することで動的システムの次の状態を予測します。状態空間とは、システムの状態を記述するすべての状態変数と、それぞれの変数が互いに関連しながら取り得る可能性の範囲を表す数学的表現です。

SSMは入力系列x(t)を受け取り、それをRNNの隠れ状態に相当する潜在状態表現h(t)にマッピングし、出力系列y(t)を予測します。あらゆるSSMの中核には2つの方程式があります。

• 状態方程式： $h^{\prime }\left(t\right)=A*h\left(t\right)+B*x\left(t\right)$ 
• 出力方程式： $y\left(t\right)=C*h\left(t\right)+D*x\left(t\right)$ 

モデルの主要なパラメーターはA、B、C、Dであり、これらは通常、重み行列の形を取ります。制御理論など、SSMが従来用いられてきた分野においては、これらの行列は固定されていると仮定されることが多く、確立されたシステムの動態を表します。そしてSSMは、望ましい出力yを導く入力xを求めるために利用されます。より現代的なSSMの概念においては、これらの行列自体が機械学習によって最適化されるパラメーターとなります。ディープラーニング・モデルにおいては、これらの行列はニューラル・ネットワークの学習可能な重みとして表現されます。

状態方程式は、状態がどのように変化するかを記述します。行列Aの値は、各状態変数が放置された場合に時間とともにどのように変化するかを決定します。行列Bの値は、テキスト系列における次のトークンのような入力が各状態変数にどのように影響するかを決定します。

言語モデリングにおいて、現在の状態はテキスト系列の文脈を表し、各トークンの後に更新されます。その役割は、TransformerモデルにおけるKVキャッシュと同等です。

出力方程式は、現在の状態が出力にどのように影響するか（行列Cを介して）、さらに入力が直接出力にどのように影響するか（行列Dを介して）を記述します。行列Dは本質的にh(t)そのもののモデル化には外部的な存在であるため、SSMの図や議論では、中核となる行列A、B、Cに焦点を当てるために省略されることが多いです。

Mamba LLMにおいては、出力方程式が次のトークンを生成するために使用されます。

従来のSSMは連続的な入力をモデル化するよう設計されていますが、テキスト系列（および現代のディープラーニングモデルで処理される他の多くのデータモダリティ）は離散的な入力です。SSMを用いて離散系列をモデル化するには、その個別の特定のタイムステップを連続信号の一部として表現する手段が必要です。

概念的には、離散化とは連続関数の値を特定の時点でサンプリングすることに相当します。これは新しいパラメーター、すなわちステップ・サイズ（∆と表記）を導入することを意味します。このパラメーターは、各離散的な時刻tにおいて、その値がどれだけの間サンプリングされ、あるいは「保持」されるかを決定します。∆の調整は、（時系列データにおける）解像度や（動画データにおける）フレーム・レートといった特性の変更に相当します。「離散化」には複数の方法がありますが、Mambaを含むほとんどの現代的なSSMの派生モデルは、単純なゼロ次ホールド（ZOH）法を使用しています。

SSMを離散化することで、系列対系列タスクにおいてRNNのように利用できるようになります。離散化されたSSMのパラメーターや方程式は、連続時間のものと区別するために、通常RNNで用いられる添字表記を使って書き換えられます。この表記では、h_tはモデルが生成する更新後の状態空間を表し、h_t-1はその前の状態、すなわち現在の状態空間を表します。

 $h_t=\overline{A}{h}_{t-1}+\overline{B}{x}_t$ 
 $y_t=\overline{C}{h}_t$

標準的な離散SSMを用いてテキスト・データをモデル化することは、RNNと共通するいくつかの欠点のために現実的ではありません。これらの欠点のうち2つは、2021年にAlbert Gu氏らによって提案された構造化状態空間系列モデル（S4モデル）によって解決されました。それは、学習の非効率性と長い系列をモデル化できないという問題です。

S4モデルの成功と、それに由来する多数の派生モデル（対角SSM〈DSS〉、対角S4〈S4D〉、H3モデルなど）の登場が、Mambaへとつながる道を直接切り開きました。

離散化されたSSMがRNNの特定の一形態と同等であることの利点は、RNNが推論において非常に高速である点にあります。しかしその一方で、RNNは学習が非常に遅いという欠点があります。

幸いなことに、離散化されたSSMには他のRNNと一線を画す重要な特性があります。それは、線形依存関係のみをモデル化するという点です。言い換えれば、単純で明快な乗算と加算の演算のみを使用します。S4論文が示すように、この単純で繰り返し相互依存する線形漸化式は、1次元畳み込みカーネルに展開でき、入力xを出力yへと直接1ステップでマッピングします。これは高速フーリエ変換を用いることで非常に効率的に計算できます。

唯一の「難点」は、入力系列全体のすべてのステップが既知である場合にのみ、これが可能になるという点です。これは推論時には不可能ですが、学習時には可能です。したがって、構造化SSMは両方の利点を享受できます。すなわち、学習時にはCNNとして非常に効率的に動作し、推論時にはRNNとして非常に効率的に動作します。

ほとんどのRNNと同様に、標準的なSSMは長距離依存関係のモデル化が本質的に苦手です。言い換えれば、段落の冒頭と末尾の単語のように系列内で離れたステップ間の関係を理解するのが苦手で、そのため長い系列全体のモデル化にも弱い傾向にあります。

この問題を解決するために、Gu氏とその共同執筆者（その1人がTri Dao氏）は、HiPPO（High-order Polynomial Projection Operatorsの略）と呼ばれる手法を用いました。これは、直交多項式から導かれた数式を用いて初期値を構造化することで、行列AとBの挙動を定義するものです。これは、モデルの重みを学習開始時にランダムに初期化するという標準的な機械学習の慣習とは対照的です。S4において、Dao氏とGu氏はルジャンドル多項式に基づく初期化手法を提案しました。彼らは続編となる論文「How to Train Your HiPPO」¹で、追加の数式を探求しました。

S4論文では、「単にSSMのランダム行列Aを\[HiPPO行列]に変更しただけで、逐次MNISTベンチマークにおける性能が60％から98％に向上し、SSMの長期記憶問題を実質的に解決した」と記されています。DSS、S5、Mambaといった構造化SSMの後続の派生モデルでは、行列AとBに対して異なる（多くの場合より単純な）初期化手法が用いられていますが、それでもHiPPOの基本原則は保持されています。すなわち、安定した更新を可能にし、行列内の各値の間にある程度の独立性を持たせる対角構造を実装しています。

Mambaアーキテクチャーの中核には2つの革新があります。1つ目は選択的状態空間モデルであり、これはMambaに対し、従来はTransformerモデルだけが持っていた重要な機能（現在の関連性に基づいて過去の入力履歴の特定部分に選択的に注目したり無視したりする能力）を与えます。もう1つはハードウェア対応の並列スキャンであり、これはGPU（グラフィックス・プロセッシング・ユニット）がモデルの計算をメモリー階層内で処理する方法を最適化し、速度と計算効率を最大化するアルゴリズムです。

Transformerにおいては、この能力はアテンション機構によって提供されます。これは現在の入力トークンとの関連性に基づき、各過去のトークンの影響を強調または減じるアテンション重みを調整する仕組みです。通常のSSMは、入力履歴全体を用いて入力を出力にマッピングするように明示的に設計されています。これは一部の系列モデリングタスクでは許容できる、あるいは望ましい場合さえありますが、先進的な言語モデリングタスクの多くにとっては大きな不利となります。

入力履歴の特定部分を動的に省略したり強調したりできないという欠点を克服するために、Dao氏とGu氏は「選択的スキャン」を備えた新しいクラスの状態空間モデルを提案しました。Mamba論文の中で、著者らは「選択的SSMをS6モデルと略すことがある。なぜならそれはS4モデルに選択メカニズムを加え、スキャンで計算するものだからだ」と述べています。また、S6ベースのアーキテクチャーに「Mamba」というニックネームを付けたのは、他の理由に加えて、それらすべてのSがヘビのシューという音のように聞こえるためです。

Mambaは、その中核に選択的状態空間モデルを備えたニューラル・ネットワーク・アーキテクチャーとして理解するのが最も適切です。わかりやすい類比で言えば、Mambaと選択的SSMの関係は、Transformerモデルとアテンション機構の関係に相当します。

従来のSSMは固定的な動態を持ちます。すなわち、隠れ状態があるステップから次のステップへとどのように変化するかを規定するモデル・パラメーターは、入力や系列内のすべてのステップにおいて同一です。この特性は線形時不変（LTI）として知られています。現在の文脈に基づいて過去の特定の情報を選択的に重視または軽視できるようにするために、Dao氏とGu氏はSSMを再構成し、主要なモデル・パラメーターの値が入力ごとに異なるようにしました。

より具体的には、選択的SSMではステップサイズ∆_tおよび行列B_tとC_tが、現在の入力トークンx_tの直接的な関数となります。これは、x_tのベクトル埋め込みを3つの並列線形射影レイヤー、言い換えれば標準的なフィードフォワード・ニューラル・ネットワークレイヤー（MLPレイヤー）に通すことで実現されます。これは、Transformerモデルにおいて並列のQuery、Key、Valueヘッドが入力に対応するQ、K、Vベクトルを生成するのと同等です。

その線形射影ネットワークにおいて、x_tのベクトル埋め込みに重みとバイアス項を掛け合わせることで、最終的に∆_t、B_t、C_tの値が得られます。線形射影レイヤー自体の重みとバイアス項は、膨大なテキストサンプルの事前学習において学習され、その後（必要に応じて）追加のファイン・チューニングによって洗練されます。

• ∆_tの値は、これまでに見た文脈に対するモデルの記憶に対してx_tが及ぼす影響の大きさ、つまり隠れ状態h_t-1からh_tへの更新がどの程度行われるかを決定します。ステップサイズ∆_tが大きいほど変化は大きくなり、状態に含まれる古い情報の減衰、言い換えれば「忘却」が速まります。逆に、ステップサイズが小さい場合、更新も小さくなります。ステップサイズが十分に小さい場合、現在の入力は隠れ状態にまったく影響を与えません。
  
  
• 行列B_kの変化は、現在の入力トークンが隠れ状態をどのように更新するかを決定します。例えば、x_tが「yesterday」という単語のトークンである場合、B_tは、その後の文脈がおそらく過去に関連することを反映するように状態を更新する形で調整される可能性があります。
  
  
• 行列C_tの変化は、この文脈情報がモデルの出力y_tにどのように反映されるかを決定します。例としてx_kが「yesterday」というトークンである場合、C_kはその後モデルが出力する動詞が過去形に活用されるような形で影響を受ける可能性があります。

注目すべきは、A行列にはこのような入力に基づく調整が一切行われないという点です。その役割はS4モデルと同じであり、過去の入力履歴全体を効率的に記憶することにあります。その履歴のどの部分を特定の時点で利用するかを決定する役割は、B行列とC行列が担います。

しかし、一度モデルが時間不変でなくなると、遷移カーネルがもはや一定ではないため、学習時に畳み込みの近道を利用できなくなります。すなわち、選択性メカニズムの要点は、h_t-1からh_tへの遷移が文脈に依存するようになるという点です。

その代わりに、Mambaは工夫された回避策を用いて、同様の並列化の利点を実現しています。SSMは乗算と加算のみを用いるため、その計算は数学でよく知られる結合則に従います。つまり、最終的な結果を変えることなく、異なる方法でまとめて計算することができます。これにより、多数の逐次計算を小さな独立したチャンクに分割し、並列プレフィックス和スキャンを通じてGPUで並列処理することが可能になります。

さらに、その結果は特定の階層的な方法で結合され、GPU上のさまざまな種類のハードウェア・メモリーを最適に活用します。これは、現在のLLMで広く用いられているFlashAttention技術（これもTri Dao氏によって開発された）と類似の原理に基づいています。

Mambaアーキテクチャーにおいて、S6モデルはより大きな「Mambaブロック」のモジュールとして機能します。これは、アテンション機構がより大きな「アテンション・ブロック」のモジュールとして機能するのと同様です。これはS6モジュールとゲーテッド・ニューラル・ネットワーク・アーキテクチャーを組み合わせたものです。Mambaモデルは通常、最終的な出力を予測する出力レイヤーの前に、複数のMambaブロック、すなわちニューラル・ネットワーク内の連続するMambaレイヤー群で構成されます。

Mambaブロックに入る前に、入力のコピーが残差接続として末尾に直接送られます。Mambaブロックの内部処理の目的は、その入力に関連する広い文脈のどの部分かを特定するだけでなく、その文脈情報が入力本来の意味をどの程度修正すべきかを決定することにあります。

Mambaブロック内では、元の入力ベクトルは次のように処理されます。

• まず、入力は入力ベクトル自体の2倍の幅を持つ線形レイヤーに通され、高次元空間へと射影されます。例えば、モデルが各入力トークンxをもともと512次元のベクトル埋め込みで表現している場合、xに線形射影レイヤーの重みを掛けることで、それを1024次元のベクトルに拡張します。
  
  
• 次に、拡張されたベクトルは2つに分割されます。そのうちの半分（x_projと呼ぶ）はSSMを通る経路に送られ、もう半分（z_projと呼ぶ）はゲーティング機構を通る別の経路に送られます。わかりやすくするために、前の拡張ステップは通常、2つの並列線形レイヤーによって実行されると図示されます。
  
  
• x_projがSSMに入る前に、1次元の畳み込みレイヤーに通されます。この畳み込みレイヤーは、隣接するトークン間の依存関係（例えば、単純な動詞と主語の組み合わせなど）といった局所的なパターンを抽出します。これにより、SSMは長距離にわたるグローバルな依存関係の文脈理解に「注目」できるようになります。
  
  
• この畳み込みレイヤーの出力は、非線形の活性化関数への入力として使用されます。非線形性の導入はすべてのニューラル・ネットワークの特徴であり、より複雑なパターンを捉えることを可能にします。Mamba論文では、シグモイド線形ユニット（SiLU）が使用されています。この結果得られるベクトルをx_actと呼びます。
  
  
• 一方、別のゲーティング機構の経路では、z_projも非線形活性化関数に入力され、z_actが得られます。
  
  
• SSM経路では、x_actが3つの並列線形射影レイヤーに入力され、それぞれ∆_x、B_x、C_xの値が生成されます。
  
  
• SSMはこれらの入力依存パラメーター（およびA行列とD行列）を用いて状態空間の更新とSSMの出力yを計算します。
  
  
• SSMの出力ベクトルyは、ゲーティング経路の出力ベクトルz_actと要素ごとに乗算されます。本質的に、z_actの各要素は音声ミキシング・コンソール上のボリューム・ノブのように機能します。つまり、z_actのある要素がゼロに近い場合、それに対応するyの部分との乗算結果はゼロに近づき、その影響は弱まります。逆に、zactのある要素が大きい場合、yとの乗算によってその文脈情報の影響が増幅されます。
  
  
• その結果得られたベクトルは元のサイズに射影し直されます。これは、元の入力ベクトルの各要素に対する重み付きの文脈的更新（あるいは非更新）のベクトルとして理解できます。
  
  
• 最後に、その更新ベクトルが、残差接続としてブロックの末尾に直接送られていた元の入力ベクトルのコピーに加えられます。
  
  
• 元の入力ベクトルは、選択的SSMによってもたらされた文脈理解を反映するように更新されました。これを次のMambaレイヤーに送ることができ、またモデルの最終レイヤーでは、完全に更新されたベクトルが語彙内の各単語に対応する確率を出力するソフトマックス関数への入力として利用されます。

元のMamba論文から1年後、Dao氏とGu氏は後続として「Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality」を発表しました。この後続論文では3つの大きな貢献が示されました。

• MambaとTransformerの理論的な関連性の探究と、両アーキテクチャー間で共有される用語体系
• Mambaモデルにおけるさまざまな設計選択に関する一連の明確化と検討
• それらの設計検討に基づいて改良されたアーキテクチャーであるMamba-2

Mamba-2アルゴリズムは、元のMambaよりも大幅に高速で実装も容易です。著者らは、選択的SSMを約25行のコードで実装した「minimal SSD」コードベースを提供しました。²この効率性により、Mamba-2はモデルを遅くすることなく、はるかに大きな隠れ状態の次元を利用でき、より大規模で強力かつ表現力豊かなモデルをこのアーキテクチャーで構築できるようになります。テストにおいて、Mamba-2モデルは同規模のMambaモデルやTransformerモデルと比較して、一連のダウンストリームタスクで確実に同等またはそれ以上の性能を示しました。

論文の序文にあるとおり、Dao氏とGu氏の「主な目的は、構造化SSMとアテンションの派生手法との間に豊富な理論的関連性を構築すること」でした。これにより両者を統合する新たな概念的枠組みが生まれ、著者らはそれを「状態空間二重性（SSD）」と呼びました。³これにより、MambaはTransformerアーキテクチャーに対して数年間にわたり行われてきた探究と最適化の成果を活用できる道が開かれました。

注目すべき1つのメリットは、マルチヘッド・アテンション（MHA）に相当するMambaの仕組みが開発されたことです。これにより、Mambaブロックを複数の「Mambaヘッド」に分割でき、Transformerにおける複数の「アテンション・ヘッド」と同様の構造を実現できます。このアプローチの一種は、グループド・クエリ・アテンションに類似するとされており、GPUにおけるテンソル並列化によってさらに高い効率性を実現できます。

Mamba-2ブロックでは、元の「逐次型」Mambaブロックに対して「並列型」Mambaブロックと呼ばれる構造において、入力依存パラメーターである∆、B、Cが初期の射影レイヤーで並列的に生成されます。特にBとCは、専用の線形レイヤーを通じてx_projを乗算するのではなく、x_projの一部を単純にコピーすることで導出されます。全体のモデル・パラメーターを簡素化し削減するだけでなく、この並列化によって大規模トレーニングを大幅に効率化できます。⁴

MambaとTransformerはいずれもそれぞれの強みを持っていますが、メモリー使用量や速度に関してはMambaベースのモデルの方が概して優れています。Mambaの論文によれば、Mambaは同等のTransformerと比べて5倍のスループットを実現しています。

Transformerは非常に精度が高く多用途に利用できますが、計算リソースの要求も非常に大きいです。事前トレーニング（およびファインチューニング）において、セルフアテンションのメモリー要件は系列長に対して二次的に拡大します。つまり、系列のコンテキスト長を2倍にすると、アテンション・メカニズムは4倍のリソースを使用します。この「二次的ボトルネック」は、コンテキスト・ウィンドウが大きくなるにつれて速度とメモリー利用可能量をますます制限します。推論時には、必要なメモリーは系列長に対して線形に拡大します。

トレーニング時において、Mambaモデルのメモリー使用量は系列長に対して線形にのみ拡大します。さらに重要なのは、推論時のメモリー使用量が一定であることです。モデルが処理したトークン数に関係なく、SSMは入力履歴の固定サイズの表現を保持します。これにより、理論的にはハードウェアの制約を除いて無制限のコンテキスト長を扱うことが可能になります。

とはいえ、Transformerのメモリー集約的で計算的に冗長な手法にも独自の利点があります。例えば、研究によれば、コンテキスト内学習（少数ショット・プロンプティングなど）、コピー、あるいは長いコンテキストを要する推論が必要なタスクにおいては、TransformerはいまだにMambaやMamba-2を上回っています。

幸いなことに、TransformerとMambaそれぞれの強みは相互に排他的なものではありません。Mamba-2の論文では、ハイブリッド・モデルは純粋なTransformerやSSMのいずれをも上回る可能性があると示唆されています。この考え方は2024年後半にNVIDIA社の研究によって正式に実証されました。⁵概して、ハイブリッド・モデルはMambaの効率性のメリットと、Transformerのよりリソース集約的なアテンション・メカニズムによって得られる微妙な表現力やコンテキスト内学習の性能を兼ね備えているようです。

これをさらに探求するために、IBM ResearchはDao氏とGu氏に加え、イリノイ大学アーバナ・シャンペーン校（UIUC）のMinjia Zhang氏と協力し、BambaおよびBamba V2に取り組みました。Bambaは、その後、IBM Granite 4.0の多くのアーキテクチャー要素に反映されています。

ハイブリッド・モデルに関する研究は、特にオープンソース・コミュニティにおいて、いまも活発に進められています。