论文：

Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer
Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

代码：https://github.com/jingyaogong/minimind/blob/master/model/model_minimind.py

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

混合专家模型 (MoE) 的一个显著优势是它们能够在远少于稠密模型所需的计算资源下进行有效的预训练。这意味着在相同的计算预算条件下，您可以显著扩大模型或数据集的规模。特别是在预训练阶段，与稠密模型相比，混合专家模型通常能够更快地达到相同的质量水平。

混合专家模型 (MoEs):

与稠密模型相比，预训练速度更快
与具有相同参数数量的模型相比，具有更快的推理速度
需要大量显存，因为所有专家系统都需要加载到内存中
在微调方面存在诸多挑战，但近期的研究表明，对混合专家模型进行指令调优具有很大的潜力。

模型架构

混合专家模型主要由两个关键部分组成:

稀疏 MoE 层: 这些MoE层代替了传统 Transformer 模型中的前馈网络 (FFN) 层。MoE 层包含若干“专家”(例如 8 个)，每个专家本身是一个独立的神经网络。
门控网络或路由: 这个部分用于决定哪些token被发送到哪个专家。一个token可以被发送到多个专家。

总结来说，在混合专家模型 (MoE) 中，将传统 Transformer 模型中的每个前馈网络 (FFN) 层替换为 MoE 层，其中 MoE 层由两个核心部分组成: 一个门控网络和若干数量的专家。

稀疏性的概念采用了条件计算的思想。在传统的稠密模型中，所有的参数都会对所有输入数据进行处理。相比之下，稀疏性允许我们仅针对整个系统的某些特定部分执行计算。这意味着并非所有参数都会在处理每个输入时被激活或使用，而是根据输入的特定特征或需求，只有部分参数集合被调用和运行。

2017 年，Shazeer 等人将MoE应用于 137B 的 LSTM，GShard首次将 MoE 引入到 Transformer 模型中，Switch Transformer相比 Gshard 等方案主要做了三点改进：

简化稀疏路由
高效稀疏路由
增强的训练和微调技巧

专家选择

路由器其实是 MoE 中最核心的组件之一。它不仅决定推理阶段该选哪些专家，也决定训练阶段哪些专家要被更新。

最基础的做法是，将输入 token 表示向量 x ，乘以一个路由权重矩阵 Wg ，对打分结果做 softmax，得到一个不同专家的概率分布G( x ) ：

V_{\pi} ( S )=\sum_{a} \pi( a | s ) Q ( s , a )

然后根据这些概率值，选出和输入最相关的专家。

最后把每个路由器的输出与对应的专家相乘，加权求和，得到该 token 的最终输出：

y=\sum_{i=1}^{n} G ( x )_{i} E_{i} ( x )

负载均衡

虽然专家选择机制看起来很简单，但实际上会出现一个问题，某些专家学得更快，导致路由器总是选它们。这会导致选择的专家不均匀，一些专家几乎不会被训练到，最终训练和推理时都容易出现偏移和性能下降的问题。为了解决这个问题，我们引入了负载均衡（Load Balancing）：它的目的是让每个专家在训练和推理中都能被公平地使用，避免某在某几个专家上过拟合。

Q：如何解决MoE架构中负载不平衡的问题？

如果所有的token都被发送到只有少数几个受欢迎的专家，那么训练效率将会降低。在通常的混合专家模型 (MoE) 训练中，门控网络往往倾向于主要激活相同的几个专家。这种情况可能会自我加强，因为受欢迎的专家训练得更快，因此它们更容易被选择。为了缓解这个问题，引入了一个辅助损失（aux_loss），旨在鼓励给予所有专家相同的重要性。这个损失确保所有专家接收到大致相等数量的训练样本，从而平衡了专家之间的选择。接下来的部分还将探讨专家容量的概念，它引入了一个关于专家可以处理多少令牌的阈值。在 transformers 库中，可以通过 aux_loss 参数来控制辅助损失。

KeepTopK 策略

一种常见的路由负载均衡方法是通过一个叫做KeepTopK的简单的扩展，引入可训练的高斯噪声，以防止模型总是选中同一批专家：

H ( x )_{i}=( x \cdot W_{g} )_{i}+S t a n d a r d N o r m a l ( ) \cdot S o f t p l u s ( ( x \cdot W_{n o i s e} )_{i} )

接下来，除了你想要激活的 Top-k 个专家（比如前 2 个），其他专家的分数全部设为 -∞：

K e e p T o p K ( v , k )_{i}=\begin{cases} {v_{i}} & {\mathrm{i f ~ v_i ~ i s ~ i n ~ t h e ~ t o p ~ k ~ e l e m e n t s ~ o f ~ v ~} .} \\ {-\infty} & {\mathrm{o t h e r w i s e} .} \\ \end{cases}

这样一来，在做 SoftMax 时，分数为 -∞ 的专家对应的概率就是 0，不会被选中：

G ( x )=S o f t m a x ( K e e p T o p K ( H ( x ) , k ) )

Token Choice

KeepTopK 策略是每个 token 会被送到少数几个专家那里处理。这种方法被称为 Token Choice 。Top-1 路由每个 token 只交给一个专家处理。Top-k 路由每个 token 被同时送给k个专家，然后再根据它们的输出加权合并。这种方法的优点是：可以根据每个专家的权重贡献，融合多个专家的知识，更灵活。

辅助损失

为了让训练过程中的专家使用更加平均，研究者在主损失之外，引入了一个辅助损失（Auxiliary Loss）或者叫负载均衡损失（Load Balancing Loss）。这个损失项的作用是强制每个专家在训练中有差不多的“重要性”。

辅助损失的第一步，是对整个 batch 中每个专家的路由概率值求和也就是统计在这个 batch 中，每个专家一共“被选中的概率”加起来是多少：

I m p o r t a n c e ( X )=\sum_{x \in X} G ( x )

这样就得到了每个专家的一个重要性分数（importance score），表示这个专家在整个 batch 中，在不考虑输入内容的情况下，有多大可能被选中。

接下来，我们可以利用这些重要性分数，计算每个专家重要性分数的变异系数（Coefficient of Variation, CV） ，表示各个专家的重要性差异有多大：

C V=\frac{\sigma} {\mu} .

其中，σ表示Importance(X)的标准差，μ表示Importance(X)平均值。

如果 CV 很高，说明某些专家总被选中，而其他专家几乎没被用；如果 CV 很低，说明所有专家被使用得差不多，这正是我们想要的“负载均衡”状态。

利用这个 CV 分数，我们可以在训练过程中不断更新辅助损失，使模型的优化目标之一就是尽可能降低 CV 值，从而让每个专家的使用重要性趋于一致：

L_{i m p o r t a n c e} ( X )=w_{i m p o r t a n c e} \cdot C V ( I m p o r t a n c e ( X ) )^{2}

其中，w_{importance}是一个常量，表示缩放因子。

最后，这个辅助损失会作为一个单独的损失项，加入到整体训练目标中一起优化。

专家容量

模型中的不平衡，不仅体现在被选中的专家不平均，也体现在发送给专家的 token 分布不均。

比如：如果大量输入 token 都被不均衡地路由到了某一个专家（而其他专家几乎不会接收到 token），那可能会导致训练不充分。为了解决这个问题，可以引入一个限制机制，限制每个专家一次最多能处理多少个 token，称之为专家容量（Expert Capacity）。一旦某个专家达到容量上限，剩下的 token 就会被路由给下一个候选专家，如果所有候选专家都已满，token 将不会被任何专家处理，而是跳过当前 MoE 层，进入下一层。这种情况被称为 token overflow（token 溢出）。

Switch Transformer允许设置专家容量因子（capacity factor），显式地控制专家容量：

\mathrm{e x p e r t ~ c a p a c i t y=} \left( \frac{\mathrm{t o k e n s ~ p e r ~ b a t c h}} {\mathrm{n u m b e r ~ o f ~ e x p e r t s}} \right) \times\mathrm{c a p a c i t y ~ f a c t o r} .

容量因子设置得越大，每个专家就能处理更多的 token。如果容量因子太大，则会浪费计算资源。相反，如果容量因子太小，则 Token 溢出会导致模型性能下降。

参考文章：

Peter's Blog

混合专家模型笔记