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

作者：Maciej Pioro´, Kamil Ciebiera, Krystian Krol´, Jan Ludziejewski, Michał Krutul, Jakub Krajewski, Szymon Antoniak, Piotr Miłos´, Marek Cygan, Sebastian Jaszczur
机构：1 IDEAS NCBR, Warsaw, Poland; 2 University of Warsaw, Poland; 3 NOMADS, University of Warsaw, Poland; 4 Polish Academy of Sciences, Poland

A1 主要贡献

本文旨在解决现有大型语言模型（LLM）架构的局限性，并探索进一步扩展语言模型的可能性。当前LLM主要依赖于Transformer架构【65, Vaswani, A. et al. Attention is all you need. 2017】，但其二次方复杂度的注意力机制在处理长序列时面临挑战。状态空间模型（SSM）作为一种有前途的替代方案，具有线性时间推理和高效并行训练的优势，其中Mamba模型【12, Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces, 2023】因其选择性机制和硬件感知设计而表现出色。另一方面，专家混合（MoE）技术【17, Jacobs, R. A. et al. Adaptive mixtures of local experts. 1991】通过稀疏激活模型参数，显著提高了模型的可扩展性，已成功应用于最先进的LLM中【18, Jiang, A. Q. et al. Mixtral of experts, 2024】。

本文的核心论点是，为了充分释放SSM的扩展潜力，应将其与MoE技术相结合。基于此，本文提出了MoE-Mamba，一个融合了Mamba和Switch MoE层【10, Fedus, W. et al. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2022】的新模型。该模型旨在结合SSM和MoE两者的效率优势，从而在性能上超越单独的Mamba和基于Transformer的MoE基线模型。

主要贡献如下：
* 提出MoE-Mamba模型：该模型将Mamba与MoE层相结合，同时利用了SSM和MoE的效率优势。实验证明，MoE-Mamba达到与Mamba相同性能所需的训练步数减少了2.35倍。
* 鲁棒性验证：通过全面的研究证实，MoE-Mamba带来的性能提升在不同模型大小、设计选择和专家数量的情况下是稳健的。
* 探索多种集成方法：探索并比较了在Mamba块内部集成MoE的多种替代方法，为未来研究提供了有价值的参考。

A3 背景知识

3.1. 预备知识

Mamba
Mamba【12, Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces, 2023】是最近提出的一种基于SSM的模型，它实现了与Transformer相当的卓越性能。通过采用一种高效工作的并行扫描算法，Mamba减轻了循环（recurrence）的顺序性带来的影响，而融合GPU操作则避免了实例化扩展状态的需求。用于反向传播的中间状态不会被保存，而是在反向传播过程中重新计算，从而降低了内存需求。与注意力机制相比，Mamba的优势在推理过程中尤其突出，因为它不仅降低了计算复杂度，而且内存使用量不依赖于上下文长度。图3展示了Mamba层的内部结构。

MoE层
在我们的工作中，我们遵循了成熟且易于实现的Switch Transformer MoE设计【10, Fedus, W. et al. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2022】，如【70, Zhao, W. X. et al. A survey of large language models, 2023a】和【40, Sanseviero, O. et al. Mixture of experts explained, 2023】所述，并将其他MoE设计的考虑留给未来的工作。我们假设有$N_{experts}$个专家$\{E_i\}_{i=1}^{N_{experts}}$，每个专家都是一个可训练的前馈网络，参数数量相同。对于给定的输入$x \in R^{d_{model}}$，一个可训练的门控网络$G(x) \in R^{N_{experts}}$计算每个专家的亲和度分数$s_i(x) = x \cdot G_i$，其中$G_i \in R^{d_{model}}$是门控网络的第$i$列。这些分数通过softmax进行归一化：

在Switch Transformer之前，为每个令牌选择$k > 1$个最合适的专家的top-k路由被认为是必要的。然而，Switch通过设置$k = 1$成功地简化了之前的MoE方法。具体来说，MoE层对输入$x$的输出由以下公式给出：

其中 $I = \text{argmax}_i p_i(x)$。在批处理执行期间（例如训练中），每个批次包含$N$个令牌。按照标准程序，如果令牌到专家的分配不完美，即某个专家$E_f$被当前批次中超过$N/N_{experts}$个令牌选中，那么多余的令牌将被丢弃且不进行更新（容量因子=1）。为了进一步鼓励令牌均匀分布到各个专家，我们将【10, Fedus, W. et al. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2022】中描述的负载均衡损失（权重$\alpha = 0.01$）添加到训练目标中。

A2 方法细节

3.2. MoE-Mamba架构

标准的Mamba架构由多个Mamba块堆叠而成，每一层的输出被加到残差流中。这在图2中有所展示。在MoE-Mamba中，我们将Mamba层与MoE层交错排列（见图2）。值得注意的是，标准的Mamba模型不使用前馈网络层。

MoE-Mamba通过这种方式分离了无条件处理和有条件处理。Mamba层对每个令牌进行无条件处理，能够有效地将整个序列上下文整合到一个内部表示中。而MoE层则进行有条件处理，可以为每个令牌应用最相关的专家（也就是参数的子集）。这种将有条件处理和无条件处理交错使用的思想在一些基于MoE的模型中也被采用，通常是通过交替使用标准的和MoE的前馈网络层【24, Lepikhin, D. et al. Gshard: Scaling giant models with conditional computation and automatic sharding, 2020; 10, Fedus, W. et al. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2022】。

3.3. 并行MoE-Mamba

除了将MoE层与Mamba层交错排列外，我们还探索了另一种设计。该设计受到【66, Wang, B. Mesh-Transformer-JAX: Model-Parallel Implementation of Transformer Language Model with JAX. 2021】和【5, Chowdhery, A. et al. Palm: Scaling language modeling with pathways. 2023】的启发，其中MoE层与Mamba层并行执行（见图3）。这种设计取得了积极的结果，但效果不如交错式的MoE-Mamba。

3.4. 修改Mamba块

除了在Mamba外部附加一个独立的MoE层，我们还进行了其他实验，通过修改Gu & Dao【12, Mamba: Linear-time sequence modeling with selective state spaces, 2023】的原始块设计来引入有条件的MoE计算。其中一些设计相比基线架构有所改进，并为未来的研究指明了有希望的方向。

A4 实验环境

• 模型架构：
  • 比较了四种解码器专用（decoder-only）模型：MoE-Mamba、Mamba、Transformer和Transformer-MoE。
  • 模型规模分为约2500万（□25M）和约1亿（□100M）两个级别。
  • 为保证公平比较，MoE-Mamba中每个专家的维度被缩小（$d_{expert} = 3 \times d_{model}$），以使得所有相似规模模型的活跃参数数量大致相同。
  • Transformer模型使用旋转位置嵌入（Rotary Position Embedding）。
• 数据集：
  • 使用C4数据集【38, Raffel, C. et al. Exploring the limits of transfer learning with a unified text-to-text transformer. 2020】进行下一词元预测任务的训练。
  • 主实验中，□25M模型训练约100亿词元，□100M模型训练约300亿词元。其他消融实验训练10亿词元。
• 硬件配置：
  • 实验在多GPU环境下进行，具体硬件未详细说明，但提到了使用PL-Grid基础设施和Entropy集群。
• 软件配置：
  • 使用PyTorch【33, Paszke, A. et al. Pytorch: An imperative style, high-performance deep learning library, 2019】框架实现。
  • 利用FSDP（Fully Sharded Data Parallel）【71, Zhao, Y. et al. Pytorch fsdp: Experiences on scaling fully sharded data parallel, 2023b】进行多GPU训练。
  • 使用GPT2分词器【37, Radford, A. et al. Language models are unsupervised multitask learners. 2019】。
  • 优化器为AdamW【27, Loshchilov, I. and Hutter, F. Decoupled weight decay regularization, 2019】。

A4 实验结果

主要结果

MoE-Mamba在两个模型规模上都显著优于原始的Mamba模型。如表1和图1所示，MoE-Mamba100M模型仅用相当于Mamba100M处理词元量的 1/2.35 的训练量，就达到了与之相当的性能。对于□25M规模的模型，性能增益较低，这可能是由于训练词元数较少。此外，MoE-Mamba的性能也优于相应的Transformer-MoE模型，这进一步印证了Mamba是Transformer的一个有力竞争者【12, Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces, 2023】的观点。

Mamba与MoE中活跃参数的最佳比例

研究在总参数固定的情况下，Mamba层和MoE层中活跃参数的最佳比例。如图5和表2所示，增加Mamba层的活跃参数可以提升性能。然而，当比例达到 3:3 后，性能增益变得微乎其微，而更高的比例由于专家数量过多会导致硬件利用率低和路由成本高，因此不切实际。本文在其他实验中默认使用此比例。

替代设计方案

• 并行MoE-Mamba：实验了一种并行设计，其中MoE层和Mamba层并行放置。如图6所示，在所有测试设置中，顺序排列的MoE-Mamba性能都优于并行变体。并行MoE-Mamba需要2到4倍的专家数量和总参数才能达到顺序变体的性能水平。
• 内部MoE：尝试在Mamba块内部用MoE层替换三个线性投影层。如表3所示，有三种设计的结果略好于原始Mamba，但没有一种能超过MoE-Mamba。
  

  

  
  

  

专家数量的影响

如图4和表4所示，MoE-Mamba的性能随着专家数量的增加而提升。当专家数量 $N_{experts} \ge 4$ 时，MoE-Mamba的性能开始超过原始Mamba。实验中，使用32个专家时效果最好，并预计更多专家会带来进一步的增益。有趣的是，专家数量较少时（例如1个），模型性能反而比原始Mamba差，这与【12, Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces, 2023】的发现一致，即Mamba与前馈网络层交错（相当于单专家MoE）效果不如纯Mamba。

准确率与困惑度

在一个小型模型（MoE-Mamba25M，32个专家）的训练过程中观察到，尽管其困惑度低于Transformer-MoE基线，但准确率却持续低于后者（见附录C，图7）。作者推测，这可能是因为基于SSM的模型将历史信息压缩到有限的隐藏状态中，从而限制了它们逐字复制词元的能力，而这正是基于注意力的模型的优势所在。

A7 补充细节

5. 未来工作与局限性

扩展
在当前工作中，我们对活跃参数少于10亿、总参数最多24亿的模型进行了实验。鉴于MoE已使Transformer能够扩展到前所未有的规模【10, Fedus, W. et al. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2022】，我们期待看到扩展对我们提出的方法产生的影响。为此，发展相应的扩展定律将是至关重要的。

将MoE集成到Mamba层中
我们的实验表明，将Mamba层与高性能的稀疏MoE前馈层交错使用，可以得到一个有前途的模型。然而，在密集设置下，不带前馈层的Mamba表现略好。这表明，在Mamba层内部集成稀疏计算可能会产生更好的结果，同时保持简单、同质的架构。我们在4.4节中详细介绍的实验带来了一些乐观情绪，我们预计这一研究方向将保持其重要性。

探索MoE-Mamba中不同类型的MoE
虽然我们的设计基于常用的Switch【10, Fedus, W. et al. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2022】，但已有许多其他MoE架构被提出。这些设计不仅可能整体表现更好，而且与SSM结合时，可能有一种不同类型的MoE会是最佳选择。这方面的可能改变包括Expert-Choice路由器【72, Zhou, Y. et al. Mixture-of-experts with expert choice routing, 2022】、完全可微的架构【35, Puigcerver, J. et al. From sparse to soft mixtures of experts, 2023; 2, Antoniak, S. et al. Mixture of tokens: Efficient llms through cross-example aggregation, 2023】、变化的专家数量及其粒度【6, Clark, A. et al. Unified scaling laws for routed language models, 2022; 21, Krajewski, J. et al. Scaling laws for fine-grained mixture of experts, 2024】以及其他修改。

蒸馏
一些工作，例如【10, Fedus, W. et al. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2022】，已经表明MoE层可以被蒸馏回前馈层。我们预计MoE-Mamba也会有类似的结果。有趣的是，【12, Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces, 2023】的研究结果表明，Mamba模块可以很好地模拟前馈层。这就提出了一个问题：MoE是否可以被蒸馏成一个vanilla Mamba模块，以及如何实现这一点。

协同效应
我们将对Mamba和MoE的协同效应进行更深入的研究留给未来的工作。我们怀疑，由于更好的硬件利用率，效率增益可能会随着上下文长度的增加而增长；至于推理，Mamba缓解了由更长上下文大小引起的计算和内存吞吐量问题，而MoE则缓解了由参数数量增加和模型中存储的知识引起的相同问题。这种协同作用可能允许语言模型在参数数量和输入/输出长度上实现前所未有的扩展。

Mamba和注意力机制
Mamba和Transformer在数据处理过程中做出了不同的权衡。这导致了它们各自不同的优缺点，例如，Mamba可以处理非常长的输入，但在需要过去输入的详细知识的任务（例如，某些复制实例）上可能会遇到困难。探索将这两种架构结合起来以实现两全其美将是很有趣的。

长上下文利用
Mamba和其他SSM因其处理长上下文的能力而受到赞誉。然而，它们能有效利用长上下文的程度以及提高利用率的技术尚未得到深入研究。为此，一些为Transformer开发的方法【42, Shi, W. et al. In-context pretraining: Language modeling beyond document boundaries, 2023; 64, Tworkowski, S. et al. Focused transformer: Contrastive ´ training for context scaling, 2023; 44, Staniszewski, K. et al. Structured packing in llm ´ training improves long context utilization, 2024】可能适用。

其他模态
这项工作探索了Mamba可以扩展的一个方向。Mamba是一种通用架构，并不局限于语言建模。我们预计可以将MoE-Mamba应用于其他任务，如【12, Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces, 2023】中提出的非文本序列建模，以及不同的模态，如视觉，初步工作已由【73, Zhu, L. et al. Vision mamba: Efficient visual representation learning with bidirectional state space model, 2024】提出。

A5 结论

本文首次将专家混合（MoE）技术与Mamba架构相结合，提出了MoE-Mamba模型。这种新方法不仅继承了Mamba在推理上的优势，而且达到相同性能所需的训练步数减少了2.35倍。我们在参数量高达24亿、训练长度达300亿词元的模型上进行实验，证实了这种性能提升对于模型大小、训练时长和专家数量都具有鲁棒性。

此外，我们探索并评估了多种在Mamba块内部集成MoE的替代设计。虽然这些变体没有超越MoE-Mamba，但这些研究有助于筛选掉无效的研究方向，并指向更有前景的路径。

我们的工作开辟了将专家混合与状态空间模型相结合的新研究方向。我们相信，这条路径将使语言模型能够更有效地扩展到更大的规模。

A6 附录

A. 超参数和训练设置

基础模型超参数（dmodel, dff, 注意力头数量, 层数）的灵感来源于BERT【7, Devlin, J. et al. Bert: Pre-training of deep bidirectional transformers for language understanding, 2019; 63, Turc, I. et al. Well-read students learn better: On the importance of pre-training compact models, 2019】。其中，□25M模型等同于BERT-MEDIUM，□100M模型则复制了BERT-BASE的配置，并将块数从12增加到16。学习率调度、权重衰减和梯度裁剪值均按照社区标准实践设定。我们使用了AdamW优化器【27, Loshchilov, I. and Hutter, F. Decoupled weight decay regularization, 2019】。我们为每个□25M模型单独调整了最大学习率，并在训练□100M对应模型时将其减半。我们使用PyTorch【33, Paszke, A. et al. Pytorch: An imperative style, high-performance deep learning library, 2019】进行模型训练，并利用FSDP【71, Zhao, Y. et al. Pytorch fsdp: Experiences on scaling fully sharded data parallel, 2023b】来支持多GPU设置。

B. 活跃参数 vs FLOPs

在本文中，我们报告的是活跃参数的数量（不包括嵌入和去嵌入层），而不是浮点运算次数（FLOPs）。这一做法遵循了【72, Zhou, Y. et al. Mixture-of-experts with expert choice routing, 2022】。这两个数值大致成正比【20, Kaplan, J. et al. Scaling laws for neural language models, 2020】，但FLOPs的计算更困难，并且对于像Mamba这样具有优化的硬件感知架构而言，尤其是在推理过程中，其相关性较低。

C. 准确率和困惑度

正如4.6节所述，我们观察到一个有趣的现象，即两个性能相近但基于不同架构的模型在度量指标上存在不一致。我们推测，这种差异暗示了Mamba及其他SSM可能存在的一种失败模式。由于历史信息被压缩到一个有限的隐藏状态中，它们进行逐字复制令牌的能力受到了限制。与之相关的能力，即在给定前缀...[A][B]...[A]的情况下预测令牌[B]（其中[A]、[B]可以是任意令牌），已由【9, Elhage, N. et al. A mathematical framework for transformer circuits. 2021】进行了机理研究，并被推测是Transformer卓越的上下文学习能力的根源【29, Olsson, C. et al. In-context learning and induction heads. 2022】。

Peng等人【34, Peng, B. et al. Rwkv: Reinventing rnns for the transformer era, 2023】提到，他们的无注意力模型RWKV在需要回忆长上下文中精确信息的任务上性能可能有限，因为其隐藏状态大小是固定的，这是Mamba和其他SSM共有的特性。然而，由于Mamba的困惑度可以与同等规模的Transformer相媲美，我们可以推测Mamba通过其他方式弥补了这种失败模式，并且在与其他任务相比时可能显示出相对优势。特别地，它可能在零样本任务中胜过Transformer，而不是那些允许少样本演示或需要上下文学习的任务。

D. 加速比与训练时间的关系

在我们的实验中，我们注意到通常随着训练的继续，MoE-Mamba相比于vanilla Mamba的加速比会增加（见图8）。也就是说，比率

随着 $l$ 的减小而增加。□25M模型的加速比在1.6到1.9之间波动，而□100M模型的加速比则稳步上升。

E. 模型参数计数

对于所有模型及其变体，我们报告的是可训练的、非嵌入参数的数量，即我们排除了输入（嵌入）和输出（去嵌入）层中的参数。这个惯例由【20, Kaplan, J. et al. Scaling laws for neural language models, 2020】提出，他们指出仅使用非嵌入参数可以使他们的扩展定律形式更清晰。嵌入参数数量对最终性能的相对较低重要性已由【22, Lan, Z. et al. Albert: A lite bert for self-supervised learning of language representations, 2020】指出。

F. 探索Mamba与MoE活跃参数的最佳比例

在异构架构中，将FLOPs和参数分配给不同组件是一个重要的设计选择。例如，在Transformer中，模型的形状已被【20, Kaplan, J. et al. Scaling laws for neural language models, 2020】广泛研究。在我们的工作中，我们研究了Mamba层中的活跃参数与MoE层中的活跃参数数量的最佳比例，见4.3节。图5可能表明增加该比例会增强性能，或许将所有活跃参数分配给Mamba会得到最佳性能（比例“6:0”）。然而，需要注意的是，所有被研究的模型都包含相同数量的总参数和每个令牌的活跃参数。上面描述的假设模型无法实现这一属性。如果我们放宽要求并将所有参数都放在Mamba中，得到的模型与具有扩展因子E=4且只有8个而不是16个Mamba层的Mamba25M相同。该模型最终的对数困惑度（3.73）略差于Mamba25M。

G. 训练集和测试集性能

在正文中，我们报告的是在训练集上获得的损失值。我们的训练过程从数据集中采样，因此即使处理的令牌数没有超过C4数据集中的总数，也可能会多次遇到相同的文档。然而，由于我们处理的令牌数不到总数的20%，训练集和测试集上的性能差异应该是可以忽略的。为了透明起见，我们也提供了测试集上的结果（图9）。由于每次评估步骤中的序列数量有限，其方差可能较高。

H. 贡献

Maciej将Mamba集成到代码库中，进行了与这项工作各方面相关的实验，并监督了项目的进程。Kamil负责了大部分实验。Krystian在Jan的帮助下探索了Mamba块的替代设计。Michał和Jakub以多种方式为项目做出了贡献，主要是通过运行实验和完善代码库。Szymon为我们的共享仓库做出了贡献，并参与了一些关于集成MoE和SSM的早期讨论。Piotr和Marek提供了高层次的科学建议。Sebastian监督了整个项目，设定了研究方向并领导了实验和分析。

I. 可复现性

用于运行实验的代码库可在 https://github.com/llm-random/llm-random 获取。