MoE load balancing recast as feedback control: PID (I) holds the busiest expert lower across a data-mixture switch
Data-mixture changes during mixture-of-experts (MoE) training disrupt the per-expert load a router had converged to, yet existing balancing methods were evaluated only on stationary data. This work formulates load balancing as feedback control: each expert's deviation from an equal share is the control error, a per-expert bias added to router scores is the control input, and the auxiliary loss and ALF-LB correspond to a proportional term and a sign-based integral term. Keeping only the integral term, PID (I) attains lower MaxVio than the compared methods before and after a mid-run mixture switch in OLMoE pretraining and in continual pretraining of OLMoE-1B-7B, at comparable perplexity.