虽然辅助损失直接激励负载均衡，但要让门控网络（即路由器）有效学习，还需要其他优化策略。路由器的主要目的是形成有用的分工，将令牌分配给最适合处理它们的专家，同时满足负载均衡机制设定的限制。若没有细致的优化，路由器可能无法学习到有用的路由策略，导致几种不理想的结果：
- 路由器崩溃： 路由器持续将大部分或所有令牌发送给少数专家（通常仅一两个），实际上忽略了辅助损失，或者找到了一个局部最小值，其中任务损失的减少超过了负载均衡的惩罚。这抵消了拥有多个专家的好处。
- 随机路由： 路由器未能学习到任何分工，几乎随机地分配令牌（受负载均衡压力的影响），导致专家成为通才而非专才。
- 训练不稳定： 路由决策的波动会使路由器和专家的学习过程变得不稳定。
以下是一些有助于稳定和有效路由器学习的策略：
初始化很重要
路由器的初始状态对训练过程的开始有很大影响。一个常见做法是初始化路由器的最终线性层权重 (weight)（如果使用偏置 (bias)，也包括偏置），使得选择任何专家的初始概率大致均匀。对于输出对数 h 的路由器，其专家概率由 g=softmax(h) 给出，将最终层的权重初始化为小值（接近零）会导致 N 个专家的概率 gi≈1/N 接近均匀分布。这避免了对特定专家的强烈初始偏好，并让负载均衡机制能及早生效。
差异化学习率
路由器和专家通常从不同的学习动态中获益。路由器的决策对哪些专家接收梯度有广泛影响，而专家根据接收到的数据学习特定功能。路由的快速变化会使专家学习不稳定。因此，通常有益的做法是，与模型的其余部分（包括专家）相比，路由参数 (parameter)使用较小的学习率。这使得路由策略能更缓慢地演变，给予专家足够的时间来适应分配给它们令牌的类型。找到主学习率和路由器学习率之间的合适比例通常需要进行实验。
路由器正则化 (regularization)和噪声
其他技术可以稳定路由器训练：
- 权重 (weight)衰减： 对路由器参数 (parameter)应用标准的 L2 正则化（权重衰减）有时有助于防止权重变得过大，从而可能提高稳定性。但是，应仔细评估其影响，因为它可能会干扰路由器在需要时做出清晰、果断决策的能力。
- 路由器噪声注入： 如架构设计（第2章）中所述，在训练过程中，于 top-k 选择之前向路由器对数添加噪声（例如，带噪声的 top-k 门控）是一种重要的优化策略。标准公式添加了由可学习权重 Wnoise 缩放的采样噪声 ϵ∼N(0,σ2)：
hnoisy=h+softplus(Wnoise)⋅ϵ
g=softmax(hnoisy)
这种注入的噪声鼓励路由决策的多样性，防止路由器过快地陷入次优分配模式。它还有助于确保在训练早期有更多专家接收到梯度，有助于形成更好的整体分工。噪声的大小（或 Wnoise 的值）通常需要调整。
梯度裁剪
流回路由器的梯度有时会变得非常大，特别是当任务损失基于路由决策突然变化，或者辅助损失施加强大压力时。大梯度可能导致较大的更新，从而使路由器学习到的策略不稳定。对路由器参数 (parameter)专门应用梯度裁剪可以缓解此问题。通过限制应用于路由器的梯度的最大范数或值，我们确保路由策略的更新更平滑。
训练期间监控路由器行为
系统性监控对于诊断路由器优化问题非常重要。重要的衡量指标包括：
- 门控概率分布： 跟踪批次中分配给每个专家的平均概率。它们是相对均衡的（由于辅助损失），还是路由器过度偏向特定专家？
- 门控概率的熵： 计算每个令牌的 softmax 输出 g 的熵，并对批次进行平均。较高的熵表明路由的不确定性更大、更均匀；较低的熵则表明路由更集中、更有分工。观察熵的变化可以表明路由器是否正在学习分工。
- 专家分配重叠： 监控路由器为相同令牌选择成对专家的频率（如果使用 k>1 的 top-k 路由）。高重叠可能表明存在冗余专家。
- 辅助损失值： 跟踪负载均衡损失的贡献。如果它保持较高，则路由器在均衡负载方面存在困难；如果它在没有相应任务改进的情况下下降过快，则路由器可能只是简单地满足了均衡，而没有形成分工。
以下图表呈现了影响路由器学习过程的不同因素：
影响训练期间路由器参数 (parameter)优化的因素。来自主任务损失和辅助负载均衡损失的梯度都提供了学习信号。噪声注入、仔细的学习率选择和正则化 (regularization)等技术有助于稳定此过程并促成有用的分工。
通过采用这些策略，您可以引导路由器学习一种稳定且有效的策略，它不仅能平衡专家之间的负载，还能发现有用的分工，最终有助于提高专家混合模型的整体性能和效率。这些技术通常结合使用，其具体配置需要根据模型架构、数据集和分布式训练设置进行仔细调整。
