$\phi$-平衡：面向混合专家训练

# φ-平衡：用于混合专家训练
来源：https://arxiv.org/html/2605.15403
Jonathan Li · Qi Wang · Runlong Liao · Shuozhe Li · Chen Liang · Ni Lao · Qiang Liu

###### 摘要

混合专家（MoE）模型依赖于专家之间的负载均衡，以充分发挥其可扩展性。然而，现有的负载均衡方法大多是启发式的，并且依赖于有噪声的小批量分配统计量，从而相对于总体层面的目标引入了偏差。我们提出φ-平衡，这是一个原则性的框架，通过最小化期望路由分布的严格凸、对称且可微的势函数，直接针对总体层面的专家平衡。利用凸对偶性，我们推导出一个等价的最小-最大公式，并通过镜像下降法获得一个简单的在线算法，从而得到一个基于指数移动平均（EMA）的高效路由调整方法，且开销可忽略不计。在大型预训练和下游微调中，φ-平衡始终优于之前的Switch风格和无损失基线，展示了更稳定和更有效的专家利用率。

混合专家，负载均衡，Transformer，稀疏性

## 1 引言

混合专家（MoE）Transformer 已成为一种有效的方法，通过为每个输入词元动态选择一小部分专家模块来扩展深度学习模型。这种策略显著增加了模型容量，同时保持计算量几乎不变（shazeer2017outrageously; fedus2022switch），使得拥有数十亿参数的大规模语言和视觉模型能够以大致恒定的FLOPs运行（lepikhin2021gshard; riquelme2021scaling; fedus2022switch）。

**算法1** 单个MoE层的φ-平衡
0: 严格凸、对称且可微的 φ, η∈(0,1], α>0,
  m←0, 路由频率
1: 计算每个词元 i 的路由概率 p_{i,e}
2: p_e ← (1/T) Σ_{i=1}^T p_{i,e} for e=1,...,E (专家负载)
3: 令 p = (p_1, ..., p_E)
4: m ← (1-η)m + ηp (负载的EMA)
5: L_aux ← { ST-MoE: Σ_{e=1}^E f_e p_e ; Ours: Σ_{e=1}^E ∇φ(m)_e p_e }
6: 使用 ∇(L_task + α·E·L_aux) 更新模型

![图1](https://arxiv.org/html/2605.15403)  
**图1：在推理和代码生成基准上的性能提升。** 我们将提出的方法（Ours）与基于Moonlight-16B-A3B-Instruct架构（LiuSY25）的ST-MoE基线进行比较。所提出的方法在所有选定任务上均优于基线，在数学推理（Math500）、通用能力（LiveBench）、代码合成（HumanEval）和逻辑（BBH）方面取得了显著提升。

MoE训练中的一个关键挑战是确保专家的均衡利用，这对于充分利用模型容量和避免性能下降至关重要。已经提出了许多方法来解决这一挑战，包括Switch风格的负载均衡损失（shazeer2017outrageously; lepikhin2021gshard; fedus2022switch）和最近的无损失均衡方法（wang2024auxiliary）。然而，一个经常未被言明的问题是，大多数现有的均衡目标本质上是启发式的，因为它们不对应于最小化一个明确定义的总体层面目标。原则上，真正的目标是在整个数据分布下实现专家使用的均衡。相比之下，广泛使用的方法（如Switch风格的MoE（ST-MoE））依赖于每小批量的统计量和实际的专家分配频率，这相对于总体层面的均匀性目标引入了系统性偏差。

我们提出φ-平衡，这是一个直接针对总体层面专家平衡的原则性负载均衡框架。我们的方法将负载均衡形式化为最小化一个应用于总体平均路由分布的严格凸、对称且可微的势函数φ。为了避免每批次近似引入的偏差，我们采用基于凸对偶性的最小-最大公式，并应用在线镜像下降法来解决由此产生的内部问题。这产生了一个简单而广泛的算法族，如算法1所示，它维护路由概率的指数移动平均（EMA），并通过镜像映射∇φ进行处理，开销可忽略不计。

实验上，我们发现φ-平衡在广泛设置中始终优于ST-MoE（图1），包括在MoE增强的Gemma模型预训练中（Kamath25; Liang25），我们在受控计算预算下系统地扩展活跃参数数量N、专家数量E和路由粒度G，并对基于EMA的负载跟踪、镜像映射φ的选择以及EMA衰减率进行了消融实验。虽然φ有多种选择，但我们推荐负熵函数作为实践中最高效的。

我们进一步在基于指令微调的MoE骨干网络（liu2024deepseek; dai2024deepseekmoe; LiuSY25）上，针对七个基准进行了逐基准的LoRA微调评估，所有实验总共使用了约40,000个NVIDIA H100 HBM3-80GB GPU小时。

## 2 混合专家背景

我们考虑一个由L层组成的标准解码器-only Transformer。在密集Transformer中，每一层通过自注意力模块处理输入序列，然后是一个共享的前馈网络（FFN）。MoE架构用稀疏模块层替代了这个密集FFN，该模块层由一个可学习路由器和一组E个专家 {FFN_1, ..., FFN_E} 组成（shazeer2017outrageously）。

令 x = (x_i)_{i=1}^T ∈ ℝ^{T×d} 表示一层的输入，其中T是序列长度，d是模型隐藏维度。对于每个词元 x_i，MoE层的输出 y_i 被计算为专家输出的路由器加权和：

y_i = Σ_{e=1}^E R(x_i)_e · FFN_e(x_i; d_ffn).  (1)

这里，每个专家被参数化为一个标准的两层MLP。遵循最新最先进的实现（shazeer2020glu; dai2024deepseekmoe; agarwal2025gpt），我们使用SwiGLU激活函数，定义如下：

FFN_e(u) = W_2^{(e)} · SwiGLU(W_1^{(e)} u),  (2)

其中 W_1^{(e)} ∈ ℝ^{d_ffn × d} 和 W_2^{(e)} ∈ ℝ^{d × d_ffn} 是专家e的独立参数。

### 2.1 稀疏路由机制

MoE的计算效率依赖于路由函数 R(·)，该函数通过将每个词元引导到一小部分k个专家（其中 k ≪ E）来强制实现稀疏性。路由器通常由一个可学习的投影矩阵 W_r ∈ ℝ^{E×d} 组成。*路由权重*通过在前k个索引上归一化投影分数来确定（shazeer2017outrageously）：

R(x) = softmax(Top-k(W_r x)).  (3)

Top-k(·) 运算符将除最大的k个元素之外的所有logits设置为 -∞。因此，对于所有未选中的专家，R(x)_e 为零，允许模型跳过大部分专家计算。如果我们只激活一个专家，那么我们将不进行softmax以避免路由器logits上的零梯度。这种条件计算将参数数量与推理成本解耦；然而，它引入了我们在第3节中解决的负载均衡挑战。

### 2.2 基线负载均衡策略

![图2](https://arxiv.org/html/2605.15403)  
**图2：受控每词元计算下的预训练扩展研究。** 我们沿三个正交的MoE扩展轴评估路由稳定性和优化，同时在每项研究中通过按需调整专家大小来保持每词元计算成本（FLOPs）大致恒定。（左）活跃参数扩展：我们使用E=16个专家和每个词元A=2个活跃专家训练模型，变化*活跃参数*数量 N∈{111M, 338M, 588M, 986M}。（中）粒度扩展：对于固定的模型大小M和激活比率A/E，我们通过将专家总数从16增加到256并相应缩小每个专家来变化粒度因子 G∈{2,4,8,16,32}，因此每词元FLOPs保持不变。（右）专家数量扩展（激活比率）：我们通过保持计算预算M、激活专家数量A=2和专家大小（粒度）固定，并变化专家总数E∈{8,16,32,64,128}，来隔离A/E的影响。

虽然路由器在总参数量中占微乎其微的比例，但它协调了模型庞大专家容量的利用。在此，我们回顾ST-MoE（fedus2022switch）使用的标准辅助负载均衡损失（LBL）。这种公式仍是训练大规模稀疏模型（包括DeepSeek（liu2024deepseek）、OlMoE（muennighoff2024olmoe）和DeepSpeed-MoE（rajbhandari2022deepspeedmoe））的主流范式。

LBL目标鼓励词元在E个专家上均匀分布。对于一个包含T个词元的小批量，令 p_e 表示分配给专家e的批量平均*前k个*路由概率，p_{i,e} 表示专家e对于词元 x_i 的路由概率，f_e 表示在前k个路由下专家e的实际路由频率：

p_e = (1/T) Σ_{i=1}^T p_{i,e}, 其中 p_{i,e} = softmax(W_r x_i)_e,  (4)
f_e = (1/(kT)) Σ_{i=1}^T I(e ∈ Top-k(W_r x_i)). 

辅助损失定义为这两个向量的点积：

L_aux = Σ_{e=1}^E f_e · p_e.  (5)

如fedus2022switch所示，最小化(5)既鼓励门控概率也鼓励离散选择趋近于均匀分布。

**无损失均衡。** 不引入会注入干扰梯度并损害任务学习的显式负载均衡损失，*无损失均衡*（wang2024auxiliary）通过直接修改路由决策来强制平衡。具体地，它在*前k个*选择*之前*向路由器logits添加一个可学习的、专家特定的偏置，并使用每个专家的近期利用率在线更新这些偏置。

![图3](https://arxiv.org/html/2605.15403)  
**图3：预训练动态和专家利用率。** 我们在10k步内比较φ-平衡（红色，实线）与ST-MoE（蓝色，虚线）。（左）验证损失和准确率显示φ-平衡（负熵）达到了可比或更优的收敛。（右）基尼系数和专家负载分析显示φ-平衡的路由不平衡显著更低。φ-平衡在最大和最小专家负载之间保持了更紧的界限，更接近完美分配线（绿色），而ST-MoE在专家容量使用上表现出更高的方差。

## 3 φ-平衡

在本节中，我们介绍φ-平衡损失。与仅在单个小批量内强制平衡的经典方法不同，我们的目标是正则化整个数据分布上的*全局*专家使用。具体地，我们通过一个严格凸、对称且可微的势函数φ来鼓励全局均匀的专家利用。

### 3.1 全局负载均衡目标

令 p(x; θ) ∈ Δ^E 表示输入词元x的预测路由概率向量，参数化为θ（即 p(x; θ) = softmax(W_r x)）。对于特定专家e，p(x; θ)_e 表示分配给该专家的概率质量。我们将*全局平均路由分布* p̄(θ) 定义为路由概率在训练语料所引导的词元分布 D 上的期望：

p̄(θ) = E_{x~D}[p(x; θ)],  (6)

满足 Σ_{e=1}^E p̄(θ)_e = 1。

#### 通过凸对偶性进行负载均衡。

我们的目标是鼓励词元总体层面的路由分布 p̄(θ) 变得均匀，以便在期望上，所有专家在数据分布上被同等利用。我们将此目标形式化为优化问题

min_θ L_bal(θ) := min_θ φ(p̄(θ)),  (7)

其中势函数 φ: ℝ^E → ℝ 被选择为严格凸、对称且可微的。

φ的严格凸性和对称性保证了 (7) 中的目标在概率单纯形上在均匀分布处达到唯一最小值，这在附录LABEL:app:proofs中的引理LABEL:lem:uniform-minimizer中形式化。φ的代表性选择总结于表1中。重要的是，φ不限于诸如 Σ_e ψ(p_e) 之类的加性或可分离形式，并且可以捕获专家之间更一般的依赖关系。

#### 估计挑战。

直接使用随机梯度下降优化 (7) 是有问题的。由于 p̄(θ) 是数据集上的期望，必须对其进行估计，并且使用小批量 B 的局部均值（记为 p̂ = (1/|B|) Σ_{x∈B} p(x; θ)）会引入显著偏差。因为 φ 是非线性的，函数的期望不等于期望的函数：

E_B[φ(p̂)] ≠ φ(E_B[p̂]) = φ(p̄(θ)).