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
Feb 23, 2026cs.LG

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

Sparse Mixture-of-Experts (SMoE) architectures are increasingly used to scale large language models efficiently, delivering strong accuracy under fixed compute budgets. However, SMoE models often suffer from severe load imbalance across experts, where a small subset of experts receives most tokens while others are underutilized. Prior work has focused mainly on training-time solutions such as routing regularization or auxiliary losses, leaving inference-time behavior, which is critical for deployment, less explored. We present a systematic analysis of expert routing during inference and identify three findings: (i) load imbalance persists and worsens with larger batch sizes, (ii) selection frequency does not reliably reflect expert importance, and (iii) overall expert workload and importance can be estimated using a small calibration set. These insights motivate inference-time mechanisms that rebalance workloads without retraining or router modification. We propose Replicate-and-Quantize (R&Q), a training-free and near-lossless framework for dynamic workload rebalancing. In each layer, heavy-hitter experts are replicated to increase parallel capacity, while less critical experts and replicas are quantized to remain within the original memory budget. We also introduce a Load-Imbalance Score (LIS) to measure routing skew by comparing heavy-hitter load to an equal allocation baseline. Experiments across representative SMoE models and benchmarks show up to 1.4x reduction in imbalance with accuracy maintained within +/-0.6%, enabling more predictable and efficient inference.

Zijie Liu, Jie Peng, Jinhao Duan +7

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

Load BalancingLLM Inference OptimizationInference WorkloadsEfficient Inference
AskPDF
May 13, 2026cs.LG

EMO: Frustratingly Easy Progressive Training of Extendable MoE

Sparse Mixture-of-Experts (MoE) models offer a powerful way to scale model size without increasing compute, as per-token FLOPs depend only on k active experts rather than the total pool of E experts. Yet, this asymmetry creates an MoE efficiency paradox in practice: adding more experts balloons memory and communication costs, making actual training inefficient. We argue that this bottleneck arises in part because current MoE training allocates too many experts from the beginning, even though early-stage data may not fully utilize such capacity. Motivated by this, we propose EMO, a simple progressive training framework that treats MoE capacity as expandable memory and grows the expert pool over the course of training. EMO explicitly models sparsity in scaling law to derive stage-wise compute-optimal token budgets for progressive expansion. Empirical results show that EMO matches the performance of a fixed-expert setup in large-scale experiments while improving wall-clock efficiency. It offers a surprisingly simple yet effective path to scalable MoE training, preserving the benefits of large expert pools while reducing both training time and GPU cost.

Linghao Jin, Chufan Shi, Huijuan Wang +4

USC-ISI · MBZUAI-IFM

Mixture-Of-ExpertsExpert DemonstrationsSparsity
AskPDF
May 13, 2026cs.LG

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

Recent frontier large language models predominantly rely on Mixture-of-Experts (MoE) architectures. Despite empirical progress, there is still no principled understanding of how hyperparameters should scale with network width NN, expert width NeN_e, number of experts MM, sparsity KK, and depth LL to ensure both stability and optimal performance at scale. We take a principled step toward resolving this gap by analyzing three different scaling regimes: (I) co-scaling N≍NeN\asymp N_e, (II) co-scaling N≍M≍KN\asymp M\asymp K, and (III) full proportional scaling of N,Ne,MN, N_e, M, and KK. For each regime, we develop a novel Dynamical Mean Field Theory (DMFT) description of the limiting training dynamics of MoEs that provides a formal foundation for our analysis. Within this framework, we derive the unique parameterization for SGD and Adam satisfying all maximal-update (μμ) desiderata. We then show that the resulting μμP prescription does not reliably induce monotonic improvement with scale or robust learning-rate transfer. We trace these pathologies to scale-dependent observables in the aggregation dynamics, which motivates a refined set of desiderata that we term maximal scale stability. Guided by this principle, we derive a Maximally Scale-Stable Parameterization (MSSP) for both SGD and Adam in all three scaling regimes, and characterize the corresponding limiting dynamics - qualitatively distinct from the μμP limit - through a separate DMFT analysis. Experiments verify that MSSP robustly recovers learning rate transfer and monotonic improvement with scale across regimes. Combined with existing depth-scaling theory, these results provide a complete scaling prescription for MoE architectures as a function of width, depth, expert width, and number of experts.

Leena Chennuru Vankadara, Moritz Haas, Luke Hayward +2

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

Mixture-Of-ExpertsScaling LawsTransfer Learning
AskPDF