Expert Load Balancing and Auxiliary Losses

A mixture-of-experts layer routes each input token to a small subset of specialist sub-networks ("experts") rather than running all of them. Left to train freely, the router tends to converge on favoring a handful of experts early on — those experts get more gradient updates because they see more data, which makes them better, which makes the router favor them even more. The result is a rich-get-richer collapse: most experts go nearly unused while a few are overloaded, wasting the whole point of having extra capacity.

An auxiliary load-balancing loss fixes this by adding a training penalty proportional to how unevenly tokens are spread across experts — typically the product of the fraction of tokens routed to each expert and the router's average probability of choosing that expert, summed over all experts. Minimizing this term alongside the main task loss pushes the router toward spreading tokens roughly evenly, without needing to hand-engineer which expert should specialize in what.

A worked example

With 8 experts and no balancing loss, training might converge to 3 experts handling 80% of tokens between them and the other 5 sharing the remaining 20%. Adding a load-balancing auxiliary loss with a small weight (commonly around 0.01 relative to the main loss) is often enough to push the distribution toward something closer to 12.5% per expert on average, letting all 8 specialize usefully instead of 5 of them going nearly dormant.