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

文章来源：机器之心

性能与 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 也许可以扩展到数百亿参数！

相关研究

状态空间模型

状态空间模型（SSM）是一类用于序列建模的架构。这些模型的思想源自控制论领域，可被看作是 RNN 和 CNN 的组合。尽管它们具有相当大的优势，但也有一些问题，因此难以成为语言建模任务的主导架构。但是，近期的一些研究突破却让深度 SSM 可以扩展到数十亿参数，同时还能维持计算效率和强大的性能表现。

Mamba

Mamba 是基于 SSM 构建的模型，能实现线性时间的推理速度（对上下文长度而言），并且其还通过硬件感知型设计实现了高效的训练流程。Mamba 采用了一种工作高效型的并行扫描方法，可以减轻循环的序列性的影响，而融合 GPU 操作则可无需实现扩展状态。反向传播所必需的中间状态不会被保存下来，而是会在反向通过过程中被重新计算，由此可以降低内存需求。Mamba 优于注意力机制的优势在推理阶段尤其显著，因为其不仅能降低计算复杂度，而且内存使用量还不会取决于上下文长度。

Mamba 能解决序列模型的效率和效果之间的根本性权衡，这就凸显了状态压缩的重要性。高效的模型必需要小状态，而有效的模型所需的状态应当包含上下文的所有关键信息。不同于其它 SSM 对时间和输入不变性的需求，Mamba 引入了一种选择机制，可以控制信息沿序列维度传播的方式。这一设计选择的灵感来自对选择性复制和归纳头等合成任务的直观理解，让模型可以分辨和保留关键信息，同时滤除无关信息。

研究发现，Mamba 有能力高效地利用更长的上下文（长达 1M token），并且随着上下文长度增长，预训练困惑度也会得到改善。Mamba 模型是由堆叠的 Mamba 块构成的，在 NLP、基因组学、音频等多个不同领域都取得了非常好的结果，其性能可以媲美和超越已有的 Transformer 模型。因此，Mamba 成为了通用序列建模骨干模型的一个有力候选模型，参阅《五倍吞吐量，性能全面包围 Transformer：新架构 Mamba 引爆 AI 圈》。

混合专家

混合专家（MoE）这类技术能极大提升模型的参数数量，同时不会影响模型推理和训练所需的 FLOPs。MoE 最早由 Jacobs et al. 于 1991 年提出，并在 2017 年由 Shazeer et al. 开始用于 NLP 任务。

MoE 有一个优势：激活很稀疏 —— 对于处理的每个 token，只会用到模型的一小部分参数。由于其计算需求，Transformer 中的前向层已经变成了多种 MoE 技术的标准目标。

研究社区已经提出了多种方法用于解决 MoE 的核心问题，即将 token 分配给专家的过程，也称路由（routing）过程。目前有两种基本的路由算法：Token Choice 和 Expert Choice。其中前者是将每个 token 路由到一定数量（K）的专家，至于后者则是路由到每个专家的 token 数量是固定的。