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

1 Introduction

• As routing becomes sparser, load imbalance can slow computation at overloaded experts and leave capacity unused at underloaded experts, limiting training efficiency.The introduction identifies effective load control as important for reliably training larger, sparser MoE models.Source: §1 · p. 2
• The unified PID view treats DeepSeek’s loss-free method as fixed-step integral control and Kimi K3’s Quantile Balancing as generalized proportional control.The comparison motivates corrections that scale with load error and respond when imbalance worsens.Source: §1 · p. 2
• ID Balancing combines a magnitude-aware integral term with a derivative term that adds correction only when imbalance worsens, accumulating corrections in expert bias.It omits the proportional term, recenters bias after updates, uses O(E) token-count feedback per layer, and adds no balancing gradient to the language-modeling objective.Source: §1 · p. 2
• In Top-3 routing, ID Balancing reduces worst-case backbone MaxVio by over 50% and training-average backbone MinVio by over 12% relative to the best baselines.The evaluation also reports nearly unchanged worst-case backbone MaxVio as total parameters rise from 18.9B to 69.9B under fixed Top-10-of-768 routing.Source: §1 · p. 2
• Evaluations span Top-10, Top-5, and Top-3 routing over 768 experts, model sizes up to 69.9B parameters, and training conditions; analyses also report fewer inactive experts at inference.ID Balancing maintains competitive language-modeling and downstream performance.Source: §1 · pp. 2–3

2.1 The PID Controller

• A PID controller uses proportional feedback from current error, integral feedback from accumulated past errors, and derivative feedback from changes in error.The generalized MoE view preserves these temporal roles while allowing nonlinear feedback transformations.Source: §2.1 · p. 3, §2.3 · p. 3

2.2 Load Balancing as a Control Problem

• In the MoE control loop, expert biases are both persistent controller state and routing input, while token counts measure output and load errors drive subsequent bias updates.Uniform load is the target; negative errors indicate overload and positive errors indicate underload.Source: §2.2 · p. 3

2.3 Existing Methods through a PID Lens

• DeepSeek’s loss-free method is fixed-step integral control, while Quantile Balancing is generalized proportional control; fixed steps ignore error magnitude, whereas batch-dependent targets can drift.These contrasts motivate corrections that scale with load error and respond more strongly when imbalance worsens.Source: §2.3 · pp. 4–5
• Auxiliary-loss balancing updates router weights through training gradients rather than a separate control bias, while fixed-step and batch-target updates have distinct response trade-offs.The fixed step can respond slowly to large errors and risk overshooting near balance; changing batch targets can produce late-stage bias drift.Source: §2.3 · p. 5

3 ID Balancing

• Adding the worsening-gated derivative lowers early load violations, while its trajectories approach the integral-only update as training progresses.The active-gate fraction falls from about 0.38 early in training to 0.25 later.Source: Figure 4 · p. 6, §3 · p. 6
• Scaling the integral update by error magnitude reduces the worst-layer MaxVio peak and sustained underload.The update grows with the magnitude of the normalized error and shrinks as an expert approaches its target load.Source: §3 · pp. 5–6
• The derivative term uses error changes but adds correction only when an existing imbalance moves farther from zero without crossing it.Its correction is gated and accumulated in the bias alongside the integral update.Source: §3 · p. 6
• Stronger EMA smoothing reduces Quantile Balancing bias drift but weakens early load control; ID Balancing achieves low drift and effective balance without EMA.The comparison includes the original Quantile Balancing update, EMA variants, DeepSeek’s loss-free method, and ID Balancing.Source: Figure 5 · p. 6, §3 · p. 7
• Re-centering removes the common bias component after each update without changing expert rankings, but does not bound differences between expert biases.The integral update has zero mean; accumulated derivative corrections can make the mean bias nonzero.Source: §3 · p. 7

4.1 Model Quality

• In Top-3, ID Balancing lowers worst-case backbone MaxVio to 15.33 from 211.19 for DeepSeek’s method and 31.11 for Quantile Balancing, and leads on key backbone load metrics across routing settings.It achieves the lowest Worst and Last1k MaxVio and Last1k and average MinVio across Top-3, Top-5, and Top-10.Source: §4.1 · p. 8
• With reduced gains, ID Balancing improves load balance during continued pretraining while keeping LM loss within 0.001 of the frozen-bias baseline in the second half.For Top-10-of-768, mean MaxVio falls from about 6 to 1.4; full gains yield the highest LM loss.Source: Figure 6 · p. 8, §4.1 · p. 9
• Under Top-10-of-768, ID Balancing raises average downstream score from 58.12 to 58.66 versus Auxiliary loss and scores higher on five of nine benchmarks.The evaluated benchmarks cover knowledge, STEM, reasoning, multilingual understanding, and code generation.Source: §4.1 · p. 9, Table 2 · p. 9

4.2 Training Stability

• At 2.3× the standard peak learning rate, ID Balancing and Quantile Balancing maintain lower backbone violations than Auxiliary loss and DeepSeek’s method, while ID Balancing has the lowest LM loss and MTP MaxVio.ID Balancing’s Last1k backbone MaxVio is 0.6437; Quantile Balancing achieves the strongest backbone balance in this test.Source: Table 3 · p. 9, §4.2 · pp. 10–11
• ID Balancing and Quantile Balancing end with gradient norms near 0.043, below Auxiliary loss at 0.050 and DeepSeek’s method at 0.049, while MoE-output magnitudes remain near 0.51–0.52 for ID Balancing.DeepSeek’s method shows stronger late-stage gradient fluctuations.Source: §4.2 · p. 11, Figure 9 · p. 11

4.3 Model Characteristics

• Across 20 layers in Top-3-of-768, ID Balancing has lower or comparable MaxVio and lower MinVio than Quantile Balancing; at inference, it also reduces inactive experts versus Auxiliary loss.For Top-10-of-768, the average inactive-expert ratio falls from 8.4% to 6.4%, and the first-layer ratio from 9.3% to 2.6%.Source: §4.3 · p. 11, Figure 10 · p. 12, Figure 11 · p. 13

4.4 Ablations

• With Kd = 0, increasing Ki from 3 × 10^-3 to 6 × 10^-3 lowers training-average backbone MaxVio from 0.9459 to 0.8003 and worst-case MaxVio from 23.26 to 14.96.Ki = 9 × 10^-3 slightly lowers the average further but raises worst-case backbone and MTP violations; Ki = 6 × 10^-3 is selected.Source: §4.4 · p. 12, Table 4 · p. 13
• With Ki = 6 × 10^-3, Kd = 6 × 10^-3 lowers mean backbone MaxVio from 2.0020 to 1.9330 over 0–1k steps and from 0.9059 to 0.8612 over 1–5k steps.It reduces sampled cumulative excess above Vt = 1 by 57.9% over 1–5k steps, but training-average MTP MaxVio rises from 0.8862 to 0.9698.Source: §4.4 · p. 13

Sparse mixture-of-experts

• Sparse MoE increases total capacity through conditional expert activation with a smaller increase in per-token computation, and this work addresses load control as the active fraction decreases.Prior routing designs include Top-1 and Top-2 methods as well as alternative token–expert assignments.Source: §5 · pp. 13–14

Balancing through an auxiliary loss

• Auxiliary-loss balancing updates router weights through gradients shared with the language-modeling objective, whereas ID Balancing uses a separate bias update without adding a balancing gradient.The shared optimization introduces a trade-off between balancing and language-modeling objectives.Source: §5 · p. 14

Auxiliary-loss-free balancing

• DeepSeek’s loss-free method is fixed-step integral control, while Quantile Balancing is generalized proportional control; ID Balancing adds magnitude-aware integral and worsening-gated derivative corrections.DeepSeek’s sign-based update discards error magnitude, whereas Quantile Balancing sets a batch-dependent target bias from router-score margins.Source: §5 · p. 14

• Experiments across routing sparsities, model sizes, and learning-rate settings show effective load control while maintaining competitive language-modeling and downstream performance.Source: §6 · p. 15
• ID Balancing combines magnitude-aware integral control with a worsening-gated derivative correction and zero-mean bias centering.It requires only O(E) token-count feedback per layer and adds no balancing gradient to the language-modeling objective.Source: §6 · p. 15
• Further analyses find consistent backbone load control across layers and fewer inactive experts on evaluation data, while reduced gains retain active load correction during continued pretraining.Source: §6 · p. 15

• The evaluation covers decoder-only MoE models with fixed Top-K routing, across several sparsities, model sizes, and training conditions.Source: Limitations · p. 15
• Scaling comparisons vary depth, width, active and total parameters, and training duration, so they assess load control across configurations rather than isolating model size.Source: Limitations · p. 15
• Evaluation on other architectures and routing methods is identified as a natural next step, and the worsening-gated derivative update is supported empirically.Source: Limitations · p. 15

• Three hybrid MoE configurations cover different expert-pool and model scales, with specifications summarized in Table A.Source: §A.0.1 · pp. 18–19
• Each model uses full attention and gated linear attention, placing one full-attention layer every four Transformer layers.Source: §A.0.1 · p. 18
• Each MoE layer contains one shared expert alongside routed experts, and the input embedding and output head are untied.Source: §A.0.1 · p. 18
• Table A's parameter counts include the input embedding and untied output head but exclude the MTP module; active counts include dense, shared, and Top-K-selected routed components.Source: §A.0.1 · p. 19