$ϕ$-Balancing for Mixture-of-Experts Training
Fuente:
arXiv
Salvato in:
| Autori principali: | , , , , , , , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2026
|
| Soggetti: | |
| Accesso online: | |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866914568013348864 |
|---|---|
| author | Chen, Lizhang Li, Jonathan Wang, Qi Liao, Runlong Li, Shuozhe Liang, Chen Lao, Ni Liu, Qiang |
| author_facet | Chen, Lizhang Li, Jonathan Wang, Qi Liao, Runlong Li, Shuozhe Liang, Chen Lao, Ni Liu, Qiang |
| contents | 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 $ϕ$-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, $ϕ$-balancing consistently outperforms prior Switch-style and loss-free baselines, demonstrating more stable and effective expert utilization. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2605_15403 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | $ϕ$-Balancing for Mixture-of-Experts Training Chen, Lizhang Li, Jonathan Wang, Qi Liao, Runlong Li, Shuozhe Liang, Chen Lao, Ni Liu, Qiang Machine Learning Optimization and Control 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 $ϕ$-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, $ϕ$-balancing consistently outperforms prior Switch-style and loss-free baselines, demonstrating more stable and effective expert utilization. |
| title | $ϕ$-Balancing for Mixture-of-Experts Training |
| topic | Machine Learning Optimization and Control |
| url | https://arxiv.org/abs/2605.15403 |