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 casts loss-free MoE balancing as PID control; its ID Balancing controller cuts worst-case expert overload from 31.1 to 15.3 at Top-3-of-768 routing with LM loss on par.
MoE grows total capacity while per-token compute stays flat: a router picks K of E experts for each token. With K fixed, adding experts makes routing sparser, and imbalance gets worse as sparsity rises. Small score changes near the Top-K boundary reshuffle token assignments; overloaded experts stall the pipeline, underloaded ones train poorly, and severe imbalance can destabilize training. At Top-3-of-768 the active fraction is 0.39%, one in 256, and that is where current balancers start to break down.
The industry default is the auxiliary-loss-free route: no balancing term in the loss, just a per-expert bias outside the gradients, tuned by routing feedback. DeepSeek-V3 and Kimi K3 are the two representatives, and the paper reads them as two fragments of one controller.
A PID controller is the workhorse of feedback control, with branches responding to current error (P), accumulated error (I), and error change (D). Set uniform load as the target, the bias vector as controller state, and token counts as the measurement, and the two existing methods fall into place. DeepSeek loss-free moves each bias by a fixed step in the direction of the error sign: magnitude is discarded, so mild and severe imbalance get equal correction, large errors are answered slowly, and the fixed step risks overshoot near balance. It is an integral-only controller. Kimi K3's Quantile Balancing computes a target bias from score-margin quantiles of the current batch and sets the bias to it in one move; the target drifts with each batch, which shows up as late-stage wobble. It behaves as a generalized proportional controller.
ID Balancing fills the gap with two branches:
Cost is O(E) token counts per layer, with no balancing gradient added to the LM objective. Defaults are Ki = Kd = 6e-3.
Main runs use the Qwen-3.8-Next architecture: 20 layers, one MTP module, 768 experts per layer, about 18.9B total and 0.87B-1.03B active parameters, 120B tokens. Top-3-of-768:
| Method | Worst MaxVio ↓ | Avg 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 |
Worst-case overload drops more than 50% versus the best baseline (Quantile Balancing); DeepSeek's method peaks at 211 at this sparsity. Scaling from 18.9B to 69.9B under Top-10-of-768 moves ID Balancing's worst MaxVio only from 8.36 to 7.79, while auxiliary loss climbs to 75.03, an 89.6% gap.
In a stress test with the learning rate held at 2.3x the standard peak, Last1k backbone MaxVio is 0.644 for ID Balancing and 0.502 for Quantile Balancing, against 1.376 for DeepSeek and 3.138 for auxiliary loss. Raising the auxiliary coefficient to 0.5 backfires: loss climbs to about 2.45, some experts receive no tokens at all, and the larger-coefficient runs die before 30k steps.
Downstream over nine benchmarks at 560B tokens: at 3.2B/69.9B, ID Balancing averages 58.66, best of the four methods, versus 58.12 for auxiliary loss; at 3.0B/24.8B it scores 55.81, just under Quantile Balancing's 55.88. In continued pretraining with gains divided by 1000, LM loss stays within 0.001 of a frozen bias while mean MaxVio drops from 6-7 to about 1.4. At inference the inactive-expert ratio falls from 8.4% to 6.4%, and from 9.3% to 2.6% in the first layer.
For teams training large sparse MoE models this is a drop-in swap of the bias-update rule: no gradients, O(E) feedback. The margin grows with sparsity, and worst-case overload that stays flat from 18.9B to 69.9B speaks directly to the worry that training stops being controllable once expert counts climb into the hundreds. The PID framing is reusable on its own: two production methods in one framework make their missing branches visible, and the same analysis extends to other feedback-driven training mechanisms. Be clear about the positioning: an incremental improvement in load control. LM loss and downstream scores are on par; the value is stability, not model quality.
The authors list three: evaluation covers one decoder-only MoE family with fixed Top-K routing; the accumulated gated-derivative form is supported empirically, with no theory; and the scaling runs change depth, width, active and total parameters, and training duration together, so model size is never isolated. The text leaves more open: the main evidence comes from 20-layer models on 120B tokens, well short of production scale; Quantile Balancing still wins training-average backbone MaxVio at Top-5 and Top-10 plus the tightest backbone balance in the high-LR stress test, so ID Balancing sweeps the board only at Top-3, and the "stable" in the title really means "better worst case"; the derivative term costs some MTP balance, with average MTP MaxVio rising from 0.886 to 0.970; and the swept gains Ki = Kd = 6e-3 (divided by 1000 for continued pretraining) have no evidence of transferring to other architectures.