Mixture of Experts (MoE)架构是一种先进的机器学习模型设计,特别适用于大语言模型(LLM)。在MoE架构中,整个模型被划分为多个专门的子网络(称为“专家Experts”),每个专家针对特定类型的数据或任务进行训练。通过一个门控网络,MoE能够动态选择和激活与输入数据最相关的专家,从而实现稀疏计算。这种方法使得在处理复杂任务时,模型能够显著减少计算成本,同时提高性能和效率。
模型规模是提升模型性能的关键因素之一。在有限的计算资源预算下,用更少的训练步数训练一个更大的模型,往往比用更多的步数训练一个较小的模型效果更佳。
Introduction
在大语言模型中,MoE架构的优势尤为明显。它允许模型在保持较大参数规模的同时,仅激活部分专家进行计算,从而降低了预训练和推理阶段的计算需求。例如,在某些实现中,MoE可以在一次前向传播中仅激活少数几个专家,这样可以在不牺牲性能的情况下,显著提高计算效率和响应速度。这使得MoE成为处理自然语言处理等领域中多样化数据和高计算需求任务的重要工具。
现有MOE模型汇总
模型 | 发布 | 规模 | 训练 | 备注 |
---|---|---|---|---|
Deepseek | DeepSeek于2024年12月10日发布并持续更新 | 16B (激活2.4B)、236B (激活22B) | 未披露 | 当前最优秀的MoE系列大模型之一 |
Qwen2.5-54B-A14B | Alibaba于2024年5月发布 | 54B (激活14B) | 包含大规模文本与代码数据 (包含中文和英文数据),采用MOE与密集层结合 | 针对聊天生成任务进行了优化,适合多种应用场景。(HF地址里写的是57B, 需要核实是否为笔误) |
Mixtral | Mistral AI于2024年1月发布 | 46.7B (8*7B) | 采用稀疏MoE架构 | 8个专家,每次选择2个。 |
Arctic | Snowflake于2024年4月发布 | 480B (128x3.66B 激活17B) | 动态数据课程,包含代码、文本等数据,强调数据质量和多样性。包含密集MoE混合变换器架构,使用1000张GPU,训练时间约3周 | 开源模型,性能指标与其他LLM相当。 |
DBRX | Databricks于2024年3月发布 | 132B (36B活跃参数) | 12T文本和代码数据,强调代码数据质量和多样性,包含多种编程语言。包含精细化MoE架构。训练使用3072张H100,约3个月时间,包括了pretraining, post-training, evaluation, red-teaming, and refining | 16个专家,每次选择4个。在多个基准测试中表现优越,尤其在编程任务上超越GPT-3.5。 |
Grok-1 | xAI于2023年11月发布 | 推测数十亿参数量 (具体未公开) | 可能包含大量社交媒体数据 (特别是X平台的数据)。推测使用了Transformer架构,具体细节未公开。训练推测使用了大量高性能GPU。 | xAI公司产品,与X(原推特)平台深度整合。 |
Grok-2 | xAI于2024年8月发布 | 推测较Grok-1有所增加 | 包括大规模文本与图像数据,改进的MoE架构,使用约10万张NVIDIA H100 GPU进行训练 | 是Grok-1的升级版本,推理、编码和多模态理解能力显著提升。 |
Grok-3 | xAI于2025年2月18日发布 | 未知 | 包括大规模文本与图像数据,改进的MoE架构,使用约10万张NVIDIA H100 GPU进行训练 | 在自然语言理解和生成方面有显著提升,具体细节尚未披露。 |
OLMoE | Allenai于2024年9月24日发布 | 7B(1B激活参数) | 包含大规模英文语料,尝试了对现有MoE算法的改进,并第一次将这个规模的MoE模型在5T tokens的语料上进行训练。使用 256 个 H100 GPU,通过 GPU 间的 NVlink 连接和节点间的InfiniBand 互连,大约进行 10 天的训练。 | 同等规模下最好的MoE模型 |
Architecture
MOE的基本组成
- 专家(Experts) 混合专家模型 (MoE) 的理念起源于 1991 年的论文 Adaptive Mixture of Local Experts。这个概念与集成学习方法相似,旨在为由多个单独网络组成的系统建立一个监管机制。在这种系统中,每个网络 (被称为“专家”) 处理训练样本的不同子集,专注于输入空间的特定区域。那么,如何选择哪个专家来处理特定的输入呢?这就是门控网络发挥作用的地方,它决定了分配给每个专家的权重。在训练过程中,这些专家和门控网络都同时接受训练,以优化它们的性能和决策能力。
在 2010 至 2015 年间,两个独立的研究领域为混合专家模型 (MoE) 的后续发展做出了显著贡献:
- 组件专家: 在传统的 MoE 设置中,整个系统由一个门控网络和多个专家组成。在支持向量机 (SVMs) 、高斯过程和其他方法的研究中,MoE 通常被视为整个模型的一部分。然而,Eigen、Ranzato 和 Ilya 的研究 探索了将 MoE 作为更深层网络的一个组件。这种方法允许将 MoE 嵌入到多层网络中的某一层,使得模型既大又高效。
- 条件计算: 传统的神经网络通过每一层处理所有输入数据。在这一时期,Yoshua Bengio 等研究人员开始探索基于输入令牌动态激活或停用网络组件的方法。
这些研究的融合促进了在自然语言处理 (NLP) 领域对混合专家模型的探索。特别是在 2017 年,Shazeer 等人 (团队包括 Geoffrey Hinton 和 Jeff Dean,后者有时被戏称为 “谷歌的 Chuck Norris”) 将这一概念应用于 137B 的 LSTM (当时被广泛应用于 NLP 的架构,由 Schmidhuber 提出)。通过引入稀疏性,这项工作在保持极高规模的同时实现了快速的推理速度。这项工作主要集中在翻译领域,但面临着如高通信成本和训练不稳定性等多种挑战。
混合专家模型 (MoE) 的引入使得训练具有数千亿甚至万亿参数的模型成为可能,如开源的 1.6 万亿参数的 Switch Transformers 等。这种技术不仅在自然语言处理 (NLP) 领域得到了广泛应用,也开始在计算机视觉领域进行探索。然而,本篇博客文章将主要聚焦于自然语言处理领域的应用和探讨。
- 门控网络(Gating Network)
这个部分用于决定哪些令牌 (token) 被发送到哪个专家。例如,在下图中,“More”这个令牌可能被发送到第二个专家,而“Parameters”这个令牌被发送到第一个专家。有时,一个令牌甚至可以被发送到多个专家。令牌的路由方式是 MoE 使用中的一个关键点,因为路由器由学习的参数组成,并且与网络的其他部分一同进行预训练。
总结来说,在混合专家模型 (MoE) 中,我们将传统 Transformer 模型中的每个前馈网络 (FFN) 层替换为 MoE 层,其中 MoE 层由两个核心部分组成: 一个门控网络和若干数量的专家。门控网络决定了哪些令牌被发送到哪个专家,而专家则负责处理这些令牌。这种方法的引入使得模型能够在保持极高规模的同时实现快速的推理速度。
- 输出组合机制(Combiner)
一般来说,组合器将所有专家模型的输出进行加权融合
MOE的工作原理
- 条件计算与稀疏激活
稀疏性的概念采用了条件计算的思想。在传统的稠密模型中,所有的参数都会对所有输入数据进行处理。相比之下,稀疏性允许我们仅针对整个系统的某些特定部分执行计算。这意味着并非所有参数都会在处理每个输入时被激活或使用,而是根据输入的特定特征或需求,只有部分参数集合被调用和运行。
让我们深入分析 Shazeer 对混合专家模型 (MoE) 在翻译应用中的贡献。条件计算的概念 (即仅在每个样本的基础上激活网络的不同部分) 使得在不增加额外计算负担的情况下扩展模型规模成为可能。这一策略在每个 MoE 层中实现了数以千计甚至更多的专家的有效利用。
这种稀疏性设置确实带来了一些挑战。例如,在混合专家模型 (MoE) 中,尽管较大的批量大小通常有利于提高性能,但当数据通过激活的专家时,实际的批量大小可能会减少。比如,假设我们的输入批量包含 10 个令牌, 可能会有五个令牌被路由到同一个专家,而剩下的五个令牌分别被路由到不同的专家。这导致了批量大小的不均匀分配和资源利用效率不高的问题。在接下来的部分中,将会讨论 让 MoE 高效运行 的其他挑战以及相应的解决方案。
那我们应该如何解决这个问题呢?一个可学习的门控网络 (G) 决定将输入的哪一部分发送给哪些专家 (E):
\[y = \sum_{i=1}^{n} G(x)_i E_i(x)\]在这种设置下,虽然所有专家都会对所有输入进行运算,但通过门控网络的输出进行加权乘法操作。但是,如果 G (门控网络的输出) 为 0 会发生什么呢?如果是这种情况,就没有必要计算相应的专家操作,因此我们可以节省计算资源。那么一个典型的门控函数是什么呢?一个典型的门控函数通常是一个带有 softmax 函数的简单的网络。这个网络将学习将输入发送给哪个专家。
\[G_\sigma(x) = Softmax(x \cdot W_g)\]Shazeer 等人的工作还探索了其他的门控机制,其中包括带噪声的 TopK 门控 (Noisy Top-K Gating)。这种门控方法引入了一些可调整的噪声,然后保留前 k 个值。具体来说:
- 添加一些噪声 \(H(x)_i=(x \cdot W_g)_i + StandardNormal() \cdot Softplus((x \cdot W_{noise})_i)\)
- 选择保留前k个值 \(KeepTopK(v, k)_i = \begin{cases} v_i, & \text{if $v_i$ is in the top $k$ elements of $v$} \\ -\infty, & \text{Otherwise} \end{cases}\)
- 应用softmax函数 \(G(x)=Softmax(KeepTopK(H(x), k))\)
这种稀疏性引入了一些有趣的特性。通过使用较低的 k 值 (例如 1 或 2),我们可以比激活多个专家时更快地进行训练和推理。为什么不仅选择最顶尖的专家呢?最初的假设是,需要将输入路由到不止一个专家,以便门控学会如何进行有效的路由选择,因此至少需要选择两个专家。Switch Transformers 就这点进行了更多的研究。
我们为什么要添加噪声呢?这是为了专家间的负载均衡!
- 混合专家模型中token的负载均衡
正如之前讨论的,如果所有的令牌都被发送到只有少数几个受欢迎的专家,那么训练效率将会降低。在通常的混合专家模型 (MoE) 训练中,门控网络往往倾向于主要激活相同的几个专家。这种情况可能会自我加强,因为受欢迎的专家训练得更快,因此它们更容易被选择。为了缓解这个问题,引入了一个 辅助损失,旨在鼓励给予所有专家相同的重要性。这个损失确保所有专家接收到大致相等数量的训练样本,从而平衡了专家之间的选择。接下来的部分还将探讨专家容量的概念,它引入了一个关于专家可以处理多少令牌的阈值。在 transformers 库中,可以通过 aux_loss 参数来控制辅助损失。