Repolode
投稿する
state-spaces

state-spaces/mamba

Mambaは、選択的状態空間モデルを基盤とするPyTorchのシーケンスモデリングアーキテクチャで、言語などのシーケンス処理においてTransformerに代わる線形時間の手法を提供します。

GitHub で見る
スター
18,939
フォーク
1,831
未解決の Issue
640
最終コミット
3 か月前

これは何？

Mambaは、選択的状態空間モデル（SSM）を基盤とするシーケンスモデリングアーキテクチャです。言語のような情報密度の高いデータを対象としており、これまでTransformerに後れを取っていた二次未満の計算量を持つモデルに対し、FlashAttentionに着想を得たハードウェアを考慮した実装を採用しています。

このリポジトリには、選択的SSMのコア演算子、Mamba、Mamba-2、Mamba-3の各ブロック、さらにブロックを積み重ねて言語モデリングヘッドを組み合わせた完全な言語モデルのラッパーが含まれています。インストール方法として、CUDA selective scanのビルドと純粋なPythonのコアパッケージの両方に対応しており、AMD ROCmのビルドについてはパッチに関する注記もあります。Hugging Faceでは、130Mから2.8Bパラメータまでの学習済みチェックポイントを公開しており、生成・ベンチマーク用スクリプトとlm-evaluation-harnessの実行手順も用意されています。

このコードは、SSMベースのシーケンスモデルを学習・実行したり、Transformerをベースラインとして比較したりしたい研究者やエンジニアを対象としています。

主な機能

  • 01ハードウェアを考慮したscan実装を備える選択的状態空間レイヤー
  • 02共通のPyTorchインターフェースを持つMamba、Mamba-2、Mamba-3のブロックモジュール
  • 03MambaブロックとLMヘッドを組み合わせた完全な言語モデルラッパー
  • 04Hugging Faceで公開されている130M〜2.8Bパラメータの学習済みチェックポイント
  • 05任意で利用できるCUDA selective scan拡張と、純粋なPythonのコアインストール
  • 06生成ベンチマーク用スクリプトとゼロショット評価の手順

ユースケース

  • Transformerに代わる手法として、長いシーケンスを扱う言語モデルを学習する
  • 状態空間モデルのアーキテクチャと効率上のトレードオフを研究する
  • 学習済みMambaチェックポイントを使ってテキスト生成を推論する
  • Transformerをベースラインとして、SSMによる生成のレイテンシとスループットをベンチマークする

クイックスタート

  1. 1

    先にPyTorchをインストールします

  2. 2

    mamba-ssmのコアパッケージをインストールします

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

使う前に知っておきたいこと

  • CUDAのselective scanを使用するには、LinuxとCUDA 11.6以降に対応したNVIDIA GPUが必要です。
  • 環境フラグで有効にしない限り、コアのインストールにはCUDAのselective_scan_cuda拡張機能は含まれません。
  • 既存のCUDA対応PyTorchが必要なため、セットアップが少し面倒な場合があります。

トピック

#state-space-model#sequence-modeling#pytorch#language-model#deep-learning#cuda#transformer-alternative

よくある質問

mamba とは？
Mambaは、選択的状態空間モデルを基盤とするPyTorchのシーケンスモデリングアーキテクチャで、言語などのシーケンス処理においてTransformerに代わる線形時間の手法を提供します。 Mambaは、選択的状態空間モデル（SSM）を基盤とするシーケンスモデリングアーキテクチャです。言語のような情報密度の高いデータを対象としており、これまでTransformerに後れを取っていた二次未満の計算量を持つモデルに対し、FlashAttentionに着想を得たハードウェアを考慮した実装を採用しています。
mamba は無料で使えますか？ライセンスは？
mamba は Apache-2.0 ライセンスのオープンソースです。商用製品で使う前にライセンス条件を確認してください。
mamba は今もメンテナンスされていますか？
半年以内に更新あり。mamba の最終コミットは 2026年7月22日 です。
mamba は何に使えますか？
Transformerに代わる手法として、長いシーケンスを扱う言語モデルを学習する 状態空間モデルのアーキテクチャと効率上のトレードオフを研究する 学習済みMambaチェックポイントを使ってテキスト生成を推論する Transformerをベースラインとして、SSMによる生成のレイテンシとスループットをベンチマークする
mamba のインストール方法は？
先にPyTorchをインストールします mamba-ssmのコアパッケージをインストールします: pip install mamba-ssm --no-build-isolation
mamba はどこで入手できますか？
ソースコードとインストール手順は GitHub にあります：https://github.com/state-spaces/mamba

概要はプロジェクトの README をもとに AI が生成しています。GitHub の数値は毎日更新されます。

関連プロジェクト