Skip to content

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

ID Balancing unifies the two leading auxiliary-loss-free MoE load balancers under a Proportional-Integral-Derivative (PID) control view: DeepSeek’s loss-free bias update is a fixed-step integral controller, and Kimi K3’s Quantile Balancing is a generalized proportional controller. From that frame, the paper proposes an Integral-Derivative controller — scale the integral term with the observed load error and activate the derivative term only when imbalance worsens — giving stronger corrections when things are off, and smaller corrections near balance. On 768-expert routing at Top-{10, 5, 3}, ID Balancing cuts worst-case backbone MaxVio by >50% and training-average backbone MinVio by >12% versus the best baselines in Top-3; scaling 18.9B → 69.9B at Top-10-of-768, its worst-case MaxVio stays nearly constant and sits ~89.6% below the auxiliary-loss baseline.

  • DeepSeek’s loss-free load balancer is equivalent to a fixed-step integral controller on expert load error; Kimi K3’s Quantile Balancing is a generalized proportional controller [§Abstract, §Method].
  • ID Balancing is an Integral-Derivative controller whose integral term scales with load error and whose derivative term fires only when imbalance is worsening, yielding larger corrections for large/growing errors and smaller updates near balance [§Abstract, §Method].
  • At 768-expert routing with Top-3, ID Balancing reduces worst-case backbone MaxVio by >50% and training-average backbone MinVio by >12% relative to the best baselines [§Abstract, §Results].
  • Scaling total parameters from 18.9B to 69.9B at Top-10-of-768, ID Balancing’s worst-case backbone MaxVio remains nearly unchanged and is ~89.6% lower than the auxiliary-loss baseline [§Abstract, §Results].
  • Advantages grow with sparsity: the gap between ID Balancing and prior balancers widens from Top-10 to Top-5 to Top-3 of 768 experts [§Abstract, §Results].
  • Language-modeling and downstream task performance remain competitive while balance improves — the method is not trading quality for balance [§Abstract].

Standard auxiliary-loss-free MoE balancing adjusts a per-expert bias added to routing logits based on observed expert load. ID Balancing reframes this bias-update rule as a controller: the “plant” is the router under current traffic, the “setpoint” is uniform load, and the “error” is per-expert load deviation. From the PID vocabulary: DeepSeek’s constant step-size update of the bias each iteration is a fixed-gain integral controller; Kimi K3’s Quantile Balancing, which sets the bias from the Top-(k+1) routing quantiles in a single forward pass, is a generalized proportional controller.

ID Balancing combines two components. The integral term is scaled by the current load error magnitude (so steps grow with error rather than being fixed). The derivative term is activated only when the error is increasing across steps — behaving as an anti-windup / anti-overshoot guard that strengthens correction when the system is diverging but stays silent near balance. There is no auxiliary loss added to the training objective; the controller operates purely on routing bias.

Evaluated across Top-10, Top-5, and Top-3 routing over 768 experts ([§Abstract]):

  • Top-3 setting: worst-case backbone MaxVio reduced >50% and training-average backbone MinVio reduced >12% vs the best baselines [§Abstract].
  • Scaling 18.9B → 69.9B at Top-10-of-768: worst-case backbone MaxVio nearly unchanged under ID Balancing and ~89.6% below the auxiliary-loss baseline [§Abstract].
  • Advantages increase monotonically with routing sparsity [§Abstract].
  • Competitive language-modeling and downstream task metrics preserved alongside the balance improvements [§Abstract].

Slots directly into the open “composition” question on MoE Routing Design: the wiki already carries DeepSeek’s bias-update (Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models, critiqued), Kimi’s Quantile Balancing (Quantile Balancing: A Hyperparameter-Free MoE Load Balancing Method, scaled to 2.8T in Kimi K3 Technical Report — Open Frontier Intelligence and validated to 67B+ by Marin in Marin scaling laws simulate the entire training trajectory of a 67B MoE run within 1%), and the symptom they both leave unaddressed at extreme sparsity (silent expert death, Mitigate Silent Expert Death In Ultra-Sparse MoE). This paper is the first to name those two methods as endpoints of a single control-theoretic design axis and to propose a filled-in point between them — making the “DeepSeek vs QB” dichotomy a parametric family rather than two separate recipes. The 768-expert / Top-3 regime also sits beyond Kimi K3’s Top-16-of-896 as the most extreme production-scale sparsity tested with aux-free balancing so far.