专家混合（MoE）解释

专家混合（MoE）使得 Transformer 模型的预训练速度显著快于密集模型，并且相对于其总参数量实现更快的推理速度。这通过将密集前馈网络（FFN）层替换为稀疏 MoE 层来实现，后者仅为给定的输入 token 激活少量的“专家”。

MoE 模型用稀疏 MoE 层与路由机制的组合取代标准的 FFN 层。该架构包括两个主要组件：

  • 稀疏 MoE 层：不再是单一的密集层，这些层包含多个“专家”，每个专家通常是一个神经网络（常为 FFN）。这些专家可以以层级结构组织，专家本身也可以是 MoE。
  • 门控网络（路由器）：一个学习得到的路由器决定将哪些 token 发送到哪个专家。路由器与网络的其他部分一起进行预训练，并且可以配置为将一个 token 发送到一个或多个专家。

虽然 MoE 提供了计算效率，但也带来了特定的权衡。它们需要大量显存，因为所有专家必须加载到内存中，即使在单次前向传播中只有少数专家是活跃的。例如，Mixtral 8x7B 需要足够的显存来容纳一个 47B 参数的模型，尽管在每个 token 使用两个专家时，其推理计算（FLOPs）更接近 12B 模型。

稀疏性与路由机制

稀疏性使模型能够在规模扩大时不成比例地增加计算量，方法是使用条件计算——仅在每个示例中激活系统的部分。

门控函数

路由由门控网络 $G$ 处理，决定激活哪些专家 $E$。在传统设置中使用 softmax 函数。更高级的方法，例如 Noisy Top-k Gating，引入可调噪声并选择前 $k$ 个值，以确保在训练期间更好的探索和负载均衡。

负载均衡与专家容量

为了防止出现“赢家通吃”的情况，即少数专家被过度使用而其他专家被忽视，MoE 采用：

  • 辅助损失：一种额外的损失函数，鼓励路由器在所有专家之间均匀分配 token。
  • 专家容量：限制单个专家可处理 token 数量的阈值。如果专家达到容量，token 可能被丢弃或通过残差连接发送到下一层。

MoE 实现的演进

GShard 与 Switch Transformers

Google 的 GShard 通过使用 top-2 门控并引入随机路由（第二个专家按其权重比例选择），将 Transformer 的规模扩展到超过 6000 亿参数。

Switch Transformers 进一步简化了这一思路，引入单专家策略。通过将每个 token 路由到唯一的专家，Switch Transformers 减少了路由器计算，将专家批量大小减半，并降低了通信成本，同时保持了质量。该架构实现了相较于 T5-XXL 的 4 倍预训练加速。

稳定性与专门化

训练的稳定性通过诸如 Router z-loss（在 ST-MoE 中引入）的技术来管理，该技术对进入门控网络的大 logits 进行惩罚，以降低舍入误差。

对专家专门化的研究表明，编码器专家往往专注于浅层概念（例如标点或专有名词），而解码器专家的专门化程度较低。在多语言设置中，由于负载均衡的需求，专家通常不会专门化到单一语言。

微调与指令微调

稀疏模型比密集模型更容易过拟合。有效的微调策略包括：

  • 更高的正则化：在稀疏层中使用更高的 dropout 率。
  • 选择性冻结：冻结 MoE 层而更新非 MoE 层，可在保持质量的同时加快训练。
  • 超参数调整：稀疏模型通常受益于更小的批量大小和更高的学习率。

最新研究表明，MoE 从 指令微调 中获益显著高于密集模型。在将 MoE 与等价的 T5 进行比较时，MoE 通过指令微调（Flan-MoE vs MoE）获得的性能提升大于密集模型的提升（Flan T5 vs T5）。

计算效率与并行性

并行策略

为处理 MoE 的海量参数，采用了不同的并行技术：

  • 专家并行：专家分布在不同的工作节点上。对于 MoE 层，token 被路由到包含相应专家的工作节点；对于非 MoE 层，则表现为数据并行。

硬件优化

  • 容量因子：在质量与通信成本之间的权衡。更高的容量因子提升质量，但会增加激活的显存占用和设备间通信。
  • 块稀疏操作：如 MegaBlocks 等项目用块稀疏操作取代传统的批量矩阵乘法，以处理不平衡的 token 分配而不丢弃 token，从而实现显著加速。

对比：稀疏 MoE 与 密集模型

特性 稀疏 MoE 密集模型
预训练计算 更高效 / 更快 效率较低
推理速度 更快（相对于总参数） 更慢（相对于总参数）
显存需求 高（必须加载所有专家） 低（相对于总参数）
最佳使用场景 高吞吐量，机器多 低吞吐量，显存受限

开源 MoE 生态系统

多个框架和模型将 MoE 引入开源社区：

  • 框架：MegaBlocks、Fairseq 和 OpenMoE。
  • 模型：Google 的 Switch Transformers（最高 1.6T 参数）、Meta 的 NLLB MoE（54B）以及 Mistral 的 Mixtral 8x7B，后者在推理速度上超越 Llama 2 70B。

Sources