State Space Models(SSM)は、Transformerのアテンション機構が持つ系列長に対する二乗計算量を線形に抑える代替アーキテクチャです。中でもMambaは、入力依存の選択メカニズムとハードウェア最適化により、SSMの実用性を一段引き上げた最も成功した実装として注目されています。本記事では、SSMの数理的基礎からMambaのブレークスルー、Transformerとの性能比較、そしてハイブリッドアーキテクチャの最新動向まで解説します。

Transformerの限界――なぜ代替アーキテクチャが必要なのか

Transformerの中核であるSelf-Attentionは、入力系列のすべてのトークン間の関係を計算します。この計算量は系列長nに対してO(n2 · d)(dは次元数)であり、系列長が増えると二次関数的に増大します。

具体的な数値で考えてみましょう。系列長1,024トークンでの計算量を1とすると、100,000トークンでは約9,500倍になります。GPT-4やClaudeが長文脈を扱えるようになった現在、この二乗コストは推論速度とメモリの両面で深刻なボトルネックです。

さらに、推論時にはKV Cache(Key-Valueキャッシュ)がトークン数に比例して増大し、GPUメモリを圧迫します。128Kトークンの文脈をGPT-4クラスのモデルで処理する場合、KV Cacheだけで数十GBに達することもあります。

一方、SSMの計算量はO(n · d)で系列長に対して線形です。この差は系列長が長くなるほど決定的になります。これが、Transformer以外のアーキテクチャが求められる根本的な理由です。

State Space Modelsの数理的基礎

SSMの理論的な基盤は、制御工学で長く使われてきた連続時間状態空間モデルです。システムの入力u(t)から出力y(t)への変換を、隠れ状態x(t)を介して表現します。

連続時間での定義式は以下のとおりです。

  • 状態方程式: x'(t) = A·x(t) + B·u(t)
  • 出力方程式: y(t) = C·x(t) + D·u(t)

ここで、Aは状態遷移行列(N×N)、Bは入力行列(N×1)、Cは出力行列(1×N)、Dはスキップ接続です。離散的なシーケンスデータを扱うため、Zero-Order Hold(ZOH)による離散化を行います。

  • 離散化後の再帰式: x[k] = A̅·x[k-1] + B̅·u[k]、y[k] = C·x[k]
  • 畳み込みカーネル表現: y = K * u (K = (C·B̅, C·A̅·B̅, C·A̅2·B̅, …))

SSMの重要な特性は、学習時には畳み込みとして並列計算でき、推論時には再帰式として逐次処理できる「二重性」にあります。

【SSMの状態遷移】

u(t) --> [B] --> [隠れ状態 x(t)] --> [C] --> y(t)
                      ^
                     [A] (状態遷移行列)

学習時:  畳み込みカーネル K で並列計算(高速)
推論時:  再帰式 x[k] = A*x[k-1] + B*u[k] で逐次処理(省メモリ)

S4(Structured State Spaces for Sequence Modeling, Gu et al., 2022)は、行列AにHiPPO(High-order Polynomial Projection Operator)初期化を導入することで、長期記憶のメカニズムを実現しました。HiPPO行列は、過去の入力をルジャンドル多項式に射影することで、理論的に最適な形で長期依存関係を保持します。これがSSMが長文脈タスクで優れた性能を示す根拠です。

Mambaのブレークスルー――選択的状態空間モデル

S4は画期的でしたが、根本的な制約がありました。パラメータA, B, Cが入力に依存せず固定されているため、「どの入力に注目し、どの入力を無視するか」を動的に判断できなかったのです。これはTransformerのAttentionが持つ、入力に応じて注目先を変える能力の欠如を意味していました。

Mamba(Gu & Dao, 2023)はこの問題を「選択メカニズム(Selection Mechanism)」で解決しました。具体的には以下の変更を加えています。

  • B, C を入力依存に: B = Linear(x)、C = Linear(x) として、入力に応じてパラメータが変化
  • 離散化ステップΔも入力依存に: Δ = softplus(Linear(x)) で、入力ごとに離散化の粒度を調整
  • ハードウェア最適化: GPUのSRAM(高速メモリ)を活用した専用カーネルで、理論的な効率性を実際の速度に変換

Δが大きいとき、モデルは現在の入力を重視して状態をリセットします。Δが小さいとき、過去の状態を保持し続けます。この選択メカニズムにより、MambaはTransformerのAttentionに近い「入力依存の情報フィルタリング」をO(n)の計算量で実現したのです。

以下は、mamba-ssmライブラリを使った推論コードの例です。

from mamba_ssm import MambaLMHeadModel
import torch

# Mamba-2.8Bモデルのロード
model = MambaLMHeadModel.from_pretrained("state-spaces/mamba-2.8b", device="cuda", dtype=torch.float16)

# トークナイザー(GPT-NeoXベース)
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")

# 推論
prompt = "State Space Models are"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to("cuda")
output = model.generate(input_ids, max_length=100, temperature=0.7)
print(tokenizer.decode(output[0]))

Mambaは同規模のTransformerと比較して、推論スループットが最大5倍、メモリ使用量が大幅に削減されることが報告されています。特に系列長が長くなるほど、その差は顕著になります。

SSM vs. Transformer――性能・効率・スケーラビリティの比較

SSMとTransformerの特性を多角的に比較します。まず、アーキテクチャレベルでの構造的な違いを整理します。

項目TransformerSSM(Mamba)ハイブリッド
計算量(学習)O(n2·d)O(n·d)O(n·d) 〜 O(n2·d)
計算量(推論)O(n) per token(KV Cache込み)O(1) per token(固定状態)アーキテクチャ依存
メモリ使用量KV Cacheが系列長に比例固定サイズの状態ベクトルAttention層分のKV Cache
長文脈性能位置エンコーディングに依存理論的に無限長対応両方の強みを活用
In-context Learning非常に高いやや劣るAttention層で補完
実装成熟度非常に高い発展途上発展途上
エコシステムHuggingFace等が充実限定的限定的
スケーリング実績数兆パラメータまで実証済数十億パラメータまで数百億パラメータまで

続いて、主要ベンチマークでのパフォーマンス比較です。

ベンチマークTransformer(同規模)Mamba備考
言語モデリング(PPL)ベースライン同等〜やや優位特に長文脈で差が顕著
常識推論(HellaSwag等)ベースライン同等パラメータ効率は Mamba が上
コード生成やや優位ベースライン正確な検索が必要なタスク
長文脈タスク品質劣化あり安定した性能Mambaの最大の強み
情報検索タスク優位やや劣るAttentionの検索能力が有利

注目すべきは、Mambaが情報検索タスク(例: 文書中の特定情報を抽出する)でTransformerに劣る点です。Attentionは「どこに何があるか」を直接参照できますが、SSMの状態ベクトルは情報を圧縮して保持するため、ピンポイントの検索が苦手です。この弱点が、後述するハイブリッドアーキテクチャの動機になっています。

ハイブリッドアーキテクチャの台頭

SSMの効率性とTransformerの検索能力を組み合わせたハイブリッドアーキテクチャが、次世代の有力候補として注目されています。

Mamba-2(Dao & Gu, 2024)

Mamba-2は、SSMとAttentionの理論的な繋がりを明示し、Structured State Space Duality(SSD)フレームワークを提案しました。SSD層はSSMとAttentionの中間的な計算構造を持ち、行列サイズの調整で両者の特性をトレードオフできます。学習速度はMamba-1比で50〜80%向上しています。

Jamba(AI21 Labs, 2024)

Jamba は、Mamba層とTransformer Attention層を交互に配置し、さらにMoE(Mixture of Experts)を組み合わせたアーキテクチャです。52Bパラメータ(活性12B)で256Kトークンの文脈長を実現し、同規模のTransformerモデルと比較してスループットが3倍に達しました。「SSMで長文脈を効率的に処理し、Attention層で精密な情報検索を補完する」という設計思想を体現しています。

Zamba(Zyphra, 2024)

Zambaは、Mamba層の間に共有Attention層を挟む独自構造を採用しました。Attention層のパラメータを複数箇所で共有することで、パラメータ効率を最大化しつつAttentionの恩恵を受けます。7Bパラメータクラスで同規模モデルを上回る性能を報告しており、小規模モデルでのハイブリッド設計の有効性を示しました。

これらのハイブリッドアーキテクチャに共通するのは、「全層をAttentionにする必要はない」という洞察です。大部分の層をSSMで処理し、要所にAttention層を配置することで、効率性と能力を両立させる設計が主流になりつつあります。

ビジネスへの示唆――SSMは実務で使えるのか

SSM/Mambaが実務で有望なユースケースは明確です。長文書処理(法律文書、論文、マニュアルの分析)、リアルタイム推論(チャットボット、音声処理)、エッジデバイス展開(固定サイズの状態ベクトルにより省メモリ)――これらはSSMの線形計算量と固定メモリが直接的なアドバンテージになる領域です。

一方で、現時点での限界も認識しておくべきです。Fine-tuningツールやサービングインフラはTransformerエコシステムに遠く及ばず、LoRAやQLoRA相当のパラメータ効率的な微調整手法もまだ発展途上です。実務での技術選定としては、「プロダクションにはTransformerベース、R&Dや特定ユースケースの検証にSSM/ハイブリッド」という棲み分けが現実的でしょう。

重要なのは、「SSMかTransformerか」という二者択一ではなく、ハイブリッドアーキテクチャの進化を含めた選択肢の拡大として捉えることです。技術トレンドへの過剰反応については「最新技術」に飛びついた会社が後悔するパターンも参考になるでしょう。

まとめ――ポストTransformer時代の設計思想を理解する

  • SSMは連続時間状態空間モデルを離散化し、畳み込み(学習)と再帰(推論)の二重性で効率的な系列処理を実現する
  • Mambaは入力依存の選択メカニズムにより、SSMに動的な情報フィルタリング能力を付与した
  • Transformerとの比較では、長文脈処理と推論効率でSSMが優位、情報検索タスクではTransformerが優位
  • ハイブリッドアーキテクチャ(Jamba, Zamba等)が両者の強みを統合する有力な方向性として台頭している
  • アーキテクチャの多様化を理解することが、適切なモデル選定の基礎になる

LLMのアーキテクチャ選定やAI基盤の設計については、DE-STKにご相談ください

FAQ

Q: MambaとTransformerの最大の違いは何ですか?

Transformerは入力系列全体に対するSelf-Attentionを計算するため系列長の二乗に比例する計算量が必要ですが、Mambaは状態空間モデルに基づく再帰構造により線形の計算量で処理できます。特に長文脈での効率性に大きな差があります。

Q: Mambaは実務でTransformerの代替になりますか?

特定のユースケース(長文書処理、リアルタイム推論など)では有望ですが、現時点ではエコシステムの成熟度やFine-tuningツールの充実度でTransformerに劣ります。ハイブリッドアーキテクチャが現実的な選択肢になりつつあります。

Q: State Space Modelの学習にはどの程度のGPUリソースが必要ですか?

同規模のTransformerと比較して学習時のメモリ効率が高いため、より少ないGPUで長い系列長を扱えます。ただし、Mambaのカスタムカーネルを活用するにはNVIDIA GPU(A100以上推奨)が必要です。