MoE 模型中的动态路由方法

24年3月来自北大等机构的论文“Harder Tasks Need More Experts: Dynamic Routing in MoE Models”。

混合专家 (MoE) 模型的动态专家选择框架，旨在通过根据输入难度调整激活专家的数量来提高计算效率和模型性能。与传统的 MoE 方法不同，传统方法依赖于固定的 Top-K 路由，无论输入的复杂性如何，都会激活预定数量的专家，而动态路由方法则根据每个输入的专家选择置信度动态选择专家。这可以更有效地利用计算资源，为需要高级推理的复杂任务激活更多专家，为简单任务激活更少专家。该模型将更多专家分派到需要复杂推理技能的任务，如 BBH（Big Bench），证实了它能够根据输入的复杂性动态分配计算资源。特别是Transformer 模型不同层所需专家数量的变化。

代码和模型开源 GitHub - ZhenweiAn/Dynamic_MoE

具体来说，首先计算选择专家的概率分布。如果某个专家的最高概率超过预定义的阈值 p，则表明置信度较高，只激活该专家。否则，逐步增加其他专家，直到选定专家的累积概率超过阈值 p。这种方法允许动态选择专家，专家数量根据输入的复杂性进行调整。

如图所示：Top-K路由机制与Top-P路由机制比较。（a）每个token以Top-K路由概率选择固定的K=2个专家。（b）在Top-P路由机制中，每个token选择路由概率较高的专家，直到累积概率超过阈值。

动态路由机制存在一个风险：它可能会为所有专家分配较低的置信度，从而激活更多专家以实现更好的性能。假设 P 是一个均匀分布，将超参数 p 设置为 0.5，那么模型将激活多达一半的专家。这违背了 MoE 框架的初衷：高效地扩展模型。

为了防止动态路由用过多的参数作弊并失去选择性选择专家的能力，对 P 引入了一个约束。希望路由机制选择一小组必要的专家，因此，目标是最小化分布 P 熵，确保每个 token 都能关注尽可能不具体的专家。 动态损失 旨在鼓励路由机制选择最小的必要专家集。

由于其他专家未得到充分利用，从而影响训练效率，因此通常希望不同专家处理的 tokens 数量大致相同。此外，先前的研究（Zuo，2022）表明，MoE 层中均匀地激活专家可以带来更好的性能。为了实现不同专家之间的平衡负载，加入负载平衡损失（Lepikhin，2021；Fedus，2022）。

该模型是一个生成模型，生成的下一个token作为训练目标。最终损失是语言模型损失、动态损失和负载平衡损失的组合。

用 RedPajama（Computer，2023）作为训练数据，它是 LLaMa 数据集的全开源实现。RedPajama 数据来自多种来源，包括常见爬虫 (CC)、C4、github、维基百科、书籍、arxiv 和 Stackexchange。在主要实验中，用 100B 个 tokens 训练所有模型。

该模型架构遵循 LLaMA（Touvron，2023）。用 Llama2 token化器，其词汇量为 32000。除非另有明确说明，将 Transformer 层数设置为 24，隐维度设置为 1024。采用多头注意机制，总共有 16 个注意头，每个头维度为 64。在 FFN 层中使用 SwiGLU（Shazeer，2020）。对于初始化，所有可学习参数都以标准差 0.006 随机初始化。每个 MoE 层有 16 个专家，具有与标准 FFN 相同的初始化参数。在这种配置下，每个密集模型大约有 374M 个参数。每个 MoE 模型总共有 3.5B 个参数。MoE-Top 1 中仅激活 374 个参数，而 MoE-Top 2 中激活 581M 个参数。