未来を形作るテクノロジーの深掘り記事。

SSM vs Transformer:実践ガイド

状態空間モデル(SSM)、Mamba-3、ハイブリッド構成を実務目線で解説。Transformerより適した場面と導入判断のポイントを紹介します。

滑らかに流れる川と、密に光るノードの格子が並ぶ画像。SSMとTransformerの対比を表現。

Transformerが唯一の主役でいられる時代は、もうそう長くは続かないでしょう。大げさに聞こえるのは承知しています。言語モデルからタンパク質の構造予測まで、Transformerアーキテクチャが長年あらゆる分野を席巻してきたのを、私たちは皆見てきました。それでも、ここ数か月、本番ワークロードでState Space Model(SSM)をTransformerのベースラインと比較検証してきた結果、この流れは本物だと確信しています。SSMが万能に優れているからではありません。そうではないからです。Transformerが根本的に解けない特定の問題を、SSMが解決できるからです。さらに最新世代は品質面の差をかなり縮めたので、無視することがもはや技術的負債の判断になりつつあります。

これは誇大広告の記事ではありません。State Space Modelが実際に何者なのか、Mamba-3が本当に改善している点、SSMがまだ苦戦する点、そしてチームが導入をどう考えるべきかを順に解説します。実務者同士の目線で話しますね。

大規模化でTransformerが壁にぶつかる理由

二次スケーリングの話は、皆さんご存じでしょう。セルフアテンションは長さNの系列に対してN×Nの行列を計算するため、コンテキスト長を2倍にすると計算量もメモリも4倍になります。長い間、これはそれほど問題になりませんでした。モデルは数千トークンで動いており、ハードウェアもそれに追いついていたからです。

その時代は終わりました。いま私たちが作っているワークロードは、10万トークン超のコンテキストを当たり前に必要とします。コードアシスタントはリポジトリ全体を見る必要がありますし、マルチモーダルのパイプラインは数時間分の動画を処理します。エージェントは何日にもわたる会話履歴を保持します。この規模になると、二次のアテンションは単に高コストなだけでなく、壁そのものです。

KVキャッシュの問題がそれをさらに悪化させます。自己回帰生成の間、Transformerの各層は、それまでに見た全トークンのKey-Valueペアを保存します。このキャッシュは層ごとに線形に増え、GPUメモリをあっという間に食いつぶします。私は7BのTransformerが128KコンテキストでKVキャッシュだけで40GBのVRAMを消費するのを見てきました。それはバッチ処理で多くのリクエストをさばくために使えるはずのメモリです。

  • 標準的なアテンションでは、メモリの二次スケーリングのせいで百万トークン級のコンテキストはほぼ不可能
  • KVキャッシュの増加がGPU1基あたりの同時ユーザー数を制限し、本番環境ではコストを直接押し上げる
  • 長いコンテキストを扱うTransformer推論の消費電力は、正当化が難しくなりつつある
  • ロボティクスやエッジAIなどのリアルタイム用途では、ミリ秒未満のトークン生成が必要だが、アテンションでは届かない
  • 「アテンションシンク」現象により、十分なメモリがあっても非常に長い系列では品質が落ちる

これらは理論上の懸念ではありません。昨年、私が一緒に仕事をした3つのチームが、代替案の本格的な検討を始めた理由がまさにこれです。

State Space Modelの仕組み

State Space Modelは制御理論が出自で、動的システムのモデル化に数十年使われてきました。中心となる考え方はシンプルです。(アテンションのように)出力を計算するたびに過去の全トークンを見る代わりに、時間とともに変化する圧縮された隠れ状態を保持します。新しいトークンが状態を更新し、状態から出力が生まれる。それだけです。

数学的には、SSMは4つの行列A、B、C、Dで定義されます。これらが入力xに応じて隠れ状態hがどう変化するかを決めます。連続時間の方程式を系列データ向けに離散化すると、推論時に計算が非常に単純な漸化式が得られます。

import torch
def ssm_step(A_bar, B_bar, C, D, h, x_t):
"""Single SSM step: O(1) memory, O(1) compute.
Compare this to attention, which needs to look at
every previous token. The SSM just updates its state.
"""
h_new = A_bar @ h + B_bar @ x_t  # Update hidden state
y_t = C @ h_new + D * x_t         # Compute output
return h_new, y_t
def ssm_generate(A_bar, B_bar, C, D, tokens, embed):
"""Autoregressive generation with constant memory.
Whether you've processed 100 tokens or 500,000,
this uses the same amount of memory.
"""
h = torch.zeros(A_bar.shape[0])
outputs = []
for t in tokens:
x_t = embed(t)
h, y_t = ssm_step(A_bar, B_bar, C, D, h, x_t)
outputs.append(y_t)
return torch.stack(outputs)

その美しさはコードを見れば一目瞭然です。推論は、系列長に関係なくトークンごとにO(1)のメモリとO(1)の計算量で済みます。KVキャッシュもなく、二次的な膨張もない。履歴全体が隠れ状態ベクトルの中に圧縮されるのです。

ただし、問題もあります。学習時にこの漸化式を逐次的に実行するとひどく遅くなります。そこで、同じ計算を畳み込みや並列スキャンとして書き直すことで、GPUで効率よく処理できるようにします。つまり、並列な学習と再帰的な推論という、両方の良いとこ取りができるわけです。

MambaからMamba-3まで:各世代が何を修正したか

S4のような初期のSSMは概念の正しさを証明しましたが、決定的な弱点がありました。コンテンツに基づく推論がうまくできなかったのです。状態遷移行列が全入力で固定されていたため、実際に読んでいる内容に応じて何を覚え何を忘れるかを決められませんでした。「3語ごとに1語書き留める」というルールでノートを取るようなもので、有用な情報は一部拾えても、何が重要かに適応することはできません。

2023年末にAlbert GuとTri Daoが発表したMambaは、エレガントなアイデアでこれを解決しました。SSMのパラメータを入力依存にしたのです。固定のA、B、C行列の代わりに、Mambaは現在のトークンの関数としてこれらを計算します。モデルは関連する情報を選択的に保持し、ノイズを捨てることを学習します。この「選択的」な仕組みによって、SSMに欠けていたコンテンツへの感度が手に入りました。

Mamba-2は、構造化SSMと線形アテンションが数学的に双対であるという理論的知見をもたらしました。State Space Duality(SSD)フレームワークです。これは単なる学術的な話ではありません。GPUのテンソルコアをより活用するハードウェア指向の実装を可能にし、学習スループットを大きく引き上げました。

デプロイの観点から特に面白くなってくるのがMamba-3です。重要な技術革新は3つあります。

  1. マルチスケールの状態追跡:モデルは複数の時間解像度で同時に状態を保持し、局所的なパターンと長距離の依存関係の両方を、どちらも犠牲にせず捉えます
  2. 適応的な状態圧縮:隠れ状態は複雑な推論が必要な箇所では動的に拡張され、予測しやすい文章では縮小します。品質を落とさずに計算量を節約します
  3. 初期化とゲーティングの改善:大規模でも学習の安定性が劇的に向上しました。数百万ドル規模の学習を回すときには、これが非常に重要になります

Mamba-3は、すべてのベンチマークでTransformerを上回るわけではありません。そこまで求める必要もありません。標準的な評価の大半で品質は同等を保ちつつ、推論計算量はごく一部で済みます。大半の本番ワークロードでは、この取引が重要なのです。

線形アテンションとSSMの収束

もう一つ理解しておく価値のある流れがあります。線形アテンションは同じ効率化の課題に、Transformerの枠組みの内側から取り組みます。標準的なアテンションは完全なN×N行列を計算します。線形アテンションはsoftmaxを分解可能なカーネル関数で置き換え、数式を並べ替えることで、その二次的な行列を実体化せずに済ませます。

# Standard attention: O(N^2 * d)
# score = softmax(Q @ K.T / sqrt(d)) @ V
# Linear attention: O(N * d^2)
# Replace softmax with kernel feature map phi()
# Rearrange: compute K^T @ V first (d×d), then multiply by Q
def linear_attention_step(q_t, running_kv, running_k, k_t, v_t, phi):
"""Incremental linear attention — runs like a recurrence.
This is why SSMs and linear attention are duals:
both compress history into a fixed-size state.
"""
k_feat = phi(k_t)
q_feat = phi(q_t)
running_kv = running_kv + k_feat.unsqueeze(-1) * v_t.unsqueeze(-2)
running_k = running_k + k_feat
y_t = (q_feat @ running_kv) / (q_feat @ running_k + 1e-6)
return y_t, running_kv, running_k

このコードを注意深く見てください。線形アテンションは、逐次的に実行すると実行中の状態を保持し、新しいトークンごとにそれを更新します。どこかで見覚えがありませんか。そのはずです。やっていることはSSMとほぼ同じなのです。SSDフレームワークはこの関係を形式的に結びつけました。これは最近の系列モデリング研究で最も重要な理論的知見の一つです。

GLA(Gated Linear Attention)やRetNetの派生アーキテクチャは、さらに一歩進めて、データ依存のゲーティングを加えました。これにより、線形アテンションと選択的SSMの境界はほとんど完全に曖昧になっています。実務上の要点は、これらを競合する手法とみなさないことです。両者は収束しつつあります。

ハイブリッドアーキテクチャ:本番で実際に勝っているのは何か

SSMへの移行を尋ねてくるチームには、いつも次のように伝えています。何か一つに純粋に振り切るのはやめましょう。いま最良の結果を出しているのは、SSM層と少数のアテンション層を組み合わせたハイブリッドです。異なる計算プリミティブは得意分野が異なり、それを無視するとパフォーマンスを取りこぼします。

SSM層は、逐次的な情報を効率よく圧縮し伝播させることに優れています。一方、アテンション層は正確で内容ベースの検索では今なお並ぶものがありません。たとえば「この質問に答える一文を47ページから探す」といった用途です。よく設計されたハイブリッドは、層の80〜90%にSSMを使い、最も効果のある場所にアテンションを散りばめます。

  • Jambaスタイルのモデル:MambaとアテンションのLayerを交互に配置し、MoEのフィードフォワードブロックを組み合わせ、効率的なSSM処理と精密なアテンションを動的にルーティングします
  • Griffinファミリーの設計:再帰的なゲート付き線形ユニットと局所的なスライディングウィンドウ・アテンションを組み合わせ、フルアテンションを最小限に抑えつつ強い結果を出します
  • Mambaとアテンションのハイブリッド:大半の層はMamba-3ブロック、戦略的な深さにフルアテンション層を挿入して大域的な情報のルーティングを担わせます
  • StripedHyenaの後継:ゲート付き畳み込み、SSM層、スパースアテンションを、NASで最適化されたパターンで交互に組み合わせます

数字がこれを裏付けています。複数の独立したグループが、同じパラメータ数で85対15(SSMとアテンションの比率)に分割すると、純粋なTransformerと同等の品質を保ちつつ、推論のFLOPsを40〜60%削減できることを示しています。長いコンテキストのワークロードでは、メモリの節約はさらに大きくなります。これは僅かな改善ではありません。GPUの請求額をほぼ半分にできるということです。

本番ベンチマーク:SSMが強い場面と弱い場面

数字について具体的に話しましょう。デプロイの判断をする人にとって、曖昧な効率化の主張は役に立たないからです。

推論スループット:8Bパラメータ規模のMamba-3ベースのモデルは、コンテキストが1Kトークンでも50万トークンでも、同じ速度でトークンを生成します。一方、同等のTransformerは、KVキャッシュが増えるにつれて徐々に遅くなります。50万トークンのコンテキストでは、SSMモデルはGPU1基あたり5〜8倍のスループットを出します。これは理論上の話ではなく、私が実際に計測した結果です。

同時ユーザー数:KVキャッシュがなければ、SSMモデルは同時リクエストをはるかに多くさばけます。A100単体で、32Kコンテキストのとき、Transformerが同時8ストリーム程度しか扱えないのに対し、同等のSSMモデルは30以上を扱えます。大規模に推論を回している人にとって、これは経済性を変える数字です。

学習速度:ここでの効果はもっと控えめです。Mamba-3は、H100クラスタ上で同等のTransformerのおよそ1.4倍のスループットで学習します。系列が長くなるほど差は広がり、32Kトークンを超えるとSSMの学習は2〜3倍速くなります。二次のアテンションを完全に避けられるからです。

ただし、限界についても正直に言わなければなりません。長いコンテキストからの正確な逐語的な再現を要するタスク、たとえば「4,382行目の正確なエラーメッセージは何だったか?」のようなものでは、純粋なSSMはまだ性能が落ちます。固定サイズの圧縮状態は非可逆な表現です。アテンションなら元のトークンを直接振り返ることができます。これこそハイブリッドアーキテクチャが機能する理由です。SSMが扱えない検索をアテンション層が担当するのです。

SSMがまだ不足している点

残っているギャップについては冷静に見ておきたいと思います。不完全な情報をもとに新しいアーキテクチャを採用すると、半年を無駄にする確実な方法だからです。

  1. インコンテキスト学習:プロンプト内のFew-shot例に基づいて振る舞いを適応させる能力は、Transformerの方がまだ優れています。SSMにもできますが、信頼性は低めです。例を使ったプロンプトエンジニアリングに強く依存するアプリケーションなら、純粋なSSMでは期待外れになるでしょう。
  2. エコシステムの成熟度:Transformerのツール群は何年もの最適化を経ています。SSM専用のカーネル、推論基盤、ファインチューニングのライブラリは急速に改善していますが、まだ同等には達していません。統合の時間は多めに見積もってください。
  3. 70B超のスケーリングの不確実性:70Bパラメータまでの Mamba-3モデルは良好なスケーリング曲線を示していますが、200B以上のフロンティア領域では確かなデータがありません。SSMのスケーリング則が極端なサイズでも成り立つかどうかは、正直なところ分かっていません。
  4. ファインチューニング手法:Transformer向けのLoRAやQLoRAはよく理解されています。SSMアーキテクチャに適用するには別のアプローチが必要で、ベストプラクティスはまだ固まっている途中です。
  5. ハードウェアとの不一致:現在のGPUは、アテンションが好む行列積に最適化されています。SSMは並列スキャンに大きく依存します。これは最新のハードウェアでも十分に動きますが、GPUが設計の前提としてきた演算ではありません。

どれも導入を断念するほどの障害ではありません。既知の解決策がある工学上の問題です。ただ、現実の問題であり、スケジュールには織り込んでおくべきです。

実践的な提言:いつ導入し、どう始めるか

複数の本番ワークロードでSSMを評価してきた結果、チームに助言するときに使っている判断の枠組みを紹介します。

長いコンテキストの推論(定常的に32K以上のトークン)、高い同時実行性の要件、あるいはレイテンシに敏感なエッジへのデプロイが関わるなら、積極的に導入すべきです。ROIは大きく、すぐに得られます。最初からSSMの純粋な構成にするよりも、JambaやGriffinファミリーのようなハイブリッドから始めましょう。リスクを抑えつつ、効率化の恩恵の大半を得られます。

ワークロードが主に短いコンテキストで、インコンテキスト学習に大きく依存しており、推論コストの圧力がないなら、様子見をおすすめします。この領域ではTransformerがまだ優位で、エコシステムの成熟度も高いです。

  • 判断の前に、実際の推論ワークロードをプロファイリングしましょう。中央値のコンテキスト長と同時ユーザー数が主要な変数です
  • 純粋なSSMではなく、ハイブリッドアーキテクチャから始めましょう。リスクが低く、それでも推論コストを40〜60%削減できます
  • 自社の具体的なタスクでベンチマークしましょう。SSMは要約や長距離の推論に優れる一方、正確な検索では遅れをとります
  • 今のうちにアーキテクチャ比較の基盤を整えましょう。精度だけでなく、レイテンシ、スループット、メモリ、クエリ単位のコストを測定する必要があります
  • SSMのツール群のエコシステムを四半期ごとに追いましょう。改善のペースが速く、今は非現実的なものが3か月後には本番投入可能になっているかもしれません

アーキテクチャの地図は分裂しつつある。それは良いことだ

一つのアーキテクチャがすべてを支配する時代は終わろうとしています。チームが計算プリミティブ(フルアテンション、線形アテンション、選択的SSM、ゲート付き畳み込み)を選び、自社の制約に応じて組み合わせる世界へと向かっています。これは成熟したエンジニアリング分野のやり方です。すべての構造物を鉄骨で作る人はいません。耐えるべき荷重に応じて材料を選ぶのです。

Transformerは死んでいません。多くのワークロードで今なお最も実績のあるアーキテクチャであり、今後何年にもわたって重要なAIシステムを支え続けるでしょう。しかし、最先端のシーケンスモデリングを独占する時代は終わりました。SSMとハイブリッドは、研究上の好奇心の対象ではなく、一級の本番ツールとしての地位を得ています。

実際のシステムを作っている私たちにとって、アーキテクチャの選択肢が増えることは、特定の問題に対してより良い道具が増えることを意味します。恐れるべき破壊ではありません。活用すべき工学的なレバレッジです。