Mixture-of-Mamba: Enhancing Multi-Modal State-Space Models with
  Modality-Aware Sparsity
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
67986cd6bdc99911a989b0a5
Weixin Liang
67986cd6bdc99911a989b0a6
Junhong Shen
admin_assigned
2025-01-28T14:18:15.019Z
6532e347b66f4bf689cf269a
/avatars/76b5dddf80a24d3ef5c68b702280da82.svg
Junhong Shen
user
sjunhongs
67986cd6bdc99911a989b0a7
Genghan Zhang
admin_assigned
2025-01-28T14:18:21.072Z
65a76ff1e504d9738d636217
/avatars/26bf5e3f19057835ee95d72c24904d77.svg
Genghan Zhang
user
Genghan
67986cd6bdc99911a989b0a8
Ning Dong
67986cd6bdc99911a989b0a9
Luke Zettlemoyer
67986cd6bdc99911a989b0aa
Lili Yu
Mixture-of-Mamba: Enhancing Multi-Modal State-Space Models with
  Modality-Aware Sparsity
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
Mixture-of-Mamba extends modality-aware sparsity to State Space Models, improving multi-modal pretraining efficiency across text, image, and speech.
State Space Models Mixture-of-Mamba modality-specific parameterization Mamba block Mixture-of-Transformers modality-aware sparsity Transfusion Chameleon three-modality framework diffusion loss text tokens continuous image tokens discrete image tokens speech tokens FLOPs decoupling projection components