阿里 Qwen 团队论文：ID Balancing——MoE 训练专家负载不均的 PID 解法

arXiv 2609.39137v1 — "ID Balancing: Stable Training of Extremely Sparse MoE via PID-Based Load Control"

问题（MoE 训练老毛病）

MoE 模型路由越稀疏，每个 token 选中的 expert 越少 → 专家分配越敏感 → 一些 expert 过载、一些无人问津。效率掉、稳定掉、参数利用率掉。

之前三大流派：

• Aux loss：加梯度惩罚（影响主 loss、收敛变慢）
• DeepSeek loss-free：用 bias 调负载（简单但负载大时跟不上）
• Kimi K3 Quantile Balancing：用分位数做平衡（平滑但反应慢）

论文核心：把平衡问题重新表述为 PID 控制

PID = 比例（P）+ 积分（I）+ 微分（D）—— 自动控制经典框架

方法	控制视角	行为
DeepSeek loss-free	纯 I（积分）	固定步长累加负载误差
Kimi K3 Quantile	纯 P（比例）	用分位数估计即时偏差
ID Balancing	I + D	累加误差 + 偏差变坏时加速响应

关键洞察：P 处理"现在有多偏"，I 处理"过去累积偏多少"，D 处理"偏差在变好还是变坏"——只在变坏时介入（derivative gate）。

公式（每层 O(E)）

e_i = (n_avg - n_i) / n_avg   // 归一化负载误差
b_i ← b_i + K_i · e_i + K_d · g_i · (e_i - e_i_prev)
g_i = 1 当 e_i 与 (e_i - e_i_prev) 同号
g_i = 0 否则
b ← b - mean(b)              // 重新中心化

工程细节：

• 只用 token-count 反馈，O(E) 每层，不引入新梯度
• 不用 aux loss——主 loss 不被打扰
• K_i = K_d = 6e-3（pretrain）；续训时降到 6e-6 修漂移

实测

方法	Worst Backbone MaxVio（Top-3-of-768，18.9B / 120B tokens）
Aux loss baseline	~150
DeepSeek loss-free	211.19
Kimi K3 Quantile	31.11
ID Balancing	15.33

跨规模 18.9B → 69.9B，89.6% 改善 vs aux loss baseline。

关键扩展

• 跨规模稳定：18.9B → 69.9B MaxVio 几乎不变
• 2.3× 学习率仍稳：aux loss 和 loss-free 在 2.3× LR 下显著退化
• 下游质量不打折：24.8B 模型 9 个 benchmark 平均 55.81；69.9B 模型 58.66

证据分层

• 211 → 15.33 MaxVio：论文正文 Table 1
• 89.6% vs aux loss：69.9B 规模 Top-10-of-768 验证
• PID 视角统一三种方法：论文 §1 + §2 公式推导
• 9 个下游 benchmark：MMLU, MMLU-Pro, SuperGPQA, MATH, GSM8K, BBH, MMMLU, EvalPlus, MultiPL-E

数据来源

• 论文：https://arxiv.org/html/2609.39137v1
• 作者：Peng Jin, Zihan Qiu, Zekun Wang, Bo Zheng, Yang Xu, Tian Xie, Xiao Li, Huaqing Zhang, Haoran Lian, Rui Men, Dayiheng Liu（Alibaba Qwen Team）