Mambaとは？Transformer以外の系列モデル

Attention matrixを作らず、tokenごとに情報を選んで固定長stateへ流すMambaの全体像
2026.09.06

Attention matrixを作らず、tokenごとに情報を選んで固定長stateへ流すMambaの全体像

Mambaは、Attentionを使わずに長いsequenceを効率よく処理するためのarchitectureです。土台にはSSM（State Space Model、状態空間モデル）があり、過去を固定長のstateへ圧縮しながら、tokenを一つずつ更新します。

ただし、単にRNNへ戻ったわけではありません。Mambaの核心は、入力に応じて「何をstateへ書き込み、何を残し、何を読み出すか」を変えるSelection mechanismと、その計算をGPU上で並列化するSelective Scanにあります。

3行で分かるMamba

  • SSMは、過去のsequenceを固定長stateへ畳み込み、現在の入力 + 前のstateから次のstateを作ります。
  • Mambaは離散化step \(\Delta\)、入力写像 \(B\)、読み出し \(C\)をtokenごとに変え、重要な情報を選択的に残します。
  • 計算量はsequence長 \(L\)に対して線形で、decode時のstateも固定長です。ただし、固定長stateへの圧縮には情報bottleneckがあります。

この記事では、Mamba原論文をもとに、SSMの式からMamba block、Selective Scan、Transformerとの違いまで順番に説明します。

SSM、S4、Selective SSM、Mambaは同じものではない

最初に用語の粒度をそろえます。

SSMを数学的な枠組み、S4を構造化SSM、Selective SSMを入力依存の状態更新、Mambaを完全なnetwork architectureとして整理した階層図

用語 粒度 役割
SSM 数学的なsystem 入力、内部state、出力の時間発展を表す
S4 neural sequence layer 構造化したstate matrixにより長いsequenceを効率よく扱う
Selective SSM / S6 入力依存のSSM layer tokenに応じて保持、忘却、読み出しを変える
Mamba end-to-end architecture Selective SSM、local convolution、gate、residualを組み合わせる

SSMはMamba固有の発明ではなく、制御工学などで使われてきた数学的表現です。S4論文は、SSMを長いsequence向けのneural layerとして効率化しました。Mambaはその流れに入力依存の選択性とhardware-awareな計算方法を加え、言語などの離散dataへ適応させました。

なぜTransformer以外の道が必要なのか

TransformerのSelf-Attentionは、queryとkeyの類似度を使い、現在のtokenが過去のどのtokenを見るかを内容に応じて決めます。このcontent-based accessは強力ですが、Dense Attentionでは長さ \(L\)のsequenceに対して概ね \(L^2\)個のtoken pairを計算します。

autoregressive decodeでは、過去tokenのKeyとValueをKV Cacheへ保存します。そのため、生成が長くなるほどcacheも増えます。

一方、recurrent modelは過去を固定長stateへ要約します。

Transformerが全tokenのKeyとValueを保持する経路と、Mambaが固定長stateを更新する経路の比較

観点 TransformerのDense Attention MambaのSelective SSM
過去の保持 過去tokenごとのKey/Value layerごとの固定長state
sequence全体の計算量 概ね \(O(L^2)\) 概ね \(O(L)\)
decode時の保持量 生成長に応じて増加 sequence長には依存しない
過去情報へのaccess tokenごとに直接参照できる 圧縮されたstateを経由する
主な強み 内容に基づく直接比較 streamingと長いsequenceの効率
主な弱み 長さに伴う計算・memory増加 固定長stateの情報bottleneck

ここでの目標は、Attentionの機能を単に削ることではありません。内容に応じた選択を残しながら、sequence長に比例する計算へ変えることです。

状態空間モデルを一つの式から理解する

連続時間のSSMは、次の二つの式で表せます。

\[\frac{d h(t)}{dt}=Ah(t)+Bx(t)\]
\[y(t)=Ch(t)+Dx(t)\]

\(x(t)\)は入力、\(h(t)\)は内部state、\(y(t)\)は出力です。matrixの役割は次のように読めます。

  • \(A\)：前のstateを時間とともにどう変化させるか
  • \(B\)：入力をstateへどう書き込むか
  • \(C\)：stateから何を読み出すか
  • \(D\)：入力をstateを経由せず出力へ渡すskip connection

入力xがBを通ってstate hへ入り、Aのfeedbackで時間発展し、CとDから出力yを作る状態空間モデル

この式は連続時間なので、token列を扱うには離散化します。入力がstep間で一定だとみなすzero-order holdでは、step幅 \(\Delta\)を使って次のparameterを作れます。

\[\bar A=\exp(\Delta A)\]
\[\bar B=(\Delta A)^{-1}\left(\exp(\Delta A)-I\right)\Delta B\]

離散化後の更新は、見慣れたrecurrent formになります。

\[h_t=\bar A h_{t-1}+\bar Bx_t\]
\[y_t=Ch_t+Dx_t\]

つまり各stepで、前のstateを \(\bar A\)で運び、現在の入力を \(\bar B\)で書き込みます。

scalar例で「過去をstateへ畳み込む」を見る

最小の例として、\(A=-1\)、\(B=C=1\)、\(D=0\)、\(\Delta=0.5\)とします。すると次のようになります。

\[\bar A=e^{-0.5}\approx0.607\]
\[\bar B=1-e^{-0.5}\approx0.393\]

したがって更新式は次のとおりです。

\[h_t=0.607h_{t-1}+0.393x_t\]

式を過去へ展開すると、現在から1 step、2 step、3 step前の入力の係数は、およそ 0.393 → 0.239 → 0.145 → 0.088 と減衰します。

A=-1、Delta=0.5のscalar SSMで過去入力の重みが0.393、0.239、0.145、0.088と減衰する具体計算

高次元のSSMでは、複数のstate成分が異なるtimescaleを担当できます。HiPPO論文は、過去のsignalを低次元stateへonlineで圧縮する理論を与え、S4へつながる長期memoryの土台になりました。

固定SSMには「再帰」と「畳み込み」の二つの顔がある

\(\bar A\)、\(\bar B\)、\(C\)が全stepで同じ、つまりtime-invariant（時不変）なら、recurrent formを展開して固定kernelを作れます。初期stateを0、簡単のため \(D=0\)とすると、長さ \(L\)のkernelは次のようになります。

\[\bar K=\left(C\bar B,\ C\bar A\bar B,\ C\bar A^2\bar B,\ldots,C\bar A^{L-1}\bar B\right)\]

出力sequence全体はconvolutionとして表せます。

\[y=x*\bar K\]

同じ固定SSMをdecodeでは左からのrecurrent update、trainingでは固定kernelとのconvolutionとして計算できる双対性

この二面性が重要です。

  • 生成時：前のstateだけを使うrecurrent formで、1 tokenずつ安く更新する
  • 学習時：sequence全体をconvolutionとして並列に計算する

S4はstate matrix \(A\)を構造化し、このkernelを長いsequenceでも効率よく計算できるようにしました。しかし、固定parameterには別の問題があります。

固定SSMは「内容を見て残す」が苦手

言語では、すべてのtokenを同じtimescaleで扱えばよいわけではありません。電話番号を後で復唱するtaskなら数字を残し、間に入る説明文はstateを大きく変えない方が有利です。

Mamba論文は、この違いをSelective Copyという人工taskで明確にしました。入力位置と、覚えるべき記号の間隔が変わると、固定した状態遷移だけでは「どのtokenを残すか」を内容に応じて切り替えにくくなります。

選択的copy taskで固定SSMが全tokenを同じ規則で混ぜる一方、Mambaが重要tokenだけをstateへ書き込んで後で読み出す様子

Attentionならqueryとkeyの内容を比較し、必要な位置へ直接重みを置けます。Mambaは同じmatrixを作る代わりに、現在のtokenを見てstate updateそのものを変える方針を採ります。

Mambaの核心は入力依存のDelta、B、C

MambaのSelective SSMでは、主に \(\Delta\)、\(B\)、\(C\)を入力 \(x_t\)の関数にします。

\[\Delta_t=s_{\Delta}(x_t),\qquad B_t=s_B(x_t),\qquad C_t=s_C(x_t)\]

それぞれの直感は次のとおりです。

parameter tokenごとに変える意味 直感
\(\Delta_t\) stateをどの程度進めるか 過去を保持するか、更新・忘却するか
\(B_t\) 現在の入力をstateへどう入れるか 何を書き込むか
\(C_t\) stateのどの成分を出力へ使うか 何を読み出すか
\(A\) stateの基本的な時間発展 Mambaでは構造化され、tokenごとの入力依存にはしない

\(\Delta_t\)はgateに近い役割を持ちます。安定なscalar systemとして \(A=-1\)を考えると、\(\bar A_t=e^{-\Delta_t}\)です。

  • \(\Delta_t=0.1\)なら \(\bar A_t\approx0.905\)：前のstateを約90.5%残す
  • \(\Delta_t=2.0\)なら \(\bar A_t\approx0.135\)：前のstateを約13.5%残し、大きく更新する

Deltaが0.1のときstate保持率0.905、2.0のとき0.135となる指数減衰曲線

実際のMambaはvector-valuedで、\(B_t\)と\(C_t\)も同時に変わります。このscalar例は、\(\Delta_t\)が時間尺度を入力に応じて変える直感を示すものです。

選択性と引き換えに固定convolutionを失う

固定SSMでは、全位置で同じ \(\bar A\)、\(\bar B\)、\(C\)を使うため、一つのkernel \(\bar K\)を先に計算できました。Mambaではこれらがtokenごとに変わるため、位置 \(t\)の係数は入力sequence自体に依存します。

したがって、全sequenceへ同じ固定kernelを適用する単純なconvolutionにはできません。ここで必要になるのがSelective Scanです。

Selective Scanは選択機構そのものではない

名前が似ていますが、二つの役割を分けると理解しやすくなります。

要素 答える問い 役割
Selection mechanism 何を残し、何を忘れるか \(\Delta_t\)、\(B_t\)、\(C_t\)を入力依存にする
Selective Scan 依存関係のある更新をどう速く計算するか recurrent updateをparallel scanで処理する

state updateは前のstateに依存しますが、affine transformationの合成として結合できます。二つのstepを (A_1, b_1)、(A_2, b_2) と書くと、合成は次の形です。

\[(A_2,b_2)\circ(A_1,b_1) =\left(A_2A_1,\ A_2b_1+b_2\right)\]

この演算は括弧の付け方を変えても結果が同じ結合則を持つため、tree状のparallel scanが使えます。全stepを完全に逐次処理する代わりに、部分区間の変換を並列に合成できます。

hardware-awareとはmemory trafficまで設計すること

理論上の演算回数だけではGPU上の速度は決まりません。Mambaでは、入力依存parameterで拡張された大きなstate tensorをHBM（GPUの大容量外部memory）へ何度も書き戻すと、memory trafficがbottleneckになります。

HBMから入力を読み、fused kernel内でparameter生成とparallel scanをSRAM上で進め、出力だけをHBMへ戻すSelective Scanのmemory経路

原論文のhardware-aware algorithmは、主に次を組み合わせます。

  1. parameter生成、離散化、scanを一つのkernelへfusionする
  2. scan中のstateを高速なon-chip SRAMへ置く
  3. 中間stateをHBMへmaterializeせず、backwardで必要な値を再計算する

recomputationは計算を少し増やしますが、HBMとの読み書きを減らせます。これはFlashAttentionと同様に、演算量だけでなくmemory hierarchyを意識する設計です。ただし、MambaはAttentionを正確に計算するalgorithmではなく、Selective SSMという別のoperatorを計算します。

Mamba blockのarchitecture

Selective SSMだけでは、完全なlanguage model blockにはなりません。original Mamba blockは、localなtoken mixing、選択的な長期state、gateを一つにまとめます。

Mamba blockで入力を二分し、左枝をLinear、depthwise Conv1D、SiLU、Selective SSMへ、右枝をLinearとSiLUのgateへ通し、積と出力projection、residualで統合するarchitecture

概念的な処理は次の順序です。

  1. normalization後の入力を二つのbranchへprojectionする
  2. main branchで短いcausal Conv1Dにより近傍tokenを混ぜる
  3. SiLU activationの後、Selective SSMで長い方向へstateを伝える
  4. gate branchをSiLUへ通し、main branchの出力と要素積を取る
  5. output projectionを通し、block入力とのresidual connectionを加える

Conv1Dは近傍pattern、Selective SSMは圧縮stateを通じた長い依存、gateは内容に応じた出力制御を担います。原論文のMamba architectureは、独立したAttention blockやMLP blockを置かず、このMamba blockを積み重ねます。

training、prefill、decodeでは計算の姿が違う

同じmodelでも、sequenceをまとめて処理する段階と、1 tokenずつ生成する段階では最適な計算方法が異なります。

Mambaがtrainingとprefillではparallel scanでsequenceを処理し、decodeでは各layerのSSM stateとConv1D stateだけを更新する時間軸図

trainingとprefill

prompt全体のtokenは既知なので、Selective Scanにより複数位置の部分計算を並列に進めます。計算量はsequence長に対して線形ですが、並列scanの実効速度はkernel、hardware、shapeに依存します。

autoregressive decode

新しいtokenが1個来るたびに、各layerのSSM stateを1回更新します。過去token全体を再入力する必要はなく、Transformerのように生成長に比例してKey/Valueを増やし続ける必要もありません。

ただし、保持量がゼロになるわけではありません。各layerはSSM stateに加え、短いcausal Conv1D用のstateも持ちます。「KV Cacheがない」は、過去tokenごとのKey/Valueを保持しないという意味です。

Transformer、RNN/LSTM、Mambaを比較する

観点 RNN / LSTM Transformer Mamba
過去の表現 固定長hidden state tokenごとのKV 固定長の構造化SSM state
内容に応じた制御 gateで更新 Attention weightで参照 入力依存の\(\Delta,B,C\)で更新・読出し
sequence全体の計算 \(O(L)\) Denseなら \(O(L^2)\) \(O(L)\)
training並列性 recurrenceが制約 token方向へ高い parallel scanで確保
decode memoryの長さ依存 なし KV Cacheが増加 なし
過去の特定tokenへの直接access できない できる できない

MambaはRNNと同じくrecurrent stateを持ちますが、長期memoryを意識した構造化SSMとparallel scanを使う点が異なります。またTransformerと同じcontent-based reasoningを、token pairの比較ではなくstate updateの選択性で近似します。

「線形時間」と「長文を完全に覚える」は別

Mambaの利点を正しく捉えるには、計算効率と情報容量を分ける必要があります。

長さ \(L\)が増えてもstateの大きさを固定すれば、計算とdecode memoryは扱いやすくなります。しかし、1,000 tokenでも100万tokenでも同じ有限stateへ圧縮する以上、過去の細部をすべて損失なく保持することはできません。

Attentionは過去tokenごとの表現を残し、必要な位置へ直接戻れます。Mambaは「何を残すべきか」を学習し、stateへ要約します。この違いから、次のtrade-offが生まれます。

  • streaming、音声、長いsignal：固定長stateと線形計算が生きやすい
  • 過去の特定位置を正確に引用するtask：直接accessできるAttentionが有利になり得る
  • 両方が必要なtask：AttentionとSSMを組み合わせるhybridも選択肢になる

前回のLong Context解説で扱った「最大context長」と「実効context長」の違いはMambaにも当てはまります。線形に計算できる長さと、その全域から情報を正確に使える長さは同じではありません。

原論文の結果はどう読むべきか

Mamba原論文は、language、audio、genomicsなどでMambaを評価しました。language modelingでは、Mamba-3Bが同規模のTransformerを上回り、2倍規模のTransformerに匹敵したと報告しています。また、論文条件ではTransformerより最大5倍高いinference throughputを報告しました。

これらは、Selective SSMが単なる長文向けの理論ではなく、汎用sequence backboneになり得ることを示す重要な結果です。一方で、次のように限定して読む必要があります。

  • 2023年の論文で比較したmodel、data、hardware、実装条件での結果である
  • 「すべてのTransformerより常に5倍速い」という意味ではない
  • model規模、training recipe、kernel対応、batch sizeで相対性能は変わる
  • benchmark上の平均性能と、特定の長文retrieval能力は分けて評価する必要がある

後続のMamba-2論文は、Structured State Space Dualityを通じてSSMとAttentionの関係を一般化し、Mamba coreを再設計しました。これは発展形であり、original Mambaと同一のarchitectureではありません。

Mambaを選ぶときの確認事項

Mamba系modelを実装やsystemへ採用するときは、次の順で確認すると判断しやすくなります。

  1. taskはstreamingか：入力が継続的に届くなら固定長stateが有利
  2. 正確な過去参照が必要か：文書中の特定箇所を直接引くtaskでは圧縮bottleneckを評価する
  3. sequence長は支配的か：短いsequenceではkernel overheadを含めた実測が必要
  4. optimized kernelを使えるか：理論上の \(O(L)\)だけでなくhardwareとsoftware stackを確認する
  5. trainingとservingを分けて測るか：prefill throughput、decode latency、memoryを別々に測る
  6. hybridが適切か：local/global AttentionとSSMを組み合わせる余地を検討する

著者の公式Mamba repositoryにはSelective ScanとMamba blockのreference implementationがあります。実装を読むときは、model定義だけでなく、optimized scan kernelとincremental inference用stateの経路も合わせて確認すると理解が深まります。

まとめ

Mambaは、SSMを入力依存にするSelection mechanismと、そのrecurrent computationをGPUで効率よく計算するSelective Scanを組み合わせたarchitectureです。

固定SSMはrecurrent formとconvolution formを使い分けられますが、全tokenを同じ規則で扱うため、離散dataの内容に応じた選択が弱いという課題があります。Mambaは \(\Delta\)、\(B\)、\(C\)をtokenごとに変え、「残す・書く・読む」を制御します。その代わり固定convolutionは使えなくなり、hardware-aware parallel scanが必要になります。

Transformerと比べた本質的なtrade-offは、過去tokenへの直接accessと固定長stateへの圧縮です。Mambaはsequence長に対して線形にscaleし、decode時の保持量も一定ですが、長い過去を有限stateへ圧縮するbottleneckは残ります。用途に応じてTransformer、Mamba、hybridを選ぶことが重要です。

参考文献

スポンサーリンク

コメント

タイトルとURLをコピーしました