MoE-Mamba: 通过混合专家实现高效选择状态空间模型.

0. TL;DR

本文介绍了MoE-Mamba，一种结合了Mamba状态空间模型和混合专家（MoE）技术的高效模型。MoE-Mamba在训练效率和性能上均优于原始的Mamba模型和Transformer-MoE模型。具体来说，MoE-Mamba在2.35倍更少的训练步数内达到了与Mamba相同的性能，同时保持了Mamba在推理性能上的优势。通过全面的研究，作者证实了MoE-Mamba的改进在不同的模型大小、设计选择和专家数量下都是稳健的。

1. 背景介绍

近年来，状态空间模型（SSMs）在序列建模领域逐渐崭露头角，成为Transformer架构的有力竞争者。SSMs的优势在于其线性时间推理、高度并行化的训练以及在长上下文处理任务中的强大性能。Mamba作为一种基于SSM的模型，通过其选择性机制和硬件感知设计，实现了令人印象深刻的性能，成为挑战注意力机制的Transformer架构的有力候选者。

与此同时，混合专家（MoE）技术在基于Transformer的大规模语言模型中取得了显著的改进。MoE通过稀疏激活，能够在增加模型参数的同时，对计算需求的影响最小化。这种特性使得MoE模型能够高效地扩展到数万亿参数。鉴于SSMs和MoE技术各自的优势，作者提出将两者结合，以释放SSMs在扩展方面的潜力。

2. MoE-Mamba 模型

（1）Mamba架构

Mamba是一种基于SSM的模型，通过采用高效的工作并行扫描，减轻了递归性质的影响。同时，通过融合GPU操作，Mamba消除了扩展状态的存储需求。在反向传播过程中，不保存中间状态，而是重新计算，从而降低了内存需求。Mamba的优势在推理过程中尤为明显，因为其计算复杂度较低，且内存使用不依赖于上下文长度。

Mamba模块如图所示。通过卷积映射和门控映射把输入特征维度扩大$E$倍，然后通过下列SSM计算：

\[\begin{aligned} \mathbf{x}^\prime(t) &= A \mathbf{x}(t) + B \mathbf{u}(t) \\ \mathbf{y}(t) &= C \mathbf{x}(t) \\ \end{aligned}\]

（2）MoE架构

作者采用了Switch Transformer MoE设计。MoE层由多个专家网络组成，每个专家网络是一个可训练的前馈网络。对于每个token嵌入，计算得分并使用softmax进行归一化。然后，选择得分最高的专家来处理该token。在批量执行中，如果某个专家被超过一定数量的token选择，则会丢弃多余的token。为了鼓励token在专家之间的均匀分布，添加了负载平衡损失。

（3）MoE-Mamba架构

MoE-Mamba通过在Mamba层之间交错插入MoE层，实现了无条件处理和条件处理的分离。Mamba层负责高效地将整个序列上下文集成到内部表示中，而MoE层则应用最相关的专家（即参数子集）来处理每个token。这种设计借鉴了一些基于MoE的模型，通过交替使用原始层和MoE层来实现条件和无条件处理的结合。

3. 实验分析

作者使用了C4数据集进行模型训练，任务是下一个token预测，使用交叉熵作为损失函数，使用了训练对数困惑度作为比较指标。

实验结果显示，MoE-Mamba在训练效率和性能上均优于原始的Mamba模型和Transformer-MoE模型。具体来说，MoE-Mamba在2.35倍更少的训练步数内达到了与Mamba相同的性能。此外，MoE-Mamba的表现优于对应的Transformer-MoE，这进一步证明了Mamba是Transformer的有竞争力的替代品。

实验结果表明，随着专家数量的增加，模型的性能单调提高。当专家数量达到32时，模型取得了最佳性能。

作者研究了Mamba层和MoE层之间活动参数的最佳比例。结果显示，增加Mamba层的活动参数可以提高性能，但在达到一定比例后，收益递减。

作者比较了不同的MoE-Mamba设计，包括交错设计和平行设计。结果显示，交错设计在所有测试设置中均优于平行设计。