Top-k门控是稀疏专家混合模型中用于路由令牌的标准方式。它的设计简单明了且作用良好,构成了许多成功MoE架构的根基。门控网络,一个简单的线性层,为每个专家生成一个对数,表示该专家处理当前输入令牌的偏好。TopK函数随后选择对数分数最高的k个专家。Top-k门控的运作方式对于给定的输入令牌表示$x$,该流程始于由权重矩阵$W_g$参数化的门控网络,该网络计算对数$h(x)$:$$ h(x) = W_g \cdot x $$这里,$h(x)$是一个向量,其中每个元素$h(x)_i$对应于专家$i$的对数。TopK操作选择$h(x)$中k个最高值的索引。我们将这组索引称为$S$。最终的门控分数$G(x)$通常通过仅对选定对数应用softmax函数来计算。这保证了所选专家的权重总和为一。$$ G(x)_i = \begin{cases} \frac{e^{h(x)i}}{\sum{j \in S} e^{h(x)_j}} & \text{如果 } i \in S \ 0 & \text{如果 } i \notin S \end{cases} $$MoE层的输出$y$是选定专家输出的加权和:$$ y = \sum_{i \in S} G(x)_i \cdot E_i(x) $$超参数$k$是一个重要的设计选择。使用$k=2$使得每个令牌可以由两个专家处理,这提供了比路由到一个专家更丰富的表示。这可以在增加计算成本的情况下提升模型质量。相比之下,像Switch Transformers所推广的那样,将$k=1$可以最大限度地减少计算和通信开销。digraph G { rankdir=TB; node [shape=record, style="rounded,filled", fillcolor="#e9ecef", fontname="sans-serif"]; edge [fontname="sans-serif"]; subgraph cluster_input { label = "输入令牌"; style = "filled"; color = "#f8f9fa"; token [label="令牌 x", shape=box, style="filled", fillcolor="#a5d8ff"]; } subgraph cluster_gating { label = "门控网络 (Top-k=2)"; style="filled"; color="#f8f9fa"; gate [label="{ <f0> 门控对数 h(x) | { E1: 0.8 | E2: 1.5 | E3: -0.2 | E4: 2.1 | ... | E8: 0.5 }}"]; topk [label="TopK 选择", shape=invhouse, style="filled", fillcolor="#ffc9c9"]; gate:f0 -> topk; } subgraph cluster_experts { label = "专家网络"; style="filled"; color="#f8f9fa"; E1 [label="专家 1", fillcolor="#e9ecef"]; E2 [label="专家 2", fillcolor="#b2f2bb"]; E3 [label="专家 3", fillcolor="#e9ecef"]; E4 [label="专家 4", fillcolor="#b2f2bb"]; E_dots [label="...", shape=plaintext]; E8 [label="专家 8", fillcolor="#e9ecef"]; } output [label="最终输出", shape=box, style="filled", fillcolor="#a5d8ff"]; token -> gate; topk -> E2 [label=" G(x)₂"]; topk -> E4 [label=" G(x)₄"]; {E2, E4} -> output [label="加权和"]; }单个令牌的门控流程。门控网络计算对数,TopK函数选择对数分数最高的两个专家(专家2和专家4)来处理该令牌。负载不平衡及其影响尽管功能简单,Top-k门控存在一个重要的运行不足:它可能导致严重的负载不平衡。在训练期间,门控网络可能偏向选择一小部分“常用”专家,而其他专家则利用不足或很少被选中。这种现象出现的原因是门控网络是一个学习函数,在没有约束的情况下,它会优化预测准确性,这可能涉及重复使用在训练早期被证明最有效的少数专家。这种不平衡带来两个主要问题:训练效率低下: 未被充分选择的专家接收到的训练信号很少,导致它们的参数优化不足。模型实际上未能充分利用其全部能力,因为其大部分权重对其性能的贡献很小。计算浪费: 在硬件层面,专家通常分布在多个设备(例如GPU)上。如果一个专家接收到过高比例的令牌,其分配到的设备就会成为瓶颈,而承载不常用专家的设备则处于闲置状态。{"layout":{"title":{"text":"专家负载不平衡示例"},"xaxis":{"title":{"text":"专家ID"},"tickmode":"array","tickvals":[0,1,2,3,4,5,6,7],"ticktext":["E1","E2","E3","E4","E5","E6","E7","E8"]},"yaxis":{"title":{"text":"批次中分配的令牌数"}},"bargap":0.2,"plot_bgcolor":"#f8f9fa","paper_bgcolor":"#ffffff"},"data":[{"type":"bar","x":[0,1,2,3,4,5,6,7],"y":[120,550,80,115,490,95,75,105],"marker":{"color":["#4dabf7","#fa5252","#4dabf7","#4dabf7","#fa5252","#4dabf7","#4dabf7","#4dabf7"],"line":{"width":0}}}]}单个训练批次中八个专家间的令牌分布概览。专家2和专家5负载过重,而其他专家则利用不足,表明负载平衡不佳。容量因子与丢弃令牌为了管理令牌流并防止任何单个专家超负荷,MoE实现引入了一个容量因子。这个超参数定义了每个专家在批次中可以处理多少令牌的缓冲区。每个专家的容量$C$定义为:$$ C = \mathrm{容量因子} \times \frac{\mathrm{批次中令牌总数}}{\mathrm{专家数量}} $$大于1.0的容量因子为轻微的不平衡提供了缓冲区。例如,1.25的值允许每个专家处理比平均值多达25%的令牌。但是,如果一个常用的专家接收到的令牌超过其容量$C$,多余的令牌就会被“丢弃”。这些丢弃的令牌不会被任何专家处理。相反,它们会完全绕过MoE层,并通过残差连接。这是一种信息丢失形式,并可能降低模型性能,因为模型无法对这些令牌应用专门的计算。找到合适的容量因子是一个权衡:较高的值会减少丢弃令牌的数量,但会增加内存分配,并且如果负载平衡,可能导致更多的计算浪费。在标准Top-k系统中应对负载不平衡的主要方式是辅助负载平衡损失,我们在第1章中已作介绍。这个损失项惩罚不平衡的路由决策,促使门控网络将令牌更均匀地分配到所有可用专家。它是稳定训练和使Top-k门控在实践中可行的一个重要组成部分。然而,正如我们在后续章节中将看到的,修改路由算法本身可以为此挑战提供更直接和有效的方法。