EMO：为涌现模块化预训练混合专家
EMO: Pretraining mixture of experts for emergent modularity
AllenAI 发布 EMO，一个 1B-active、14B-total-parameter、128-expert MoE，在 1T tokens 上训练。其用文档边界约束 router 共享 expert 池，并全局 load balancing；保留 25% experts 性能约降 1%，12.5% 约降 3%，同时发布模型、baseline、代码和可视化。
](https://huggingface.co/Ai2Comms)
🧠 Models:https://huggingface.co/collections/allenai/emo | 📄 Tech report:https://allenai.org/papers/emo | 💻 Code:https://github.com/allenai/EMO | 📊 Visualization:https://emovisualization.netlify.app/
今天我们发布 EMO，这是一个新的 mixture-of-experts (MoE) model，采用端到端预训练，使模块化结构直接从数据中涌现，而不依赖人工定义的先验。对于给定任务，EMO 允许你只使用一小部分 experts——总量的 12.5%——同时保持接近完整模型的性能；当所有 experts 一起使用时，它仍然是一个强大的通用模型。
大型语言模型通常作为单体系统进行训练和部署：一个模型被初始化、预训练、微调，并作为一个统一实体提供服务。但应用通常只需要能力的一个子集，例如代码生成、数学推理或特定领域知识。随着前沿语言模型经常达到万亿级参数，使用和适配完整模型对大多数用户来说变得不切实际，并且会为托管那些可能根本不需要的参数带来不必要的计算成本和内存开销。
Mixture-of-experts (MoE) models 看起来像是一种自然的方式，可以放松这一约束。MoE 不在每一层使用一个大型前馈网络，而是包含许多较小的网络，称为 experts，并且只为每个输入 token 激活其中一小部分。原则上，一个只需要某一种能力的任务可以只加载相关 experts。
然而在实践中，现有 MoE 仍然需要完整模型才能良好工作。即使在单个输入内部，不同 token 也经常激活不同 experts，因此一个任务在生成过程中最终可能会用到所有 experts。正如我们在论文中所展示的，这种情况部分是因为标准 MoE 中的 experts 通常专门化于介词或标点等低层词汇模式，而不是更高层的领域或能力。因此，小型 expert 子集并不能可靠地单独使用。
我们希望得到的是 experts 能够组织成一致性群组、并可被选择性使用和组合的 MoE models。
在预训练期间鼓励这一点的一种方式，是根据预定义的语义领域（例如数学、生物或代码）将 token 路由到 experts。BTX 以及我们的 FlexOlmo 项目等先前工作尝试过这种做法。然而，预定义领域有重要限制。它们要求在预训练语料中提供领域标签，而这些标签可能含糊且获取成本高；它们还可能把过多的人类偏见注入模型被允许组织自身的方式。更重要的是，预先固定领域也会固定模型的模块化结构：如果在推理时出现新的领域或能力，很难明确应该使用哪些 experts。
这就是 EMO 的切入点。
我们展示了 EMO——一个在 1 trillion tokens 上训练的 1B-active、14B-total-parameter（8-expert active、128-expert total）MoE——支持选择性使用 expert：对于给定任务或领域，我们可以只使用一小部分 experts（仅占总 experts 的 12.5%），同时保留接近完整模型的性能。与此同时，当所有 experts 一起使用时，EMO 仍然是一个强大的通用模型。相比之下，一个在相同数据上训练、架构相同的标准 MoE，在选择性使用其 expert 子集时会出现严重性能下降。
EMO 是一个以 modularity 作为一等目标训练的 MoE。对于给定领域（例如 math、code、biomedical），用户可以选择任意大小的一小部分 experts，并保留接近完整模型的性能。这把单个模型变成一种可组合架构，为大型稀疏 MoE 带来更灵活的部署方式，并改善内存-准确率权衡。
我们如何让 modularity 涌现？
在 MoE 中，一个称为 router 的小型网络决定每个 token 激活哪些 experts。我们希望 router 学会：来自相似领域的 token 应该激活相似的 expert 子集。我们的关键观察是：同一文档中的 token 通常来自同一领域。因此，我们使用文档边界作为弱监督信号：在训练期间，一个文档中的所有 token 都被限制为从一个共享 expert 池中选择其 active experts。
标准 MoE 与 EMO 的训练对比（k = 2, n = 10，为简洁起见省略 shared experts）。（左）在标准 MoE 中，每个 token 独立选择其 top-k experts。跨 token 来看，所有 experts 都会被使用。（右）在 EMO 中，router 首先为每个文档选择一个 expert 子集，所有 token 都被约束在该子集内路由。这会强制文档内部保持一致的 expert 使用方式，从而鼓励 expert 群组形成领域专门化。
例如，在一个总共有 10 个 experts、每个 token 有 2 个 active experts 的 MoE 中，如上图所示，一个文档中的所有 token 都被限制在同一个包含 4 个 experts 的池内路由。这个池由 router 自身选择：我们对文档中所有 token 的 router expert 偏好取平均，然后选择使用最多的 experts 作为该文档的共享池。不同文档可以使用不同的池，从而让反复出现的 expert 群组直接从训练数据中涌现。
实现该系统时需要考虑几点：
Load balancing。 一个技术挑战是 load balancing。在标准 MoE 训练中，load-balancing 目标用于防止模型坍缩到只使用少数 experts。乍看之下，这似乎与 EMO 的训练目标冲突：我们明确限制每个文档只使用 expert 的一个子集。
冲突来自通常应用 load balancing 的尺度。在许多 MoE 实现中，load balancing 是局部计算的，通常在只包含少量文档的 micro-batch 内进行。这个局部目标可能会推动同一文档内的 token 分散到许多 experts 上，这与 EMO 希望在文档内保持 expert 使用一致的目标直接相反。
为解决这一点，我们在许多文档的全局范围内应用 load balancing。在这个更大的尺度上，两个目标变得互补：EMO 鼓励同一文档内的 token 使用一个一致的 expert 池，而全局 load balancing 鼓励不同文档整体覆盖所有 experts。在实践中，我们发现全局 load-balancing 对稳定训练很重要。
Document pool size。 document pool size 控制 modularity 约束的严格程度。较小的池会迫使同一文档中的 token 共享更紧凑的一组 experts，从而鼓励更强的 modularity；较大的池给模型更多灵活性，但会削弱约束。
我们没有固定一个池大小，而是在训练期间随机采样它。这可以防止 EMO 过拟合到单一子集大小，并使其在推理时支持不同的 expert 子集大小。
Benchmark 结果
在通用 benchmark 上，EMO 与标准 MoE model 的性能相当，表明 modularity 目标并不以完整模型性能为代价。然而，更重要的问题是：当我们只保留 expert 子集时，模型是否仍然能工作。在这种设置下，我们通过在少量任务验证数据上根据路由使用情况对 experts 排序，构建任务特定的 expert 子集，保留使用最多的 experts，并丢弃其余部分。
下图显示，EMO 在选择性使用 expert 时仍然稳健。当我们只保留 25% 的 experts（32 expert 子集）时，EMO 在所有 benchmark 上只损失约 1% 的绝对性能；即使只保留 12.5% 的 experts（16 expert 子集），整体下降也只有约 3%。这一点在 fine-tuning 前后都成立。相比之下，匹配的标准 MoE 会随着 expert 子集变小而急剧退化，在最小 expert 子集设置中经常接近或低于随机性能。
此外，我们表明，为任务选择正确 experts 的成本低得出人意料：一个带有 few-shot demonstrations 的单个样例就足以识别出一个模块，其表现与使用完整验证集选择的模块相当。而且 EMO 不绑定任何特定选择方法：它可以很好地配合 Easy-EP 等现有 expert-pruning 方法，两者相互补充。
较小的 130B-token 设置。在不同内存预算下，对 16 个 MMLU 类别的性能取平均。EMO expert 子集推进了内存-准确率权衡的 Pareto frontier，优于标准 MoE，甚至优于从头训练的固定预算模型。
expert 子集在专门化到什么？
为了观察 EMO 在训练后实际学到了什么，我们对 12K 个预训练文档中前 100 个 token 的 router activations 进行了聚类。它与标准 MoE 的差异非常明显。
EMO 的 token 聚类对应的是 Health, Medical & Wellness、News Reporting、US Politics & Elections 和 Film & Music 等内容。标准 MoE 产生的聚类则像是 Prepositions、Proper Names、Copula Verbs 或 Definite Articles。在 EMO 中，来自同一文档的 token 大多落在同一个聚类中；在标准 MoE 中，它们会分散到多个聚类中。
这种对比在单个例子上最容易看清。以一篇健康文章为例：在 EMO 中，几乎每个 token 都会路由到 Health, Medical & Wellness 聚类。在标准 MoE 中，顶部聚类是 Possessives & Definite Articles；模型会把这篇文章与其他任何恰好使用 the 或 your 的文本归为一组，而不管那些文本的主题是什么。
在 1T tokens 上训练的 MoE 上，对预训练数据进行 token 聚类。EMO 聚类对应语义上有意义的领域，来自同一文档的 token 大体被分到一起。标准 MoE 训练产生的是表层或句法特征聚类，文档 token 分散在多个聚类中。
由于 EMO 形成的模块映射到语义领域而非表层特征，你可以选择一个小型 expert 子集，同时仍然得到一个能工作的模型：这个群组对应一种真实能力。
你可以在我们的交互式可视化中自行查看聚类结果。
我们发布的内容
我们发布了完整的 EMO 训练模型、一个在相同数据上训练的匹配标准 MoE baseline，以及训练代码。我们希望这些 artifacts 能对其他研究 MoE 中 emergent modularity 的团队有所帮助。
还有更多工作要做。EMO 是朝着让大型稀疏模型更加模块化迈出的早期一步，但仍有许多问题：如何更好地选择和组合 expert 子集，如何在不扰动完整模型的情况下更新模块，以及如何利用模块化结构提升可解释性和控制能力。发布这些模型应能帮助社区研究这些问题，并推动构建更易部署、适配、检查和组合的模块化语言模型。