MoE++: 颜水成团队与北大袁粒团队推出新一代 MoE 架构，专家推理速度倍增，性能全面提升

作者：新消费日报

2024-10-21
云南
本文字数：2891 字
阅读完需：约 9 分钟

颜水成领衔的昆仑万维 2050 研究院与北大袁粒团队联合提出了 MoE++，相较于传统的 MoE，具有更快的推理速度和更高的性能。

MoE++是一个通用框架，能够与任何现有的 MoE 模型无缝集成。其核心创新在于引入了“零计算量专家”，这一设计带来了三大优势：

(1) 降低计算成本：MoE++允许每个 Token 使用可变数量的 FFN 专家，甚至可以完全跳过当前的 MoE 层。

(2) 提升性能：通过减少简单 Token 所需的 FFN 专家数量，MoE++使更多专家能够专注于复杂的 Token，释放出比传统 MoE 更大的性能潜力。

(3) 易于部署：零计算量专家的参数极小，可以在每个 GPU 上同时部署所有的零计算量专家，避免了分布式 FFN 专家部署带来的通信开销和专家负载不均的问题。

此外，MoE++还让每个 Token 在选择专家时参考前一层的路由路径。

实验结果表明，在 0.6B 到 7B 参数规模的 LLMs 上，MoE++在相同模型大小的情况下，相比传统 MoE，性能更优，同时实现了 1.1 到 2.1 倍的专家吞吐速度。该模型权重已开源。

● 论文地址： https://arxiv.org/abs/2410.07348

● GitHub 地址：https://github.com/SkyworkAI/MoE-plus-plus

● Huggingface 地址：https://huggingface.co/Chat-UniVi/MoE-Plus-Plus-7B

1729505202140276.png

图 1 丨 MoE++和普通 MoE 的对比

MoE++是如何做到的？

现有的大多数混合专家（MoE）方法通常为所有 Token 激活固定数量的 FFN 专家。

在许多研究中，每个 Token 会选择 Top-2 FFN 专家，并将它们的输出加权合成为下一层的输入。然而，并非所有 Token 的预测难度都是相同的。

例如，像逗号等简单符号可能只需要一个 FFN 专家来处理。甚至在某些特殊情况下，某些 Token 如果与当前 MoE 层的专家不匹配，绕过该层而不选择 Top-2 FFN 专家反而可能更为高效。

基于这一见解，研究团队认为，现有 MoE 使用的固定混合机制可能导致训练和推理效率下降，同时限制模型性能。

为了在提升速度的同时增强性能，研究团队提出了一种通用的异构 MoE 框架，称为 MoE++。

具体来说，团队引入了三种零计算量专家：

Zero 专家，输出空向量；

Copy 专家，将输入直接作为输出；

Constant 专家，用可训练的向量替代输入作为输出。

如图 1 所示，与传统 MoE 方法不同，MoE++允许每个 Token 使用可变数量的 FFN 专家，接受恒定向量的替换，甚至完全跳过当前的 MoE++层。

这种异构结构通过扩大网络的组合空间，提升了模型的拟合能力，并显著降低了计算成本。此外，研究团队还将前一层的路由分数整合到当前层的专家选择中，使 Token 在选择专家时能够参考其先前的路由路径，从而实现更稳定的专家分配。

研究团队认为，新设计的 MoE 架构应满足以下标准：

(1) 设计应尽量简化，以高效处理简单的 Token；

(2) 为了确保公平比较，新增参数应保持在可忽略的范围。

在这些原则的指导下，研究团队引入了零计算量专家，每个专家仅执行最基本的操作。如图 2（a）所示，团队设计了三种类型的零计算量专家：Zero 专家、Copy 专家和 Constant 专家，分别对应丢弃、跳过和替换操作。

1729505213189537.png

图 2 丨 MoE++的核心组成部分

Zero 专家。最简单的零计算量专家是丢弃当前输入的 Zero 专家。本质上，Zero 专家的存在可以将 Top-2 MoE++层降级为 Top-1 MoE++层。具体来说，当 Zero 专家被激活时，Top-2 MoE++层的输出将等同于另一个专家的单独输出。这样，Zero 专家的引入提升了模型在处理简单 Token 和复杂 Token 时的灵活性。

Copy 专家。Copy 专家直接将输入作为输出，直观上相当于跳过当前 MoE++层。具体而言，当输入 Token 与现有专家的匹配较差时，选择绕过 MoE++层可能更为有利。

Constant 专家 Constant 专家通过可训练向量替换输入 Token。然而，完全替换会导致输入 Token 信息的丢失。为此，研究团队引入了可训练的权重矩阵，用于动态预测替换的比例。由于 Constant 专家的计算开销极小，因此仍被归类为零计算量专家。

路由分数残差。由于 MoE++包含异构专家，因此与普通 MoE 相比，路由器的设计变得更加关键。为此，如图 2（b）所示，研究团队提出了一种路径感知路由器，它在选择合适的专家时考虑了前一层所采用的路径。具体来说，MoE++将前一层的路由分数通过一个可学习的转换矩阵合并到当前层的专家选择中。这些路由分数残差使每个 Token 在选择专家时可以考虑其先前的路由路径。

为什么 MoE++比 MoE 更好（Why）？

（1）灵活的计算量分配。MoE++通过为简单 Token 分配较少的 FFN 专家，优化了计算资源的分配，从而使更多 FFN 专家能够专注于处理更具挑战性的 Token。

正如图 3 所示，研究团队发现，在 MoE++中，动词激活的 FFN 专家数量最多，其次是名词，而拆分后的词片激活的 FFN 数量最少。这表明，MoE++能够让语义较少的 Token 使用更少的 FFN 专家，从而释放更多专家来处理语义更丰富的 Token。因此，MoE++不仅减少了计算开销，还提升了整体性能。

1729505229154190.png

图 3 丨在 MoE++中不同 Token 所需要的平均 FFN 专家数量

（2）稳定的路由。MoE++将前一层的路由分数合并到当前层的专家选择中。这些路由分数残差使每个 Token 在选择专家时考虑其先前的路由路径。

如图 4 所示，路由分数残差有效地建立了不同 MoE++层之间的联系，减小了路由分数的方差。同时，路由分数残差不改变路由分数的均值和取值范围。因此，路由分数残差有助于在 MoE++中实现异构专家架构的稳定路由。

1729505240778653.png

图 4 丨路由分数残差对路由分数分布的影响

（3）更低的计算复杂度。如下表所示，MoE++具有比普通 MoE 更低的理论计算复杂度。

1729505244675979.png

实验结果

从 0.6B 的参数量逐渐扩展到 7B 参数量的大量实验结果表明，MoE++方法明显优于普通 MoE 方法。

与相同大小的普通 MoE 模型相比，MoE++的专家吞吐量提高了约 15% ~ 111%，同时具有更高的性能。

1729505250226893.png

1729505257783992.png

现有的 LLMs 模型通常需要大量的训练预算，比如 OpenMoE-8B/32E 使用 1.1T Tokens，TinyLlama-1.1B 使用 3T Tokens。

研究人员也将 MoE++模型的训练预算扩展到 1T Tokens。研究人员发现 MoE++模的性能与具有 2 到 3 倍激活参数的稠密模型相当。

值得注意的是，MoE++优于 OpenMoE-8B/32E，这是一个从零开始训练的更大的 MoE 模型，使用更多的训练 Tokens。这些结果表明，MoE++框架是一种很有前途的 LLMs 框架方案。

1729505263612320.png

任务级专家负载分布的可视化

研究人员探索了 MoE++模型中跨不同任务的专家负载分。这些可视化揭示了几个有趣的发现:

（1）专家负载在不同层之间存在相关性，特别是在相邻层之间。例如，当第 j 层激活很大比例的 FFN 专家时，第 j + 1 层很可能也会以同样大的比例激活 FFN 专家。

（2）与中间层相比，浅层和最后一层的专家分配模式在不同任务之间的差异更大。这表明该模型主要通过其浅层和最终层而不是中间层来适应不同的任务。未来的工作可以集中在这些层中设计更复杂的结构，以增强模型对不同任务的适应性。

（3）不同任务中每个 Token 激活的 FFN 专家数量存在显著差异，但并不一定是更简单的任务激活更少的 FFN 专家。例如，ARC Challenge 任务通常比 ARC Easy 任务激活更多的 FFN 专家。这些结果表明，MoE++模型根据知识内容和 Token 级别的复杂性来分配专家，而不是根据整体任务难度来分配专家。

（4）在所有专家类型中，Zero 专家的平均激活次数最高，更简单的任务显示出更高的平均激活次数。例如，ARC Easy 任务比 ARC Challenge 任务激活更多的零专家。这表明 Zero 专家的激活水平可能可以作为模型任务难度的一个指标。

（5）在 MoE++模型的所有层中，不同任务主题的专家分配差异显著，这表明 MoE++模型通过采用不同的专家分配模式来处理不同主题的任务。