🤖 machine learning

Hierarchical Mixture-of-Experts with Two-Stage Optimization

本文介绍了 Hi-MoE，这是一种分层混合专家框架，它采用两阶段优化策略，同时实现组间负载均衡与组内专家专业化，从而在解决路由稳定性与多样性之间权衡的同时，相较于现有基线实现了显著的性能提升。

原作者： Gleb Molodtsov, Alexander Miasnikov, Aleksandr Beznosikov

发布于 2026-05-12

📖 1 分钟阅读☕ 轻松阅读

原作者： Gleb Molodtsov, Alexander Miasnikov, Aleksandr Beznosikov

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在运营一个庞大且高科技的呼叫中心。你拥有数千名专家坐席（即“专家”），随时准备帮助客户（即“令牌”或数据点）。为了保持高效与快速，你并不希望每位坐席都监听每一通电话。相反，你希望一位智能调度员（即“路由器”）将每通电话仅转接给少数几位最适合解决该特定问题的专家。

这就是人工智能中**混合专家模型（MoE）**的工作原理。它们极其强大，但存在一个著名的缺陷：调度员会变得懒惰或困惑。

问题所在：“明星坐席”与“幽灵坐席”

在标准设置中，调度员倾向于将几乎所有电话都转给相同的几位“明星坐席”，因为他们在初期看起来表现不错。与此同时，数百名其他坐席却闲置无事，无所事事。

结果： 你为那些闲置坐席浪费了资金，而你的“明星”们则不堪重负。在人工智能术语中，这被称为路由崩溃。模型停止学习多样化的技能，转而仅仅依赖其大脑中一个微小且过度工作的子集。

一些先前的解决方案试图通过强制调度员向每个团队发送相同数量的电话来解决这一问题。但这引发了一个新问题：“群体思维”效应。 如果你强制每个团队处理完全相同的电话组合，不同团队中的坐席就会开始做完全相同的事情。他们变成了彼此冗余的复制品，失去了独特的专长。

解决方案：Hi-MoE（两级管理体系）

本文作者提出了Hi-MoE，这是一种组织呼叫中心的新方式。他们不再使用扁平的坐席列表，而是将坐席组织成组（类似于不同的部门），并采用两阶段管理策略来保持平衡与多样性。

这就像一家拥有四个专科院区（组）的大型医院。

第一阶段：院区经理（组间平衡）

目标： 确保没有单个院区不堪重负，而另一个院区却空无一人。
运作方式： 系统确保发送给院区 A、院区 B、院区 C 和院区 D 的患者总数大致相等。
类比： 如果医院收到 1,000 名患者，“院区经理”会确保每个院区各接收 250 名。这防止了“掉队者”问题，即某个 GPU（计算机芯片）被过多的工作卡住，而其他芯片则在等待。

第二阶段：科室主任（组内专业化）

目标： 确保每个院区内部的医生不会都做完全相同的事情。
问题： 如果院区 A 接收了 250 名患者，而院区 A 内的所有医生都被迫治疗完全相同类型的患者，他们都会变成千篇一律的通才。
Hi-MoE 的修正： 系统鼓励院区 A 内部的医生进行 specialization（专业化）。也许史密斯医生只处理骨折，琼斯医生只处理发烧，而李医生只处理过敏症。
类比： “科室主任”告诉医生们：“你们在同一个院区，所以你们必须分担工作量，但你们也必须彼此不同。不要都试图去治同一条断腿！”

为何有效（“两阶段优化”）

本文认为，你需要这两个层级协同工作：

层级 1（组间）： 让硬件保持愉悦。它确保计算机芯片（GPU）不会互相等待。
层级 2（组内）： 让人工智能保持智能。它迫使专家学习不同且互补的技能，而不是仅仅互相模仿。

结果：他们发现了什么？

作者在三种不同场景中测试了这种“两级管理”系统：

视觉（Tiny ImageNet）： 他们训练了一个 AI 来识别图像。Hi-MO 在识别物体（如手与背景）方面表现更好，并保持计算机芯片的工作均匀。
语言（nanoGPT）： 他们训练了一个 AI 来撰写文本。Hi-MoE 写出了更好的文本，并且没有让“明星坐席”独占所有注意力。
大规模（OLMoE-7B）： 他们测试了一个拥有 70 亿参数的巨型模型。
- 质量： 该模型在从数学到维基百科文章的各种主题上，犯错更少（即“困惑度”更低）。
- 平衡性： 工作负载比标准模型平衡了 40%。这意味着计算机硬件的利用效率要高得多。

核心结论

Hi-MoE 就像雇佣了一位优秀的经理，他懂得两件事：

公平性： 每个人都能分到公平的工作量，这样没人会精疲力竭。
多样性： 每个人都能成为独特的专家，这样整个团队就能解决更复杂的问题。

通过将管理划分为“组间”和“组内”，AI 获得了双赢：它在硬件上运行得更快，并学会了更聪明、更多样化的技能。

技术摘要：具有两阶段优化的分层混合专家模型（Hi-MoE）

1. 问题陈述

稀疏混合专家（MoE）模型通过将令牌路由到一小部分专家网络来扩展容量。然而，训练这些模型面临着负载均衡与专家专业化之间的根本权衡：

路由崩溃：激进的多样性机制往往导致“路由崩溃”，即少数专家主导令牌选择，而其他专家保持空闲。这导致容量利用率低和设备间负载不平衡（拖尾现象）。
分组困境：最近的方法通过将专家分组并强制每组执行固定选择（例如，从 $M$ 个组中的每组选择 $K/M$ 个专家）来缓解不平衡。虽然这确保了设备级计算的均匀性，但它限制了路由自由度。这种约束往往导致冗余，即不同组中的专家学习相似的功能，因为它们是在相似的令牌混合体上优化的，这实际上用组间的“多样性崩溃”取代了专家级的崩溃。

现有的基于损失的平衡方法（如 GShard）经常与任务梯度冲突，减缓收敛，而无损失方法可能无法防止冗余。因此，需要一个能够同时强制硬件对齐的平衡并促进有意义、互补的专业化的框架。

2. 方法：Hi-MoE

作者提出了Hi-MoE，这是一个分层框架，将路由控制分解为两个耦合的层级：组间平衡和组内专业化。

架构

分组结构：专家被划分为 $M$ 个不相交的组（通常与 GPU 分片对齐）。
路由约束：对于每个令牌，路由器从每个组 $m$ 中选择固定数量的专家（ $K_m$ ）。
两级优化：该框架为分层的每个层级引入了特定的正则化项。

关键组件

组间正则化（ $R_{inter}$ ）：
- 目标：强制组间公平的流量分布，以确保平衡的设备利用率。
- 机制：最小化跨组的选择后路由权重的 $\ell_2$ 集中度。通过惩罚特定组上集中的路由质量，它作为最小化组负载变异系数（CV）的代理。
- 理论基础：定理 4.1 证明，最小化该项提供了组级不平衡的认证上界。
组内正则化（ $R_{intra}$ ）：
- 目标：促进组内的互补专业化，防止专家学习冗余功能。
- 机制：最大化选择前路由分布的 $\ell_2$ 集中度（实际上是一个抗重叠项）。这鼓励路由器在组内对特定专家做出果断、稀疏的承诺，而不是以低置信度混合许多专家。
- 理论基础：定理 4.2 表明，减少路由重叠可以降低专家之间的梯度耦合，使它们能够遵循不同的优化轨迹。定理 4.3 证明，在平衡边际下，增加路由集中度可以最大化输入与所选专家之间的互信息，这是专业化的代理指标。
稳定机制：
- 为了防止组内多样性信号漂移到退化状态，作者采用基于路由器 Logits 指数移动平均（EMA）的偏差校正路由计算。这在不添加单独平衡损失的情况下，抑制了对历史上首选专家的持续过度选择。

目标函数

总训练目标结合了任务损失、标准负载均衡损失以及提出的分层项：
$\mathcal{L} = \mathcal{L}_{task} + \mathcal{L}_{load} + R_{intra} + R_{inter}$
作者将此框架化为一个约束优化问题，其中系数 $\lambda_{intra}$ 和 $\lambda_{inter}$ 分别充当多样性和平衡约束的拉格朗日乘子。

3. 主要贡献

张力识别：该论文指出，通过严格的组约束来强制设备级平衡可以稳定利用率，但会增加组间的冗余。
分层框架：Hi-MoE 引入了一种模块化的、可即插即用的目标，解耦了这些关注点：
- 用于硬件对齐平衡的组间目标。
- 用于功能多样性和专业化的组内目标。
理论保证：该论文提供了形式化证明，将正则化项与组不平衡界限、梯度解耦以及路由信息（专业化）的增加联系起来。
实证验证：在视觉和语言领域均取得了一致的改进，证明了 Hi-MoE 扩展了困惑度/准确率与负载均衡（CV）之间的帕累托前沿。

4. 实验结果

作者在三种不同的预训练设置中评估了 Hi-MoE：

A. 计算机视觉（Swin-MoE on Tiny ImageNet）

设置：10 亿参数模型，包含 64 个专家（4 个组，每组 16 个）。
结果：与 Vanilla GShard、无损失平衡和 MoGE 相比，Hi-MoE 实现了最高的 Top-1 准确率（59.13%）和最佳的负载均衡（CV 0.29）。
观察：注意力图的可视化显示，Hi-MoE 的组发展出了独特的视觉专业化（例如，一个组专注于身体部位，另一个组专注于背景表面），而基线模型则显示出更多的冗余。

B. 语言建模（nanoGPT on OpenWebText）

设置：10.6 亿参数模型，包含 8 个专家（4 个组，每组 2 个）。
结果：Hi-MoE 实现了最低的困惑度（2.947）和最佳的平衡（CV 0.23）。
消融实验：移除 $R_{inter}$ 会损害平衡；移除 $R_{intra}$ 会损害质量。偏差校正起到了稳定器的作用。
效率：分层路由器增加了极小的开销（每层 1.3 毫秒至 1.6 毫秒），同时将同步开销减少了约 3 倍（与扁平路由相比）。

C. 大规模预训练（OLMoE-7B vs. Hi-MoE-7B）

设置：70 亿参数模型在 585.4 亿个令牌（Dolma 混合体）上进行训练。
困惑度：与标准 OLMoE-7B 基线相比，Hi-MoE-7B 在 11 个验证域上的平均困惑度降低了5.6%。
下游任务：Hi-MoE 在 MMLU 切片（例如，STEM 领域 +22.2%，人文学科 +11.9%）和常识推理任务上显示出显著增益。
负载均衡：Hi-MoE 将专家平衡（CV）提高了40%（0.153 对比 0.257），直接解决了大规模训练中的资源效率瓶颈。

5. 意义与主张

该论文声称，Hi-MoE 解决了稀疏 MoE 训练中固有的“平衡–多样性”张力。通过显式建模路由的分层结构（组与专家），该方法使模型能够：

高效扩展：在不牺牲模型条件容量的情况下保持均匀的 GPU 利用率。
学习专业化表示：防止“多样性崩溃”，即分组专家变得功能冗余。
提升质量：在视觉和语言任务中实现困惑度和下游任务性能的一致提升。

作者将 Hi-MoE 定位为现有稳定机制的补充，而非替代品，它是一种模块化的增强，使专家分组成为一种原则性的优化机制，而不仅仅是系统约束。这项工作表明，分层协调是解锁万亿参数模型可靠质量提升的一条可行途径。