EMO：预训练混合专家模型以实现涌现模块化

来源：https://huggingface.co/blog/allenai/emo 返回文章列表（https://huggingface.co/blog）

Kyle Wiggers 的头像（https://huggingface.co/Ai2Comms）

Ryan Wang 的头像（https://huggingface.co/ryanyxw）

• 如何让模块化涌现？（https://huggingface.co/blog/allenai/emo#how-do-we-get-modularity-to-emerge）
• 基准测试结果（https://huggingface.co/blog/allenai/emo#benchmark-results）
• 专家子集专精于什么？（https://huggingface.co/blog/allenai/emo#what-are-expert-subsets-specializing-to）
• 我们发布的内容（https://huggingface.co/blog/allenai/emo#what-were-releasing）

🧠**模型：**https://huggingface.co/collections/allenai/emo | 📄**技术报告：**https://allenai.org/papers/emo | 💻**代码：**https://github.com/allenai/EMO | 📊**可视化：**https://emovisualization.netlify.app/

EMO 博客文章草稿 ryan - Google Docs-image-1 (1) (https://cdn-uploads.huggingface.co/production/uploads/638e39b249de7ae552d977b5/812CTj_1vPuk8Xrgt5hfi.png)

今天，我们发布 EMO（https://huggingface.co/collections/allenai/emo），一种新的混合专家（MoE）模型，经过端到端预训练，使得模块化结构直接从数据中涌现，而无需依赖人工定义的先验知识。EMO 允许你仅使用一小部分专家（仅占总专家的 12.5%）来完成特定任务，同时保持接近全模型的性能；当所有专家一起使用时，它仍然是一个强大的通用模型。

大型语言模型通常作为整体系统进行训练和部署：一个单一模型被初始化、预训练、微调，并作为一个统一实体提供服务。但应用通常只需要一部分能力，例如代码生成、数学推理或特定领域的知识。随着前沿语言模型的参数规模常达到数万亿，使用和适配全模型对大多数用户来说变得不切实际，并且会带来不必要的计算成本和内存来承载可能根本不需要的参数。

混合专家（MoE）模型似乎是放宽这一约束的自然方式。MoE 不是在每一层使用一个大型前馈网络，而是包含许多较小的网络，称为专家，并且每个输入 token 仅激活一小部分专家。原则上，一个只需要单一能力的任务可以只加载相关的专家。

然而在实践中，现有的 MoE 仍然需要全模型才能良好工作。即使在单个输入中，不同的 token 也常常激活不同的专家，因此一个任务在生成过程中可能会使用所有专家。正如我们在论文中所展示的，这在一定程度上是因为标准 MoE 中的专家往往专精于低层次的词法模式（如介词或标点符号），而非更高层次的领域或能力。结果，小规模的专家子集无法可靠地独立使用。

相反，我们希望 MoE 模型的专家能够组织成连贯的组，这些组可以被选择性地使用和组合。

在预训练中鼓励这一点的一种方法是根据预定义的语义领域（如数学、生物学或代码）将 token 路由到专家。先前的工作，如 BTX 和我们的 FlexOlmo 项目，已经尝试过这种方法。然而，预定义的领域具有重要的局限性。它们需要在预训练语料库中标注领域标签，这可能是模糊且昂贵的，并且可能会向模型的组织方式注入过多的人类偏见。更重要的是，预先固定领域也固定了模型的模块化结构：如果在推理时出现新的领域或能力，则不清楚应该使用哪些专家。

这就是 EMO 的用武之地。

我们展示了 EMO——一个拥有 10 亿活跃参数、总计 140 亿参数（每层 8 个专家活跃，总共 128 个专家）的 MoE 模型，在 1 万亿 token 上训练——支持选择性专家使用：对于给定的任务或领域，我们可以仅使用一小部分专家（仅占总专家的 12.5%），同时保持接近全模型的性能。同时，当所有专家一起使用时，EMO 仍然是一个强大的通用模型。相比之下，同等架构的标准 MoE 在相同数据上训练，当选择性使用其专家子集时，性能会严重下降。

EMO 博客文章草稿 ryan - Google Docs-image-2 (1) (https://cdn-uploads.huggingface.co/production/uploads/638e39b249de7ae552d977b5/CPWUSB64LhBEjMI0Rgg6L.png)

EMO 是一种将模块化作为首要目标进行训练的 MoE。对于给定的领域（例如数学、代码、生物医学），用户可以选择任意大小的小规模专家子集，并保持接近全模型的性能。这便将单一模型转变为可组合的架构，从而实现了灵活部署，并在大型稀疏 MoE 中改善了内存与准确率的权衡。

https://huggingface.co/blog/allenai/emo#how-do-we-get-modularity-to-emerge 如何让模块化涌现？

在 MoE 中，一个被称为路由器的简单网络决定每个 token 激活哪些专家。我们希望路由器学会，来自相似领域的 token 应激活相似的专家子集。我们的关键观察是：来自同一文档的 token 通常来自同一领域。因此，我们利用文档边界作为弱监督信号：在训练过程中，同一文档中的所有 token 被限制从共享的专家池中选择其活跃专家。

EMO 博客文章草稿 ryan - Google Docs-image-3 (1) (https://cdn-uploads.huggingface.co/production/uploads/638e39b249de7ae552d977b5/slqCFgfncHGGK1lErZNvl.png)

标准 MoE 与 EMO 的训练对比（k = 2, n = 10，为简洁起见省略共享专家）。(左) 在标准 MoE 中，每个 token 独立选择其 top-k 专家。跨 token 来看，所有专家都被使用。(右) 在 EMO 中，路由器首先为每个文档选择一个专家子集，所有 token 都被限制在该子集内进行路由。这强制了文档内专家使用的一致性，从而鼓励专家群组形成领域专精。

例如，在一个拥有 10 个专家、每个 token 有 2 个活跃专家的 MoE 中，文档中的所有 token 被限制在同一个 4 专家池内进行路由，如上图所示。该池由路由器自身选择：我们平均整个文档中所有 token 的路由器专家偏好，然后选择最常用的专家作为文档的共享池。不同的文档可以使用不同的池，从而允许重复出现的专家群组直接从训练数据中涌现。

实现该系统时需要考虑几个方面：

**负载均衡。**一个技术挑战是负载均衡。在标准 MoE 训练中，负载均衡目标用于防止模型崩溃到仅使用少数几个专家。乍一看，这似乎与 EMO 的训练目标相冲突：我们明确限制每个文档仅使用一部分专家。

这种冲突源于负载均衡通常应用的尺度。在许多 MoE 实现中，负载均衡是局部计算的，通常在一个仅包含少量文档的微批次内。这种局部目标可能促使同一文档内的 token 分散到许多专家中，直接与 EMO 保持文档内专家使用一致的目标相悖。

为了解决这个问题，我们在跨许多文档的全局尺度上应用负载均衡。在这个更大的尺度上，这两个目标变得互补：EMO 鼓励同一文档内的 token 使用一致的专家池，而全局负载均衡则鼓励不同文档共同覆盖所有专家。在实践中，我们发现全局负载均衡对于稳定训练至关重要。

**文档池大小。**文档池大小控制模块化约束的严格程度。较小的池迫使同一文档内的 token 共享更紧密的专家集，从而鼓励更强的模块化；较大的池则给模型更多灵活性，但会削弱约束。

我们并没有固定一个池大小，而是在训练过程中随机采样。这可以防止 EMO 过拟合到单一的专家子集大小，并使其在推理时支持不同的专家子集大小。

https://huggingface.co/blog/allenai/emo#benchmark-results 基准测试结果

在通用基准测试上，EMO 与标准 MoE 模型的性能相当，这表明模块化目标并未以牺牲全模型性能为代价。然而，更重要的问题是，当我们只保留一部分专家时，模型是否仍然有效？在这种设置下，我们通过根据少量任务验证数据上的路由使用情况对专家进行排序，从而构建任务特定的专家子集，保留最常用的专家并丢弃其余部分。

下图显示，EMO 在选择性专家使用下保持稳健。当我们仅保留 25% 的专家（32 个专家子集）时，EMO 在所有基准测试上的绝对性能仅下降约 1%；即使仅保留 12.5% 的专家（16 个专家子集），总体下降也仅约 3%。这在微调前后都成立。相比之下，匹配的标准 MoE 随着专家子集变小而急剧下降，在最小专家子集设置下往往接近或低于随机性能。

EMO 博客文章草稿 ryan - Google Docs-image-4 (https://cdn-uploads.huggingface.co/production/uploads/638e39b249de7ae552d977b5/ki3pHaOktjGM1qI4JBeYG.png)

此外，我们展示了为任务选择正确专家出奇地廉价：一个带有少样本演示的示例就足以识别出一个性能与使用完整验证集选择的模块相当的模块。而且 EMO 并不局限于任何特定的选择方法：它与现有的专家剪枝方法（如 Easy-EP）配合良好，两者相互补充。

EMO 博客文章草稿 ryan - Google Docs-image-5 (1) (https://cdn-uploads.huggingface.co/production/uploads/638e39b249de7ae552d977b5/NMSuavox_S24mxIBovCMd.png)

较小的 130B token 设置。在 16 个 MMLU 类别上，不同内存预算下的平均性能。EMO 专家子集推动了内存-准确率权衡的帕累托前沿，优于标准 MoE 甚至从头训练的固定预算模型。

https://huggingface.co/blog/allenai/emo#what-are-expert-subsets-specializing-to 专家子集专精于什么？

为了了解 EMO 在训练后实际学到了什么，我们对 12K 预训练文档的前 100 个 token 的路由器激活进行了聚类。与标准 MoE 的差异非常明显。

EMO 的 token 聚类对应诸如“健康、医疗与保健”、“新闻报道”、“美国政治与选举”以及“电影与音乐”等内容。标准 MoE 则产生诸如“介词”、“专有名词”、“系词动词”或“定冠词”等聚类。在 EMO 中，来自同一文档的 token 大多落在同一个聚类中；而在标准 MoE 中，它们则分散在许多聚类中。

这种对比在一个例子中最容易看出。以一篇健康文章为例：在 EMO 中，几乎每个 token 都会路由到“健康、医疗与保健”聚类。在标准 MoE 中，最顶部的聚类是“所有格与定冠词”；模型会将这篇文章与所有其他恰好使用了单词“the”或“your”的文本归为一类，而不管这些文本的内容是什么。

EMO 博客文章草稿 ryan - Google Docs-image-6 (1) (https://cdn-uploads.huggingface.co/production/uploads/638e39b249de7ae552d977b5/FG27xp8oJXUZW9cTju235.png)

在 1T token 上训练的 MoE 的预训练数据 token 聚类图。EMO 的聚类对应语义上有意义的领域，同一文档的 token 大部分聚合在一起。标准 MoE 训练产生基于表层或句法特征的聚类，文档 token 分散在多个聚类中。

由于 EMO 形成的模块映射到语义领域而非表层特征，你可以选择一个小型专家子集并仍然拥有一个功能完整的模型：该群组对应着真实的能力。

你可以通过我们的交互式可视化（https://emovisualization.netlify.app/）亲自探索聚类结果。

https://huggingface.co/blog/allenai/emo#what-were-releasing 我们发布的内容

我们正在发布完整的 EMO 训练模型（https://huggingface.co/collections/allenai/emo）、一个在相同数据上训练的匹配标准 MoE 基线模型，以及训练代码（https://github.com/allenai/EMO）。我们希望这些成果对其他研究 MoE 中涌现模块化的团队有所帮助。

还有更多工作要做。EMO 是朝向使大型稀疏模型更加模块化迈出的早期一步，但仍有许多问题悬而未决：如何更好地选择和组合专家子集，如何更新模块而不破坏全模型，以及如何利用模块化结构实现更好的可解释性和控制。发布这些模型应该有助于社区研究这些问题，并朝着更易于部署、适配、检查和组合的模块化语言模型前进。