ID Balancing: PID Control Stabilizes Extremely Sparse MoE with 768 Experts

Mixture-of-Experts (MoE) models scale parameter count by routing each token to only a few experts. As expert counts grow and the number of activated experts per token (Top-K) shrinks, load imbalance becomes severe: some experts receive excessive tokens while others stay underutilized. Traditional auxiliary loss methods use a fixed objective that cannot adapt to the dynamic load shifts during training. Recent loss-free approaches — DeepSeek's Loss-Free and Kimi K3's Quantile Balancing — avoid auxiliary losses but still struggle at extreme sparsity.

Framing Load Balancing as a PID Control Problem

The paper ID Balancing: Stable Training of Extremely Sparse MoE via PID-Based Load Control by Dayiheng Liu et al. (Qwen team) observes that both loss-free methods fit into an incomplete PID control framework. DeepSeek Loss-Free acts like a fixed-step integral controller, while Quantile Balancing resembles a generalized proportional controller. Building on this insight, the authors introduce ID Balancing (Integral-Derivative Balancing) , which adds both integral and derivative terms to the router's bias updates.

Integral term : continuously corrects based on current load error; larger error triggers larger adjustment.

Derivative term : monitors the trend of load error — if imbalance is worsening, it intervenes earlier; if load is stabilizing, it reduces extra correction.

In effect, the router knows not only which experts are busy now, but also which are becoming busier.

Experimental Setup

Experiments use a 18.9B parameter model with 768 experts, trained on 120B tokens. Three routing sparsity levels are tested: Top-10, Top-5, and Top-3 of 768. Metrics include MaxVio (worst-case expert overload) and MinVio (average expert underload). Comparisons are made against auxiliary loss, DeepSeek Loss-Free, and Quantile Balancing.

Key Results

Extreme Sparsity (Top-3-of-768)

ID Balancing maintains stable load control as sparsity increases. At Top-3:

Worst Backbone MaxVio: 15.33 (ID Balancing) vs 31.11 (Quantile Balancing) vs 211.19 (DeepSeek Loss-Free).

Avg. MinVio: 0.5480 (ID Balancing) vs 0.6232 (auxiliary loss) vs 0.7441 (DeepSeek Loss-Free) vs 0.6373 (Quantile Balancing).

MaxVio reduced >50% and MinVio reduced >12% compared to baselines.

Language modeling loss: 1.7426 , competitive with other methods.

Scaling Model Size (18.9B → 69.9B, Top-10-of-768)

As total parameters grow from 18.9B to 69.9B (activated params 1.03B → 3.2B), ID Balancing's worst Backbone MaxVio barely degrades. Compared to auxiliary loss, worst-case load deviation drops ~89.6% .

Training Dynamics and Robustness

Training curves (Figure 2) show ID Balancing suppresses transient expert overload and reduces chronic underload throughout training. Additional tests with higher learning rates confirm the method's stability across training configurations.

Downstream Task Evaluation

ID Balancing is evaluated on MMLU, MMLU-Pro, SuperGPQA, MATH, GSM8K, BBH, MMMLU, EvalPlus, and MultiPL-E. It retains competitive downstream performance while achieving superior load balance.

Context: Qwen's MoE Research Trajectory

This work follows Qwen's earlier Global Load Balancing , which addressed load distribution at a broader scope to balance compute efficiency with expert specialization. ID Balancing tackles a different layer: keeping the routing system itself stable as MoE becomes extremely sparse (more experts, lower activation ratio). The trend points toward future models with even more experts and lower activation ratios, where dynamic load control will be essential for scalable training.