让状态空间模型按模态智能分配计算资源，提速降耗且性能不降。
Mixture-of-Mamba: Enhancing Multi-Modal State-Space Models with Modality-Aware Sparsity
- 通过模态专属参数化实现稀疏计算，保留状态空间模型高效性。
- 在多模态预训练中减少35%~65%的算力消耗，保持相同损失水平。
- 适合追求高效多模态模型的开发者，尤其关注推理成本优化者。
状态空间模型（SSMs）作为序列建模的高效替代方案，但在多模态预训练中因无法充分利用模态特异性特征而受限。本文提出Mixture-of-Mamba，通过在Mamba块中引入模态感知稀疏性，实现模态专属参数化。基于Mixture-of-Transformers的思路，将模态感知稀疏性扩展至SSMs，同时保持其计算效率。我们在三种多模态预训练设置下评估：Transfusion（文本与连续图像令牌交错，扩散损失）、Chameleon（文本与离散图像令牌交错），以及包含语音的三模态扩展框架。Mixture-of-Mamba在更早训练步数达到相同损失值，且显著降低计算开销。在1.4B规模下，Transfusion中仅需34.76%的训练FLOPs即可达同等图像损失；Chameleon中图像损失仅需42.50%的FLOPs，文本损失仅需65.40%的FLOPs；三模态场景下，语音损失仅需24.80%的FLOPs。消融实验表明，投影组件解耦具有协同增益，联合解耦优于单一修改。结果证明模态感知稀疏性是跨模型通用且有效的设计原则，从Transformer延伸至SSMs，树立多模态预训练新基准。代码见：https://github.com/Weixin-Liang/Mixture-of-Mamba。
原文摘要 · Abstract (English)
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
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。