MoE load balancing recast as feedback control: PID (I) holds the busiest expert lower across a data-mixture switch

MoE load balancing is recast as a feedback control loop, so the balancer no longer acts only as a penalty added to the training objective but adjusts routing from each expert's deviation from an equal share of the batch. The two dominant families, an auxiliary balancing loss and the loss-free bias update of ALF-LB, were largely developed and evaluated under a stationary training distribution, where per-expert loads drift slowly and a balancer can settle and hold. On an OLMoE backbone trained from scratch with 16 to 64 experts per layer and a mid-run mixture switch, PID (I) holds balance before the shift and restores it more tightly afterwards, with the lowest MaxVio in every setting and perplexity comparable to the auxiliary loss and ALF-LB.

The control input is a per-expert bias added to router scores, which changes only expert selection and leaves the gate weights and their gradients untouched. The auxiliary loss folds balancing pressure into the same backward pass that updates router weights, whereas bias-based methods move that pressure out of the gradient and act only at selection time. Reading the router as a static monotone map with a one-step delay, only the integral term both removes persistent imbalance and averages out step-to-step load fluctuation, while the proportional and derivative terms pass or amplify it into the bias. In continual pretraining of OLMoE-1B-7B, PID (I) reaches postMV of 0.90 cold-start and 0.87 warm-start, against 1.01 and 1.03 for ALF-LB.

PID (I) modifies neither model weights nor the training objective, so it can be inserted directly into an existing pretrained model. Continual pretraining starts from the released OLMoE-1B-7B checkpoint, first trains on the annealing corpus until the balancer reaches steady state, then switches to a code corpus, testing the balancer as a drop-in component. The cold-start checkpoint's router has never seen these balancers, yet its postMV comes close to the warm-start setting; PID (I) has the lowest postMV in both settings, and held-out perplexity on the new domain is nearly the same across the four methods.

The demonstrated scope is an OLMoE backbone: pretraining from scratch with 16 to 64 experts per layer, one eighth of experts activated per token, and a single instantaneous domain switch mid-training; continual pretraining starts from the OLMoE-1B-7B-0125 checkpoint on a code corpus in cold-start or warm-start form. The distribution change studied is one hard switch between two domains, and the largest model has 6.9B total and 1.3B active parameters.

The authors note that the balancer's margin over ALF-LB widens with expert count, but three scales are too few to turn this trend into a firm conclusion, and whether it continues at production expert counts remains open. Only a single hard switch is covered, not gradual drift or repeated switches among more than two domains; results come from one backbone, and transfer to expert-choice routing or routers with capacity limits is left to future work. In continual pretraining the bias-based balancers forget more within the same training budget, with higher held-out perplexity on the annealing corpus than Vanilla and the auxiliary loss.