{"id":"arxiv:2408.15664","kind":"preprint","title":"Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts","authors":["Lean Wang","Huazuo Gao","Chenggang Zhao","Xu Sun","Damai Dai"],"author_count":5,"year":2024,"published":"2024-08-28","venue":null,"cited_by":null,"doi":null,"arxiv_id":"2408.15664","nct_id":null,"source_url":"https://arxiv.org/abs/2408.15664","archive_url":"https://space-kitty.com/papers/arxiv:2408.15664","updated":"2024-08-28","pmid":null,"pmcid":null,"openalex_id":null,"pdf_url":"https://arxiv.org/pdf/2408.15664","abstract":"For Mixture-of-Experts (MoE) models, an unbalanced expert load will lead to routing collapse or increased computational overhead. Existing methods commonly employ an auxiliary loss to encourage load balance, but a large auxiliary loss will introduce non-negligible interference gradients into training and thus impair the model performance. In order to control load balance while not producing undesired gradients during training, we propose Loss-Free Balancing, featured by an auxiliary-loss-free load balancing strategy. To be specific, before the top-K routing decision, Loss-Free Balancing will first apply an expert-wise bias to the routing scores of each expert. By dynamically updating the bias of each expert according to its recent load, Loss-Free Balancing can consistently maintain a balanced distribution of expert load. In addition, since Loss-Free Balancing does not produce any interference gradients, it also elevates the upper bound of model performance gained from MoE training. We validate the performance of Loss-Free Balancing on MoE models with up to 3B parameters trained on up to 200B tokens. Experimental results show that Loss-Free Balancing achieves both better performance and better load balance compared with traditional auxiliary-loss-controlled load balancing strategies.","abstract_license":"CC0-1.0 (arXiv descriptive metadata)","license":null,"is_open_access":false,"retracted":false,"peer_reviewed":false,"categories":["cs.LG","cs.CL"],"details":null,"provenance":"arxiv","attribution":"Abstract: arXiv descriptive metadata (CC0). Thank you to arXiv for use of its open access interoperability. This service was not reviewed or approved by, nor does it necessarily express or reflect the policies or opinions of, arXiv.","topics":[{"slug":"llm-architectures","name":"LLM Architectures & Pretraining","domain_slug":"ai"}],"cited_in_dossiers":[{"slug":"llm-architectures","title":"LLM Architectures & Pretraining: State of the Field (October 2026)","url":"https://space-kitty.com/dossiers/llm-architectures"}]}