路由优化策略

虽然辅助损失直接激励负载均衡，但要让门控网络（即路由器）有效学习，还需要其他优化策略。路由器的主要目的是形成有用的分工，将令牌分配给最适合处理它们的专家，同时满足负载均衡机制设定的限制。若没有细致的优化，路由器可能无法学习到有用的路由策略，导致几种不理想的结果：

• 路由器崩溃： 路由器持续将大部分或所有令牌发送给少数专家（通常仅一两个），实际上忽略了辅助损失，或者找到了一个局部最小值，其中任务损失的减少超过了负载均衡的惩罚。这抵消了拥有多个专家的好处。
• 随机路由： 路由器未能学习到任何分工，几乎随机地分配令牌（受负载均衡压力的影响），导致专家成为通才而非专才。
• 训练不稳定： 路由决策的波动会使路由器和专家的学习过程变得不稳定。

以下是一些有助于稳定和有效路由器学习的策略：

初始化很重要

路由器的初始状态对训练过程的开始有很大影响。一个常见做法是初始化路由器的最终线性层权重（如果使用偏置，也包括偏置），使得选择任何专家的初始概率大致均匀。对于输出对数 $h$ 的路由器，其专家概率由 $g = \text{softmax}(h)$ 给出，将最终层的权重初始化为小值（接近零）会导致 $N$ 个专家的概率 $g_i \approx 1/N$ 接近均匀分布。这避免了对特定专家的强烈初始偏好，并让负载均衡机制能及早生效。

差异化学习率

路由器和专家通常从不同的学习动态中获益。路由器的决策对哪些专家接收梯度有广泛影响，而专家根据接收到的数据学习特定功能。路由的快速变化会使专家学习不稳定。因此，通常有益的做法是，与模型的其余部分（包括专家）相比，路由参数使用较小的学习率。这使得路由策略能更缓慢地演变，给予专家足够的时间来适应分配给它们令牌的类型。找到主学习率和路由器学习率之间的合适比例通常需要进行实验。

路由器正则化和噪声

其他技术可以稳定路由器训练：

1. 权重衰减： 对路由器参数应用标准的 $L_2$ 正则化（权重衰减）有时有助于防止权重变得过大，从而可能提高稳定性。但是，应仔细评估其影响，因为它可能会干扰路由器在需要时做出清晰、果断决策的能力。
2. 路由器噪声注入： 如架构设计（第2章）中所述，在训练过程中，于 top-k 选择之前向路由器对数添加噪声（例如，带噪声的 top-k 门控）是一种重要的优化策略。标准公式添加了由可学习权重 $W_{\text{noise}}$ 缩放的采样噪声 $\epsilon \sim \mathcal{N}(0, \sigma^2)$： $$h_{\text{noisy}} = h + \text{softplus}(W_{\text{noise}}) \cdot \epsilon$$ $$g = \text{softmax}(h_{\text{noisy}})$$ 这种注入的噪声鼓励路由决策的多样性，防止路由器过快地陷入次优分配模式。它还有助于确保在训练早期有更多专家接收到梯度，有助于形成更好的整体分工。噪声的大小（或 $W_{\text{noise}}$ 的值）通常需要调整。

梯度裁剪

流回路由器的梯度有时会变得非常大，特别是当任务损失基于路由决策突然变化，或者辅助损失施加强大压力时。大梯度可能导致较大的更新，从而使路由器学习到的策略不稳定。对路由器参数专门应用梯度裁剪可以缓解此问题。通过限制应用于路由器的梯度的最大范数或值，我们确保路由策略的更新更平滑。

训练期间监控路由器行为

系统性监控对于诊断路由器优化问题非常重要。重要的衡量指标包括：

• 门控概率分布： 跟踪批次中分配给每个专家的平均概率。它们是相对均衡的（由于辅助损失），还是路由器过度偏向特定专家？
• 门控概率的熵： 计算每个令牌的 softmax 输出 $g$ 的熵，并对批次进行平均。较高的熵表明路由的不确定性更大、更均匀；较低的熵则表明路由更集中、更有分工。观察熵的变化可以表明路由器是否正在学习分工。
• 专家分配重叠： 监控路由器为相同令牌选择成对专家的频率（如果使用 $k > 1$ 的 top-k 路由）。高重叠可能表明存在冗余专家。
• 辅助损失值： 跟踪负载均衡损失的贡献。如果它保持较高，则路由器在均衡负载方面存在困难；如果它在没有相应任务改进的情况下下降过快，则路由器可能只是简单地满足了均衡，而没有形成分工。

以下图表呈现了影响路由器学习过程的不同因素：

影响训练期间路由器参数优化的因素。来自主任务损失和辅助负载均衡损失的梯度都提供了学习信号。噪声注入、仔细的学习率选择和正则化等技术有助于稳定此过程并促成有用的分工。

通过采用这些策略，您可以引导路由器学习一种稳定且有效的策略，它不仅能平衡专家之间的负载，还能发现有用的分工，最终有助于提高专家混合模型的整体性能和效率。这些技术通常结合使用，其具体配置需要根据模型架构、数据集和分布式训练设置进行仔细调整。