MoE-Mamba: Efficient Selective State Space Models with Mixture of Experts

TL;DR

MoE-Mamba combines State Space Models with Mixture of Experts to significantly improve sequential modeling efficiency and performance, achieving comparable results to Mamba with fewer training steps and outperforming Transformer-based models.

Contribution

This paper introduces MoE-Mamba, a novel integration of SSMs and MoE that enhances scaling and efficiency in sequential modeling tasks.

Findings

MoE-Mamba outperforms Mamba and Transformer-MoE in benchmarks.

Achieves same performance as Mamba in 2.35x fewer training steps.

Maintains inference performance gains of Mamba over Transformer.

Abstract

State Space Models (SSMs) have become serious contenders in the field of sequential modeling, challenging the dominance of Transformers. At the same time, Mixture of Experts (MoE) has significantly improved Transformer-based Large Language Models, including recent state-of-the-art open models. We propose that to unlock the potential of SSMs for scaling, they should be combined with MoE. We showcase this on Mamba, a recent SSM-based model that achieves remarkable performance. Our model, MoE-Mamba, outperforms both Mamba and baseline Transformer-MoE. In particular, MoE-Mamba reaches the same performance as Mamba in $2.35\times$ fewer training steps while preserving the inference performance gains of Mamba against Transformer.