title: Mixture-of-Mamba: Enhancing Multi-Modal State-Space Models with
    Modality-Aware Sparsity
summary: 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
paper_authors._id: 67986cd6bdc99911a989b0a5
paper_authors.name: Weixin Liang
paper_authors._id: 67986cd6bdc99911a989b0a6
paper_authors.name: Junhong Shen
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2025-01-28T14:18:15.019Z
paper_authors._id: 6532e347b66f4bf689cf269a
paper_authors.avatarUrl: /avatars/76b5dddf80a24d3ef5c68b702280da82.svg
paper_authors.fullname: Junhong Shen
paper_authors.type: user
paper_authors.user: sjunhongs
paper_authors._id: 67986cd6bdc99911a989b0a7
paper_authors.name: Genghan Zhang
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2025-01-28T14:18:21.072Z
paper_authors._id: 65a76ff1e504d9738d636217
paper_authors.avatarUrl: /avatars/26bf5e3f19057835ee95d72c24904d77.svg
paper_authors.fullname: Genghan Zhang
paper_authors.type: user
paper_authors.user: Genghan
paper_authors._id: 67986cd6bdc99911a989b0a8
paper_authors.name: Ning Dong
paper_authors._id: 67986cd6bdc99911a989b0a9
paper_authors.name: Luke Zettlemoyer
paper_authors._id: 67986cd6bdc99911a989b0aa
paper_authors.name: Lili Yu
paper_title: Mixture-of-Mamba: Enhancing Multi-Modal State-Space Models with
    Modality-Aware Sparsity
paper_summary: 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
paper_ai_keywords: State Space Models (SSMs) Mixture-of-Mamba modality-aware sparsity Mamba block Mixture-of-Transformers Transfusion Chameleon diffusion loss interleaved text continuous image tokens discrete image tokens three-modality framework speech FLOPs decoupling projection components ablation study modality-specific parameterization