Mixture-of-Mamba: Enhancing Multi-Modal State-Space Models with Modality-Aware Sparsity

要約

状態空間モデル（SSM）は、シーケンシャルモデリングのための変圧器の効率的な代替品として浮上していますが、モダリティ固有の機能を活用できないと、マルチモーダルな事前トレーニングでのパフォーマンスが制限されます。
ここでは、MAMBAブロックのモダリティ固有のパラメーター化を通じてモダリティに対応するスパースを導入する新しいSSMアーキテクチャであるMambaの混合を提案します。
混合物の変換者（W. Liang etal。Arxiv：2411.04996; 2024）に基づいて、計算効率を維持しながら、モダリティを意識したスパースの利点をSSMに拡張します。
3つのマルチモーダルな事前トレーニング設定にわたるマンバの混合物を評価します：輸血（拡散損失を伴うインターリーブテキストと連続画像トークン）、カメレオン（インターリーブテキストと離散画像トークン）、およびスピーチを含む3つのモダリティフレームワーク。
MAMBAの混合は、早期のトレーニングステップで一貫して同じ損失値に達し、計算コストが大幅に削減されました。
輸血の設定では、混合物は、1.4Bスケールでトレーニングフロップの34.76％のみを使用して、同等の画像損失を達成します。
カメレオンの設定では、混合物は、1.4Bスケールでフロップのわずか42.50％で同様の画像損失に達し、フロップのわずか65.40％で同様のテキスト損失に達します。
3つのモダリティの設定では、Momは1.4Bスケールでフロップの24.80％で音声喪失を一致させます。
私たちのアブレーション研究は、共同分離が個々の修正よりも大きな利益をもたらすという分離投影成分の相乗効果を強調しています。
これらの結果は、モダリティを認識したスパースを多用途で効果的な設計原則として確立し、トランスからSSMへの影響を拡大し、マルチモーダルな事前トレーニングで新しいベンチマークを設定します。
私たちのコードは、https：//github.com/weixin-liang/mixture-of-mambaでアクセスできます

要約(オリジナル)

State Space Models (SSMs) have emerged as efficient alternatives to Transformers for sequential modeling, but their inability to leverage modality-specific features limits their performance in multi-modal pretraining. Here, we propose Mixture-of-Mamba, a novel SSM architecture that introduces modality-aware sparsity through modality-specific parameterization of the Mamba block. Building on Mixture-of-Transformers (W. Liang et al. arXiv:2411.04996; 2024), we extend the benefits of modality-aware sparsity to SSMs while preserving their computational efficiency. We evaluate Mixture-of-Mamba across three multi-modal pretraining settings: Transfusion (interleaved text and continuous image tokens with diffusion loss), Chameleon (interleaved text and discrete image tokens), and an extended three-modality framework incorporating speech. Mixture-of-Mamba consistently reaches the same loss values at earlier training steps with significantly reduced computational costs. In the Transfusion setting, Mixture-of-Mamba achieves equivalent image loss using only 34.76% of the training FLOPs at the 1.4B scale. In the Chameleon setting, Mixture-of-Mamba reaches similar image loss with just 42.50% of the FLOPs at the 1.4B scale, and similar text loss with just 65.40% of the FLOPs. In the three-modality setting, MoM matches speech loss at 24.80% of the FLOPs at the 1.4B scale. Our ablation study highlights the synergistic effects of decoupling projection components, where joint decoupling yields greater gains than individual modifications. These results establish modality-aware sparsity as a versatile and effective design principle, extending its impact from Transformers to SSMs and setting new benchmarks in multi-modal pretraining. Our code can be accessed at https://github.com/Weixin-Liang/Mixture-of-Mamba

arxiv情報

提供元, 利用サービス

arxiv.jp, Google