MoE与Mamba强强联合，将状态空间模型扩展到数百亿参数

转载

311 天前

1911

机器之心

文章转载来源：机器之心

文章来源：机器之心

性能与 Mamba 一样，但所需训练步骤数却少 2.2 倍。

图片来源：由无界 AI生成

状态空间模型（SSM）是近来一种备受关注的 Transformer 替代技术，其优势是能在长上下文任务上实现线性时间的推理、并行化训练和强大的性能。而基于选择性 SSM 和硬件感知型设计的 Mamba 更是表现出色，成为了基于注意力的 Transformer 架构的一大有力替代架构。

近期也有一些研究者在探索将 SSM 和 Mamba 与其它方法组合起来创造更强大的架构，比如机器之心曾报告过《Mamba 可以替代 Transformer，但它们也能组合起来使用》。

近日，波兰一个研究团队发现，如果将 SSM 与混合专家系统（MoE/Mixture of Experts）组合起来，可望让 SSM 实现大规模扩展。MoE 是目前常用于扩展 Transformer 的技术，比如近期的 Mixtral 模型就使用了这一技术，参阅机器之心文章。

这个波兰研究团队给出的研究成果是 MoE-Mamba，即将 Mamba 和混合专家层组合起来的模型。

论文地址：https://arxiv.org/pdf/2401.04081.pdf

MoE-Mamba 能同时提升 SSM 和 MoE 的效率。而且该团队还发现，当专家的数量发生变化时，MoE-Mamba 的行为是可预测的。

该团队也进行了实验论证，如图 1 所示，结果表明：相比于 Mamba，MoE-Mamba 达到同等性能时所需的训练步骤数少 2.2 倍，这彰显了新方法相较于 Transformer 和 Transformer-MoE 的潜在优势。这些初步结果也指出了一个颇具潜力的研究方向：SSM 也许可以扩展到数百亿参数！

模型架构

尽管 Mamba 的主要底层机制与 Transformer 中使用的注意力机制大不相同，但 Mamba 保留了 Transformer 模型的高层级、基于模块的结构。使用这一范式，由相同模块构成的一层或多层会彼此堆叠在一起，而每一层的输出都会被添加到残差流（residual stream）中，见图 2。之后，这个残差流的最终值会被用于预测语言建模任务的下一个 token。

MoE-Mamba 利用了这两种架构的兼容能力。如图 2 所示，在 MoE-Mamba 中，每间隔一个 Mamba 层就会被替换成一个基于 Switch 的 MoE 前馈层。

不过该团队也注意到这一设计和《Mamba: Linear-time sequence modeling with selective state spaces》的设计有些相似；后者交替堆叠了 Mamba 层和前馈层，但得到的模型相比于单纯的 Mamba 还略有不及。该设计在图 1 中被记为 Mamba-MLP。

MoE-Mamba 分开了 Mamba 层执行的每个 token 的无条件处理和 MoE 层执行的有条件处理；其中的无条件处理可高效地将序列的整个上下文整合到一个内部表征中，而有条件处理可为每个 token 使用最相关的专家。这种将有条件处理和无条件处理交替起来的思路在一些基于 MoE 的模型中已经得到了应用，不过它们通常是交替基本的和 MoE 的前馈层。

主要结果

训练设置

该团队比较了 5 种不同设置：基本 Transformer、Mamba、Mamba-MLP、MoE 和 MoE-Mamba。

在大多数 Transformers 中，前馈层包含 8dm² 个参数，而 Mamba 论文中则让 Mamba 更小一些（约 6dm²），这样两个 Mamba 层的参数数量与一个前馈层和一个注意力层加起来差不多。为了让 Mamba 和新模型中每个 token 的活动参数数量大致一样，该团队将每个专家前向层的大小缩小到了 6dm²。除了嵌入层和解除嵌入（unembedding）层，所有模型都是每个 token 使用大约 2600 万参数。训练过程使用了 65 亿个 token，训练步骤数为 100k。

训练使用的数据集是 English C4 数据集，任务是预测下一 token。文本的 token 化使用了 GPT2 tokenizer。表 3 给出了超参数的完整列表。