做MoE模型训练或路由优化的研究者,这篇论文解释了路由崩溃和负载均衡损失的底层机制,看完会对如何设计更有效的路由策略有直接启发。
稀疏混合专家模型(SMoE)在扩展语言模型时面临路由崩溃和负载均衡损失导致专业化下降的问题。本文揭示了路由器与其对应专家之间的几何耦合机制:对于给定token,所选专家的路由器权重和专家权重沿相同输入方向接收梯度,仅标量系数不同,因此匹配的路由器-专家方向累积相同的路由历史。实验表明,从零训练的1B SMoE中,更高的路由器分数预测更强的专家神经元激活,路由决策在所选专家内部被镜像。辅助负载均衡损失会破坏这种几何结构,使不同路由器方向相似度增加近三倍。最后,作者提出无参数在线K-Means路由器,通过维护专家隐藏状态运行平均值并基于余弦相似度分配token,在最低负载不平衡和适度困惑度增加下验证了几何耦合对有效路由的核心作用。
Routers Learn the Geometry of Their Experts: Geometric Coupling in Sparse Mixture-of-Experts
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.