ID Balancing: Stable Training of Extremely Sparse MoE via PID-Based Load Control
Peng Jin, Zihan Qiu, Zekun Wang, Bo Zheng, Yang Xu, Tian Xie, Xiao Li, Huaqing Zhang, Haoran Lian, Rui Men, Dayiheng Liu
cs.LG
2026-09-30
Qwen 把 loss-free MoE 均衡统一成 PID 控制问题,新方法在 256 倍稀疏路由下把最差过载从 31 降到 15,LM loss 持平。
MoE 的扩展方式是总参数涨、单 token 计算量几乎不变:router 每个 token 只从 E 个专家里选 K 个。专家数加大、K 不动,路由越来越稀疏,负载失衡也随之加剧:Top-K 边界附近的微小分数变化就能改掉 token 归属,过载专家拖慢流水线,欠载专家训练不足,严重时威胁训练稳定性。Top-3-of-768 的活跃比例只有 0.39%(1/256),现有均衡方法在这个稀疏度下开始失灵。
工业界主推 auxiliary-loss-free 路线:不加均衡损失,给每个专家挂一个不参与梯度更新的 bias,用路由反馈去调它。DeepSeek-V3 和 Kimi K3 是两个代表,论文把它们读成同一个控制器的两块碎片。
PID 控制器是过程控制的标准件,P、I、D 分支分别响应当前误差、累积误差、误差变化率。把均匀负载设为目标、专家 bias 当控制状态、token 计数当测量输出,两个已有方法就能对号入座:DeepSeek loss-free 让 bias 按误差符号走固定步长,误差大小被扔掉,大失衡和小失衡拿到同样的修正,响应大误差慢、接近均衡时容易过冲,是只有 I 项的固定步长积分控制;Kimi K3 的 Quantile Balancing 从当前 batch 分数分布的分位数直接算出目标 bias,一步设过去,目标随 batch 漂移、后期抖动大,近似只有 P 项的广义比例控制。
ID Balancing 补上缺口,只用 I 和 D 两项:
开销只有每层 O(E) 的 token 计数反馈,不给语言建模目标加任何均衡梯度,默认 Ki = Kd = 6×10⁻³。
主实验用 Qwen-3.8-Next 架构,20 层加一个 MTP 模块,每层 768 个专家,约 18.9B 总参、0.87B–1.03B 活跃,120B token。Top-3-of-768 一档:
| 方法 | Worst MaxVio ↓ | 平均 MinVio ↓ | LM loss |
| Auxiliary loss | 70.76 | 0.623 | 1.7527 |
| DeepSeek loss-free | 211.19 | 0.744 | 1.7426 |
| Quantile Balancing | 31.11 | 0.637 | 1.7432 |
| ID Balancing | 15.33 | 0.548 | 1.7426 |
最差过载比最优基线(Quantile Balancing)低 50% 以上;DeepSeek 的方法在这一档最差过载冲到 211。规模从 18.9B 加到 69.9B(Top-10-of-768),ID Balancing 的最差 MaxVio 只从 8.36 挪到 7.79,auxiliary loss 涨到 75.03,差距 89.6%。
学习率恒定在标准峰值 2.3 倍的压力测试里,ID Balancing 和 Quantile Balancing 的 Last1k backbone MaxVio 为 0.644 和 0.502,DeepSeek 1.376,auxiliary loss 3.138;后者把系数 α 加到 0.5 反而崩掉,loss 升到约 2.45,有专家接不到任何 token,加大系数的运行都在 30k 步内提前终止。
下游九个基准(560B token):3.2B/69.9B 档 ID Balancing 平均 58.66,四方法最高,auxiliary loss 58.12;3.0B/24.8B 档 55.81,略低于 Quantile Balancing 的 55.88。持续预训练把增益除以 1000,LM loss 与冻结 bias 的差距不到 0.001,平均 MaxVio 从 6–7 降到约 1.4;推理时未激活专家占比从 8.4% 降到 6.4%,第一层从 9.3% 降到 2.6%。
对训练大稀疏 MoE 的团队,这是即插即用的 bias 更新规则替换:不加梯度、开销 O(E)。稀疏度越高优势越大,而 18.9B 到 69.9B 最差过载不随规模恶化,直接回应了专家数上到几百个之后训练还能不能稳住的担忧。把两个工业级方法放进同一个 PID 框架、看清各自缺哪个分支,这套分析也能搬到其他反馈式训练机制上。定位要说清:这是负载控制层面的渐进改进,LM loss 和下游分数与基线持平,价值在稳定性,不在模型质量。
作者自述三条:只验证了 fixed Top-K 的 decoder-only MoE 家族;门控微分的累积形式只有经验支撑、没有理论;scaling 实验里深度、宽度、活跃与总参数、训练时长同时变化,效果无法单独归因给规模。另有几点论文没有正面处理:主实验只有 20 层、120B token,离产线规模尚远;Quantile Balancing 在 Top-5/Top-10 的训练平均 MaxVio 和高学习率下的 backbone 均衡上仍占上风,ID Balancing 只在最稀疏的 Top-3 全面领先,标题里的「稳定」更接近「最差情况更好」;加 D 项让 MTP 模块平均 MaxVio 从 0.886 升到 0.970;Ki、Kd 靠扫描选定,持续预训练还要再除以 1000,换架构后是否仍成立没有验证。