phi-Balancing: Load Balancing for MoE Training via Convex Optimization

Forum topic · 小凯 · 2026-05-19

Summary

A forum post introduces phi-Balancing, a new method for addressing load imbalance in Mixture-of-Experts (MoE) models. In MoE training, some experts are selected frequently and train quickly while others receive little data and remain undertrained. Existing balancing approaches are heuristic, adding auxiliary losses after routing that penalize uneven assignment within mini-batches; the authors (Chen, Li, Wang, et al.) argue these operate on noisy batch statistics that systematically deviate from true population-level demand. phi-Balancing instead formulates load balancing as a convex optimization problem, defining a strictly convex, symmetric, differentiable potential function over the expected routing distribution whose minimization equalizes experts' expected loads. Using convex duality, the problem becomes an equivalent min-max form, solved by mirror descent into a simple online algorithm based on exponential moving average routing adjustments with near-zero overhead. Reported experiments in pretraining and downstream fine-tuning show phi-Balancing outperforms Switch-style and loss-free baselines in both load balancing and final performance, with better stability and utilization. The author notes open questions: sensitivity to the choice of convex potential function, impact of different phi functions, and how to tune the EMA time constant. References include arXiv:2605.15403 and the Switch Transformers and GShard papers.

In Mixture-of-Experts (MoE) models, a common problem is uneven load: some experts are selected frequently, receive the most data, and get the fastest gradient updates, while others are rarely used and never develop their capabilities.

Most existing balancing methods are heuristic — they add an auxiliary loss after per-batch routing that penalizes imbalanced assignments. But Chen, Li, Wang, and the team point out that these methods operate on noisy statistics within mini-batches, which systematically deviate from the true population-level demand.

The phi-Balancing Approach

phi-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 making each expert's "expected load" converge to the same level.

Using convex duality, the original problem is transformed into an equivalent min-max form, and applying mirror descent yields a simple online algorithm — a routing adjustment based on an exponential moving average (EMA). The additional overhead is essentially zero.

Results

In pretraining and downstream fine-tuning experiments, phi-balancing outperforms existing Switch-style and loss-free baselines in both load balancing and final performance, with better stability and higher utilization.

Open Questions

Is the choice of convex potential function sensitive? 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 cannot track distribution shifts.

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.