Mambaとは
Selective State Space Model
選択的状態空間モデルに基づき、系列長に線形の計算量で長い系列を処理できるアーキテクチャ
ひとことで言うと
Transformerの「文が長くなるほど計算が急増する」弱点を避け、長い文章を軽い計算で処理できる新しい仕組み。
概要
Mambaとは、選択的状態空間モデル（Selective SSM）に基づく系列モデリングのアーキテクチャを指す。 Transformerの自己注意機構は系列長の2乗に比例する計算量を要するのに対し、Mambaは系列長に線形の計算量で処理でき、長い系列を効率よく扱える。 入力に応じて状態の更新を変える選択機構の導入により、従来の状態空間モデルの弱点だった内容依存の処理を可能にした。 Transformerの代替や、注意機構と組み合わせたハイブリッド構成の要素として、言語・音声・ゲノムなどの長系列タスクで採用が広がっている。
歴史
2023年12月にAlbert GuとTri Daoが論文「Mamba: Linear-Time Sequence Modeling with Selective State Spaces」で発表した。2024年5月には、状態空間モデルと注意機構の理論的な関係を示し学習効率を高めた後継のMamba-2が発表された。
アーキテクチャ
状態空間モデル（SSM）は入力を潜在状態へ逐次的に圧縮しながら系列を処理する仕組みだが、従来は状態遷移のパラメータが入力によらず固定だった。Mambaはこの遷移パラメータ（選択機構）を入力トークンごとに変化させることで、内容に応じて重要な情報を状態に残したり忘れたりできるようにした。この選択機構は、GPU上で効率的に計算できるハードウェア対応の並列スキャンアルゴリズムと組み合わせて実装される。
ワークフロー
入力トークン列を受け取る → 各トークンごとに入力依存の選択機構がSSMのパラメータを計算 → 並列スキャンにより状態を系列全体で効率的に更新 → 各時点の状態から出力を計算する。
コード例
mamba-ssmライブラリでのMambaブロックの利用
import torch
from mamba_ssm import Mamba
model = Mamba(d_model=256, d_state=16, d_conv=4, expand=2).to("cuda")
x = torch.randn(1, 128, 256).to("cuda")
y = model(x)利点
- 系列長に対して線形の計算量で処理でき、長い系列でもTransformerより効率的に扱える
- 推論時の状態サイズが系列長に依存せず一定であり、長いコンテキストでもメモリ使用量が抑えられる
- 選択機構により、入力内容に応じて必要な情報を状態に保持し不要な情報を忘れられる
欠点
- Transformerの自己注意のような、任意の2トークン間を直接参照する仕組みを持たないため、精密な照合を要するタスクでは注意機構に劣る場合がある
- 登場から日が浅く、大規模モデルでの学習・運用のノウハウはTransformerほど蓄積されていない
- 既存のTransformer向け推論エンジンやツールをそのまま流用しにくい
比較
- Transformer — Transformerが系列長の2乗に比例する自己注意で全トークン間の関係を捉えるのに対し、Mambaは線形計算量の状態空間モデルで系列を逐次的に圧縮しながら処理する
- State Space Model — 状態空間モデルの枠組みに、入力内容に応じて状態の更新を変える選択機構を加えたものがMamba
関連用語
よくある質問
MambaはTransformerを置き換える？
長い系列の処理効率では優位性を示す一方、精密な照合能力ではTransformerに一日の長があるとされ、両者を組み合わせたハイブリッド構成を採る研究・モデルも増えている。