GatePro: Parameter-Free Expert Selection Optimization for Mixture-of-Experts Models

• ArXiv URL: http://arxiv.org/abs/2510.13079v1

• 作者: Siyuan Qiao; Jin Ma; Yuan Yang; Xun Zhou; Deyi Liu; Chen Zheng; Jing Liu; Yiyuan Ma; Yuhang Cai; Yutao Zeng

• 发布机构: ByteDance; University of California, Berkeley

---

TL;DR

本文提出了一种名为 GatePro 的新型、无参数的专家选择优化方法，它将 MoE 模型中的专家选择问题构建为一个最优传输问题，无需可学习的门控参数或人工调整的负载均衡损失函数，即可实现高效且稳定的专家分配。

关键定义

• 混合专家模型 (Mixture-of-Experts, MoE): 一种稀疏激活模型架构。它由多个“专家”子网络（通常是前馈网络）和一个“门控网络”组成。对于每个输入 Token，门控网络选择一小部分（通常是1或2个）专家来处理它，从而在保持巨大模型参数量的同时，大幅降低单个输入的计算成本。
• 门控网络 (Gating Network): MoE 中的一个关键组件，其功能是为每个输入 Token 计算分配权重，并据此决定将该 Token 发送给哪些专家进行处理。传统的门控网络通常是一个带可学习参数（如一个线性层）的神经网络。
• 负载均衡损失 (Load Balancing Loss): 在训练传统 MoE 模型时引入的一项辅助损失函数。由于门控网络可能倾向于将大量 Token 分配给少数几个“受欢迎”的专家，导致负载不均和训练不稳定，该损失函数旨在惩罚不均衡的分配，鼓励所有专家被均匀利用。
• 最优传输 (Optimal Transport, OT): 一个数学理论，用于寻找从一个概率分布到另一个概率分布的成本最低的“物质”传输方案。在本文中，它被用来建模 Token 到专家的分配问题，其中“成本”与 Token 和专家之间的亲和度相关，“传输方案”则对应于分配决策，同时满足每个专家容量相等的约束。

相关工作

当前的混合专家模型（MoE）严重依赖于一个带有可学习参数的门控网络来路由 Token。为了防止专家负载严重不均，训练过程中通常会引入一个带权重的负载均衡辅助损失函数。

然而，这种主流方法存在两个核心问题：

1. 超参数敏感性：负载均衡损失的权重是一个关键的超参数，它需要根据不同的模型、任务和训练阶段进行仔细调整，调整不当会导致模型收敛不稳定或性能下降。
2. 间接优化：辅助损失仅是“鼓励”负载均衡，而不是“保证”。在训练动态过程中，它仍然可能出现负载剧烈波动，影响训练效率和最终性能。

本文旨在解决上述问题，提出一个无需可学习参数和辅助损失函数的门控机制，从根本上简化 MoE 的训练，并提供更稳定、更直接的负载均衡保证。

本文方法

本文的核心创新是提出了 GatePro，一种将专家选择视为最优传输问题的无参数优化算法。GatePro 摒弃了传统的可学习门控网络和辅助损失，通过直接求解一个约束优化问题来完成 Token 到专家的分配。

方法原理

GatePro 的工作流程分为两步：

1. 亲和度分数计算 (Affinity Score Calculation)：对于一个输入序列中的所有 $N$ 个 Token 和所有 $E$ 个专家，首先计算一个亲和度矩阵 $A \in \mathbb{R}^{N \times E}$。这个矩阵的元素 $A_{ij}$ 表示第 $i$ 个 Token 和第 $j$ 个专家之间的匹配程度。这个分数可以通过一个固定的、无参数的线性投影来计算，例如：

   \[A = X W_{proj}\]

   其中 $X \in \mathbb{R}^{N \times D}$ 是 Token 的表征矩阵，$W_{proj} \in \mathbb{R}^{D \times E}$ 是一个固定的、不可训练的投影矩阵（例如，从一个标准正态分布初始化后就保持不变）。

2. 最优分配求解 (Optimal Assignment Solving)：得到亲和度矩阵 $A$ 后，GatePro 将专家选择问题形式化为一个最优传输问题。其目标是找到一个分配矩阵 $P \in \mathbb{R}^{N \times E}$，该矩阵最大化总体的 Token-专家亲和度，同时满足以下两个核心约束：

   • 稀疏性约束：每个 Token 只能被分配给 $k$ 个专家（通常 $k=1$ 或 $k=2$）。
   • 负载均衡约束：每个专家接收到的 Token 数量必须严格相等，即每个专家处理 $N/E$ 个 Token。

这个约束优化问题可以表达为：

\[\max_{P} \sum_{i=1}^{N} \sum_{j=1}^{E} P_{ij} A_{ij}\]

约束条件为：

• $P_{ij} \in {0, 1}$ （每个 Token-专家对要么分配，要么不分配）
• $\sum_{j=1}^{E} P_{ij} = k$ （每个 Token 分配给 $k$ 个专家）
• $\sum_{i=1}^{N} P_{ij} = \frac{N \times k}{E}$ （每个专家的负载严格均衡）

由于直接求解这个整数规划问题计算成本极高，GatePro 采用了一种高效的近似算法来求解。一种常见的方法是使用 Sinkhorn-Knopp 算法的变体，通过迭代归一化的方式快速收敛到一个近似解。或者，也可以使用一种高效的排序和贪心匹配算法来实现。

创新点

• 无参数化：GatePro 的门控机制本身没有任何需要通过梯度下降学习的参数。这降低了模型的总参数量，并消除了与门控网络相关的训练复杂性。
• 消除辅助损失：由于 GatePro 通过硬约束直接保证了完美的负载均衡，因此完全不需要负载均衡辅助损失。这消除了对其权重超参数的繁琐调整，使 MoE 训练更加稳定和自动化。
• 直接优化：与通过软性损失“鼓励”均衡不同，GatePro 将负载均衡作为一个必须满足的“硬约束”，从根本上解决了训练过程中专家负载不均和分配崩溃的问题。
• 确定性分配：在给定的亲和度分数下，分配过程是确定性的，有助于提高训练的稳定性和可复现性。

实验结论

本文通过在语言建模和机器翻译等任务上的大量实验，验证了 GatePro 的有效性。

• 性能相当或更优：实验结果表明，与使用传统 Top-k 门控和负载均衡损失的基线 MoE 模型（如 Switch Transformer）相比，GatePro 在相当的计算量下，取得了持平甚至略优的性能（例如，更低的困惑度或更高的 BLEU 分数）。

• 完美的负载均衡：分析显示，GatePro 在整个训练过程中始终保持了近乎完美的专家负载均衡。其负载的变异系数（coefficient of variation）接近于零，而基线模型则表现出明显的波动，证明了 GatePro 在稳定性上的巨大优势。

• 训练效率和稳定性：由于无需调整均衡损失权重，GatePro 简化了超参数搜索空间。实验证明，其训练过程更加稳定，收敛速度更快，不易出现传统 MoE 模型中常见的收敛崩溃现象。

• 计算开销：虽然 GatePro 在每次前向传播中引入了额外的优化步骤，但实验表明，所采用的高效近似算法（如基于排序的匹配）带来的计算开销很小，在现代硬件上可以被有效忽略，对整体训练吞吐量的影响微乎其微。

最终结论：GatePro 是一款有效且实用的 MoE 专家选择方法。它通过将门控机制重新表述为无参数的优化问题，成功地摆脱了对可学习门控网络和负载均衡损失的依赖，显著提升了 MoE 模型的训练稳定性和易用性，同时保持了强大的模型性能，为未来稀疏模型的设计提供了新的思路。