8. MoE

介绍混合专家（Mixture of Experts，MoE）架构及其在大语言模型中的应用

• 参考资料

  • 帖子

    • A Visual Guide to Mixture of Experts (MoE)

    • Mixture-of-Experts (MoE) LLMs - by Cameron R. Wolfe, Ph.D.

    • nanoMoE: Mixture-of-Experts (MoE) LLMs from Scratch in PyTorch

    • MoE 相关成果：Mixture-of-Experts (MoE): The Birth and Rise of Conditional Computation

    • MoE 发展时间线与 MoE 并行：Distributed Mixture-of-Experts and Expert Parallelism | Bruno Magalhaes

  • 论文

    • 综述：A Review of Sparse Expert Models in Deep Learning

    • 综述：A Survey on Mixture of Experts in Large Language Models

    • Towards Understanding Mixture of Experts in Deep Learning

    • 专家路由选择：Mixture-of-Experts with Expert Choice Routing

什么是 MoE

MoE 的基本架构

MoE Abstract View

• Dense 架构

  • Dense 架构是传统的大语言模型架构，其中所有的参数都参与每个输入的前向计算

  • 模型的每一层对输入数据进行完整的计算，参数是全局共享的

  • 尽管每个前向计算使用了所有参数，但对于不同的任务或样本，很多参数的贡献可能较低，导致计算资源利用率不高

• 混合专家模型（Mixture of Experts，MoE）的优势

  • 模型规模是提升模型性能的关键因素之一，在有限的计算资源预算下，用更少的训练步数训练一个更大的模型，往往比用更多的步数训练一个较小的模型效果更佳

  • MoE 的一个显著优势是它能够在远少于 Dense 模型所需的计算资源下进行有效的预训练，这意味着在相同的计算预算条件下，可以显著扩大模型或数据集的规模

  • 特别是在预训练阶段，与 Dense 模型相比，混合专家模型通常能够更快地达到相同的质量水平

  • 关键技术：专家网络、门控网络、Top-K激活、负载均衡机制、稀疏激活机制

• MoE 模型主要由两个关键部分组成

  - 稀疏 MoE 层 - MoE 层代替了传统 Transformer 模型中的前馈网络 (FFN) 层 - MoE 层包含若干“专家”（例如 8 个），每个专家本身是一个独立的神经网络 - 在实际应用中，这些专家通常是前馈神经网络，也可以是更复杂的网络结构，甚至可以是 MoE 层本身，从而形成层级式的 MoE 结构 - GPU / TPU 的计算和内存访问优化主要针对 Dense 计算，稀疏激活可能导致硬件利用率降低 - 门控网络或路由（Gating Network） - 这个部分用于动态决定哪些 token 被发送到哪个专家，从而决定哪些专家模块参与计算 - 例如，在上图中，“More” 可能被发送到第二个专家，而 “Parameters” 被发送到第一个专家；有时，一个 token 可以被发送到多个专家 - token 的路由方式是 MoE 使用中的一个关键点，因为门控网络由可学习的参数组成，并且与网络的其他部分一同进行预训练 - 门控网络可能引入优化难题，例如专家负载不均（Load Balancing）

• 训练挑战：虽然 MoE 能够实现更高效的计算预训练，但在微调阶段往往面临泛化能力不足的问题，长期以来易于引发过拟合现象

• 推理挑战

  • MoE 模型虽然可能拥有大量参数，但在推理过程中只使用其中的一部分，这使得推理速度快于具有相同数量参数的 Dense 模型

  • 然而，MoE 模型需要将所有参数加载到内存中，因此对内存的需求非常高

  • 以 Mixtral 8x7B 这样的 MoE 为例，需要足够的 VRAM 来容纳一个 47B 参数的稠密模型。之所以是 47B 而不是 8 x 7B = 56B，是因为在 MoE 模型中，只有 FFN 层被视为独立的专家，而模型的其他参数是共享的

  • 此外，假设每个 token 只使用两个专家，那么推理速度（以 FLOPs 计算）接近于 12B 模型（而不是 14B 模型），因为虽然它进行了 2x7B 的矩阵乘法计算，但某些层是共享的

MoE 的本质

• 专家不是在领域知识层面的，而是语法句法层面的

• encoder 中的 MoE 专家并非专注于“心理学”或“生物学”等特定领域，而是呈现出明显的 token / 概念级分工，例如：标点符号、数字、格式化 token、专有名词、实体、形态学或浅层句法线索，更擅长在特定语境下处理特定 tokens

  • encoder 的目标是构建输入表示，而不是做决策，每个 token 的表示是“并行”的，不需要为最终输出承担直接责任

  • encoder 中的路由看到的是当前 token 的 embedding，梯度主要来自 token 级重建 / attention 对齐

  • 浅层统计模式在 encoder 中梯度最稳定，标点、大小写、子词边界在所有任务中高度一致，这些模式对 loss 的贡献密集、噪声小

  • 因此，encoder 中非常容易形成这种结构：

    • 专家 A：高频结构 token

    • 专家 B：实体类 token

    • 专家 C：函数词、连接词

• decoder 中的专家往往表现为语义混合、专精边界模糊、很难指认“这个专家只做某一类 token”，原因恰恰来自 decoder 的角色不同：

  • decoder 的损失是序列级、决策级，每一步预测都会影响未来 token，梯度高度耦合、非局部

  • decoder token 不再是“被解释的对象”，而是“要被生成的决策”，不同类型 token 的预测共享相同的上下文依赖

  • 从优化角度看，一个 decoder 专家如果只擅长“某一类 token”，那么在其他时间步就会变成低效参数，门控网络会自然避免这种高度专一但覆盖率低的专家

  • 结果是 decoder 专家更像“通用语言建模子网络”，差异体现在风格或高阶特征，而不是 token 类型

• 为什么多语言训练中不会出现“一专家一语言”

  • MoE 路由器优化的目标是：降低 task loss，同时满足负载均衡约束，而不是最大化语言间可分性或最小化跨语言干扰

  • MoE 的路由粒度是 token 级，而不是句子、段落或语言标签级

  • 在多语言文本中同一语言内部 token 类型高度多样，不同语言之间 token 类型反而大量重叠

  • 于是自然形成的是“跨语言的功能专家”，而不是“单语言专家”

  • 从模型能力角度看，这种现象反而是 MoE 的优势，专家学的是“可迁移的语言结构”，而不是“语言标签本身”，多语言泛化能力更强，低资源语言可以共享高资源语言的专家能力

MoE 发展历史

• 混合专家模型的概念最早可以追溯到 1991 年的工作 Adaptive Mixture of Local Experts

  - 这一思想与集成学习方法在动机上较为相似，其目标是为由多个独立子网络组成的系统引入一种统一的调度与决策机制 - 在该框架下，每个子网络被视为一个“专家”，负责处理训练数据的不同子集，并专注于输入空间中的特定区域 - 专家之间的选择与分配由门控网络完成，门控网络为每个输入样本分配对应的专家权重 - 在训练过程中，专家网络与门控网络是联合优化的，从而同时提升专家的表达能力与路由决策的有效性

• 在 2010 至 2015 年期间，两个相对独立的研究方向为混合专家模型的进一步发展奠定了基础：

  • 作为网络组件的专家模型

    • 在早期的 MoE 设定中，系统通常由一个门控网络和多个专家整体构成

    • 在支持向量机、高斯过程等模型中，MoE 往往被视为完整模型的一种形式

    • Eigen、Ranzato 和 Ilya 的研究 则提出将 MoE 作为深层神经网络中的一个组成模块

    • 这种设计使得 MoE 可以嵌入到多层网络的特定位置，从而在显著提升模型容量的同时保持计算效率

  • 条件计算

    • 传统神经网络在前向传播时会对每一层、每一个参数执行完整计算

    • 这一时期，Yoshua Bengio 等研究者开始探索根据输入动态激活或跳过部分网络结构的条件计算机制

    • 条件计算为在不成比例增加计算量的前提下扩展模型规模提供了新的可能性

  • 上述两个研究方向的逐步融合，直接推动了混合专家模型在自然语言处理领域中的应用与发展

• 在 2017 年，Shazeer等人 将混合专家模型引入到大规模序列建模中，并将其应用于一个参数规模达到 137B 的 LSTM 模型

  • LSTM 是当时 NLP 领域的主流架构之一，最早由 Schmidhuber 提出

  • 通过引入稀疏激活机制，该工作在保持极大模型规模的同时，实现了相对高效的推理速度

  • 尽管该模型主要应用于机器翻译任务，但仍面临通信开销较高、训练过程不稳定等实际问题

• 混合专家模型的提出与持续演进，使得参数规模达到数千亿乃至万亿级别的模型在工程上成为可能，例如公开发布的 1.6 万亿参数 Switch Transformers

Dense vs. Sparse

• 标准的 decoder-only Transformer 架构会在层归一化之后使用前馈神经网络

• 为了重新学习 token 中的复杂关系，FFN 通常会对输入进行扩展，隐藏层维度很大

  

• 上述前馈神经网络被称为 Dense 模型，因为所有参数（其权重和偏置）都会被激活，没有任何参数被遗漏，所有参数都会用于计算输出

• 相比之下，稀疏模型仅激活总参数的一部分，举个例子，可以将 Dense 模型分割成多个部分（即所谓的专家），对其重新训练，并且在特定时间只激活一部分专家：

• 其核心思想是，每个专家在训练过程中会学习不同的信息，在进行推理时，只会使用那些与特定任务最相关的专家

• 当接收到一个输入时，可以据此选择最适合特定任务的专家：

  

• 实际上，这里的专家学到的是更细粒度的信息，而非整个领域的知识，因此称之为“专家”是有一定误导性的，如 ST-MoE 论文的统计：

  

• Mixtral 的论文展示了每个 token 被分配到的第一个专家，也表明专家们倾向于关注语法而非特定领域

  

• 从理论上来讲，稀疏性的核心思想来源于条件计算（Conditional Computation）：

  • 在传统的 Dense 模型中，所有参数都会参与对每一个输入样本的计算过程，而稀疏模型允许仅对部分输入执行计算

  • 具体而言，并非所有参数在处理每个输入时都会被激活，而是根据输入的特征，仅激活与之相关的一小部分参数

  • 这种“按样本激活不同子网络”的条件计算机制，使得模型参数规模可以显著扩展，而不会线性增加计算开销

  • 在 MoE 层中，该策略使得数以千计甚至更多专家能够被有效整合到同一个模型中

• 在上面的说明中，我们将专家描述为由一个 Dense 模型分割为不同的子模块，但在实际工程中，专家本身通常是完整的前馈神经网络

  

• 由于存在多个 Decoder Block，因此一般需要经过多个专家的处理

  

• 由于不同的专家在不同的 token 处理中存在差异，因此会有多条流向路径产生

  

• Dense 模型与稀疏模型的对比

Router / Gating Network

• 既然现在有了一组专家，那么如何选择使用哪些专家？

• 为此，引入一个可学习的门控网络 $G$，用于决定输入应被分配给哪些专家 $E$，并基于此解决负载均衡问题：

  $$y = \sum_{i=1}^{n} G(x)_i E_i(x)$$

  • 在这一形式化描述中，输出由各专家的结果加权求和得到

  • 当某个专家对应的门控权重 $G(x)_i = 0$ 时，该专家的计算可以被跳过，从而节省计算资源

  • 实际应用中，门控网络通常采用一个简单的前馈神经网络加 Softmax，用于学习输入到专家的分配关系：

    $$G_\sigma(x) = \text{Softmax}(x \cdot W_g)$$
• 门控网络的示意图如下：

  

• 这样路由和一组专家组成了 MoE 层

  

• MoE 层也可以分为 Dense 和 sparse 两种，两者都使用路由器来选择专家，但 Sparse MoE 模型只选择少数专家，而 Dense MoE 模型则会选择所有专家，只是可能分布不同，一般来说，MoE 指的是 Sparse MoE 模型

• 路由的工作原理

  

• 负载均衡问题

  • 过于简单的路由网络往往会导致路由器选择同一个专家，不仅选定专家的分布会不均衡，而且有些专家几乎不会得到任何训练，而这会在训练和推理阶段都引发问题

  • 即尽管较大的 batch size 通常有助于提升训练效率，但在 MoE 中，token 在专家之间的动态路由会导致有效 batch size 的不均匀分布

  • 例如，对于包含 10 个 token 的输入批量，可能有 5 个 token 被路由到同一个专家，而剩余 token 分散到多个不同专家

  • 这会造成专家负载不平衡，从而降低计算资源的整体利用率

Load Balancing

KeepTopK

• Shazeer 等人的工作进一步探索了更复杂的门控策略，其中最具代表性的是带噪声的 Top-K 门控（Noisy Top-K Gating），该方法在门控网络中引入可学习的高斯噪声，并仅保留得分最高的 $k$ 个专家

  • 首先，在门控 logits 上加入噪声：

    $$H(x)_i = (x \cdot W_g)_i + \epsilon \cdot \text{Softplus}((x \cdot W_{\text{noise}})_i)$$

    • 其中，$\epsilon \sim \mathcal{N}(0,1)$是每个令牌采样的标准高斯噪声，$W_{noise}$是另一个可训练的权重矩阵，用于根据输入对噪声进行缩放

    • softplus 函数确保噪声缩放因子为正，随后使用这些含噪声的分数进行 top-k 选择

    • 而后续用于加权专家输出的概率通常仍从所选专家的原始、非噪声对数几率中得出，以保持输出计算的稳定性，即噪声仅仅用于选择专家

  • 示意图如下：

    

  • 随后，仅保留前 $k$ 个值，其余位置置为 $-\infty$：

    $$\text{KeepTopK}(v, k)_i = \begin{cases} v_i & \text{if } v_i \text{ is in the top } k \text{ elements of } v, \\ -\infty & \text{otherwise.} \end{cases}$$
  • 示意图如下：

    

  • 最后对结果应用 Softmax，得到最终门控权重：

    $$G(x) = \text{Softmax}(\text{KeepTopK}(H(x), k))$$
  • 示意图如下：

    

• KeepTopK 策略将每个 token 路由到几个选定的专家，这种方法被称为 token 选择，并允许将给定的 token 发送到一个专家（如 top-1）

  

• 也可以发送到多个专家（如 top-2）

  

• KeepTopK 的特性：

  • 通过选取较小的 $k$（如 1 或 2），可以在保证模型容量的同时显著提升训练和推理速度

  • 若仅始终选择单一最优专家，门控网络难以学习有效的路由策略，因此通常至少需要激活两个专家

  • Switch Transformers 对这一现象及其工程影响进行了系统分析

• 在 KeepTopK 路由策略中，引入噪声的核心目的在于缓解专家负载不均衡，防止少数专家在训练过程中长期占据大部分计算资源

  • 在 MoE 训练中，路由网络往往容易出现负载坍塌：大量 token 被持续分配给少数专家，而其余专家几乎得不到训练信号。这不仅导致专家利用率严重失衡，也会降低整体吞吐率与训练效率，削弱 MoE 所期望的容量扩展优势

  • 为缓解这一问题，KeepTopK 在执行 top-k 专家选择之前，对路由器输出的门控 logits 注入随机噪声，并在加入噪声后的 logits 上进行 top-k 操作。该机制在 logits 接近的专家之间引入随机扰动，打破确定性的路由决策，使原本处于边缘的专家获得被选中的机会，从而逐步缓解路由偏置

  • 从优化角度看，噪声并非用于提升单次路由决策的准确性，而是一种探索机制

  • 在训练初期和中期，较高幅度的噪声能够提升路由选择的多样性，避免“冷启动”问题，使所有专家都能接收到足够的 token 并获得有效梯度

  • 随着训练的推进，噪声强度通常会逐步衰减，路由决策从探索阶段过渡到利用阶段，此时，路由器可以基于已经形成的专家专业化结构，做出更加稳定和确定的选择

  • 从动态系统的角度看，噪声的引入有助于打破早期的确定性反馈循环，防止专家仅因初始化或微小的早期优势而持续占据主导地位，使专家分工能够在更公平的条件下逐步形成

Auxiliary Loss

• KeepTopK 中的噪声

  • 噪声的作用本质是在 logits 接近的专家之间打破确定性排序，给“边缘专家”一个被偶然选中的机会

  • 噪声提供的是一种局部层面、样本级别的探索机制，作用于单个 token 的路由决策

  • 噪声（如 Gaussian）在期望上是零均值的，如果某个专家在平均意义上 logit 明显更大，那么即使加噪声，它仍然会在大多数样本中胜出

• 但负载不均衡不是一个局部现象，而是一个跨 token、跨 batch、跨 step 的统计问题：

  • 专家 $e_i$ 在一个训练窗口中接收的 token 数

  • 是否接近均匀分布

  • 是否超过容量、或长期低于容量

• 为缓解这一现象，MoE 模型通常引入辅助损失（auxiliary loss）

  • 又称为负载均衡损耗（load balancing loss），用于鼓励各个专家在整体上具有相近的重要性

  • 该损失项通过约束路由分布，使每个专家接收到的大致训练样本数量保持均衡，从而改善负载分配与训练稳定性

  • 在 transformers 库的实现中，可以通过 aux_loss 参数调节上述辅助损失项

• 该损耗的第一个组成部分是指定专家对整个批次中每个 token 的路由值的和，这个结果表示每个专家的重要性分数，该分数表示无论输入如何，某个特定专家被选中的可能性

  

• 然后利用重要性分数来计算 coefficient variation（CV），衡量专家重要性差异

  $$CV=\frac{标准差\,\sigma}{平均值\,\mu}$$
• 如果重要性分数存在很大差异，那么 CV 将会很高

  

• 如果重要性分数差异较小，那么 CV 较低，也正是我们所希望得到的

  

• 利用这个 CV 值，可以在训练过程中更新辅助损失，使其尽可能降低 CV，从而给予每个专家同等的重要性

  

• 辅助损失被作为一个单独的损失项添加进损失函数中，以在训练过程中进行优化

• CV 为什么能约束负载均衡

  • 设第 $i$ 个专家的负载为 $x_i$，均值为

    $$\mu = \frac{1}{E} \sum_{j=1}^E x_j$$
  • 采用 $CV^2$ 作为辅助损失：

    $$\mathcal{L}_{cv} = \frac{1}{\mu^2}\frac{1}{E}\sum_{j=1}^E (x_j-\mu)^2$$
  • 对 $x_i$ 求偏导，注意到 $\mu=\frac{1}{E}\sum_j x_j$，但

    $$\sum_{j=1}^E (x_j-\mu)=0$$
  • 因此均值项对一阶梯度的贡献相互抵消，得到

    $$\frac{\partial \mathcal{L}_{cv}}{\partial x_i} = \frac{2}{E\mu^2}(x_i-\mu)$$
  • 去掉正的常数因子，方向上等价于

    $$\frac{\partial \mathcal{L}_{cv}}{\partial x_i} \propto (x_i - \mu)$$
  • 由于

    $$x_i = \sum_t p_{t,i}, \quad p_{t,i} = \mathrm{softmax}(z_t)_i$$
  • 梯度继续传到 router logit $z_{t,i}$，其方向保持一致：

    • 若 $x_i > \mu$：

      $$\frac{\partial \mathcal{L}_{cv}}{\partial x_i} > 0 \Rightarrow z_{t,i} \downarrow$$
    • 若 $x_i < \mu$：

      $$\frac{\partial \mathcal{L}_{cv}}{\partial x_i} < 0 \Rightarrow z_{t,i} \uparrow$$

Expert Capacity

• Auxiliary Loss 解决的是“平均负载”，是一种长期、统计意义上的均衡机制

• 而专家容量（Expert Capacity）作用在前向路由阶段中，规定了每个专家在一个 batch 中最多接收的 token 数

• 其通过硬截断来防止某个专家在单个 step 内被压爆，限制最坏情况的显存和计算量

• 当一个专家达到容量时，多余的 token 将被发送给下一个专家：

  

• token 溢出：如果两位专家都已达到其容量，token 将不会被任何专家处理，而是被发送到下一层

  

参数对比

• 尽管 MoE 模型推理时需要加载更多参数（稀疏参数），但实际激活的参数更少，因为在推理过程中只使用部分专家（活跃参数）

  

• 在推理时，仍然需要将整个模型（包括所有专家）加载到设备上（稀疏参数），但在进行推理时，只需要使用其中的一个子集（激活参数）

• 混合专家模型需要更多的显存来加载所有专家，但在推理过程中运行速度更快

• Mixtral 8x7B 的模型参数计算

  - 单个专家规模为 5.6B（非 7B），共 8 个专家； - 加载时需处理 “8×5.6B=46.7B” 参数（含所有共享参数）； - 推理时仅使用 “2×5.6B=12.8B” 参数，体现 “高显存占用但快推理速度” 的特点

• Mixtral 参数示意

  

GShard

• 以往对 Transformer 的研究表明，模型性能通常会随着参数规模的增加而持续提升

• 基于这一经验，Google 提出了 GShard，尝试将 Transformer 的参数量扩展到 6000 亿以上，从而探索超大规模模型的可行性与效率边界

• GShard 将 Encoder 和 Decoder 中的每一个前馈神经网络（FFN）层，统一替换为采用 Top-2 门控机制的 MoE 层

• 上图展示的是 Encoder 侧的整体结构，这种设计在大规模分布式环境下尤为高效：

  • 当模型扩展到多张 GPU 或多台设备时，MoE 层中的专家参数在不同设备之间进行分片与共享

  • 其余非 MoE 层（如注意力层、嵌入层等）则在每个设备上完整复制

  • 这种划分方式在保证模型容量的同时，有效控制了通信与计算开销

• 为了在超大规模训练中维持负载均衡与训练稳定性，GShard 在引入前文所述的辅助损失之外，还对 MoE 的路由与容量机制进行了两个关键改进：

  • 随机路由（Random Routing）

    • 在 Top-2 路由设置中，第一个专家始终选择门控得分最高的专家

    • 第二个专家不再固定为次高得分，而是按照门控权重比例进行随机采样

    • 这种方式可以缓解专家负载过度集中在少数专家上的问题，从而提升整体训练效率

  • 专家容量（Expert Capacity）

    • 为每个专家设置一个容量上限，用于限制其在一次前向计算中最多可处理的 token 数

    • 当一个 token 被分配到的两个专家均已达到容量上限时，该 token 会发生溢出

    • 溢出的 token 通常通过残差连接直接传递到下一层，在部分实现中也可能被直接丢弃

• 除了模型结构本身，GShard 还探讨了 MoE 并行计算，这部分内容可以在分布式训练中的专家并行（Expert Parallelism）章节查看

Switch Transformers

• Switch Transformers 是首批基于 Transformer 的 MoE 模型之一，它解决了 MoE 的训练不稳定性问题（如负载均衡），并简化了大部分架构和训练流程，同时提高了训练的稳定性

• 该工作同时在 Hugging Face 上发布了一个 1.6 万亿参数的 MoE，其中包含 2048 个专家，展示了该架构在极端规模下的可行性

• 模型架构

  • 主要目标：在保持 T5 任务范式（如微调、多任务）的前提下，通过 MoE 提升参数规模与效率

  • Switching Layer：

    • 基于 T5 Encoder-Decoder 架构，用 “稀疏 MoE 层” 替代传统 FFN 层

    • 采用 “Top-1 路由”，该简化显著降低了门控网络的计算与通信开销，每个专家所需处理的有效 batch size 至少减少一半，在降低系统复杂度的同时，模型质量得以保持

  • Capacity factor：

    $$\text{Expert Capacity} = \left( \frac{\text{tokens per batch}}{\text{number of experts}} \right) \times \text{capacity factor}$$

    • 决定了专家可以处理的 token 数量，如果提高容量因子，每个专家将能够处理更多的 token

    • 当容量因子大于 1 时，可以在一定程度上缓解 token 在专家之间分配不均的问题

    • 容量因子过大，会显著增加跨设备通信开销，浪费计算资源

    • 容量因子过小，模型性能会因 token 溢出而下降，因此需要在性能与效率之间进行权衡

    • 实验表明，Switch Transformers 在较低的容量因子（如 1 至 1.25）条件下即可获得良好性能

  • 容量因子与 token 溢出

    

  • 容量因子与专家容量的计算关系

    

  • 简化的辅助损失

    - 无需计算 coefficient variation（CV），而是根据每个专家的路由概率占比来权衡分配的 token 占比 - 由于目标是在 N 个专家之间实现 token 的统一路由，所以希望向量 P 和 f 的值为 1/N - 会倾向于降低“高 $f_i$ + 高 $P_i$”的专家，鼓励概率质量向低负载专家扩散 - 超参数 α 用于微调该损失在训练中的重要性，值过高会超过主要的损失函数，值过低则对负载均衡几乎没有作用

  • 辅助损失的公式表示

    

  • 混合精度训练策略：

    • 专家网络可采用 bfloat16 精度训练，而其余计算保持全精度，从而降低计算、通信与激活存储成本

    • 实验发现，当专家与门控网络同时使用 bfloat16 时，训练过程容易出现不稳定

    • 不稳定性主要来源于路由计算中涉及的指数等对数值精度敏感的操作

    • 因此，为保证路由决策的稳定性与准确性，门控与路由相关计算仍采用全精度实现

• 模型对比

  • 在预训练效率方面，Switch Transformers 相较于 T5-XXL 实现了约 4 倍的训练速度提升

  • 与 GShard 类似，Switch Transformers 通过将 Transformer 中的前馈神经网络层替换为 MoE 层，从而显著提升模型容量

  • 在结构设计上，Switch Transformers 提出了专门的 Switch Transformer 层，该层在概念上接收多个输入 token，并在一个包含多个专家的 MoE 模块中完成路由与计算

GLaM

• 参考论文：GLaM: Efficient Scaling of Language Models with Mixture-of-Experts

• 架构：decoder-only（更接近 GPT 系列）

• 目标：在语言建模质量上达到 GPT-3 水平，但训练能耗只用大约 1/3

• 核心动机：MoE 在每个 token 上只激活少量 expert，有效 FLOPs 显著低于同参数量的 Dense 模型，因此能显著降低训练计算量

• 评测方式：侧重 few-shot / one-shot，而不是下游任务微调（这点和 GPT-3 一致）

• 路由：Top-2 routing，每个 token 同时发给两个 expert，负载更均衡、稳定性更好，但计算量和通信略高

• 使用更大的 capacity factor：capacity factor 控制每个 expert 能接收的 token 上限，对于大 CF，更少 token 被丢弃，更好的质量，更高的计算和显存开销

• GLaM 把 capacity factor 当成一个训练 / 推理时可调的算力–性能权衡参数

ST-MoE

• 参考论文：ST-MoE: Designing Stable and Transferable Sparse Expert Models

• Switch-Tranformer 模型存在训练不稳定性，尽管 Switch Transformer 在预训练速度上超过了 Dense 模型，但在 SuperGLUE 等主流基准上进行微调时，其性能落后于之前的最先进模型

• 这些基于 MoE 的语言模型只有在数据集非常大时才能表现良好，否则会因过拟合而导致性能下降

• ST-MoE 系统性地解决了训练不稳定的问题

  • 通过 dropout 或输入抖动（将路由器的 logits 乘以 [0.99, 1.01] 范围内的随机值）向模型注入噪声

  • 在反向传播过程中，基于固定阈值裁剪梯度或权重更新，以解决梯度爆炸问题

• 在稀疏 MoE 中，训练不稳定往往不是来自专家本身，而是来自路由网络：

  • 路由网络通常输出 logits，然后通过 softmax 或 Top-K + softmax 决定专家分配

    

  • 随着训练进行，门控网络很容易产生绝对值越来越大的 logits，一旦 logits 的幅值变大，softmax 中的指数项数值范围急剧扩大，不同 expert 之间的概率差异被极度放大，浮点舍入误差和数值不稳定开始主导梯度

  • 这会导致几类典型问题：梯度突然爆炸或消失、路由分布高度尖锐，专家塌缩、混合精度下训练直接发散

  • 传统的负载均衡损失只能约束「概率分布形态」，但无法限制 logits 的数值尺度本身，而简单地对路由输出到 logits 进行截断可能会损害模型性能

• ST-MoE 引入了 Router z-loss，来促使路由输出的值更小——路由网络的正则化技术

  • Router z-loss 的目标非常明确：不去干预 token 被送到哪个专家，而是约束「路由网络输出 logits 的数值规模」，这是一种纯数值稳定性正则项，而非结构性或语义性约束

  • 实际上，softmax 对 logits 的相对差异敏感，但对整体平移不敏感，因此 Router z-loss 实际上在做的是抑制 logits 的无意义整体放大，强迫模型通过“相对差异”而非“绝对数值”来表达路由偏好

  • 这使得 softmax 输入始终处于安全数值区间，指数计算不容易溢出或产生严重舍入误差，从而使混合精度下梯度更加平滑

  • 设路由对某个 token 输出的 logits 为 $z \in \mathbb{R}^E$，其中 $E$ 是专家数量

  • Router z-loss 定义为：

    $$ \mathcal{L}_{z} = \lambda_z \cdot \left( \log \sum_{i=1}^{E} e^{z_i} \right)^2 $$

  • z-loss 常被称为对 log-sum-exp 的平方惩罚，当 logits 整体平移或幅值增大时，该项会迅速变大，平方使得对“大值”惩罚更强

  • 该损失会被加到总训练损失中：

    

• 对比几种常见方法：

  • dropout：稳定，但直接破坏表示能力，在 MoE 中容易导致专家学习不足

  • 减小模型容量 / 减少专家数：稳定，但直接降低上限性能

  • 强行平滑路由概率：稳定，但限制模型做“尖锐选择”

  • 而 Router z-loss 不限制 Top-1 / Top-2 的选择，不干预专家间的相对竞争关系，只约束一个数值冗余自由度

• 同时，尽管稀疏模型在大规模数据集上表现良好，但对这些模型进行微调可能存在困难

  • 稀疏模型在小型/大型数据集这两种情况下都比 Dense 模型收敛得更快，然而在数据较少的下游任务中，稀疏模型的验证质量远不如 Dense 模型——换句话说，稀疏模型由于缺乏足够的数据而发生过拟合

  • 为了减轻过拟合，ST-MoE 首先尝试只对 MoE 模型的某些部分参数进行微调，但仅微调 MoE 层会导致性能急剧下降，而选择性地微调其他参数组的效果则与微调所有参数相当，这表明在微调阶段，非 MoE 层对任务适配起着关键作用

  • 在训练稀疏模型时，采用更强的内部正则化通常是有益的，例如，可以对 Dense 层使用较低的 dropout 率，而对稀疏（MoE）层使用更高的 dropout 率，以平衡模型容量与泛化能力

  • ST-MoE 的实验表明，即使关闭辅助损失并导致高达约 11% 的 token 被丢弃，模型质量也未出现明显下降，这表明 token 丢弃本身可能起到了一种正则化作用，有助于缓解过拟合

  • 相较于 Dense 模型，稀疏模型通常更适合使用较小的批量大小，同时采用更高的学习率往往可以获得更好的微调效果

• MoEs Meets Instruction Tuning 对微调的发现

  • Flan-MoE 相对于原始 MoE 的性能增益，明显大于 Flan-T5 相对于原始 T5 的提升幅度，表明 MoE 模型可能比 Dense 模型从指令微调中获得更大的收益

  • MoE 在多任务学习场景下具有更强的优势

  • 保留辅助损失反而有助于抑制过拟合，从而提升整体泛化性能

MoE 的稳定性训练

• 熵正则化

  • 与侧重于平衡各专家接收 token 总量的负载均衡损失不同，熵正则化直接作用于路由器对单个 token 的专家分配概率分布

  • 其目标是鼓励 softmax 输出在专家维度上更加平坦，即具有更高的熵，从而避免路由器过早做出过于确定的决策

  • 熵正则化项通常定义为：

    $$\mathcal{L}_{\text{entropy}} = -\lambda_e \cdot \frac{1}{N \cdot T} \sum_{i=1}^{N \cdot T} \sum_{j=1}^E g_{i,j} \log (g_{i,j} + \epsilon)$$
  • 其中，$g_{i,j}$ 表示 token $i$ 被分配给专家 $j$ 的门控概率，$E$ 为专家数量，$\epsilon$ 为数值稳定项，$\lambda_e$ 为熵正则化系数

  • 该损失项通过惩罚低熵分布，促使路由器在专家选择上保持一定的不确定性

  • 这在训练早期尤为重要，有助于提高路由探索性，避免因初始化偏置或早期噪声导致专家利用率迅速失衡

  • 与路由噪声类似，熵正则化的权重 $\lambda_e$​ 通常会在训练过程中逐步衰减，使路由策略从探索阶段平滑过渡到更加确定的利用阶段

• 至此，已经引入了五种有助于 MoE 负载均衡和稳定性训练的技术

  机制	作用对象	主要解决的问题	核心机制	是否参与反向传播	典型使用阶段
  熵正则化	单个 token 的专家概率分布 $g_{i,\cdot}$	路由过早变得确定，softmax 分布过于尖锐	惩罚低熵分布，鼓励概率更平坦	是	训练早期为主，权重常随时间衰减
  Noise	router logits（softmax 之前）	确定性 top-k 导致的路径依赖与早期偏置	在 logits 上注入随机扰动，提升探索性	否	训练早期与中期，噪声幅度逐步衰减
  auxiliary loss	跨 token 的专家负载统计 $x_i$	专家负载不均衡、表示坍塌	惩罚负载的相对方差，拉平专家使用率	是	全训练过程，通常保持恒定或弱衰减
  Expert Capacity	单个专家可接收的 token 数	专家过载、显存和吞吐失控	对路由结果施加硬上限，超额 token 被丢弃或重路由	否	全程启用，工程层面约束
  Router z-loss	router logits 的整体尺度	logits 过大导致 softmax 饱和与梯度消失	惩罚 $\log\sum_j e^{z_j}$ 的幅值，稳定数值尺度	是	全训练过程，通常权重较小

• 完整的路由参数决定因素

MoE 高效推理

• token 级分组与置换（Token-level Grouping & Permutation），其核心思想是：根据每个 token 被分配到的专家，对批次内的 token 进行主动重排，从而把“稀疏、不规则”的专家访问，转化为“致密、规则”的批处理计算

• 整体工作流程通常如下：

  • 门控计算

    • 门控网络对输入批次中的所有 token 进行处理

    • 根据门控输出，确定每个 token 应被路由到的专家

    • 推理阶段通常采用 Top-1 路由以降低计算与通信开销

    • 在部分设置中也可能使用 Top-2 路由以提升模型质量

  • 路由决策与排序

    • 为每个 token 确定其目标专家 ID

    • 根据专家 ID 对 token 进行排序或建立索引

    • 排序后的目标是所有分配给同一专家的 token 在逻辑或物理上相邻，为后续的批量专家计算做准备

  • token 置换（Gather）

    • 根据排序结果，对 token 表示进行内存中的重新排列，使分配给专家 0 的 token 连续，接着是专家 1、专家 2，依此类推

    • 如果专家分布在不同设备上，该阶段会触发 All-to-All 通信，每个设备发送属于远程专家的 token，同时接收分配给本地专家的 token

    • 高效实现通常依赖优化的集体通信库（如 NCCL）或专门针对 MoE 优化的框架（如 Tutel）

  • 批处理专家计算

    • 每个专家（或承载该专家的设备）接收到一个只包含“自己负责 token ”的致密小批次

    • 专家网络（通常是 FFN）对该批次执行高效的向量化计算

    • 这一阶段充分利用 GPU 对致密矩阵运算的高吞吐能力

  • 令牌反置换（Scatter）

    • 专家计算完成后，输出 token 的顺序仍然是“按专家分组”的

    • 需要执行与置换相反的操作，将 token 恢复到原始批次与序列顺序

    • 在分布式设置下，通常涉及第二次 All-to-All 通信，将处理后的 token 结果发回其原始设备

  • 组合与继续计算（Combine & Continue）

    • 来自专家的输出按照原始 token 顺序重新对齐

    • 根据门控 logits 对专家输出进行组合，即使推理时使用 Top-1 计算，组合阶段仍可能保留加权形式以匹配训练行为

    • 得到的 token 表示进入模型的下一层，如后续 FFN、LayerNorm、残差连接等

• “分组与置换”将 MoE 的不规则路由问题转化为专家内部的规则批处理计算

• All-to-All 通信是跨设备专家部署时的关键成本，该机制是 MoE 在推理阶段实现高吞吐与可扩展性的核心基础

开源 MoE 模型

• 有多种开源项目可用于训练 MoE 模型：

  • Megablocks

    • 提供高效的稀疏 MoE 训练实现

    • https://github.com/stanford-futuredata/megablocks

  • Fairseq

    • 在其语言模型示例中支持 MoE 训练

    • https://github.com/facebookresearch/fairseq/tree/main/examples/moe_lm

  • OpenMoE

    • 面向社区的 MoE 训练与实验框架

    • https://github.com/XueFuzhao/OpenMoE

• 开源的 MoE 模型：

  • Switch Transformers（Google）

    • 基于 T5 架构的 MoE 模型系列

    • 专家数量从 8 到 2048 不等

    • 最大模型规模达到 1.6 万亿参数

    • https://huggingface.co/collections/google/switch-transformers-release-6548c35c6507968374b56d1f

  • NLLB MoE（Meta）

    • NLLB 翻译模型的混合专家版本

    • https://huggingface.co/facebook/nllb-moe-54b

  • OpenMoE

    • 社区对基于 LLaMA 架构的 MoE 模型探索

    • https://huggingface.co/fuzhao

  • Mixtral 8×7B（Mistral）

    • 高质量混合专家模型，在多项基准上性能超过 LLaMA 2 70B

    • 同时具备更快的推理速度

    • 提供指令微调版本

    • https://mistral.ai/news/mixtral-of-experts/