发现专家路由的几何耦合机制,揭示路由如何与专家协同学习。
Routers Learn the Geometry of Their Experts: Geometric Coupling in Sparse Mixture-of-Experts

- 路由器与专家在梯度方向上保持几何同步,共享相同输入方向。
- 训练中高路由得分对应强专家激活,证明路由决策被专家内部体现。
- 无辅助损失的K-Means路由实现最低负载不均,凸显几何耦合重要性。
稀疏混合专家(SMoE)模型可高效扩展语言模型,但训练仍具挑战,因路由可能集中于少数专家,且辅助负载均衡损失会削弱专家专属性。我们研究了SMoE中路由决策的形成机制。首先,揭示路由器与其对应专家间存在几何耦合:对特定标记,路由器权重与专家权重在相同输入方向上接收梯度,仅缩放系数不同。因此,匹配的路由器-专家方向累积相同的路由历史。该理论耦合在10亿参数的SMoE中得到实证。其次,分析辅助负载均衡的影响,发现其破坏几何结构,使不同路由器方向相似度提升近三倍。最后,通过无参在线K-Means路由器验证其核心作用:每个专家维护被路由隐藏状态的运行平均,按余弦相似度分配。相比带辅助损失和无损失均衡方法,该方案实现最低负载不均,仅小幅增加困惑度,表明几何耦合捕获了路由器学习的关键部分。整体结果解释了路由如何形成支持有效分工的分配几何。
原文摘要 · Abstract (English)
Sparse Mixture-of-Experts (SMoE) models enable scaling language models efficiently, but training them remains challenging, as routing can collapse onto few experts and auxiliary load-balancing losses can reduce specialization. Motivated by these hurdles, we study how routing decisions in SMoEs are formed mechanistically. First, we reveal a geometric coupling between routers and their corresponding experts. For a given token, the router weights for the selected expert and the expert weights processing it receive gradients along the same input direction, differing only in scalar coefficients. Thus, matched router--expert directions accumulate the same routed token history. This theoretical coupling also appears empirically in routing dynamics. In a $1$B SMoE trained from scratch, higher router scores predict stronger expert neuron activations, showing that routing decisions are mirrored inside the selected expert. Next, we analyze the effects of auxiliary load balancing on the router--expert geometric coupling, showing that such losses break this structure by spreading input-directed gradients across router weights, making distinct router directions nearly three times more similar to each other. Last, we demonstrate the centrality of geometric coupling for effective routing with a parameter-free online K-Means router, in which each expert maintains a running average of the hidden states routed to it and tokens are assigned based on cosine similarity. Compared with auxiliary-loss and loss-free balancing, this router achieves the lowest load imbalance with only a modest perplexity increase, indicating that geometric coupling captures a substantial part of what the router learns. Overall, our results explain how routers form assignment geometry that supports an effective division of labor.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。