2024/01/24 11:00

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

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

状态空间模型（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 数量是固定的。

Fedus et al. 在 2022 年的论文《Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity》中提出的 Switch 是一种 Token Choice 架构，其是将每个 token 路由到单个专家（K=1），而他们使用该方法将 Transformer 的参数规模成功扩增至了 1.6 万亿。波兰的这个团队在实验中也采用了这种 MoE 设计。

最近，MoE 也开始进入开源社区，比如 OpenMoE。

项目地址：https://github.com/XueFuzhao/OpenMoE

尤其值得一提的是 Mistral 开源的 Mixtral 8×7B，其性能可比肩 LLaMa 2 70B，同时所需的推理计算预算只有后者的约六分之一。

模型架构

尽管 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 给出了超参数的完整列表。

结果

表 1 给出了训练结果。MoE-Mamba 的表现显著优于普通 Mamba 模型。

值得注意的是，MoE-Mamba 只用仅仅 46% 的训练步骤就达到了与普通 Mamba 同等的结果水平。由于学习率是针对普通 Mamba 进行调整的，因此可以预计，如果再针对 MoE-Mamba 对训练流程进行一番优化，MoE-Mamba 的表现还会更好。

消融研究

为了评估 Mamba 是否能随专家数量的增长而很好地扩展，研究者比较了使用不同数量专家的模型。

图 3 展示了使用不同数量的专家时的训练运行步骤情况。

表 2 给出了 100k 步骤后的结果。

这些结果表明新提出的方法能随专家数量而很好地扩展。如果专家的数量为 8 或更多，新模型的最终性能优于普通 Mamba。由于 Mamba-MLP 比普通 Mamba 差，可以预见使用少量专家的 MoE-Mamba 的性能表现会比 Mamba 差。当专家数为 32 时，新方法得到了最佳结果。

产业MoE-MambaTransformer状态空间模型

相关数据

感知技术

知觉或感知是外界刺激作用于感官时，脑对外界的整体的看法和理解，为我们对外界的感官信息进行组织和解释。在认知科学中，也可看作一组程序，包括获取信息、理解信息、筛选信息、组织信息。与感觉不同，知觉反映的是由对象的各样属性及关系构成的整体。

来源：维基百科

参数技术

在数学和统计学裡，参数（英语：parameter）是使用通用变量来建立函数和变量之间关系（当这种关系很难用方程来阐述时）的一个数量。

来源：维基百科

学习率技术

在使用不同优化器（例如随机梯度下降，Adam）神经网络相关训练中，学习速率作为一个超参数控制了权重更新的幅度，以及训练的速度和精度。学习速率太大容易导致目标（代价）函数波动较大从而难以找到最优，而弱学习速率设置太小，则会导致收敛过慢耗时太长

来源：Liu, T. Y. (2009). Learning to rank for information retrieval. Foundations and Trends® in Information Retrieval, 3(3), 225-331. Wikipedia

超参数技术

在机器学习中，超参数是在学习过程开始之前设置其值的参数。相反，其他参数的值是通过训练得出的。不同的模型训练算法需要不同的超参数，一些简单的算法（如普通最小二乘回归）不需要。给定这些超参数，训练算法从数据中学习参数。相同种类的机器学习模型可能需要不同的超参数来适应不同的数据模式，并且必须对其进行调整以便模型能够最优地解决机器学习问题。在实际应用中一般需要对超参数进行优化，以找到一个超参数元组（tuple），由这些超参数元组形成一个最优化模型，该模型可以将在给定的独立数据上预定义的损失函数最小化。

来源：Wikipedia

注意力机制技术

我们可以粗略地把神经注意机制类比成一个可以专注于输入内容的某一子集（或特征）的神经网络. 注意力机制最早是由 DeepMind 为图像分类提出的，这让「神经网络在执行预测任务时可以更多关注输入中的相关部分，更少关注不相关的部分」。当解码器生成一个用于构成目标句子的词时，源句子中仅有少部分是相关的；因此，可以应用一个基于内容的注意力机制来根据源句子动态地生成一个（加权的）语境向量（context vector）, 然后网络会根据这个语境向量而不是某个固定长度的向量来预测词。

来源：机器之心

混合专家系统技术

混合专家系统是指一种机器学习技术，其中多名专家（学习器）被用来将问题空间分为均匀的区域。计算机视觉领域的一个例子是将用于人类检测的神经网络模型与用于姿态估计的神经网络模型相结合。

来源：Goker Erdogan Wikipedia

控制论技术

控制论是一门跨学科研究，它用于研究控制系统的结构，局限和发展。在21世纪，控制论的定义变得更加宽泛，主要用于指代“对任何使用科学技术的系统的控制”。由于这一定义过于宽泛，许多相关人士不再使用“控制论”一词。控制论与对系统的研究有关，如自动化系统、物理系统、生物系统、认知系统、以及社会系统等等。

来源：维基百科

堆叠技术

堆叠泛化是一种用于最小化一个或多个泛化器的泛化误差率的方法。它通过推导泛化器相对于所提供的学习集的偏差来发挥其作用。这个推导的过程包括：在第二层中将第一层的原始泛化器对部分学习集的猜测进行泛化，以及尝试对学习集的剩余部分进行猜测，并且输出正确的结果。当与多个泛化器一起使用时，堆叠泛化可以被看作是一个交叉验证的复杂版本，利用比交叉验证更为复杂的策略来组合各个泛化器。当与单个泛化器一起使用时，堆叠泛化是一种用于估计（然后纠正）泛化器的错误的方法，该泛化器已经在特定学习集上进行了训练并被询问了特定问题。

来源：Wolpert, D. H. (1992). Stacked generalization. Neural networks, 5(2), 241-259

机器之心机构

机器之心，成立于2014年，是国内最具影响力、最专业、唯一用于国际品牌的人工智能信息服务与产业服务平台。目前机器之心已经建立起涵盖媒体、数据、活动、研究及咨询、线下物理空间于一体的业务体系，为各类人工智能从业者提供综合信息服务和产业服务。

https://www.jiqizhixin.com/