この記事はWIP（Work in Progress）です。
概要
Mambaは時系列データのような系列データの処理に特化した深層学習モデルである。すでに、RNNやLSTM、GRUなどの再帰型ニューラルネットワークや、Transformerのような自己注意機構（Scaled dot-product Attention）を用いたモデルがある。これらに対して、Mambaは線形な状態空間モデル（State Space Model; SSM）を用いたモデルである。これは、隠れ状態が入力と隠れ状態自身によって変化するものとし、その隠れ状態から出力を計算するような定式化であり、再帰型ニューラルネットワークと似ている一方、このモデルは中で計算される隠れ状態の変化が線形であることが特徴である。Mambaではこのモデルのパラメータを入力依存にすることで非線形な活性を外部に出したSelective SSMを用いている。以下は、Mamba-1 のブロック図である。
Mamba-2ブロックではさらなる制約を加え、Selective SSMを行列積で表現する。行列積は現代のGPUで非常に効率的に計算できる。結果としてScaled dot-product Attentionに非常によく似た計算（これをSSDという）になる。Mamba-2ブロックの図は以下の通りである。
この図において、
はAttentionにおけるQuery, Key, Valueに相当する。
離散時間の状態空間モデル（SSM）
ここではMambaの基礎である状態空間モデル（State Space Model; SSM）による時系列データのモデリングについて説明する。離散時間に入る前に、まずは連続時間の状態空間モデルについて導入し、離散化する。
まず、ここでは連続的な変化を行う隠れ状態 と入力 、出力 を持つ次のようなモデルを考える。
つまり、状態は今の状態と入力の線形変換で決まり、出力は状態と入力の線形変換で決まる。ここで、 は状態遷移行列、 は入力行列、 は出力行列、 は直接伝達行列と呼ばれ、ここでは何にも依存しない定数であるとする。また、ここではそれぞれのベクトル、行列の次元は適切であるとする。
ここで、 上の分割 を考える。そして、 の元で であると仮定する。 とおく。このとき、
が成り立つ。天下り的ではあるが、 をかけた の微分を考えた。この両辺を について から まで積分すると、
となる。右辺の積分について、 と置換すると、 であり、 である。よって、
となる。これを先ほどの式に代入すると、
となる。上のように と を定義した。これらは を刻み幅 で離散化した行列で、形は元と同じである。添字は刻み幅そのものを表す。刻み幅が別の値になれば、添字もその値になる。文献ではこれを と書くことも多いが、本章では上線付きの記号を後の節で別の対象に使うので、離散化した行列は刻み幅を添字にして表すことにする。隠れ状態を 、出力を とおくと、離散時間の状態空間モデルは次のように表される。
この条件では は閉形式で計算できる。
さて、一般に深層学習では入力 は時系列データの観測値を処理した後の隠れ状態を とすることが多い。なので とおくと、次のように書き換えられる。
なお、ここで出力は入力を処理した後の隠れ状態 から計算すると決めた。数列 を漸化式を繰り返し用いて次のように書き下してみる。なお、 とする。
よって、 の一般項は次のように表される。
そして、出力 は次のように表される。
ここで、
とすると、
と表せ、総和部分は と の一次元の畳み込みであることがわかる。よって、離散時間の状態空間モデルは、入力 と出力 の間に線形な畳み込み関係があることがわかる。 から までの入力が一度に得られる学習フェーズでは、畳み込みとして並列計算ができる。
Selective SSM
上記の離散時間の状態空間モデルで入力 が出力 に対する重みは である。 の定義を思い出すと、これは にのみ依存する。つまり、入力 が出力 に対する重みは、入力と出力の距離 にのみ依存する。しかし、実際は入力 に依るはずである。たとえば、時系列データ上のノイズの影響を受ける入力 は、距離が近くても出力 に対する重みは小さくなるべきである。そこで、Selective SSMでは、入力 に依存する重みを導入する。ある関数 を用いて、 を入力に依存するようにする。すなわち、
である。ここで、 は入力に依存しない定数であるとする。 は 上の分割 を決める。先ほどと同様に の元で であると仮定する。すると、これらも離散化され、次のように表される。
これらは で一定である。よって、SSMの離散化の議論をそのまま適用できる。すなわち、次のように表される。
ここで は、先ほどの定義の刻み幅に を入れただけのものである。すなわち
である。刻み幅が時刻ごとに違うので、離散化した行列も時刻の数だけ存在することになる。 さて、同様に、数列 を漸化式を繰り返し用いて次のように書き下してみる。なお、 とする。
よって、 の一般項は次のように表される。
そして、出力 は次のように表される。
Mamba-1
入力
Mamba ブロックでは得られた入力を少し処理した後、Selective SSM に入力する。今までの説明では入力はベクトルであったが、Mambaブロックへの入力はそれを束ねたシーケンスである。
この入力に線形変換を行い、特徴量方向の次元を変化させる。その際、 と を作る。
各シーケンスの特徴量方向に線形射影を行って、最終的にそれを2つに分けて、初めの 次元を 、後ろの 次元を としている。Selective SSM に対する入力 は以下のようにして与えられる。
ここで、CausalConv1Dは過去の情報のみを使うシーケンス方向のみに対する（Depthwiseな）畳み込みで、重み を用いて以下のように計算される。
こうして得た入力 が、先ほどのSelective SSMの一つの入力 と以下のように対応する。
Mamba-1におけるSelective SSM
Selective SSMの漸化式は以下のようなものであった。
Mamba-1 では、この漸化式を入力の feature-wise （つまり異なる特徴量の状態を直接混合しないよう）に計算する。ここで記号の形が一段変わるので、先に断っておく。
を対角行列に制限すると、 も対角行列になる。対角行列は対角成分だけで決まるので、 の行列を持ち回る必要はなく、 次元のベクトルで足りる。同じことが にも起きる。また は、状態が 次元で入力と出力がスカラーになるため、行列ではなく 次元ベクトルになる。
そこで、行列から退化したこれらには別の記号を割り当てる。ベクトルであることを示すために小文字にし、離散化したものには上線を付ける。
上付きの は特徴量の番号である。 と は特徴量ごとに別なので を持つが、 は全特徴量で共有するので持たない。第 番目の特徴量について、以下のように計算される。
入力はスカラーのまま、その特徴量に対して得られる隠れ状態はベクトルになる。下にそれぞれの次元をまとめる。
| 変数 | 次元 | 定義 |
|---|---|---|
| スカラ | 時刻 の第 番目の特徴量の入力 | |
| 次元ベクトル | 上記 | |
| スカラ | 出力 | |
| 次元ベクトル | に依存しない。学習可能な定数 | |
| スカラ | の第 成分 | |
| 次元ベクトル | （ に依存） | |
| 次元ベクトル | ||
| 次元ベクトル | ||
| 次元ベクトル | ||
| スカラ | 学習可能な定数 |
連続時間の状態空間モデルでの を対角行列 にしている。これによって隠れ状態の各次元が独立に計算される。
これを各 について計算するが、これは巨大な配列に集約することができる。 本の 次元ベクトルを縦に積んだ の配列であり、状態に掛ける行列ではなく要素ごとの積 で作用する。行列と区別するため、この種の配列には大文字の上に線を引いた記号を割り当てる。次のように決める。
すると、
と表される。ここで、 は 次元の全ての成分が1のベクトルである。初期状態として とする。
値の並び方を確認しておく。 は、第 行に特徴量 の状態 を横倒しにして置いた行列である。成分で書けば
となる。行が特徴量、列が状態の次元である。 と も同じ並びなので、漸化式の は同じ位置の成分どうしを掛けるだけである。
は、 次元の列ベクトルと 次元の行ベクトルの積であるから
となる。各行に同じ値を 回並べただけの行列であり、 と形を揃えるためだけに置いている。出力側の は、 の行列に 次元ベクトルを掛けて 次元ベクトルを得る、通常の行列とベクトルの積である。図にすると次のようになる。
また、 とおくと、以下のように書き下せる。
なので、一般項は
となる。しかし、これを各 について計算するのは、計算量が となり、長いシーケンスに対しては計算量が大きくなる。そこで、Mamba-1では、これを効率的に計算するためのアルゴリズムを導入する。まず、以下の行列の組み（順序対）を定義する。
これは、それぞれ にしか依存しないので、前の状態 によらず独立に計算できる。つまり並行計算が可能である。この時、 を を用いて書き下すと、次のようになる。
ここで について次のような写像 を定義する。
すると、定義から
である。また、先ほどの式から とすると、
でもある。そうすると、試しに以下のような演算 を定義する。
すると、一般に以下が成り立つ。
主張
証明
実際、
つぎに、この演算は結合的であることを証明する。
主張
は結合的である。つまり、任意の について、
証明
実際、
次は連続した結合から隠れ状態を得られることを証明する。ここで として以下のような記号を定義する。
主張
証明
のとき、 となり、これは定義から成り立つ。 ある のとき、次が成り立つと仮定する。
すると、
ここで、 であるので、結論が得られる。
ここで、特に としたとき、 とおくと、次のように書き下せる。
なお、 である。ここで、 は の順に結合したものである。つまり、 は の順に結合した結果から得られる。結合的であることからどんな順に結合してもいいが、
- を得る。これの2要素目は である。
- を得る。これの2要素目は である。
-
を得る。これの2要素目は
である。
- を得る。これの2要素目は である。
という順番で結合すれば、 を順番に得られる。これは並列時間計算量が である。
また、最後の を得るだけであれば、 を二分木の葉に置き、親ノードは子ノードの2つの要素を結合するようにすれば、並列時間計算量は である。
prefix scan
先ほど紹介した順番に結合していく方法よりも、効率的な方法をここでは議論する。なお、途中の隠れ状態すべてを生成できるようなアルゴリズムを prefix scan と呼ぶ。prefix scan は、与えられたシーケンスの各要素に対して、ある二項演算を順番に適用した結果を返すものである。ここでは、二項演算として を用いる。
Hillis-Steele型の prefix scan
まず、 と定義する。これを用いて以下のルールで再帰的に を定義する。
これを繰り返し、 まで計算する。すると、最終的に が得られる。
以下に のときの例を示す。
のとき、
のとき、
のとき、
過去の計算結果（部分的な連続結合）を賢く使うことで、並列計算量は で、計算量（結合回数）は である。
Blelloch scan
まず、結合演算の単位元を定義する。単位元は、結合演算において要素を変化させないような要素である。ここでは、単位元を とする。これは実際単位元である。
Blelloch scanでは以前紹介した二分木構造を用いた結合を行う。ただこれだけだと、木の根以外にノードを使ったとしても全ての隠れ状態は得られない。そこで、木の根から逆方向に足りない結合を行うことで、全ての隠れ状態を得ることができる。前者の木の根までの結合をUp-sweep、後者の逆方向の結合をDown-sweepと呼ぶ。
まず、Up-sweepを行う。これは単純に から までの結合を二分木構造で行う。以下では の木構造の図を示す。
ここまでで、 が得られているがそれ以外は得られていない。そこで、Down-sweepを行う。Down-sweepでは、Up-sweepで得られた結合結果を使って、全ての を計算する。具体的にはまずそれぞれのノードに を割り当て、ルートノードから始めて以下の操作を再帰的に行う。なお、 はノードの の添字と対応する。ルートノードでは を割り当てる。
Up-sweepで の結合が であったとする。
このとき、
とする。これを全てのノードに対して行う。 のとき、各ノードの は以下のようになる。
このように、Down-sweepを行うことで、全ての を計算することができる。これにより、全ての隠れ状態 を効率的に計算することが可能となる。
この方法の良いところは、Up-sweepとDown-sweepの両方を並列に計算できるため、全体の計算時間を大幅に短縮できる点である。特に、長いシーケンスに対しても、結合回数は に抑えられ、並列計算を活用することで実際の計算時間は まで削減可能である。これは後で数える。加えて、このアルゴリズムはほとんど長さ のメモリがあれば済むところである。つまりin-placeでできる。Up-sweepではノードの結合はその右側のノードに上書きする。つまり、以前の例で図示してみると以下のようになる。
このようにすると途中の状態が消えてしまうのではないかという懸念があるかもしれないが、Down-sweepでは、Up-sweepで上書きされたセル、つまり右側の子供のノードは使わない。これはアルゴリズムの式に が現れないことからもわかる。Up-sweepで左側のノードになってしまったセルはその後上書きされることはないため、必要な情報は全て残っている。具体的にはDown-sweepでは、 の結果を のセルに上書きする。 は、 のセルに上書きする。なお、最初に を割り当てる。
以下に、 のときのDown-sweepの様子を示す。
Down-sweepが終わった時点で、第 セルには
が入っている。つまり、自分より前だけを結合したものが並んでいる状態である。このように自分自身を含まないprefixを返すものを排他的（exclusive）なprefix scanと呼ぶ。
一方、我々が欲しかったのは 、すなわち の第2要素であった。これは自分自身を含む結合であり、包含的（inclusive）なprefixである。両者の違いは最後に を結合するかどうかだけであるから、各 について
を計算すればよい。これは ごとに独立なので、 個すべてを1段で同時に計算できる。図示すると以下のようになる。
ここで一つ問題がある。この最後の結合には元の が必要であるが、in-placeで進めてきた以上、それはもう残っていない。Up-sweepの時点で右側の子にあたるセルが上書きされ、Down-sweepではさらに全てのセルが で埋められている。直前の図の最終行を見ればわかるとおり、 の8つのセルに元の は1つも残っていない。
対処は2通りある。1つは を別の配列に取っておくことで、この場合メモリは 個分必要になる。もう1つは、必要になった時点で を作り直すことである。 は のみから決まるのであった。 はブロックの入力側で計算されており、いずれにせよメモリ上に存在する。したがって、 は捨ててしまっても必要になったときに作り直せる。この考え方は後で改めて扱う。
さて、Blelloch scanの結合回数を数えておこう。Up-sweepでは、レベル で 回の結合を行うので、合計は
回である。これは二分木の内部ノードの数に一致する。Down-sweepでも内部ノード1つにつき1回結合するので、同じく 回である。最後の結合が 回であるから、全体では
回となる。並列段数は、Up-sweepが 段、Down-sweepが 段、最後の結合が1段で、合計 段である。
ここまでに出てきた3つの方法をまとめる。
| 方法 | 結合回数 | 並列段数 |
|---|---|---|
| 順に結合 | ||
| Hillis-Steele | ||
| Blelloch |
Hillis-Steeleは並列段数を まで縮めたが、結合回数は であり、順に結合する場合よりも仕事量そのものは増えている。プロセッサが無限にあるならこれでよいが、実際には有限である。Blelloch scanは並列段数を に保ったまま結合回数を に抑えており、この性質を仕事効率的（work-efficient）であるという。その代わり木を2度たどるので、段数の係数はHillis-Steeleの2倍になる。
メモリ階層を意識した実装
ここまでは結合回数と並列段数だけを見てきた。しかし実際のGPU上では、もう一つ別の制約のほうが強く効いてくる。メモリの読み書きである。Mamba-1の実装が「ハードウェアを意識している」と言われるのは、この部分を指している。
GPUのメモリは階層になっている。演算器のすぐ近くには小さく高速なSRAM（共有メモリ）があり、その外側に大きく低速なHBM（一般に「GPUのメモリ容量」と言われるもの）がある。おおよその規模感は次の通りである。
| 階層 | 容量 | 帯域 |
|---|---|---|
| SRAM（共有メモリ） | スレッドブロックあたり数十〜数百KB | 非常に高速 |
| HBM | 数十GB | SRAMより1桁以上遅い |
したがって、ある計算が速いかどうかは浮動小数点演算の回数だけでは決まらない。「HBMから1つ値を読むごとに何回演算するか」が効いてくる。この比を演算強度（arithmetic intensity）と呼び、これが小さい計算はメモリ帯域で頭打ちになる。この状態をメモリ律速（memory-bound）であるという。
Selective SSMのscanがまさにこれである。結合演算 の中身は
であり、 個の要素それぞれについて積2回と和1回を行うだけである。読み書きするデータ量も 個であるから、演算強度は定数である。行列積のように「一度読んだ値を何度も使い回す」構造が無い。
さらに悪いことに、素直に実装すると途中の状態 を全時刻分HBMに書き出すことになる。 が 個であるから 個の値であり、入力 の 個に比べて 倍大きい。 なら16倍である。つまり、状態を素直に持つと、入力を読む何倍ものデータをHBMとやりとりすることになってしまう。
カーネル融合とチャンク分割
そこでMamba-1では、離散化からscan、出力の計算までを1つのGPUカーネルにまとめる。これをカーネル融合（kernel fusion）という。加えて、系列を長さ のチャンクに区切り、チャンク単位でSRAMに載せて処理する。 は必要な作業領域がSRAMに収まるように選ぶ。
具体的な手順は次の通りである。 から始め、チャンク について順に行う。
- HBMから、チャンク に対応する時刻の と、時刻に依らない をSRAMへ読み込む。
- SRAM上で離散化を行い、 と を作る。すなわち を作る。
- SRAM上でチャンク内のprefix scanを実行する。このとき、直前のチャンクから引き継いだ を初期値とする。 が得られる。
- 同じくSRAM上で を計算する。
- HBMへ書き出すのは だけである。 と は捨てる。ただしチャンク末尾の だけは次のチャンクへ引き継ぐ。
3のチャンク内scanには、前節までに述べたBlelloch scanをそのまま使う。チャンク間は逐次だが、チャンクは 個しかない。またバッチ方向と特徴量方向 には完全な並列性があるので、GPUのスレッドブロックはそちらで埋まる。図にすると以下のようになる。
これでどれだけ減るのか、1つのMambaブロックについてHBMとやりとりする値の個数を数えてみる。素直に中間結果を書き出す実装では、 と がそれぞれ 個あるので、全体で 個の読み書きが発生する。一方、融合カーネルがHBMとやりとりするのは、
- 読み込み: と が各 個、 が各 個
- 書き出し: が 個
だけであり、合計は 個である。 のとき、比はおよそ 倍になる。 であれば、HBMとの通信量が 程度まで減る計算である。演算回数は変わっていないが、律速していたのはメモリのほうなので、これがそのまま速度に効く。
逆伝播と推論
順伝播の話は以上だが、逆伝播では が必要になる。 であるから、 についての勾配を求めるには そのものが要るからである。しかし順伝播でそれを捨ててしまっている。
素直に考えれば順伝播で をHBMに保存しておけばよいが、それはちょうど 個の書き出しと読み込みであり、いま避けたばかりのものである。そこで、逆伝播の時点で から を作り直し、もう一度scanして を得る。この「保存せずに作り直す」戦略を再計算（recomputation）と呼ぶ。先ほど、Blelloch scanの最後の結合で を作り直せると述べたのと同じ発想である。
計算量は順伝播1回分増えるが、演算はメモリアクセスに比べて桁違いに安いので、全体としては速くなる。なお、 についての勾配は時刻の逆向きに伝わるので、逆伝播では時刻を逆にたどるscanも併せて行うことになる。
一方、1トークンずつ生成する推論時には、そもそもscanは要らない。漸化式を1ステップ回すだけである。
このとき次の時刻へ持ち越すのは だけである。ただし実装上はもう一つ注意点があり、ブロック入口のCausalConv1Dが直近 時刻の入力を参照するので、 の末尾 時刻分も一緒に保持しておく必要がある。この2つを合わせたものが、Transformerで言うところのKVキャッシュに相当する。違いは、こちらは系列長 に依存しない固定サイズであるという点である。
出力
ここまでで、Selective SSMの出力 が得られた。これを入力と同じように縦に束ねて
とおく。
さて、入力のところで と一緒に作った がまだ使われていなかった。Mamba-1ブロックでは、この をゲートとして使う。
ここで は入力のところでも使った活性化関数で、
で定義される。 も も要素ごとに作用する。
ゲートの役割は、SSMが出した特徴量のうちどれを通すかを入力自身に決めさせることである。 は の線形射影であったから、 の各成分は入力に応じて 付近から大きな値までを取り、 に近い成分は出力にほとんど現れない。Selective SSMが「時間方向にどの入力を残すか」を選んでいたのに対し、こちらは「特徴量方向にどの成分を残すか」を選んでいることになる。
最後に、特徴量方向の次元を入力と同じ に戻す。
ここで である。以上をまとめると、Mamba-1ブロックは次の図のようになる。
Mamba-1ブロックの計算手順
ここまでで部品が出揃ったので、入力から出力までを順に並べておく。 を入力とし、行列の第 行を と書く。
1. 入力の射影 — 特徴量方向を 次元に広げ、半分ずつに分ける。
2. 畳み込みと活性化 — だけがSSMへ向かう。
3. 時刻ごとのパラメータ — として、1つの射影から3つをまとめて作り、切り分ける。
4. 離散化 — として、
とする。 は を 列ぶん横に並べた配列、 は 次元ベクトルと 次元ベクトルの外積である。2つ目の式は、 に一次近似 を入れて整理したものである。
5. scan — から始めて
を全時刻について求める。学習時はprefix scanで並列に、推論時は1ステップずつ回す。
6. SSMの出力
7. ゲートと出力の射影 — を縦に束ねて とし、
3から6は時刻 ごとに書いたが、3と4と6は について完全に独立なので一度に計算できる。逐次性があるのは5だけである。
ブロックの積み重ねとモデル全体
ここまでで1つのMambaブロックの中身が分かった。ここからは、それを何層も積み重ねて言語モデルに仕立てるところまでを見る。
Mambaブロックは、単体で使うのではなく、残差接続と正規化を伴って何層も積み重ねる。第 層は
である。図にすると次のようになる。
RMSNorm
ここで使っているRMSNormについて、何をしている層なのか、なぜ必要なのか、そしてなぜLayerNormではないのかを順に見ていく。
RMSNormは、各トークンの特徴量ベクトル に対して
で定義される。分母に現れる
は各成分の二乗平均の平方根であり、二乗平均平方根（root mean square）と呼ばれる。要するにベクトルの大きさを次元数で均したものである。 は零除算を避けるための小さな定数、 は学習可能なスケールで、通常は全成分 で初期化する。
を で割るということは、ベクトルの向きだけを残して大きさを揃えるということである。実際、 かつ のとき、出力のノルムは
となり、入力の大きさによらず一定である。そのうえで によって特徴量ごとに大きさを付け直す。つまりRMSNormは「大きさをいったん捨てて、望ましい大きさを学習し直す」層である。
では、なぜこれが必要なのか。理由は2つある。1つ目は、残差接続によって値がどんどん大きくなるのを抑えるためである。 という形をしているので、層を通るたびに何かが足され続ける。各層の出力が独立に近ければ、分散はおおむね層数に比例して増えていく。ところがブロックの中には や といった、入力の大きさに敏感な非線形関数が入っている。入力の大きさが層ごとに変わってしまうと、浅い層と深い層で の効き方も の値域も変わってしまう。ブロックの手前で大きさを揃えておけば、どの層も同じ条件で入力を受け取れる。
2つ目は、勾配が深いところまで届くようにするためである。ここで正規化を残差の経路の外、すなわちブロックの手前に置いていることが効いてくる。この置き方をpre-normという。式が
の形なので、 から へは、 を通らない恒等の経路が残る。したがって勾配は、層をいくつ跨いでも減衰しない経路を1本持つことになる。逆に、正規化を残差の後ろに置くと（post-norm）、この経路にも正規化が挟まってしまい、層を深くしたときに学習が難しくなる。
最後に、平均を引かない点について触れておく。9章で扱ったBatch Normalizationや、Transformerで使われるLayer Normalizationは、平均を引いてから標準偏差で割る。
これに対しRMSNormは平均を引かず、単に大きさで割るだけである。上で見たとおり、正規化に期待している役割は「大きさを揃えること」であって、「中心を に寄せること」ではない。実際、平均を引かなくても精度はほとんど変わらないことが知られている。
そして省ける計算が意外に大きい。平均を引く操作は特徴量方向の総和を取ってから全成分に減算を配る必要があり、これも演算に対して読み書きの比率が高い、メモリ律速な処理である。メモリ階層の節で見たとおりMambaはメモリ帯域を節約する設計になっているので、正規化だけ重いのでは筋が通らない。RMSNormはその方針に合っている。
言語モデルとして動かす
ここまででMambaブロックの中身と積み重ね方が分かった。では、これだけで言語モデルとして動くのだろうか。答えは「動かない。ただし足りないのは入口と出口だけで、それはTransformerと全く同じものである」となる。
12章のTransformerによる言語モデルは、
という構成であった。Mamba-1はこのうち真ん中だけを差し替えたものである。すなわち、Attentionとフィードフォワードネットワークという2種類のサブブロックの繰り返しが、Mambaブロック1種類の繰り返しに置き換わる。前後は変わらない。
| 部分 | Transformer | Mamba-1 |
|---|---|---|
| 入口 | トークン埋め込み ＋ 位置エンコーディング | トークン埋め込みのみ |
| 中身 | （Self-Attention ＋ FFN） | Mambaブロック |
| 正規化の位置 | 各サブブロックの手前（pre-norm） | 各ブロックの手前（pre-norm） |
| 出口 | 正規化 線形層 softmax | 同じ |
入口で違うのは位置エンコーディングが要らないことである。Attentionは全トークンを対称に扱うので順序の情報を外から与える必要があったが、Mambaは を時刻順に更新していく漸化式であり、順序そのものが計算の構造に入っている。
全体を図にすると以下のようになる。
破線のブロックが、この節の初めに図示した「RMSNorm Mambaブロック 残差の加算」を 回繰り返す部分である。最後にもう一度RMSNormを置くのは、残差を足し続けた結果の大きさを出力層の手前で揃えるためである。出力層 は語彙数 次元へ射影し、softmaxを通して次トークンの確率分布を得る。埋め込み行列と出力層の重みを共有する（weight tying）ことも多い。
離散化の一次近似
Selective SSMの表では、 を積分の形で定義していた。ところが実際の実装では、この積分を厳密に計算せず、一次近似で済ませることが多い。ここでは、その近似がどこから出てくるのか、どれくらいずれるのか、そしてなぜそれが許されるのかを順に導出する。
積分を厳密に計算する
まず、そもそも厳密な値がどうなるかを求める。連続時間の を対角行列にしたことで、 次元の各成分は独立に動く。したがって成分ごとに書き下せる。
記号を軽くするため、特徴量 と時刻 を固定して
とおく。定義から、 の第 成分は
である。 は に依らないので積分の外に出せる。残った積分は、 と置換すると、 であり、 が から まで動くとき は から まで動くので、
となる。 のとき、
である。よって厳密な値は
となる。 の第 成分が であったことを思い出せば、これは
とも書ける。すなわち、 さえ計算してあれば、 は で割るだけで得られる。
一次近似
ここから近似を導く。上の式で を括り出してみる。分母と分子に を掛けて
とすると、右側の因子は だけの関数になっている。そこで
とおくと、厳密な値は
と、 に補正係数 が掛かった形に整理できる。
この は で になるが、これは除去可能な特異点である。実際、 のマクローリン展開
から
であり、両辺を で割ると
を得る。この級数は全ての で収束するので、 は を含めて解析的であり、 である。
したがって、 が小さいところでは
であり、 であれば 、すなわち
と近似できる。ベクトルの形に戻すと
となる。これが実装で使われている形である。
この近似には別の見方もある。連続時間の式 を、区間 で前進Euler法により離散化すると
となる。つまりEuler法では 、 である。Mamba-1が使っているのは、 はZOHの厳密な のまま、 だけEuler法の に取り替えた混ざった形である、と言える。
近似が許される理由
を近似するなら も近似してよさそうに思えるが、この2つは役割が違う。
は状態に繰り返し掛かる。 の一般項に という積が現れていたことを思い出すと、 の誤差は掛け算で積み上がる。長さ の系列では最大 回掛かるので、1回あたりが小さな誤差でも系列全体では効いてしまう。
さらに重要なのが安定性である。 として を保証しているので、厳密な は がどれだけ大きくなっても必ず に入り、状態は必ず減衰する。ところがEuler法の は、 で負になり、 では絶対値が を超えて発散する。 は入力依存で、しかも学習で変わる量であるから、この保証が失われるのは困る。
一方 は各時刻の入力に1回だけ掛かる。誤差が積み上がる構造にはなっていないし、発散にも関わらない。
では実際どれくらいずれるのか。 かつ なので であり、このとき である。つまり近似 は真の値を過大に見積もる。いくつかの値を挙げる。
| 真の値からのずれ | ||
|---|---|---|
は が 程度になるよう初期化され、 の初期値は であった。 とすると は最大でも である。つまり、 が小さい成分や が小さい成分ではほぼ厳密だが、両方が大きい成分では2倍近くずれうる。
ただし、大きくずれる領域は 、すなわち1ステップで状態の8割が消える成分である。そこに書き込まれた情報はどのみち数ステップしか残らないので、その寄与の大きさが多少ずれても出力への影響は限定的である。加えて は学習で決まるので、系統的なずれはある程度 の側で吸収されうる。
もう一つ、近似が乱暴でないことの根拠を挙げておく。入力を一定値 に保ったときに状態が落ち着く先、すなわち定常状態を比べてみる。漸化式 の不動点 は
である。厳密な値を代入すると、分子が であるから
となる。これは連続時間の式で とおいた の解と完全に一致する。ZOHによる離散化は定常状態を厳密に保存するのである。
近似 を使うとどうなるか。
であり、分母は
と展開できるので、
となる。定常状態も一次の精度では保たれている。
最後に、実装上の利点を2つ挙げておく。1つは計算が減ることである。厳密な形は
に加えて
による除算を必要とするが、近似では
を掛けるだけで済む。
は
の計算でどのみち必要なので、差分は時刻あたり
回の除算がまるごと消えることになる。もう1つは数値的な安全性である。
は
で
になるため、素直に書くと桁落ちする。避けるにはexpm1や場合分けが要る。
にはその問題がなく、学習の途中で
が
に近づいても壊れない。
なお繰り返しになるが、 についてはこの近似を使わず、 をそのまま計算する。
実装のためのまとめ
ここで、実装に必要な量を一覧にしておく。 をモデルの次元、 をブロック内部の次元、 を状態の次元、 を畳み込みの幅、 を の低ランク射影の階数とする。代表的な設定は 、 、 、 である。
は、1つの線形射影でまとめて作ってから切り分ける。
は時間の刻み幅であるから正でなければならない。それを保証するために を通している。また だけ 次元を経由する低ランクの射影になっているのは、 が 次元すべてに対して必要である一方、 のように 次元で済むものと比べてパラメータ数が大きくなりすぎるためである。
| 記号 | 形 | 役割 | 初期化の要点 |
|---|---|---|---|
| トークン埋め込み | 通常の乱数 | ||
| への射影 | 通常の乱数 | ||
| 深さ方向の因果的畳み込み | 通常の乱数 | ||
| への射影 | 通常の乱数 | ||
| , | への射影 | が 程度になるよう調整 | |
| スキップ結合の係数 | 全成分 | ||
| ブロック出力の射影 | 通常の乱数 | ||
| RMSNormのスケール | 全成分 | ||
| 出力層 | 埋め込みと共有することも多い |
の扱いには理由がある。 を直接学習可能なパラメータにするのではなく、 と置いて のほうを学習する。こうすると は必ず負であり、したがって
が保証される。 であることと合わせて、状態は時刻を追うごとに必ず減衰する。もし が を超えると、漸化式を回すうちに が発散してしまう。初期値を 、すなわち 番目の状態ほど速く減衰するようにしておくと、状態の各次元が異なる時間スケールを担当することになり、短期の情報と長期の情報を同時に保持しやすくなる。
については、前節で導いたとおり一次近似の を使えばよい。 のほうは をそのまま計算する。
Transformerとの比較
最後に、12章で扱ったTransformerと比較しておく。 を系列長、 をトークン表現の次元、 をMambaブロック内部の次元、 を状態の次元とする。
| 項目 | Transformer（Self-Attention） | Mamba-1 |
|---|---|---|
| 学習時の計算量 | ||
| 学習時の並列段数 | ||
| 推論時の1トークンあたりの計算量 | ||
| 推論時に持ち越す状態 | KVキャッシュ | と畳み込みの直近 時刻分 |
| 位置の情報 | 位置エンコーディングが必要 | 漸化式が順序を保持する |
Attentionは全てのトークン対について内積を取るため、系列長に対して二乗で計算量が増える。その代わり、全ての対を独立に計算できるので並列段数は定数である。Mambaは漸化式であるから本来は逐次的だが、prefix scanによって 段まで並列化でき、計算量は について線形に収まる。
推論時の差はさらに大きい。Transformerは過去の全トークンのキーとバリューを保持し続ける必要があり、生成が進むほどメモリも1トークンあたりの計算量も増えていく。一方、Mambaが次の時刻へ持ち越すのは だけであり、これは に依存しない。長い系列を扱うほどこの差は効いてくる。
ただし、Attentionが持つ「任意の過去のトークンを直接参照できる」という性質は、固定サイズの状態 には無い。 は過去を 個の値に圧縮したものであり、そこに入り切らない情報は失われる。Selective SSMにおいて を入力依存にしたことは、この限られた容量に何を残すかを入力自身に選ばせるための仕組みであった、と見ることもできる。
Mamba-2
Mamba-1によって長い系列を線形の計算量で扱えるようになったが、速度の面では不満が残っていた。原因はscanにある。結合演算 の中身は要素ごとの積と和だけであり、行列積の形をしていない。ところが現代のGPUは行列積のための専用の演算器を積んでおり、そこを通る計算とそうでない計算とでは最大で16倍ほどの速度差がある。scanはその恩恵をまったく受けられない。状態次元 を までしか大きくできなかったのも、突き詰めればここに理由がある。
Mamba-2はこの壁を理論の側から崩した。Selective SSMの計算を行列ひとつの積として書き直すと、それが注意機構とほとんど同じ形になる。この対応を状態空間双対性、あるいはSSDと呼ぶ。同じ計算に、漸化式として時刻順に回す見方と、行列積として一度に片づける見方の2通りがあることになり、後者を選べば行列積の演算器が使える。この節では、まずその書き直しを導き、次にそれを実際の計算手順に落とす。
半分離可能行列としてのSSM
Selective SSMの一般項を思い出す。特徴量 を固定し、記号から上付きの を落として書けば、
であった。出力は であるから、これを代入して
を得る。スキップ結合の項は以降の議論に関わらないので省いた。
この式は、出力 が入力 の線形結合であることを言っている。そこで の係数だけを取り出して
と定義する。 のとき としたのは、未来の入力が出力に影響しないためである。すると、出力を縦に並べた と入力を縦に並べた の間には
という関係が成り立つ。 は下三角行列である。漸化式が、行列をひとつ掛けるだけの計算に化けた。
どういう値がどこに並んでいるのかを見ておこう。 として書き下すと
となる。第 行は出力 を、第 列は入力 を担当する。ひとつ下の行へ移るごとに がひとつ増え、左へ一列進むごとにも がひとつ増える。つまり、対角から左下へ離れるほど減衰が積み重なる。上三角が なのは、未来の入力を見ないという因果性そのものである。
減衰が積み重なる様子は、もう少し長く取ったほうが見えやすい。 として、各成分に何回 が掛かるかだけを描くと次のようになる。
図の は、その成分に減衰が 回掛かっていることを表す記号として使った。実際に掛かるのは と時刻ごとに違うベクトルであり、同じものの累乗ではない。
この には強い構造がある。それを見るために、時刻 から までの減衰をまとめて
とおく。 のとき と書ける。 は要素ごとの除算を表す。これを代入すると
となる。2つ目の等号では、要素ごとの積を内積のどちら側に寄せてもよいことを使った。そこで
とおけば、対角より下の成分はすべて
という 次元ベクトル同士の内積で書けることになる。ここから次が従う。
主張
の対角線より真に下の領域に完全に含まれる部分行列は、そのランクが 以下である。
証明
そのような部分行列の成分はすべて を満たすので、上で得た がそのまま使える。行の添字を 、列の添字を とすると、この部分行列は
と、 行列と 行列の積に分解される。行列の積のランクは各因子のランクを超えないから、この部分行列のランクは 以下である。
対角より下がこのように低ランクになっている行列を半分離可能行列と呼ぶ。ランクの上限が であることを強調して -半分離可能行列ということもある。Selective SSMとは、要するに半分離可能行列を掛ける操作だったわけである。
ただし、この分解をそのまま計算に使うことはできない。 は減衰を積み重ねたものであり、 が大きくなると急速に へ近づく。それで割る は数値として発散してしまう。分解の存在は理論として重要だが、実際の計算では後で述べるチャンク分割を使う。
注意機構との双対性
ここでMamba-2は に制限を課す。 個の成分がすべて等しいとするのである。すなわちスカラー を用いて
と書けるとする。連続時間に戻せば、対角行列 の対角成分がすべて同じ値であること、つまり がスカラーと単位行列の積であることに当たる。Mamba-1では状態の 個の次元がそれぞれ別の速さで減衰していたが、その自由度を捨てて、ひとつの速さに揃える。
この制限のもとでは
となり、減衰がスカラーとして内積の外に出せる。したがって
である。右辺が2つの部分に分かれたことが重要である。 は と の両方に依存するが、 次元ベクトル同士の内積でしかない。一方 はスカラーであり、時刻の情報だけを持つ。この2つをそれぞれ の行列に集約する。
まず内積のほうを扱う。 と を、今度は時刻の方向に積んで の配列 を作る。Mamba-1の が特徴量の方向に積んだ配列で時刻ごとに1つずつ存在したのに対し、こちらは系列全体で1つである。積む方向が違えば別の量なので、記号も分けて 、 と書く。
とおく。 の第 行が 、 の第 行が である。このとき
となり、 成分がちょうど欲しかった内積になる。 の行列と の行列を掛けて を作る形であり、12章の と同じ計算である。
次にスカラーのほうを
とまとめて行列 を作る。原論文はこれを と書くが、ここでは系列長 と紛れるので とした。 で書き下すと
である。対角は積が空なので 、左下へ進むごとに因子がひとつずつ増えていく。
この2つを要素ごとに掛ければ に戻るから、
と書ける。 は の半分離可能行列になっているので、1-半分離可能行列と呼ばれる。
この形まで来たものがSSDである。冒頭で状態空間双対性という名前だけ挙げておいたが、実体はこの式である。以降、 をスカラーに制限したSelective SSMのことをSSDと呼び、その計算の仕方を詰めていく。
それぞれの行列の形を並べると次のようになる。
ここで12章の注意機構を並べてみる。因果マスクを掛けたAttentionの出力は
であった。一方、SSDは
と計算される。 がクエリ、 がキー、 がバリューに対応している。違いは2つしかない。
| 因果マスク付きAttention | SSD | |
|---|---|---|
| 類似度 | ||
| 正規化 | softmaxで行ごとに和を にする | 何もしない |
| マスク | 未来を にするだけ | 未来を にし、過去には を掛ける |
Attentionのほうも同じ構図で描いてみると、違いが目で見える。
1段目はSSDとまったく同じ形をしている。細長い2つの配列を掛けて の正方行列を作るところまでは共通である。
分かれるのは2段目以降である。SSDは を要素ごとに掛けて終わる。Attentionはまずマスク を足す。これは未来にあたる上三角の位置に を置いた行列で、足された成分は になる。そのうえで行ごとにsoftmaxを通すと、 より上三角がちょうど になり、残った下三角の和が に揃う。マスクを掛け算ではなく足し算で入れるのは、softmaxの前段だからである。指数を取った後の世界では、足し算は掛け算に化ける。
出来上がる重み行列の中身も違う。SSDの は対角から左下へ離れるほど薄くなる。減衰 が掛かる回数が増えるからで、遠い過去ほど自動的に効かなくなる。一方Attentionの重みは、softmaxが行ごとに和を に揃えるだけなので、どの過去に濃さを置くかは完全に自由である。遠く離れたトークンに重みのほとんどを集めることもできる。
この差が、そのまま両者の性格の違いになっている。SSDは「近い過去を見る」という偏りを構造として持ち込むかわりに、 を掛けるだけで済むので線形形式でも計算できる。Attentionは偏りを持たないかわりに、 の行列を実際に作らないとsoftmaxが計算できない。
以上をまとめると、SSDとAttentionの違いはsoftmaxが無いことと、マスクが か かではなく過去へ遡るほど小さくなる重みになっていることの2点だけである。逆に言えば、Attentionのほうも「マスクを と に固定し、softmaxを入れたSSD」として見ることができる。
こうして、SSDにはひとつの計算に2つの顔があることがわかった。漸化式として時刻順に回せば結合回数は で済むが、逐次であるから並列に進められない。行列 を作って掛ければ の仕事量になるが、中身はすべて行列積である。前者を線形形式、後者を二次形式と呼ぶ。SSDがこの2つの形を同時に持っているという事実が、状態空間双対性である。
チャンク分割による計算
SSDを実際にどう計算するかを決めよう。二次形式は行列積だけで書けるものの、 が大きいと の行列を作る時点で破綻する。線形形式は仕事量こそ小さいが逐次である。そこでMamba-2は、系列を長さ のチャンクに区切り、チャンクの中では二次形式を、チャンクの間では線形形式を使う。SSDが2つの顔を持っていることを、そのまま計算に使うわけである。
をチャンクごとのブロックに分けて考える。チャンク が時刻 を担当するとし、行がチャンク 、列がチャンク に属する部分行列を と書く。 のブロックは因果性からすべて である。 の対角ブロックは、チャンクの中で閉じた添字だけを含むので、前節の二次形式をそのままチャンクに適用したものになる。大きさは であり、 を手頃に取れば行列積で扱える。
問題は のブロックである。ここで の積を分解する。 はチャンク 、 はチャンク の時刻であるから、区間 はチャンク の残り、その間のチャンク全体、そしてチャンク の先頭から までの3つに分けられる。すなわち
である。 はチャンク の入口から までの減衰、 は からチャンク の出口までの減衰、 はその間をまたぐ減衰を表す。これを に戻すと
となる。 を行に並べた 行列と、 を行に並べた 行列があれば、非対角ブロックはその積のスカラー倍として書ける。ランクは 以下である。 の行列を作る必要はない。
図にすると次のようになる。
非対角ブロックが低ランクであることから、チャンク より前のすべての入力が に与える寄与は、ひとまとめにできる。実際、
であり、下線を引いた だけ持ち回ればよい。これはチャンク の入口における状態にほかならない。しかも 自身が、チャンクをまたぐ減衰 を用いた漸化式
を満たす。これは時刻についての漸化式と同じ形であり、長さは しかない。
以上をまとめると、計算は次の4段階になる。
- 対角ブロック — 各チャンクについて、入口の状態を と思ったときの出力を二次形式で計算する。チャンクごとに独立なので、行列積をまとめて実行できる。
- 各チャンクが生む状態 — を計算する。これもチャンクごとに独立な行列積である。
- チャンク間の受け渡し — 上の漸化式を回して、各チャンクの入口の状態 を求める。ここだけがscanだが、長さは に縮んでいる。
- 入口の状態からの寄与 — を各時刻に足す。再び行列積である。
1と2と4は完全に行列積であり、行列積の演算器がそのまま使える。3のscanには前節までに述べたBlelloch scanをそのまま使えるが、対象がチャンク境界の状態だけに減っているので、全体に占める割合は小さい。仕事量も、二次形式をそのまま使ったときの から について線形なところまで戻る。Mamba-2がMamba-1の2倍から8倍速いと報告されているのは、この置き換えによる。以降、この4段階の手順をまとめてSSDの計算と呼ぶ。ブロックの図で「SSD」と書かれた箱の中身がこれである。
の選び方には両側から圧力がかかる。 を大きくすると対角ブロックの の計算が増え、小さくするとscanの回数が増える。実際には から 程度が使われる。
アーキテクチャの変更
双対性を使うために をスカラーに落としたので、Mamba-1のブロックもそれに合わせて作り替える必要がある。変更は3点ある。
状態をヘッドにまとめる
がスカラーになったということは、Mamba-1で特徴量ごとに持っていた 個の減衰率が1つに減ったということである。そのままでは表現力が落ちる。そこでMamba-2は、 個の特徴量を 個ずつの組に分け、組の中では同じ を共有することにした。この組をヘッドと呼ぶ。 をヘッド次元といい、 が使われる。ヘッドの数は である。
これは12章のマルチヘッドAttentionと同じ構造である。Mamba-1は 、すなわちすべての特徴量が独立なSSMを持っていたことになる。ヘッドにまとめると自由度は下がるが、その代わりに状態次元 を大きくできる。行列積が使えるようになったことで まで引き上げられ、差し引きでは表現力が増している。
なお と はヘッド間で共有する。Attentionで言えば、クエリとキーを共有したまま複数のバリューを持つ形に当たる。論文はこれを多入力SSMと呼び、multi-value attentionに対応づけている。
パラメータを並列に作る
Mamba-1では、 を作るのに、まずCausalConv1DとSiLUを通した後の が必要だった。つまり入力の射影とパラメータの生成が直列に並んでいた。Mamba-2では、ブロックの入口の線形射影ひとつで 、 、 、 、 をまとめて作る。これらの間に依存関係が無くなるので、複数のGPUに重みを分けて配置するテンソル並列がやりやすくなる。
畳み込みの掛け方も変わる。Mamba-1は だけに掛けていたが、Mamba-2は と と を横に連結したものに掛ける。
出力の直前に正規化を置く
Mamba-1は、SSMの出力に を掛けたら、そのまま出力側の射影に渡していた。Mamba-2はその間に正規化を挟む。ゲートを掛けるところまでは同じで、
としたあとに
を通してから に渡す。大規模なモデルで学習が不安定になる問題への対策であり、Transformerの派生であるNormFormerが同じ位置に正規化を置いたのを引き継いでいる。
なお実装では、このゲートと正規化はひとつの層にまとめられている。順序を入れ替えて とすることもできるようになっているが、既定は上の順序である。
以上をまとめると、Mamba-2のブロックは次のようになる。
Mamba-1のブロック図と見比べてほしい。骨格はそのままで、違いは2つだけである。ひとつは、左の列に流れるものが ひとつではなく の3つになっていること。もうひとつは、 で合流した後に がひとつ挟まっていることである。
は図では省いたが、 たちと同じく入口の射影から作られ、SSDの中で使われる。ブロックの外側も変わらない。RMSNormを手前に置いて残差で繋ぎ、それを 層積み重ね、最後に正規化と出力層を置く構成はMamba-1と同じである。
Mamba-2ブロックの計算手順
Mamba-1と同じように、入力から出力までを並べておく。ヘッド次元を 、ヘッド数を とし、 と は全ヘッドで共有する場合を書く。
1. 入力の射影
必要なものを1つの射影ですべて作り、切り分ける。
それぞれ 、 、 である。Mamba-1と違い、 のための低ランクの射影は無く、ここで直接 次元ぶん作られる。
2. 畳み込みと活性化 — を横に連結したものに掛け、掛けたら元の3つに戻す。
3. 刻み幅と減衰 — として、ヘッドごとにスカラーを作る。
4. ヘッドごとのSSD — の第 行を 次元ずつ 個に分け、ヘッド の入力を とする。各ヘッドについて、減衰 、 、 を使ってSSDを計算する。中身はチャンク分割の節で述べた4段階そのままである。状態は 、出力は である。
5. ヘッドの連結とスキップ結合 — 個の出力を縦に繋いで 次元に戻す。
6. ゲート、正規化、出力の射影 — を縦に束ねて とし、
Mamba-1の手順と見比べると、1で全部まとめて作るようになったこと、2の畳み込みが にも掛かること、3の がヘッドあたり1つのスカラーに減ったこと、4が行列積で書けるSSDになったこと、6に正規化が入ったことの5点が違う。
Mamba-1との比較
最後に、2つの違いを一覧にする。
| 項目 | Mamba-1 | Mamba-2 |
|---|---|---|
| の構造 | 対角。特徴量ごとに 個の減衰率 | スカラーと単位行列の積。ヘッドごとに1個 |
| ヘッド次元 | ||
| 状態次元 | ||
| と の生成 | 畳み込みと活性化の後に直列で | 入口の射影でまとめて並列に |
| 畳み込みの対象 | のみ | を連結したもの |
| 出力側 | ゲートのみ | ゲートの後に正規化 |
| 主計算 | prefix scan | チャンク分割した行列積 |
| 行列積の演算器 | 使えない | 使える |
をスカラーに制限したことだけを見れば、Mamba-2はMamba-1より表現力の低いモデルである。それでも性能が上がるのは、その制限と引き換えに状態次元を8倍にでき、しかも学習が速くなるからである。制約を課すことでハードウェアに合う形へ持ち込み、浮いた分を別の次元に注ぎ込む。この構図は、Mamba-1がメモリ階層に合わせてカーネルを設計したのと同じ発想の延長にある。