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

TL;DR

The paper introduces $oldsymbol{ ext{ extphi}}$-balancing, a new principled framework for improving expert load balancing in mixture-of-experts models, leading to more stable and effective utilization.

Contribution

It proposes a convex, population-level balancing method with an efficient online algorithm, outperforming prior heuristics in large-scale MoE training.

Findings

$oldsymbol{ ext{ extphi}}$-balancing outperforms previous methods in large-scale experiments.

The method achieves more stable expert utilization during training.

It introduces an efficient EMA-based routing adjustment with negligible overhead.

Abstract

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.