$\phi$-Balancing for Mixture-of-Experts Training

Mixture-of-Experts (MoE) models rely on balanced expert utilization to fully realize their scalability. However, existing load-balancing methods are largely heuristic and operate on noisy mini-batch assignment statistics, introducing bias relative to population-level objectives. We propose **$\phi$-balancing**, a principled framework that directly targets population-level expert balance by minimizing a strictly convex, symmetric, and differentiable potential of the expected routing distribution. Using convex duality, we derive an equivalent min-max formulation and obtain a simple online algorithm via mirror descent, yielding an efficient EMA-based routing adjustment with negligible overhead. Across large-scale pretraining and downstream fine-tuning, $\phi$-balancing consistently outperforms prior Switch-style and loss-free baselines, demonstrating more stable and effective expert utilization.

Modern Mixture-of-Experts (MoE) models can scale to very large sizes by dividing work across many specialized "experts." Many of today’s most capable AI models rely on this architecture because it enables dramatic increases in model capabilities without a matching increase in computation. However, these systems only work well when the workload is distributed evenly across the experts. In practice, some experts often become overloaded while others are underused, which reduces efficiency and degrades performance. This paper introduces a simple and efficient algorithm, $\phi$-balancing, for keeping expert usage balanced during training. Unlike existing, largely heuristic approaches for addressing this challenge, our method is grounded in a principled mathematical framework. Large-scale experiments across diverse tasks show that $\phi$-balancing achieves better performance than prior balancing methods, affirming the theoretical motivation.