EMO：为涌现模块化预训练 Mixture of Experts | Ai2

今天，我们发布 EMO，这是一个新的 mixture-of-experts（MoE）模型，以端到端方式预训练，使模块化结构直接从数据中涌现，而不依赖人工定义的先验。对于给定任务，EMO 只需使用一小部分 experts——总量的 12.5%——就能保持接近完整模型的性能；当所有 experts 一起使用时，它仍然是一个强大的通用模型。

大型语言模型通常作为单体系统进行训练和部署：一个模型被初始化、预训练、fine-tune，并作为统一实体提供服务。但应用往往只需要一部分能力，例如 code generation、mathematical reasoning 或领域特定知识。随着 frontier language models 的参数规模经常达到数万亿级，使用和适配完整模型对大多数用户来说变得不切实际，并且会带来不必要的计算成本和内存开销，用于承载可能根本不需要的参数。

Mixture-of-experts（MoE）模型似乎是放松这一约束的自然方式。MoE 不是在每一层使用一个大型 feedforward network，而是包含许多较小的网络，称为 experts，并且只为每个输入 token 激活其中一小部分。原则上，一个只需要某种能力的任务可以只加载相关 experts。

然而在实践中，现有 MoE 仍然需要完整模型才能表现良好。即使在单个输入内部，不同 token 也常常会激活不同 experts，因此一个任务在生成过程中最终可能使用所有 experts。正如我们在论文中所示，这种情况部分是因为标准 MoE 中的 experts 往往专门处理介词或标点等低层级词汇模式，而不是更高层级的领域或能力。因此，较小的 expert 子集本身并不总是可靠可用。

我们希望得到的是这样一种 MoE 模型：它的 experts 会组织成连贯的组，并且这些组可以被选择性使用和组合。

在预训练期间鼓励这一点的一种方法，是根据预定义的语义领域（如 math、biology 或 code）将 token 路由到 experts。BTX 以及我们的 FlexOlmo 项目等先前工作尝试过这种方法。然而，预定义领域有重要局限。它们要求在整个预训练语料中提供领域标签，而这些标签可能含糊且获取成本高；它们还可能把过多人工偏见注入模型被允许组织自身的方式。更重要的是，预先固定领域也会固定模型的模块化结构：如果在 inference 时出现一个新的领域或能力，应该使用哪些 experts 并不显然。

这正是 EMO 的用武之处。

我们表明，EMO——一个 1B-active、14B-total-parameter（8-expert active、128-expert total）的 MoE，在 1 trillion tokens 上训练——支持选择性使用 experts：对于给定任务或领域，我们只需使用一小部分 experts（仅占总 experts 的 12.5%），就能保留接近完整模型的性能。同时，当所有 experts 一起使用时，EMO 仍然是一个强大的通用模型。相比之下，在相同数据上训练、架构相同的标准 MoE，在选择性使用其 expert 子集时性能会严重下降。

我们如何让模块性涌现？

在 MoE 中，一个称为 router 的小型网络决定每个 token 激活哪些 experts。我们希望 router 学会：来自相似领域的 token 应该激活相似的 expert 子集。我们的关键观察是，同一文档中的 token 通常来自同一领域。因此，我们使用文档边界作为弱监督信号：在训练期间，一个文档中的所有 token 都被限制为从共享的 expert pool 中选择其 active experts。

例如，在一个总共有 10 个 experts、每个 token 激活 2 个 experts 的 MoE 中，同一文档中的所有 token 都被限制在同一个包含 4 个 experts 的 pool 内路由，如上图所示。这个 pool 由 router 自己选择：我们对文档中所有 token 的 router expert 偏好取平均，然后选择使用最多的 experts 作为该文档的共享 pool。不同文档可以使用不同的 pool，从而让反复出现的 expert 组直接从训练数据中涌现出来。

实现该系统时有几个需要考虑的点：

Load balancing。 一个技术挑战是 load balancing。在标准 MoE 训练中，load-balancing objective 用于防止模型坍缩到只使用少数 experts。乍看之下，这似乎与 EMO 的训练目标冲突：我们明确限制每个文档只使用 experts 的一个子集。

冲突来自通常应用 load balancing 的尺度。在许多 MoE 实现中，load balancing 是局部计算的，往往在只包含少量文档的 micro-batch 内完成。这种局部目标可能推动同一文档内的 token 分散到许多 experts 上，直接对抗 EMO 让一个文档内 expert 使用保持一致的目标。

为了解决这个问题，我们在大量文档上全局应用 load balancing。在这个更大的尺度上，两个目标变得互补：EMO 鼓励同一文档中的 token 使用连贯的 expert pool，而全局 load balancing 鼓励不同文档合起来覆盖所有 experts。实践中，我们发现全局 load-balancing 对稳定训练很重要。

Document pool size：document pool size 控制模块化约束的严格程度。较小的 pool 会迫使同一文档中的 token 共享更紧的 experts 集合，从而鼓励更强的模块性；较大的 pool 则给模型更多灵活性，但会削弱约束。

我们并不固定一个 pool size，而是在训练期间随机采样。这样可以防止 EMO 过拟合到单一子集大小，并使其能够在 inference 时支持不同大小的 expert 子集。

Benchmark 结果

在通用 benchmark 上，EMO 的性能与标准 MoE 模型相当，说明模块化目标并没有以完整模型性能为代价。然而，更重要的问题是，当我们只保留一个 expert 子集时，模型是否仍能工作。在这种设置下，我们通过在少量任务验证数据上根据 experts 的路由使用情况对其排序，保留使用最多的 experts，并丢弃其余 experts，来构建任务特定的 expert 子集。

下图显示，EMO 在选择性使用 experts 时仍然稳健。当我们只保留 25% 的 experts（32 expert subset）时，EMO 在所有 benchmarks 上的绝对性能只下降约 1%；即使只保留 12.5% 的 experts（16 expert subset），整体下降也只有约 3%。这一点在 fine-tuning 前后都成立。相比之下，匹配的标准 MoE 随着 expert 子集变小而急剧退化，在最小 expert 子集设置中，性能往往接近或低于随机水平。

此外，我们表明，为任务选择正确 experts 的成本出人意料地低——一个带有 few-shot demonstrations 的样例就足以识别出一个模块，其表现可与使用完整验证集选择的模块相当。而且 EMO 并不绑定任何特定选择方法：它可以与 Easy-EP 等现有 expert-pruning 方法良好配合，两者相互补充。

Expert 子集专门化到了什么？

为了观察 EMO 在训练后实际学到了什么，我们对 12K 个预训练文档中前 100 个 token 的 router activations 进行了聚类。它与标准 MoE 的差异非常明显。

EMO 的 token clusters 对应的是 Health, Medical & Wellness、News Reporting、美国 Politics & Elections、Film & Music 等类别。标准 MoE 产生的 clusters 则是 Prepositions、Proper Names、Copula Verbs 或 Definite Articles 这类类别。在 EMO 中，给定文档中的 token 大多落入同一个 cluster；在标准 MoE 中，它们最终分散到许多 cluster 中。

这种差异在单个示例上最容易看出。以一篇健康类文章为例——在 EMO 中，几乎每个 token 都会被路由到 Health, Medical & Wellness cluster。在标准 MoE 中，顶部 cluster 是 Possessives & Definite Articles；模型会把这篇文章与所有其他碰巧使用 the 或 your 的文本归为一组，而不管那些文本的主题是什么。

因为 EMO 形成的模块映射到语义领域，而不是表层特征，你可以选择一个小的 expert 子集，仍然得到一个可工作的模型——这个组对应着真实能力。

你可以在我们的交互式可视化中自行探索聚类结果。

我们发布了什么

我们发布了 完整的 EMO 训练模型、一个在相同数据上训练的匹配标准 MoE baseline，以及 training code。我们希望这些产物能对其他研究 MoE 中 emergent modularity 的团队有所帮助。

还有更多工作要做。EMO 是让大型 sparse models 更模块化的早期一步，但仍有许多问题有待回答：如何更好地选择和组合 expert 子集，如何在不扰动完整模型的情况下更新模块，以及如何利用模块化结构提升 interpretability 和 control。发布这些模型应能帮助社区研究这些问题，并推动构建更易部署、适配、检查和组合的模块化语言模型。

订阅以接收关于最新 Ai2 新闻的月度更新。