mamba-architecture

Mamba - 選択的状態空間モデル

クイックスタート

Mamba は、シーケンスモデリングにおいて O(n) の線形計算量を実現する状態空間モデルのアーキテクチャです。

インストール：

# Install causal-conv1d (optional, for efficiency)
pip install causal-conv1d>=1.4.0

# Install Mamba
pip install mamba-ssm
# Or both together
pip install mamba-ssm[causal-conv1d]

前提条件：Linux、NVIDIA GPU、PyTorch 1.12 以降、CUDA 11.6 以降

基本的な使い方（Mamba ブロック）：

import torch
from mamba_ssm import Mamba

batch, length, dim = 2, 64, 16
x = torch.randn(batch, length, dim).to("cuda")

model = Mamba(
    d_model=dim,      # Model dimension
    d_state=16,       # SSM state dimension
    d_conv=4,         # Conv1d kernel size
    expand=2          # Expansion factor
).to("cuda")

y = model(x)  # O(n) complexity!
assert y.shape == x.shape

一般的なワークフロー

ワークフロー 1：Mamba-2 による言語モデル

生成機能を備えた完全な LM：

from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
from mamba_ssm.models.config_mamba import MambaConfig
import torch

# Configure Mamba-2 LM
config = MambaConfig(
    d_model=1024,           # Hidden dimension
    n_layer=24,             # Number of layers
    vocab_size=50277,       # Vocabulary size
    ssm_cfg=dict(
        layer="Mamba2",     # Use Mamba-2
        d_state=128,        # Larger state for Mamba-2
        headdim=64,         # Head dimension
        ngroups=1           # Number of groups
    )
)

model = MambaLMHeadModel(config, device="cuda", dtype=torch.float16)

# Generate text
input_ids = torch.randint(0, 1000, (1, 20), device="cuda", dtype=torch.long)
output = model.generate(
    input_ids=input_ids,
    max_length=100,
    temperature=0.7,
    top_p=0.9
)

ワークフロー 2：事前学習済み Mamba モデルを使用する

HuggingFace から読み込む：

from transformers import AutoTokenizer
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel

# Load pretrained model
model_name = "state-spaces/mamba-2.8b"
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")  # Use compatible tokenizer
model = MambaLMHeadModel.from_pretrained(model_name, device="cuda", dtype=torch.float16)

# Generate
prompt = "The future of AI is"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to("cuda")
output_ids = model.generate(
    input_ids=input_ids,
    max_length=200,
    temperature=0.7,
    top_p=0.9,
    repetition_penalty=1.2
)
generated_text = tokenizer.decode(output_ids[0])
print(generated_text)

利用可能なモデル：

• state-spaces/mamba-130m
• state-spaces/mamba-370m
• state-spaces/mamba-790m
• state-spaces/mamba-1.4b
• state-spaces/mamba-2.8b

ワークフロー 3：Mamba-1 と Mamba-2 の比較

Mamba-1（小さな状態）：

from mamba_ssm import Mamba

model = Mamba(
    d_model=256,
    d_state=16,      # Smaller state dimension
    d_conv=4,
    expand=2
).to("cuda")

Mamba-2（マルチヘッド、大きな状態）：

from mamba_ssm import Mamba2

model = Mamba2(
    d_model=256,
    d_state=128,     # Larger state dimension
    d_conv=4,
    expand=2,
    headdim=64,      # Head dimension for multi-head
    ngroups=1        # Parallel groups
).to("cuda")

主な違い：

• 状態サイズ：Mamba-1（d_state=16）と Mamba-2（d_state=128）
• アーキテクチャ：Mamba-2 はマルチヘッド構造を採用しています
• 正規化：Mamba-2 は RMSNorm を使用します
• 分散処理：Mamba-2 はテンソル並列処理をサポートします

ワークフロー 4：Transformers とのベンチマーク比較

生成速度の比較：

# Benchmark Mamba
python benchmarks/benchmark_generation_mamba_simple.py \
  --model-name "state-spaces/mamba-2.8b" \
  --prompt "The future of machine learning is" \
  --topp 0.9 --temperature 0.7 --repetition-penalty 1.2

# Benchmark Transformer
python benchmarks/benchmark_generation_mamba_simple.py \
  --model-name "EleutherAI/pythia-2.8b" \
  --prompt "The future of machine learning is" \
  --topp 0.9 --temperature 0.7 --repetition-penalty 1.2

期待される結果：

• Mamba：推論が 5 倍高速
• メモリ：KV キャッシュが不要
• スケーリング：シーケンス長に対して線形

代替手段と比べて Mamba を使う場合

Mamba を使うケース：

• 長いシーケンス（100K 以上のトークン）が必要な場合
• Transformers より高速な推論を求める場合
• メモリに制約がある場合（KV キャッシュなし）
• ストリーミングアプリケーションを構築する場合
• 線形スケーリングが重要な場合

利点：

• O(n) の計算量：二次ではなく線形
• 推論が 5 倍高速：アテンションのオーバーヘッドなし
• KV キャッシュ不要：メモリ使用量を削減
• 100 万トークンのシーケンス：ハードウェア効率が高い
• ストリーミング：トークンごとのメモリ使用量が一定

代わりに別の手段を使うケース：

• Transformers：最高水準の性能が必要で、計算リソースに余裕がある場合
• RWKV：RNN と Transformer のハイブリッドを求める場合
• RetNet：保持ベースのアーキテクチャが必要な場合
• Hyena：畳み込みベースのアプローチを求める場合

よくある問題

問題：CUDA のメモリ不足

バッチサイズを小さくするか、勾配チェックポイントを使用してください：

model = MambaLMHeadModel(config, device="cuda", dtype=torch.float16)
model.gradient_checkpointing_enable()  # Enable checkpointing

問題：インストールが遅い

バイナリホイールをインストールしてください（ソースからではなく）：

pip install mamba-ssm --no-build-isolation

問題：causal-conv1d が見つからない

個別にインストールしてください：

pip install causal-conv1d>=1.4.0

問題：HuggingFace からモデルを読み込めない

AutoModel ではなく、MambaLMHeadModel.from_pretrained を使用してください：

from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
model = MambaLMHeadModel.from_pretrained("state-spaces/mamba-2.8b")

発展的なトピック

選択的 SSM：数式による定式化、状態空間方程式、および選択性によって O(n) の計算量が実現される仕組みについては、references/selective-ssm.md を参照してください。

Mamba-2 のアーキテクチャ：マルチヘッド構造、テンソル並列処理、および分散学習の設定については、references/mamba2-details.md を参照してください。

性能の最適化：ハードウェアを考慮した設計、CUDA カーネル、メモリ効率の高い手法については、references/performance.md を参照してください。

ハードウェア要件

• GPU：CUDA 11.6 以降を搭載した NVIDIA GPU
• VRAM：
  • 130M モデル：2GB
  • 370M モデル：4GB
  • 790M モデル：8GB
  • 1.4B モデル：14GB
  • 2.8B モデル：28GB（FP16）
• 推論：Transformers より 5 倍高速
• メモリ：KV キャッシュなし（Transformers より少ない）

性能（Transformers との比較）：

• 速度：推論が 5 倍高速
• メモリ：50％削減（KV キャッシュなし）
• スケーリング：二次ではなく線形

リソース

• 論文（Mamba-1）：https://arxiv.org/abs/2312.00752（2023 年 12 月）
• 論文（Mamba-2）：https://arxiv.org/abs/2405.21060（2024 年 5 月）
• GitHub：https://github.com/state-spaces/mamba ⭐ 13,000+
• モデル：https://huggingface.co/state-spaces
• ドキュメント：リポジトリの README と Wiki

エージェントの作業手順

1. 環境を確認します。 GPU の種類、メモリ、ドライバー／CUDA のバージョン（nvidia-smi）、フレームワークのバージョン、データセットの場所とサイズを確認します。
2. 入力を特定します。 データまたはユーザーから、形式、識別子、パラメーターを確認します。結果に影響する値は推測せず、質問します。
3. まず小規模なバージョンを実行します。 スモークテストとして、小さなモデルまたはサブセットを少ないステップ数で実行し、損失が減少することとチェックポイントが保存されることを確認します。
4. 上記の手順と参考資料に従って、タスク全体を実行します。
5. 結果を検証します。 保留データで指標を追跡し、ベースラインと比較して、シード、設定、ハードウェアを記録します。
6. 報告します。 実行した内容（バージョン、コマンド、パラメーター）、確認した内容、未確定の点を明記します。

発生した問題	対応方法
CUDA のメモリ不足	バッチサイズを小さくする、勾配累積／チェックポイントまたは混合精度を有効にする、あるいはモデルを分割します。
損失が NaN になる、または発散する	学習率を下げ、データに無効な値がないか確認し、勾配クリッピングを有効にします。
インストール済みバージョンに、この手順で指定された関数、フラグ、エンドポイントがない	インストール済みバージョンのドキュメント（help()、--help、公式ドキュメント）を確認して適応し、ユーザーに伝えます。API をでっち上げてはいけません。
必須の入力、識別子、パラメーターが曖昧	ユーザーに確認するか、実行前に仮定を明示します。

誠実性に関するルール

• 結果、パラメーター、識別子、引用、統計情報を捏造してはいけません。実行または検証できない場合は、その旨を明確に伝えます。
• ログなしに学習結果を主張してはいけません。大規模なジョブを開始する前に、計算コストを見積もります。
• ここに記載されているバージョン固有の情報は古くなっている可能性があるものとして扱い、インストール済みバージョンの公式ドキュメントで確認します。
• 費用が発生する操作、共有 GPU やクラウドのクォータを消費する操作、個人データまたは患者データに触れる操作、取り消せない操作を行う前に、確認を求めます。

関連スキル

• rwkv-architecture：RWKV（Receptance Weighted Key Value）アーキテクチャについて説明します。RWKV は、O(n) の推論と KV キャッシュ不要を特徴とする RNN／Transformer ハイブリッドで、RWKV…
• optimizing-attention-flash：PyTorch ネイティブの scaled_dot_product_attention（PyTorch 2.2 以降）または flash-attn ライブラリを使用して、Transformer モデルで Flash Attention を有効にします…