💻 computer science

Selective Sinkhorn Routing for Improved Sparse Mixture of Experts

本文介绍了选择性 Sinkhorn 路由（Selective Sinkhorn Routing, SSR），这是一种轻量级机制，它将 Token 到专家（token-to-expert）的分配建模为一个受约束的最优传输问题，旨在不依赖辅助平衡损失或复杂可训练组件的情况下，实现均衡的专家利用率并提升模型性能。

原作者： Duc Anh Nguyen, Huu Binh Ta, Nhuan Le Duc, Tan Minh Nguyen, Toan Tran

发布于 2026-06-05

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Duc Anh Nguyen, Huu Binh Ta, Nhuan Le Duc, Tan Minh Nguyen, Toan Tran

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

想象一下你经营着一家规模宏大、高科技的呼叫中心。你有成千上万的进线电话（Token）和一支由 100 名专业坐席（专家）组成的团队。你的目标是将每一通电话路由到最合适的坐席，以便快速解决问题。

在标准设置中，你会使用一个简单的规则：“将电话发送给当前看起来最资深的坐席。”这就像是一个 Softmax 路由。问题在于？同样的几个“超级坐席”会接到所有的电话，而其他 90 名坐席却处于闲置状态。这被称为路由崩溃（Routing Collapse）。由于你没有充分利用整个团队，呼叫中心变得效率低下。

为了解决这个问题，之前的方法尝试通过强制平衡来干预。他们增加了一个“经理”，不断地向系统唠叨：“嘿，5 号坐席已经一个小时没接到电话了，给他派一个！”或者“1 号坐席太忙了，别再给他发电话了！”这些就是文中提到的辅助损失（Auxiliary Losses）。虽然它们有所帮助，但非常笨重，增加了计算机的额外工作量，有时还会让系统对它真正要学习的目标感到困惑。

新的想法： “完美平衡”的任务分配

作者们提出了一种更聪明的方法来分配电话，使用了数学中的一个概念——最优传输（Optimal Transport）（具体来说是 Sinkhorn 算法）。

请不要将其视为一个不停唠叨坐席的经理，而应将其视为一场完美的编舞舞蹈。

目标： 随着时间的推移，每位坐席必须获得完全相同数量的电话，并且每位来电者都必须被匹配到擅长其业务领域的坐席。
方法： 系统不再只是为每通电话挑选“最佳”坐席，而是计算一个全局图谱。它同时观察所有的电话和所有的坐席，并计算出一种最高效的方式来分配工作，确保没有人过载，也没有人无事可做。

“完美舞蹈”的问题

这里有一个陷阱。如果你在每一通电话进来时都强行进行这种完美的平衡，系统就会陷入混乱。它可能会为了维持数字平衡，把一通关于“编程”的电话发给一位擅长“烹饪”的坐席。这会损害性能。

论文的突破点在于 选择性 Sinkhorn 路由（Selective Sinkhorn Routing, SSR）。

SSR 如何运作：“混合”策略

SSR 并没有对每一通电话都使用复杂的“完美舞蹈”，而是使用了一种聪明的混合方式：

绝大多数时间（99%+）： 它使用标准的、快速的方法（Soft 函数）来路由电话。这让系统能够学习坐席实际擅长什么。
极少数时间（0.1% 到 1%）： 它会暂停并运行“完美的舞蹈”（Sinkhorn 算法）。
- 原因： 这微小的“完美平衡”就像一次温柔的推力。它提醒系统：“别忘了其他坐席！”但又不会在每一次通话中都强行进行错误的匹配。
- 结果： 系统学会了自然地实现自我平衡，而不需要一个唠叨的经理（辅助损失），也不需要庞大的额外计算能力。

秘诀：加入一点噪声

论文还建议在训练过程中加入一点随机噪声（就像收音机里的静电干扰）。

类比： 想象坐席们有点醉酒，或者电话线路有点模糊。系统无法 100% 确定谁才是“最佳”坐席，所以它会尝试尝试不同的对象。
益处： 这防止了系统陷入一种死板的套路，即总是选择那前 3 名顶尖坐席。它迫使系统去探索和发现，从而意识到其他坐席其实也相当出色。
重要提示： 论文强调你在系统实际工作（推理）时要关闭这种噪声。你肯定不希望你的呼叫中心在处理客户时是随机的；你希望它是快速且确定的。

他们的发现

作者在两个主要任务上测试了该方法：

语言建模（写作）： 他们在 WikiText-103 等数据集上进行了测试。
- 结果： 相比于之前的方法，他们的方案（SSR）写出的文本更好（具有更低的“困惑度/Perplexity”，这是衡量 AI 有多困惑的分数）。
- 速度： 它的训练速度更快，因为它不需要沉重的“唠叨型”损失函数。它只在极小比例的时间内使用复杂的数学运算。
图像分类（视觉）： 他们在 ImageNet 上测试了图像识别。
- 结果： 他们的识别准确率更高，并且能更好地处理异常或“对抗性”图像（即旨在欺骗 AI 的图像）。

核心结论

论文声称，选择性 Sinkhorn 路由（SSR） 是解决稀疏混合专家模型（SMoE）中“路由崩溃”问题的一种轻量级且高效的方法。

旧方法： 使用沉重、复杂的数学运算或强制性的惩罚项来强制实现平衡。（速度慢，有时不稳定。）
新方法（SSR）： 仅在需要时偶尔使用复杂的数学运算来引导系统，并加入一点随机性以保持探索性。
结果： 你得到了一个更聪明、更平衡的 AI，它训练得更快、表现更好，且没有额外的负担。

至关重要的一点是，论文强调“完美平衡”和“噪声”仅用于训练阶段。当模型在现实世界中投入使用时，它会切换回标准的、快速的、确定性的模式。这确保了最终产品既高质量又高效。

技术摘要：用于改进稀疏混合专家模型的选择性 Sinkhorn 路由

问题陈述

稀疏混合专家（SMoE）模型通过为每个 token 仅激活一部分专家，提供了一种在增加模型容量的同时保持计算效率的可扩展方法。然而，依赖于基于 Softmax 路由的传统 SMoE 模型通常面临“路由崩溃”问题，即一小部分专家被过度使用，而其他专家则利用不足。

现有的缓解策略通常依赖于：

辅助目标： 负载均衡损失或 z-loss，用于惩罚不均匀的专家使用情况。这些方法可能会引入目标失配、训练不稳定以及增加复杂度。
噪声门控： 注入噪声以鼓励探索，这会增加可训练参数或超参数。
基于 Sinkhorn 的路由： 使用 Sinkhorn 算法进行最优传输（OT）的前人方法通常仅将传输映射用于专家选择（top-k），而非权重分配。此外，由于门控矩阵无法以与 Softmax 方法相同的方式通过基于梯度的学习进行直接优化，它们可能会降低路由的灵活性；或者它们会产生巨大的训练开销，因为 Sinkhorn 更新具有迭代性质。

方法论

本文提出了一种轻量级机制——选择性 Sinkhorn 路由（Selective Sinkhorn Routing, SSR），将 token 到专家的分配重新表述为一个熵正则化的最优传输问题。

核心公式

作者将 $m$ 个 token 到 $n$ 个专家的分配建模为一个带有熵正则化的最大成本最优传输问题。其目标是找到一个传输计划 $\hat{\Pi}$ ，在满足特定约束的同时最大化兼容性：

正定性： $\Pi > 0$ 。
行和： 每个 token 将其全部质量路由至专家（ $\Pi \mathbf{1}_n = \mathbf{1}_m$ ）。
列和： 每个专家接收相等的预期总负载（ $\Pi^\top \mathbf{1}_m = \frac{m}{n} \mathbf{1}_n$ ）。

成本矩阵 $C$ 由门控分数 $S$ 推导而来。作者提出了两种变体：

线性成本： $C = S$ 。
Softmax 成本： $C_{i,:} = \text{softmax}(S_{i,:})$ ，用于防止 Sinkhorn 迭代过程中的数值溢出。

与以往仅将传输计划用于选择的 OT 方法不同，SSR 直接使用 $\hat{\Pi}$ 中的值来计算路由权重。对于每个 token，根据 $\hat{\Pi}$ 对应行中最高 entries 的专家进行 top- $k$ 选择，并从这些 entries 中归一化权重。

选择性策略与噪声注入

一个关键挑战在于，Sinkhorn 路由不会更新门控权重矩阵 $W_g$ ，因为 $W_g$ 与 OT 目标是解耦的。为了解决这个问题，SSR 采用了混合训练策略：

选择性应用： 在训练期间，每个 MoE 模块以较小的概率 $p$ （例如 0.001）使用 Sinkhorn 路由，否则使用标准的 Softmax 门控。这使得 Softmax 门控能够有效地更新 $W_g$ ，同时 Sinkhorn 路由能提供平衡的专家利用信号。
噪声注入： 为了防止专家崩溃并鼓励探索，在成本矩阵中加入高斯噪声（ $\tilde{C} = C + \alpha_{noise} \cdot \epsilon$ ）。理论分析（命题 4.3）表明，这确保了每个专家都有非零的选择概率。

推理行为

论文认为，在推理阶段对单个输入强制执行平衡约束（如相等的列和）是适得其反的，因为这会迫使无论输入的特定兼容性如何都进行均匀路由。因此，在推理期间，Sinkhorn 路由和噪声注入均被禁用，恢复为确定性的 Softmax 路由，以确保一致且依赖于输入的预测。

核心贡献

新型路由框架： 一种将熵正则化最优传输与随机噪声注入相结合的方法，旨在促进平衡的专家利用，而不使用辅助的负载均衡损失。
理论见解： 证明了 Sinkhorn 路由和噪声注入通过鼓励探索和平衡来辅助训练，但应在推理时禁用以避免扭曲的分配。
效率与性能： 广泛的评估表明，与最先进的基准相比，SSR 提高了训练效率、准确性和鲁棒性。

实验结果

作者在语言建模（WikiText-103, Enwik-8）和图像分类（ImageNet-1K, ImageNet-A, ImageNet-O, ImageNet-R）上评估了 SSR。

语言建模：
- 在 WikiText-103 上，SSR 变体始终优于 Vanilla SMoE 以及使用负载均衡损失、z-loss 或噪声的基准模型。
- SSR-L w/ noise 实现了 34.367 的测试困惑度（PPL），相比 Vanilla SMoE（35.550）降低了 1.183，显著超过了次优的基准模型。
- 至关重要的是，SSR 的训练开销极小（0.33% – 0.65%），而全量 Sinkhorn 路由 SMoE 的开销为 72.47%。
- 在 Momentum 设置下，SSR 保留了 Vanilla SMoE 特有的范数演化特征，同时提升了性能，而其他平衡方法则降低了性能。
字节级建模（Enwik-8）：
- SSR-S w/ noise 实现了最佳的测试每字符比特数（BPC）1.128，比 Vanilla SMoE 提升了 0.010 BPC。
- 其速度比全量 Sinkhorn 路由 SMoE 快 2.28倍。
视觉任务：
- 在 ImageNet-1K 上，SSR-L w/ noise 实现了 77.420 的 Top-1 准确率，比 Vanilla SMoE 提升了 +2.368。
- 它在 ImageNet-O 和 ImageNet-R 上也表现出强大的鲁棒性。
消融研究：
- 应用 Sinkhorn 路由的概率 $p$ 至关重要；极小的值（如 $10^{-3}$ 或 $10^{-4}$ ）足以提供平衡信号，而不会过度约束路由器。
- 噪声注入在各种设置下一致地提高了性能。
- 在推理时禁用平衡机制产生了最佳结果，证实了理论主张。

意义与主张

论文声称，选择性 Sinkhorn 路由为 SMoE 设计提供了一种实用且有效的解决方案。通过使用高效的、间歇性的 Sinkhorn 优化来取代复杂的辅助损失，SSR 实现了：

提高训练稳定性： 通过内置的专家平衡机制，避免了目标失配。
更高的准确率： 在文本和视觉基准测试中均得到验证。
鲁棒性： 增强了针对输入损坏和对抗性样本的性能。
高效性： 与辅助损失方法和全量 Sinkhorn 路由相比，计算开销极低。

作者强调，该方法是轻量级的，（在噪声变体中）不需要额外的可训练参数，并且在保持专家选择灵活性的同时，确保了训练阶段的平衡利用。