💬 NLP

Routers Learn the Geometry of Their Experts: Geometric Coupling in Sparse Mixture-of-Experts

本文揭示了稀疏混合专家模型中路由器和专家之间的一种基本几何耦合，即匹配的方向会累积共享的令牌历史，并表明这种耦合虽会被辅助负载均衡损失所破坏，但可通过一种无参数的 K-Means 路由器有效复现，从而实现更优的负载均衡。

原作者： Sagi Ahrac, Noya Hochwald, Mor Geva

发布于 2026-05-13

📖 1 分钟阅读☕ 轻松阅读

原作者： Sagi Ahrac, Noya Hochwald, Mor Geva

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象你正在运营一个超大规模、高速运转的图书馆，每秒都有成千上万本书（数据）需要被分类和处理。为了应对如此庞大的工作量，你并没有配备一位巨型图书管理员，而是拥有一位路由器（一位聪明的交通协管员）和一支由专家（专业图书管理员）组成的团队。

在“稀疏混合专家”（SMoE）模型中，路由器会审视一段文本，并决定从 64 位图书管理员中选出最合适的 6 位来处理它。只有这 6 位开始工作，其余的则休息。这节省了能源并提升了系统速度。

然而，训练这些系统颇具挑战性。有时，路由器会变得懒惰，将所有内容都派发给仅一两位图书管理员，导致其他人无所事事、毫无用处。为了解决这个问题，工程师通常会添加一个“惩罚分数”（辅助损失），以迫使路由器均匀地分配工作。

本文深入探究了该系统“大脑”内部实际发生的情况。以下是通俗易懂的解析：

1. 秘密握手：“几何耦合”

作者发现了路由器与专家之间存在一种隐藏的、自然的联系。

类比：想象路由器和专家是两位正在共同学习一套舞步的舞者。每当路由器将特定类型的书发送给特定的图书管理员时，他们都会在同一时间、针对完全同一本书进行学习。
发现：从数学上讲，路由器对某本书的“记忆”与该图书管理员对同一本书的“记忆”，其增长方向完全一致。本质上，他们正在构建相同的、关于共同所见书籍的心理地图。
证据：研究人员检查了一个真实模型，发现当路由器表示“我确实认为图书管理员 A 应该处理这个”时，图书管理员 A 内部的齿轮实际上会以比平时更快、更强的速度运转。路由器的决策在图书管理员的“大脑”内部得到了物理层面的镜像反映。

2. “强行拟合”平衡的问题

为了防止路由器偷懒，工程师们通常使用前述的“惩罚分数”。本文认为，这种惩罚实际上破坏了自然的舞蹈。

类比：想象这个惩罚分数是一位严厉的管理员，他向每一位图书管理员大喊，即使是那些没有拿到书的管理员，也会听到：“你也需要关注这本书！”
结果：由于每位图书管理员都被迫从每一本书中学习（仅仅是为了保持分数平衡），他们开始变得看起来和思考方式完全一样。每位图书管理员独特的“个性”被冲刷殆尽。
数据：研究人员发现，在使用这种惩罚的情况下，路由器针对不同专家的方向变得彼此相似度高出近三倍。这种使系统高效运行的独特“专业化”，正被强行平衡的尝试所抹除。

3. 解决方案：“质心”路由器

如果路由器和专家自然学会总结他们处理的书籍，那我们为什么还需要一个复杂的、可训练的路由器呢？

类比：与其让一位聪明的交通协管员去尝试学习复杂的规则，不如想象每位图书管理员只需维护一个他们通常收到的书籍的运行平均值（即“质心”）。当新书到来时，系统只需询问：“这本书最像哪位图书管理员的平均藏书？”
实验：作者构建了一个系统，其中的路由器拥有零个可训练参数。它并非以传统意义上的方式“学习”，而只是更新它所见书籍的简单平均值。
结果：这个简单、“笨拙”的路由器实际上比那些复杂的、受惩罚的系统更好地平衡了工作量。它几乎毫无困惑地完美分配了工作。唯一的缺点是在系统理解文本的能力（困惑度）上出现了微小且可忽略的下降。

核心结论

本文得出结论：这些 AI 模型的魔力不仅仅在于我们强加给它们学习的复杂数学。它们成功的一大部分源于一种自然的几何联系，即路由器和专家共同成长，学习相同的历史。

当我们试图用沉重的惩罚来强行平衡时，我们破坏了这种自然联系。相反，如果我们让路由器仅仅追踪每位专家处理内容的“平均值”，系统就能几乎同样高效地工作，保持工作平衡，并且需要远为简单的训练。

简而言之：路由器和专家是天然的伙伴。不要试图强迫他们变得完美；只需让他们记住共同做过的事，他们就会找到最佳的工作方式。

技术摘要：路由器学习其专家的几何结构

问题陈述

稀疏混合专家（SMoE）模型通过在每个 token 上仅激活部分参数，实现了语言模型的高效扩展。然而，训练这些模型面临重大挑战。若无干预，路由机制往往会发生崩溃，导致流量集中在不断缩小的专家子集上（表示崩溃）。为缓解这一问题，从业者采用辅助负载均衡损失来强制实现均匀的专家利用率。尽管这些损失在平衡负载方面行之有效，但已知会降低专家的专业化程度。本文解决的核心问题是：缺乏对 SMoE 中路由决策如何形成，以及常见训练干预措施（如辅助损失）如何改变路由器与专家之间潜在动态的机制性理解。

方法论

作者通过对 SMoE 架构内梯度动态的几何分析来探讨该问题。其方法论包含三个主要阶段：

理论梯度推导：作者推导了路由器权重（ $r_i$ ）和专家输入侧权重（ $W^{gate}_i$ ）的反向传播更新。他们分析了应用于 SMoE 层的链式法则，以确定匹配的路由器 - 专家对与未选专家之间的梯度更新方向。
实证验证：使用在 OLMoE-mix-0924 数据集（约 500 亿 token）上从头训练的 10 亿参数 SMoE 模型，作者测量了路由器评分与所选专家内部激活之间的相关性。他们特别检查了“门控神经元”激活（ $\sigma(W^{gate}_i x)$ ），以观察其是否反映了路由器的偏好。
消融与干预：
- 辅助损失影响：他们比较了两个架构相同但平衡规则不同的模型：一个使用标准辅助负载均衡损失（如 Switch Transformer 风格），另一个使用“无损失”的基于偏置的平衡方法。他们测量路由器权重向量之间的余弦相似度，以评估几何崩溃。
- 质心路由器：为了检验路由器本质上是在学习总结输入分布这一假设，他们实现了一个无参数的在线 K-Means 路由器。在此设计中，路由器权重被替换为不可学习的质心（路由隐藏状态的指数移动平均），token 根据余弦相似度进行分配。该路由器无需可训练的路由参数或平衡梯度即可运行。

主要贡献

几何耦合的发现：本文揭示了路由器与其对应专家之间存在根本的几何耦合。对于路由到专家 $i$ 的 token，路由器权重向量 $r_i$ 和专家的输入侧权重均接收与同一输入方向 $x$ 成正比的梯度更新。因此，匹配的路由器 - 专家对会演变为它们共享的路由 token 历史的耦合累加器。
耦合的实证确认：在训练好的 10 亿参数 SMoE 中，作者证明，特定专家的较高路由器评分与该专家门控神经元的较强激活之间存在强相关性。这证实了路由决策不仅仅是外部分配，而是反映在被选专家的内部计算中。
几何崩溃的机制：研究表明，辅助负载均衡损失会破坏这种自然耦合。通过向每个 token 上的每个路由器权重向量发送输入导向的梯度（无论是否被选中），这些损失导致不同的路由器方向彼此之间的相似度增加了近三倍，有效地统一了特定专家的方向并侵蚀了专业化。
基于质心的路由：作者实例化了一个简单的、无参数的 K-Means 路由器，该路由器显式跟踪路由隐藏状态的质心。该路由器在所有测试变体（包括无损失和辅助损失基线）中实现了最低的负载不平衡，且困惑度仅略有增加。

结果

梯度对齐：理论推导证实，对于选中的 token，路由器和专家权重沿相同的输入方向 $x$ 更新，仅相差标量系数。
激活相关性：在 10 亿参数模型中，路由器评分与专家门控神经元激活显示出单调关系（ $\rho = 0.43$ ），验证了路由器的几何偏好反映在专家的计算中。
辅助损失效应：与无损失模型（ $\mu \approx 0.13$ 至 $0.32 $）相比，使用辅助负载均衡损失训练的模型表现出显著更高的路由器权重向量间成对余弦相似度（$ \mu \approx 0.63$）。这表明辅助损失压缩了路由器的几何多样性。
K-Means 路由器性能：
- 负载不平衡：K-Means 路由器的 MaxVio（最大相对过载）为 0.037，显著低于无损失基线（0.084）和辅助损失基线（0.526）。
- 困惑度：与学习的无损失路由器相比，K-Means 路由器的困惑度增加适度（训练 PPL：15.40 对比 15.01），代表约 2.6% 的成本。
- 参数：K-Means 路由器使用 0 个可学习的路由参数，完全依赖运行统计量（质心和偏置）。

意义与主张

本文主张，几何耦合是梯度训练的路由器所学内容的重要组成部分，而非仅仅是副产品。研究结果表明：

机制性洞察：路由器并非学习任意的分配规则；它们与专家共同演化，形成支持有效分工的共享几何结构。
对当前实践的批判：常见的辅助负载均衡损失虽然在平衡负载，但通过注入统一路由器方向的“干扰梯度”，主动破坏了使专家专业化成为可能的几何结构。
设计启示：未来的路由方法可能受益于保留这种自然的路由器 - 专家几何结构。无参数 K-Means 路由器的成功表明，显式的质心跟踪可以在不需要复杂基于梯度的路由权重或平衡损失的情况下，恢复路由器的许多功能。

作者总结道，他们的工作解释了路由器如何形成分配几何结构，并指出保留路由器与专家之间的自然耦合对于有效的 SMoE 训练至关重要。他们承认局限性，指出其实证结果基于单一的 10 亿参数配置，且 K-Means 路由器虽然是一个建设性的概念验证，但由于仍存在困惑度差距，尚未成为学习式路由器的实用替代品。