cs.LGSep 30, 2026

ID Balancing: Stable Training of Extremely Sparse MoE via PID-Based Load Control

Authors: Peng Jin, Zihan Qiu, Zekun Wang, Bo Zheng, Yang Xu, Tian Xie, Xiao Li, Huaqing Zhang, +3 more

Organizations: Qwen Team, Alibaba Token Hub, Alibaba Group

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Scaling Large Language Models (LLMs) via Mixture-of-Experts (MoE) enables massive parameter growth with nearly constant per-token computation. However, further scaling the parameter count requires increasingly sparse routing, where expert load imbalance becomes more severe. This imbalance reduces parameter utilization and training efficiency, and can undermine training stability, becoming a bottleneck to reliable scaling. In this work, we unify two representative auxiliary-loss-free methods as incomplete Proportional-Integral-Derivative (PID) controllers: DeepSeek's loss-free method acts as a fixed-step integral controller, while Kimi K3's Quantile Balancing functions as a generalized proportional controller. Building on this control perspective, we propose ID Balancing, an Integral-Derivative controller. It scales its integral term with load error and activates its derivative term only when imbalance worsens, enabling stronger corrections for large or worsening errors and smaller updates near balance. Evaluated across Top-1010, Top-55, and Top-33 routing over 768768 experts, ID Balancing reduces worst-case backbone MaxVio and training-average backbone MinVio by over 50%50\% and 12%12\%, respectively, relative to the best baselines in the Top-33 setting. When the total parameter count increases from 18.918.9B to 69.969.9B (Top-1010-of-768768), ID Balancing's worst-case backbone MaxVio remains nearly unchanged and is approximately 89.6%89.6\% lower than that of the auxiliary-loss baseline. ID Balancing also maintains competitive language-modeling and downstream performance. The advantages of ID Balancing grow as sparsity increases, making it a promising solution for scaling larger, sparser MoE models.

Figures & tables

Figure 1: ID Balancing maintains load control as routing becomes sparser and model capacity grows. (a) Worst backbone MaxVio for Top- 1010 , Top- 55 , and Top- 33 routing over 768768 experts. (b) Training-average backbone MinVio, expressed as a decimal ratio, across the same settings. Models in (a) and (b) have 18.918.9 B total parameters and are trained on 120120 B tokens. (c) Worst backbone MaxVio as total parameters increase from 18.918.9 B to 69.969.9 B ( 1.031.03 B to 3.23.2 B active) under fixed Top- 1010 -of- 768768 routing, using the first 3030 k and 5050 k steps of the small and large models, respectively.
Figure 2: ID Balancing limits transient overload and sustained underload in highly sparse MoE training. Results use Top- 33 -of- 768768 routing in an 18.918.9 B-parameter model trained on 120120 B tokens, with metrics averaged over 2020 backbone layers. (a) Mean MaxVio on a logarithmic scale ( ↓\downarrow ). (b) Mean MinVio ( ↓\downarrow ). ID Balancing achieves the lowest training-average underload.
Figure 3: Magnitude-aware integral control reduces early overload and sustained underload. We compare the integral-only update with DeepSeek’s loss-free method in an 18.918.9 B model using Top- 33 -of- 768768 routing and 120120 B training tokens. (a) Maximum MaxVio across the 2020 backbone layers, shown on a logarithmic scale ( ↓\downarrow ). (b) Mean MinVio across the same layers ( ↓\downarrow ).
Figure 4: The worsening-gated derivative term reduces early overload and expert concentration. We compare magnitude-aware integral control with and without the derivative term in an 18.918.9 B model using Top- 33 -of- 768768 routing and 120120 B training tokens. Metrics are averaged over 2020 backbone layers. (a) gˉ(t)\bar{g}^{(t)} is the fraction of active gates averaged over the 20 backbone layers. (b) Mean MaxVio ( ↓\downarrow ). (c) Share of token assignments routed to the busiest 10%10\% of experts ( →0.10\rightarrow 0.10 ), with a balanced target of 0.100.10 .
Figure 5: EMA smoothing reduces bias drift but slows load correction in Quantile Balancing. We compare the original update ( ρ=0\rho=0 ), its EMA variants, DeepSeek’s loss-free method, and ID Balancing in an 18.918.9 B model using Top- 33 -of- 768768 routing and 120120 B training tokens. All other comparisons use Quantile Balancing without EMA. (a) Mean absolute bias change ∣Δb∣|\Delta b| between checkpoints 11 k steps apart during the final 55 k steps ( ↓\downarrow ). (b) Mean MaxVio on a logarithmic scale ( ↓\downarrow ). (c) Mean MinVio ( ↓\downarrow ).
  BackboneMTP Module
MethodAct./TotalLMMaxVio ↓\downarrowMinVio ↓\downarrowMaxVio ↓\downarrowMinVio ↓\downarrow
  lossLast1kAvg.WorstLast1kAvg.Last1kAvg.WorstLast1kAvg.
Top-3 of 768 experts, 120B tokens
Auxiliary loss0.87B / 18.9B1.75271.48091.723970.760.54620.62323.89524.463943.380.65800.7502
DeepSeek loss-free1.74260.73141.9356211.190.74870.74410.74770.893017.320.68680.6533
Quantile Balancing1.74320.80930.798531.110.65040.63731.87771.207130.890.71060.7083
Table 1: ID Balancing improves backbone load control while maintaining competitive LM loss. We compare Top- 33 , Top- 55 , and Top- 1010 routing over 768768 experts in 18.918.9 B models trained on 120120 B tokens. All summaries use the same 3030 k-step window. Backbone Last1k and Avg. metrics are averaged over 2020 layers, while Worst MaxVio is the maximum across layers and steps. MTP metrics are reported separately.
Figure 6: Reduced ID Balancing gains limit load drift while retaining competitive LM loss during continued pretraining. We compare full gains ( Ki=Kd=6×10−3K_{i}=K_{d}=6\times 10^{-3} ), gains reduced by 1000×1000\times ( 6×10−66\times 10^{-6} ), and a frozen bias in 2828 -layer models at a constant learning rate of 3×10−53\times 10^{-5} . The frozen setting matches the continued-pretraining baselines for DeepSeek’s loss-free method and Quantile Balancing. The top and bottom rows use Top- 88 -of- 256256 and Top- 1010 -of- 768768 routing, respectively. (a, d) LM loss with 201201 -step smoothing ( ↓\downarrow ). (b, e) Mean MaxVio ( ↓\downarrow ). (c, f) Mean MinVio ( ↓\downarrow ).
MethodKnowledgeSTEMReasoningMultilingualCodeAvg.
MMLUMMLU-ProSuperGPQAMATHGSM8KBBHMMMLUEvalPlusMultiPL-E
Top-8 of 256 experts, 3.0B / 24.8B, 560B tokens
Auxiliary loss68.0047.1027.3647.1274.9469.5759.8251.5146.2254.63
DeepSeek loss-free70.0847.2327.4345.8476.5471.0960.6153.9646.6655.49
Quantile Balancing70.4948.0627.6148.5474.3769.9460.1752.2742.4654.88
ID Balancing69.5848.1927.5247.9274.4972.9160.7856.4844.4055.81
Table 2: ID Balancing maintains competitive downstream performance. Results use Top- 88 -of- 256256 and Top- 1010 -of- 768768 routing. Avg. is the mean across nine benchmarks covering knowledge, STEM, reasoning, multilingual understanding, and code generation. Block headers list active and total parameter counts. Higher scores are better, and bold marks the best result in each column within a block.
 BackboneMTP Module
MethodLMMaxVio ↓\downarrowMinVio ↓\downarrowMaxVio ↓\downarrowMinVio ↓\downarrow
 lossLast1kAvg.WorstLast1kAvg.Last1kAvg.WorstLast1kAvg.
Auxiliary loss2.00993.13832.837920.890.84230.84583.56403.05334.830.65690.7130
DeepSeek loss-free2.00311.37551.781172.540.87300.82690.50330.712417.980.61050.5514
Quantile Balancing2.00090.50210.49792.960.40830.39380.62430.583610.770.43460.4308
ID Balancing2.00070.64370.727411.970.48280.49500.46420.50687.200.39400.3927
Table 3: ID Balancing and Quantile Balancing maintain load control at a higher learning rate. Results use a 1.031.03 B-active/ 18.918.9 B-total model with Top- 1010 -of- 768768 routing, trained on 120120 B tokens at a constant learning rate of 5.86×10−35.86\times 10^{-3} ( 2.3×2.3\times the standard peak). Both methods yield lower backbone load violations than Auxiliary loss and DeepSeek’s loss-free method.
Figure 7: ID Balancing and Quantile Balancing maintain lower backbone overload and underload at a higher learning rate. Results use Top- 1010 -of- 768768 routing and 120120 B training tokens at a constant learning rate of 5.86×10−35.86\times 10^{-3} ( 2.3×2.3\times the standard peak). Metrics are averaged over 2020 backbone layers. (a) Mean MaxVio on a logarithmic scale ( ↓\downarrow ), showing a large early spike for DeepSeek’s loss-free method. (b) Mean MinVio ( ↓\downarrow ), showing more severe underload for Auxiliary loss and DeepSeek’s loss-free method.
Figure 8: Larger auxiliary-loss coefficients worsen LM loss and load balance at a higher learning rate. We compare α∈{0.05,0.10,0.50}\alpha\in\{0.05,0.10,0.50\} in a 1.031.03 B-active/ 18.918.9 B-total model with Top- 1010 -of- 768768 routing at a constant learning rate of 5.86×10−35.86\times 10^{-3} ( 2.3×2.3\times the standard peak). Training targets 120120 B tokens, but both larger- α\alpha runs terminate before 3030 k steps. Load metrics are averaged over 2020 backbone layers. (a) LM loss ( ↓\downarrow ), with higher values at larger coefficients. (b) Mean MaxVio on a logarithmic scale ( ↓\downarrow ), exceeding 5050 for α=0.50\alpha=0.50 . (c) Mean MinVio ( ↓\downarrow ), reaching 1.01.0 for both larger coefficients.
Figure 9: Load-balancing methods show distinct gradient and activation dynamics. Results use a 3.03.0 B-active/ 24.824.8 B-total model with Top- 88 -of- 256256 routing over 5050 k steps. (a) Global gradient norm. ID Balancing and Quantile Balancing have lower, smoother trajectories, while DeepSeek’s loss-free method shows stronger late-stage fluctuations. (b) MoE-output magnitude averaged over the 2828 backbone layers.
Figure 10: ID Balancing maintains consistent overload and underload control across layers. Results use Top- 33 -of- 768768 routing over 3030 k steps and cover all 2020 backbone layers. The top and bottom rows show the training average and the final 11 k-step average, respectively. (a, c) MaxVio ( ↓\downarrow ). (b, d) MinVio ( ↓\downarrow ). ID Balancing achieves lower or comparable MaxVio and lower MinVio than Quantile Balancing across all layers. MinVio generally increases with depth for all four methods.
Figure 11: ID Balancing reduces expert inactivity relative to Auxiliary loss. Results are averaged by layer over nine downstream benchmarks. The top and bottom rows use Top- 88 -of- 256256 and Top- 1010 -of- 768768 routing, respectively. (a, d) Inactive-expert ratio, the fraction of experts receiving no evaluation token ( ↓\downarrow ). (b, e) Top- KK gating-score entropy, where a larger value means a more even mixture over selected experts. (c, f) Full-pool selection-score entropy, where a larger value means a flatter selection-score distribution.
 BackboneMTP module
SettingLMMaxVio ↓\downarrowMinVio ↓\downarrowMaxVio ↓\downarrowMinVio ↓\downarrow
 lossLast1kAvg.WorstLast1kAvg.Last1kAvg.WorstLast1kAvg.
Integral gain KiK_{i} , integral term only, Kd=0K_{d}{=}0
Ki=3×10−3K_{i}=3\times 10^{-3}1.74220.67590.945923.260.51900.57770.63750.79919.140.56030.5775
Ki=6×10−3K_{i}=6\times 10^{-3}1.74310.69240.800314.960.52330.55270.85570.886215.180.58480.5714
Ki=9×10−3K_{i}=9\times 10^{-3}1.74290.68310.787220.130.52630.54360.92221.038835.560.62790.5800
Table 4: An integral gain of Ki=6×10−3K_{i}=6\times 10^{-3} yields the lowest worst-case backbone MaxVio among the tested gains. Results use a 0.870.87 B-active/ 18.918.9 B-total model with Top- 33 -of- 768768 routing over 3030 k steps. We fix Kd=0K_{d}=0 and vary Ki∈{3,6,9}×10−3K_{i}\in\{3,6,9\}\times 10^{-3} . Increasing KiK_{i} to 9×10−39\times 10^{-3} further reduces training-average backbone violations but worsens MTP balance. Lower values are better for all metrics.
 BackboneMTP module
SettingLM00 – 11 k steps11 – 55 k stepsAvg.00 – 55 k stepsAvg.
 lossMaxVio ↓\downarrowMinVio ↓\downarrowMaxVio ↓\downarrowMinVio ↓\downarrowMaxVio ↓\downarrowMinVio ↓\downarrowMaxVio ↓\downarrowMinVio ↓\downarrowMaxVio ↓\downarrowMinVio ↓\downarrow
Derivative gain KdK_{d} , fixed Ki=6×10−3K_{i}=6\times 10^{-3}
Kd=0K_{d}=0 (I only)1.74312.00200.79550.90590.59890.80030.55271.49650.62220.88620.5714
Kd=3×10−3K_{d}=3\times 10^{-3}1.74281.89000.78530.88690.59920.79100.54851.74540.62010.94590.5799
Kd=6×10−3K_{d}=6\times 10^{-3}1.74261.93300.77910.86120.58990.79310.54801.77240.63360.96980.5782
Table 5: Larger derivative gains reduce backbone overload over 11 – 55 k steps but increase MTP overload. Results use a 0.870.87 B-active/ 18.918.9 B-total model with Top- 33 -of- 768768 routing over 3030 k steps. We fix Ki=6×10−3K_{i}=6\times 10^{-3} and vary Kd∈{0,3,6,12}×10−3K_{d}\in\{0,3,6,12\}\times 10^{-3} , with Kd=0K_{d}=0 as the integral-only baseline. Load metrics summarize the indicated early intervals and the full training window (Avg.).
Figure 12: A small integral gain responds slowly to the early imbalance, while larger gains converge to similar backbone trajectories. Results use Top- 33 -of- 768768 routing over 3030 k steps with Kd=0K_{d}=0 . We compare Ki∈{3,6,9}×10−3K_{i}\in\{3,6,9\}\times 10^{-3} . (a) Mean backbone MaxVio ( ↓\downarrow ). (b) Mean backbone MinVio ( ↓\downarrow ).
Figure 13: The derivative term reduces early backbone overload. Results use Top- 33 -of- 768768 routing with Ki=6×10−3K_{i}=6\times 10^{-3} and Kd∈{0,3,6,12}×10−3K_{d}\in\{0,3,6,12\}\times 10^{-3} , where Kd=0K_{d}=0 is the integral-only baseline. Let VtV_{t} denote mean MaxVio over the 2020 backbone layers. (a) VtV_{t} over the first 1010 k steps, smoothed with a five-point moving average. (b) Cumulative excess above Vt=1V_{t}=1 over 11 – 55 k steps.
Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

SpecificationSmall MoEMedium MoELarge MoE
Transformer layers202828
Hidden size102420482048
Vocabulary size248320248320248320
Softmax-attention heads161616
Query groups222
Softmax-attention head dimension256256256
Appendix
Table A: Architectural specifications for the evaluated MoE models. Parameter counts include the input embedding and untied output head, and exclude the MTP module. Active parameters include all dense and shared components together with the routed experts selected by Top- KK .

Explore similar work

CardsList

  1. A Replicate-and-Quantize Strategy for Plug-and-Play Load Balancing of Sparse Mixture-of-Experts LLMs

    Feb 23, 2026Zijie Liu, Jie Peng, Jinhao Duan +7Load BalancingLLM Inference Optimization

    University of North Carolina at Chapel Hill · University of Minnesota Twin Cities · North Carolina State University +2

  2. EMO: Frustratingly Easy Progressive Training of Extendable MoE

    May 13, 2026Linghao Jin, Chufan Shi, Huijuan Wang +4Mixture-Of-ExpertsExpert Demonstrations

    USC-ISI · MBZUAI-IFM

  3. How to Scale Mixture-of-Experts: From muP to the Maximally Scale-Stable Parameterization

    May 13, 2026Leena Chennuru Vankadara, Moritz Haas, Luke Hayward +2Mixture-Of-ExpertsScaling Laws

    1Gatsby Computational Neuroscience Unit, University College London · 2Amazon · University of Tübingen, Tübingen AI Center