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.


翻译:稀疏混合专家(Sparse Mixture-of-Experts, SMoE)模型能够高效扩展语言模型规模,但其训练仍面临挑战:路由可能坍缩到少数专家,而辅助负载均衡损失可能降低专家专业化程度。受这些难题启发,我们研究了SMoE中路由决策的机制形成过程。首先,我们揭示了路由器与对应专家之间存在一种几何耦合关系。对于给定词元,所选专家的路由权重与处理该词元的专家权重会沿相同输入方向获得梯度,二者仅在标量系数上存在差异。因此,匹配的路由-专家方向会累积相同的路由词元历史信息。这种理论耦合在路由动态中也得到经验验证:在从零训练的1B规模SMoE中,更高的路由分数预测更强的专家神经元激活,表明路由决策会镜像反映在所选的专家内部。其次,我们分析了辅助负载均衡对路由-专家几何耦合的影响,表明此类损失通过将输入方向梯度分散到路由权重上,破坏了这种结构,使得不同路由方向之间的相似性提高近三倍。最后,我们通过一个无参数在线K-Means路由器证明了几何耦合对有效路由的核心作用——该路由器让每个专家维护路由至其隐藏状态的滑动平均,并基于余弦相似度分配词元。与辅助负载均衡和无损失均衡方法相比,该路由器在仅适度增加困惑度的前提下实现了最低的负载不均衡度,表明几何耦合捕获了路由器学习内容中相当重要的部分。总体而言,我们的结果揭示了路由器如何形成支持有效分工的分配几何结构。

0
下载
关闭预览

相关内容

混合专家模型简述
专知会员服务
18+阅读 · 2025年5月30日
《混合专家模型推理优化技术综述》
专知会员服务
46+阅读 · 2024年12月21日
混合专家模型在大模型微调领域进展
专知会员服务
48+阅读 · 2024年9月23日
Jeff Dean署名《深度学习稀疏专家模型》综述论文
专知会员服务
39+阅读 · 2022年10月4日
稀疏大模型简述:从MoE、Sparse Attention到GLaM
夕小瑶的卖萌屋
14+阅读 · 2022年3月22日
稀疏性的3个优势 -《稀疏统计学习及其应用》
遇见数学
15+阅读 · 2018年10月24日
国家自然科学基金
0+阅读 · 2017年12月31日
国家自然科学基金
6+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
3+阅读 · 2014年12月31日
国家自然科学基金
1+阅读 · 2014年12月31日
VIP会员
最新内容
学习数据的几何:形状空间分析数学综述
专知会员服务
4+阅读 · 6月17日
定向能反无人机系统最新发展动态
专知会员服务
6+阅读 · 6月17日
从燃煤战舰到算法战争:水面指挥的永恒要求
专知会员服务
3+阅读 · 6月17日
相关基金
国家自然科学基金
0+阅读 · 2017年12月31日
国家自然科学基金
6+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
3+阅读 · 2014年12月31日
国家自然科学基金
1+阅读 · 2014年12月31日
Top
微信扫码咨询专知VIP会员