【译】稀疏混合专家模型的崛起： Switch Transformers

原作： 塞缪尔·弗兰德

引言：为最强大的语言模型铺平道路的核心技术

使用 Dall-E 生成的图像

稀疏混合专家模型 (MoE) 已成为最新一代 LLMs 的核心技术，例如 OpenAI 的 GPT-4、Mistral AI 的 Mixtral-8x7 等。简而言之，稀疏MoE是一种非常强大的技术，因为理论上，它允许我们以 O(1) 的计算复杂度扩展任何模型的容量！

然而，正如通常的情况一样，问题在于细节，要让稀疏的 MoE 正常工作就需要确保这些细节完全正确。

在本文中，我们将深入探讨稀疏MoE领域的一个核心贡献，即Switch Transformer（Fedus等人，2022年），它首次展示了利用这项技术实现了令人印象深刻的扩展特性，在Transformer模型的训练中实现了7倍以上的加速。我们将讨论以下内容：

• 硬路由：仅针对每个token执行单个专家模型所带来的有利扩展特性，
• Switch Transformer 架构：MoE 如何适应更广泛的 Transformer 架构背景，
• token动态路由：如何使用容量因子来权衡计算效率和建模精度，以及
• 实验结果：Switch Transformer的令人印象深刻的扩展特性。

 让我们开始吧。

 硬路由

回顾一下，MoE 的关键思想是使用专家的线性组合 E(x) 对给定输入 x 的输出 y 进行建模，每个专家的权重由门控 G(x) 控制，

其中门控是输入x与可学习的权重矩阵W相乘后的softmax函数：

因此，在训练 MoE 模型时，学习目标有两个：

1. 专家将学习将给定的输入处理成最佳的输出（即预测）
2. 门控将通过学习矩阵W来将正确的训练样本分配给正确的专家。

MoE 的这种原始表述可以追溯到30多年前 Geoffrey Hinton 在 90 年代的原创作品，如今已被称为软路由（soft routing）。之所以称为“软”，是因为尽管不同的专家可能具有非常不同的权重（其中一些几乎不明显），我们仍然将所有专家的输出组合在最终结果中，无论他们的贡献多么小。

相比之下，在硬路由（hard routing）中，我们只对由门控决定的单个最合适的专家进行向前传递，也就是说，我们近似地

其中I是使G取得最大值的索引。

这种方法的动机是在计算成本上进行建模精度的权衡：如果一个专家的权重只有0.01，那么是否真的值得对该专家进行向前传递呢？

顺便说一句，硬路由实际上是 top-k 路由的一种特殊情况，最初由 Shazeer 等人 (2017) 提出，其中 k=1。虽然Shazeer等人假设k=1在实践中可能效果不佳，因为这会严重限制通过专家反向传播的梯度，但Switch Transformer证明并非如此。

（注意，“gate”和“router”这些术语通常可以互换使用——在稀疏MoE的背景下，它们实际上意味着相同的东西。）

具有 5 个专家和 3 个token的硬路由插图，每个token都有一个 3 维嵌入。此处，路由器将token 1 分配给专家 1，将token 2 分配给专家 4，将token3 分配给专家 2。图片来自 Fedus et al 2022。

Switch Transformer架构

在常规的（密集）Transformer块中，我们将一个前馈神经网络（FFN）层堆叠在一个自注意力层之上，并在它们之间添加残差连接。在高层次上，自注意力层选择模型关注的输入，而FFN层处理该输入并将输出传递到下一个块，依此类推。例如，BERT-Large是一个将24个这些Transformer块堆叠起来的Transformer模型。

这是这个密集的 Transformer 块的可视化效果：

密集的变压器块。图片来源：Fedus 等人 2022。

在 Switch Transformer 中，我们用多个 FFN“专家”替换单个 FFN 模块，并让硬路由器向专家分配token，如下所示：

Switch Transformer 中的稀疏Transformer块。图片来源：Fedus 等人 2022。

最重要的是，这两种架构具有相同的计算复杂度，但与前者相比，后者实际上具有4倍的建模容量（即神经元数量）！这要归功于稀疏 MoE 的硬路由：虽然我们现在有4个FFN而不是一个，但对于任何给定的token，只有其中一个会处于活动状态。

token动态路由和容量因子

到目前为止，我们只考虑了路由器如何在专家之间分配token，但没有具体说明这些专家位于何处。在实践中，我们通常将专家分布在多台机器上，这是一种模型并行化形式，也称为专家并行化。

专家并行化有一个重要的物理限制：我们不能将比专家所在机器的内存允许的更多token路由到该专家！

token动态路由作为容量因子的函数。如果因子为 1，我们就必须丢弃一个token（红色箭头）。在容量因子为1.5的情况下，我们需要添加过多的填充（白色矩形）。来自 Fedus 等人 2022

在这种背景下，我们将专家容量定义为

capacity = f x T/E

其中T是token数量，E是专家总数，f是一个被称为容量因子的自由超参数。（如果我们使用k>1的top-k路由，则还会在方程的右侧添加k的因子——然而，在这里我们考虑的是使用硬路由的Switch Transformer，所以我们将使用k=1。）

例如，对于T=6，E=3和f=1，我们允许路由器向每个专家发送最多 2 个token，如上图左侧面板所示。如果我们发送更多token，我们就需要丢弃额外的token，如果我们发送较少token，我们就需要对输入进行填充，以确保一致性——毕竟，在GPU上可以运行的计算图必须是静态的，不能是动态的。

因此，容量因子f引入了一个权衡：如果太大，我们会通过过多的填充浪费计算资源（即上图中的白色矩形）；如果太小，我们会因为token丢失而牺牲模型性能（如上图中的红色箭头所示）。

在Switch Transformer论文中，作者发现较低的容量因子可以获得更好的性能，例如，将 f 从 2.0 降低到 1.0，在 100k 步后，对数困惑度从 -1.554 提高到 -1.561。这表明我们因丢弃token而遭受的惩罚并不像计算资源利用不足所带来的惩罚那么严重。或者更简单地说：最好优化资源利用率，即使以丢弃token为代价。

Switch Transformer的可扩展性特征

最重要的是，只需增加专家数量，Switch Transformer可以以近似恒定的计算复杂度进行扩展。这是因为更多的专家不会导致更多的专家向前传递——得益于硬路由——而且与整个模型架构相比，添加更多专家时需要由门控完成的额外计算可以忽略不计，特别是与计算量大的 Transformer 块相比。

下图显示了通过用128个稀疏专家替换单个FFN模块，作者能够达到与T5语言模型相同的性能，但速度快了7倍！所有这些模型都在相同的机器上进行训练，使用32个TPUv3核心，并且每个训练样本具有相等的FLOPs。

Switch Transformer通过硬路由实现了7倍的加速，同时保持了恒定的计算复杂度。

乍一看，这似乎像是魔法。在相同的FLOPs下，怎么可能更快呢？这就好像一辆汽车以相同的速度行驶，却比原来快7倍到达目的地！

答案就是，Switch Transformer 利用了稀疏性，特别是稀疏 MoE。虽然我们通过增加专家数量增加了建模容量，但由于硬路由，我们保持了恒定的FLOPs，也就是说，在每个训练迭代中我们并没有执行所有的专家，而是只选择适合每个token的最佳专家进行计算。

最后

 让我们回顾一下：

• 稀疏MoE是一项突破性的技术，因为它使我们能够以近似O(1)的计算复杂度扩展建模容量，并实现了Google的Switch Transformer、OpenAI的GPT-4、Mistral AI的Mixtral-8x7b等重大突破。
• 硬路由意味着我们只对最合适的专家运行向前传递，而不是所有专家，从而节省了FLOPs。
• Switch Transformer 用使用硬路由的稀疏 MoE 层替换了 T5 Transformer 中的 FFN 层。
• 容量因子决定了每个专家允许的token数量，是在token丢弃和机器利用率之间权衡的一个手段。实验证明，即使以牺牲token为代价，最大化机器利用率也是正确的目标。
• 由于其良好的可扩展性特征，与T5（密集）Transformer模型相比，Switch Transformer在训练速度上实现了7倍的加速。

虽然Switch Transformer已经在语言模型领域取得了突破，但我相信我们才刚刚开始看到它对行业的全面影响。稀疏 MoE 带来的强大扩展特性有可能为跨领域的应用带来巨大的建模改进，而不仅仅局限于LLMs。这是机器学习领域激动人心的时刻！