Load imbalance is a well-known pain point in Mixture-of-Experts (MoE) models: some experts are routed constantly, receive the most data, and update fastest, while others are almost never selected and never get trained. Existing balancing methods are mostly heuristic—an auxiliary loss is added after each batch's routing to penalize uneven assignments. Chen, Li, Wang, and their team point out that these methods operate on noisy statistics within mini-batches, which systematically deviate from the true population-level demand.
φ-balancing turns load balancing into a convex optimization problem. The objective is a strictly convex, symmetric, differentiable potential function over the expected routing distribution. Minimizing this potential is equivalent to driving each expert's expected load toward equality.
Using convex duality, the primal problem is transformed into an equivalent min-max form, and applying mirror descent yields a simple online algorithm—routing adjustments based on an exponential moving average. The additional overhead is nearly zero.
In pretraining and downstream fine-tuning experiments, φ-balancing beats existing Switch-style and loss-free baselines on both load balancing and final performance, with better stability and higher utilization.
Unclear points: How sensitive is the choice of the convex potential function? Do different φ functions affect final performance? How should the EMA time constant be set—too short and it gets swayed by batch noise; too long and it lags behind distribution changes?
References
1. Chen, L., Li, J., Wang, Q., et al. (2026). *φ-Balancing for Mixture-of-Experts Training*. arXiv:2605.15403 [cs.LG]. 2. Fedus, W., et al. (2022). *Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity*. JMLR. 3. Lepikhin, D., et al. (2021). *GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding*. ICLR.