分层MoE结构

标准专家混合架构，通过单一门控网络将令牌路由到一组扁平的专家，尽管提供了很大的扩展优势，但复杂数据集和任务可能从更结构化的专业化中获益。分层专家混合（HMoE）引入了多层路由，创建树状结构，使得令牌逐渐被导向更专业的专家。

设想一个情况，一个标准MoE层将令牌 $x$ 路由到 $N$ 个专家中的一个。在一个简单的两层HMoE中，该过程变得更细致：

1. 第一层路由： 输入令牌 $x$ 首先由顶层门控网络 $G_1$ 处理。该路由器选择，通常使用 $top-k$ 机制（例如，$k=1$ 或 $k=2$），从总共 $M$ 个组中选择一个或多个专家组。
2. 第二层路由： 对于每个选定的组 $m$，令牌 $x$（或从中衍生的表示）然后由该组特有的第二层门控网络 $G_{2,m}$ 处理。该路由器从组 $m$ 内部的 $N_m$ 个专家中选择一个或多个专家。
3. 专家计算： 令牌 $x$ 由最终选定的专家 $E_{m,n}$ 处理。
4. 输出组合： 所选专家的输出被组合，由两层门控概率的乘积加权。

例如，如果 $G_1$ 以概率 $p_{1,m}$ 将令牌 $x$ 路由到组 $m$（为简化起见，假设第一层采用 $top-1$ 路由），并且 $G_{2,m}$ 以概率 $p_{2,m,n}$ 将其路由到该组内的专家 $n$，则专家 $E_{m,n}$ 对令牌 $x$ 输出的最终贡献将由 $p_{1,m} \cdot p_{2,m,n}$ 加权。

架构变体与设计考量

设计选择包括：

• 层数与分支因子： 确定层级的深度以及每层组/专家的数量（$M$, $N_m$ 等）。这会影响专业化的细致程度和总参数 (parameter)量。
• 路由器架构： 不同层级的路由器（$G_1$, $G_{2,m}$）可以采用多种设计（线性、MLP、基于注意力），类似于扁平MoE。它们可以是独立的、部分共享参数，或者根据前一层的表示进行条件化。
• 路由机制： $Top-k$ 路由在每个层级都很常见。$k$ 的值在不同层级可能不同。噪声路由也可以分层应用。
• 专家同质性： 同一组内或不同组之间的专家在架构上可以相同（同质）或不同（异质），尽管同质性简化了实现和并行化。

一个两层HMoE结构路由单个令牌的视图：

令牌通过两层分层MoE的流程。第一个门控网络（$G_1$）选择一个专家组，该组内的第二层门控网络（$G_{2,m}$）选择最终的专家（$E_{m,n}$）。

分层结构的好处

分层MoE相比于扁平MoE有几个潜在的好处：

1. 更细致的专业化： 通过多阶段路由令牌，HMoE可以学习到更具体的专业化。较深层级的专家可以专注于由高层路由器识别出的数据分布中高度特定的子集。
2. 参数 (parameter)效率： 对于大量所需的细致专业化（叶子专家），分层结构可能比扁平MoE更具参数效率。扁平MoE需要一个大型的单一门控网络来区分所有叶子专家，而HMoE将路由决策分配到较小的、层级特定的路由器中。
3. 结构化知识： 该层级可以隐式地促进知识的结构化表示，可能反映数据本身中的层级。

分层MoE中的挑战

尽管有潜在好处，实现和训练HMoE带来了特定的挑战：

• 训练复杂性与稳定性： 训练多个相互依赖的门控网络可能不如训练单个门控网络稳定。确保所有层级都能学习有意义的路由策略而不会崩溃，需要仔细调整并可能使用专门的正则化 (regularization)技术。
• 负载均衡： 负载均衡变得更复杂。不仅要在最终专家（叶子）之间，还要在高层级的中间组之间平衡令牌分配。辅助损失可能需要在层级的多个层级上进行调整或应用。例如，可以根据 $G_1$ 的组分配应用负载均衡损失，并针对 $G_{2,m}$ 在每个组内的专家分配应用单独的损失。
• 计算成本： 路由决策的顺序性会增加训练和推理 (inference)期间的计算延迟，相比于路由在一个步骤中发生的扁平MoE。
• 分布式训练： 在分布式环境中高效地实现HMoE会增加复杂性。尽管仍然可以应用专家并行性（例如，将组内专家分布到不同设备上），但管理多层路由和潜在的通信模式需要仔细考量，特别是当层级的不同层存在于不同的设备集上时。框架可能需要对分层路由模式和通信的特定支持。

分层MoE代表一种先进的架构模式，用于推动大型模型专业化的边界。尽管它们有望实现更细致的专业性以及可能更好的参数 (parameter)效率，但它们需要仔细考量训练稳定性、跨多个层级的负载均衡以及增加的实现复杂性，特别是在分布式环境中。当任务或数据表现出模型可以从中获益的强烈固有分层结构时，其应用通常是合理的。