state-spaces/mamba

何を解決するか

Mambaは、情報密度の高いデータを処理する際のTransformerの限界に対処し、特に長文シーケンス処理における二次スケーリングの問題を克服します。言語モデルなどのタスクにおいても高いパフォーマンスを維持しつつ、二次未満の線形時間のシーケンスモデリングの代替手段を提供します。

動作原理

Mambaは選択的状態空間モデル（SSM）アーキテクチャに基づいています。ハードウェアに配慮した設計と効率的なCUDA実装（FlashAttentionの精神に近い）を採用し、パフォーマンスを最適化しています。このプロジェクトは以下の段階を経て進化してきました：

• Mamba-1：線形時間のシーケンスモデリングのための選択的状態空間を導入。
• Mamba-2：構造化状態空間双対性（SSD）を導入し、TransformerをSSMとして定式化することで、より効率的なアルゴリズムを構築。
• Mamba-3：状態空間の原理をさらに活用し、推論を最優先としたアプローチでシーケンスモデリングを改善。

対象ユーザー

大規模なシーケンスモデルや言語モデルを開発するAI研究者や開発者向けです。Transformerに典型的に見られるパフォーマンスを損なうことなく、線形時間スケーリングの効率性を必要とする方々に最適です。

主な特徴

• 線形時間スケーリング：Transformerの二次複雑性に対するより効率的な代替手段を提供。
• ハードウェアに配慮した設計：高性能GPU実行を可能にする最適化されたCUDA拡張（selective_scan_cuda）を含む。
• 事前学習済みモデル：PileおよびSlimPajamaデータセットで学習された、130M～2.8Bパラメータのさまざまな事前学習済みベースモデルを提供。
• 柔軟なインターフェース：選択的SSMレイヤーから完全なMambaブロック、および完全な言語モデルの例まで、複数レベルのAPIを提供。