mamba-architecture

Mamba - Selective State Space Models

クイックスタート

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 を使うべき場合:

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

利点:

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

代わりに使うべきもの:

• Transformers: 最高水準の性能が必要で、計算資源がある場合
• RWKV: RNN と Transformer のハイブリッドが欲しい場合
• RetNet: retention ベースのアーキテクチャが必要な場合
• Hyena: 畳み込みベースのアプローチが欲しい場合

よくある問題

問題: CUDA out of memory

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

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")

発展的なトピック

Selective 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