路由器学习专家的几何：稀疏专家混合中的几何耦合

稀疏专家混合 (SMoE) 模型可以有效地扩展语言模型，但训练它们仍然具有挑战性，因为路由可能会崩溃到少数专家身上，而辅助负载平衡损失可能会减少专业化。受这些障碍的推动，我们研究了 SMoE 中的路由决策是如何机械形成的。首先，我们揭示了路由器及其相应专家之间的几何耦合。对于给定的词元，所选专家的路由器权重和处理它的专家权重沿着相同的输入方向接收梯度，仅标量系数不同。因此，匹配的路由器专家方向会累积相同的路由词元历史记录。这种理论上的耦合也以经验的形式出现在路由动态中。在从头开始训练的 $1$B SMoE 中，较高的路由器分数预示着更强的专家神经元激活，这表明路由决策反映在所选专家内部。接下来，我们分析了辅助负载平衡对路由器-专家几何耦合的影响，表明这种损失通过在路由器权重上传播输入导向的梯度来打破这种结构，使不同的路由器方向彼此相似近三倍。最后，我们证明了几何耦合对于使用无参数在线 K-Means 路由器进行有效路由的中心性，其中每个专家维护路由到它的隐藏状态的运行平均值，并根据余弦相似度分配词元。与辅助损耗和无损耗平衡相比，该路由器实现了最低的负载不平衡，而困惑度仅略有增加，这表明几何耦合捕获了路由器学习内容的很大一部分。总的来说，我们的结果解释了路由器如何形成支持有效分工的分配几何结构。