MoE-Mamba: Efficient Selective State Space Models with Mixture of Experts (2401.04081v2)

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 LLMs, 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.

Introduction to MoE-Mamba

In the domain of machine learning, new models are continually introduced, enhancing both performance and efficiency. A compelling advancement has been made through the integration of Mixture of Experts (MoE) with State Space Models (SSMs), particularly in creating MoE-Mamba. This combination addresses key challenges in scaling SSMs, offering a method to retain the benefits of Mamba—a notable SSM—and MoE, resulting in significant gains in training efficiency without compromising on inference performance.

Breaking Down the Concepts

To understand MoE-Mamba's contributions, it's essential to grasp the foundational concepts. SSMs are a sequence modeling paradigm, offering a fusion of recurrent neural networks (RNNs) and convolutional neural networks (CNNs). They have made rapid progress, allowing for the scaling of deep SSMs with maintained computational efficiency and strong performance. Mamba, an SSM-based model, offers benefits like linear-time inference and parallelizable training, making it a robust alternative to attention-based Transformers.

Conversely, MoE has revolutionized the scaling of Transformers by enabling a dramatic increase in model parameters without proportionate increases in FLOPs required for training and inference. This is achieved through sparse activations, where only a subset of the model's parameters is engaged for each token processed. MoE thus provides a pathway to handle vastly large-scale models.

Architectural Insights

MoE-Mamba's architecture extends the existing Mamba framework by replacing every other Mamba layer with a MoE feed-forward layer. This arrangement benefits from interleaving unconditional processing (by the Mamba layer, integrating the entire sequence context) and conditional processing (by the MoE layer, applying the most pertinent expert for each token). This hybrid design ensures processing each token effectively, making the most of Mamba's strength to manage long contexts and the efficiency of MoE's selective activation.

Evaluating MoE-Mamba's Performance

Upon evaluation and testing, MoE-Mamba has demonstrated an ability to achieve equivalent performance levels to the original Mamba with 2.2x fewer training steps. Not only does this model advance over Transformer-MoE structures, but it also scales effectively with the increase in the number of experts. The empirical results place MoE-Mamba as a particularly promising model for efficient and powerful sequence modeling, potentially catalyzing further developments in machine learning architectures.

In conclusion, MoE-Mamba represents a notable development in sequence modeling, indicating a future where the agility of SSMs can be enhanced greatly. This model's efficiencies suggest an exciting direction for the advancements of LLMs, offering a glimpse into the next generation of machine learning innovation.