[Paper Review] BlackMamba: Mixture of Experts for State-Space Models

Quentin Anthony, Yury Tokpanov+2 more|arXiv (Cornell University)|Feb 1, 2024

Simulation Techniques and Applications10 citations

TL;DR

BlackMamba combines Mamba state-space blocks with Mixture-of-Experts routing to achieve linear-time generation and reduced training/inference FLOPs, with open-source 340M/1.5B and 630M/2.8B models trained on 300B tokens.

ABSTRACT

State-space models (SSMs) have recently demonstrated competitive performance to transformers at large-scale language modeling benchmarks while achieving linear time and memory complexity as a function of sequence length. Mamba, a recently released SSM model, shows impressive performance in both language modeling and long sequence processing tasks. Simultaneously, mixture-of-expert (MoE) models have shown remarkable performance while significantly reducing the compute and latency costs of inference at the expense of a larger memory footprint. In this paper, we present BlackMamba, a novel architecture that combines the Mamba SSM with MoE to obtain the benefits of both. We demonstrate that BlackMamba performs competitively against both Mamba and transformer baselines, and outperforms in inference and training FLOPs. We fully train and open-source 340M/1.5B and 630M/2.8B BlackMamba models on 300B tokens of a custom dataset. We show that BlackMamba inherits and combines both of the benefits of SSM and MoE architectures, combining linear-complexity generation from SSM with cheap and fast inference from MoE. We release all weights, checkpoints, and inference code open-source. Inference code at: https://github.com/Zyphra/BlackMamba

Motivation & Objective

• Motivate combining state-space models (SSMs) with mixture-of-experts (MoE) to improve efficiency and scalability over dense transformers.
• Design and implement BlackMamba by substituting Mamba blocks for attention and integrating MoE routing for the MLP components.
• Empirically evaluate BlackMamba against Mamba and transformer baselines in terms of performance, training FLOPs, and inference efficiency.
• Demonstrate scalable training of large BlackMamba variants and release model weights and inference code for community use.

Proposed method

• Introduce BlackMamba as a stacked architecture of alternating attention-free Mamba blocks and routed MoE layers.
• Use input-dependent gating to make SSM matrices A, B, C in Eq. (4) depend on x(t).
• Apply a top-1 routing with a Sinkhorn-based optimizer to balance load across 8 MoE experts per layer.
• Train two model sizes (340M/1.5B and 630M/2.8B forward/total parameters) on a 1.8–2.0 trillion-token mixture dataset.
• Train with bf16 precision using Megatron-LM framework; activate 8 experts per MoE block; disable biases; SwiGLU in expert MLPs.
• Provide inference code and open-source checkpoints under Apache 2.0.

(c) Transformer-MoE

Experimental results

Research questions

• RQ1 Does combining Mamba-like state-space blocks with MoE routing yield competitive language modeling performance at scale while reducing training and inference FLOPs?
• RQ2 How do BlackMamba models compare to dense transformers and pure Mamba baselines in generation latency and long-context processing?
• RQ3 How does MoE routing behave in BlackMamba across depth and over training time, and what is the impact of the Sinkhorn routing initialization on convergence?
• RQ4 What is the data and parameter efficiency when training BlackMamba variants on a large, multi-dataset pretraining corpus?

Key findings

• BlackMamba achieves competitive evaluation performance with significantly lower training FLOPs and faster inference than dense transformers and pure Mamba baselines.
• Inference latency for BlackMamba, particularly at longer sequence lengths, is substantially faster than canonical transformers and comparable MoE transformers, due to the combination of linear-time SSM generation and sparse MoE routing.
• Most layers show well-balanced expert utilization with the Sinkhorn-based routing, though the final layers exhibit emerging specialization/imbalance patterns that warrant further study.
• Two open-source BlackMamba configurations (340M/1.5B and 630M/2.8B) were trained on 300B tokens and released with weights and inference code for community use.

(d) Mamba-MoE