mlsanityData to 7 Oct 2026

Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models

arXiv 2501.11873 ↗First posted 21 Jan 2025cs.LG cs.CLcited by 53

On the map: Mixture-of-experts → Efficient MoE training and inference systems

Abstract (arXiv)

This paper revisits the implementation of Load-balancing Loss (LBL) when training Mixture-of-Experts (MoEs) models. Specifically, LBL for MoEs is defined as NE sumi=1NE fi pi, where NE is the total number of experts, fi represents the frequency of expert i being selected, and pi denotes the average gating score of the expert i. Existing MoE training frameworks usually employ the parallel training strategy so that fi and the LBL are calculated within a micro-batch and then averaged across parallel groups. In essence, a micro-batch for training billion-scale LLMs normally contains very few sequences. So, the micro-batch LBL is almost at the sequence level, and the router is pushed to distribute the token evenly within each sequence. Under this strict constraint, even tokens from a domain-specific sequence (textite.g., code) are uniformly routed to all experts, thereby inhibiting expert specialization. In this work, we propose calculating LBL using a global-batch to loose this constraint. Because a global-batch contains much more diverse sequences than a micro-batch, which will encourage load balance at the corpus level. Specifically, we introduce an extra communication step to synchronize fi across micro-batches and then use it to calculate the LBL. Through experiments on training MoEs-based LLMs (up to 42.8B total parameters and 400B tokens), we surprisingly find that the global-batch LBL strategy yields excellent performance gains in both pre-training perplexity and downstream tasks. Our analysis reveals that the global-batch LBL also greatly improves the domain specialization of MoE experts.

Key facts. Cited by 53 papers in the map (data through 7 Oct 2026). Placed in Efficient MoE training and inference systems, one of 12 directions in Mixture-of-experts. Builds on 27 of its 30 references that are in the map. First posted on arXiv 21 Jan 2025.

Builds on

27 of 30 references.

Closest ancestors · 3

Same direction, newest first

  1. 2022MegaBlocks: Efficient Sparse Training with Mixture-of-Experts29 Nov
  2. 2022Tutel: Adaptive Mixture-of-Experts at Scale7 Jun
  3. 2022DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale14 Jan

Other references · 24

  1. 2017Attention Is All You Needcited by 59,932
  2. 2021Training Verifiers to Solve Math Word Problemscited by 9,743
  3. 2020Measuring Massive Multitask Language Understandingcited by 7,768
  4. 2024Qwen2.5 Technical Reportcited by 4,713
  5. 2019HellaSwag: Can a Machine Really Finish Your Sentence?cited by 4,018
  6. 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,413
  7. 2021Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsitycited by 3,014
  8. 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,829
  9. 2024Mixtral of Expertscited by 1,785
  10. 2024DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Modelscited by 834
  11. 2023C-Eval: A Multi-Level Multi-Discipline Chinese Evaluation Suite for Foundation Modelscited by 694
  12. 2022Mixture-of-Experts with Expert Choice Routingcited by 611
12 more
  1. 2022ST-MoE: Designing Stable and Transferable Sparse Expert Modelscited by 391
  2. 2024MoE-LLaVA: Mixture of Experts for Large Vision-Language Modelscited by 298
  3. 2024OLMoE: Open Mixture-of-Experts Language Modelscited by 291
  4. 2024Mixture-of-Depths: Dynamically allocating compute in transformer-based language modelscited by 242
  5. 2024Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Expertscited by 191
  6. 2024OpenMoE: An Early Effort on Open Mixture-of-Experts Language Modelscited by 185
  7. 2025MiniMax-01: Scaling Foundation Models with Lightning Attentioncited by 157
  8. 2024Lory: Fully Differentiable Mixture-of-Experts for Autoregressive Language Model Pre-trainingcited by 37
  9. 2024Fewer Truncations Improve Language Modelingcited by 33
  10. 2024Layerwise Recurrent Router for Mixture-of-Expertscited by 21
  11. 2023Unlocking Emergent Modularity in Large Language Modelscited by 17
  12. 2024GW-MoE: Resolving Uncertainty in MoE Router with Global Workspace Theorycited by 8

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2025: 15; 2026: 38.

53 papers.

Sort:

By area:

  1. 2025Qwen3 Technical Reportcited by 8,195
  2. 2025Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Freecited by 266
  3. 2026Controlled LLM Training on Spectral Spherecited by 28
  4. 2025DiffMoE: Dynamic Token Selection for Scalable Diffusion Transformerscited by 23
  5. 2025Load Balancing Mixture of Experts with Similarity Preserving Routerscited by 22
  6. 2026A Unified View of Attention and Residual Sinks: Outlier-Driven Rescaling is Essential for Transformer Trainingcited by 21
  7. 2026On the Design of Qwen3.8-Next Architecture: Evaluation, Efficiency, and Training Stabilitycited by 20
  8. 2025MoRE: 3D Visual Geometry Reconstruction Meets Mixture-of-Expertscited by 19
  9. 2025Pangu Ultra MoE: How to Train Your Big MoE on Ascend NPUscited by 12
  10. 2025ProPhy: Progressive Physical Alignment for Dynamic World Simulationcited by 11
  11. 2025Coupling Experts and Routers in Mixture-of-Experts via an Auxiliary Losscited by 11
  12. 2026STEM: Scaling Transformers with Embedding Modulescited by 10
38 more
  1. 2026The Myth of Expert Specialization in MoEs: Why Routing Reflects Geometry, Not Necessarily Domain Expertisecited by 9
  2. 2025Dense Backpropagation Improves Training for Sparse Mixture-of-Expertscited by 8
  3. 2026Least-Loaded Expert Parallelism: Load Balancing An Imbalanced Mixture-of-Expertscited by 8
  4. 2026Seeing but Not Thinking: Routing Distraction in Multimodal Mixture-of-Expertscited by 8
  5. 2025ERMoE: Eigen-Reparameterized Mixture-of-Experts for Stable Routing and Interpretable Specializationcited by 6
  6. 2026Knowledge Localization in Mixture-of-Experts LLMs Using Cross-Lingual Inconsistencycited by 6
  7. 2026PROBE: Co-Balancing Computation and Communication in MoE Inference via Real-Time Predictive Prefetchingcited by 4
  8. 2026SlimQwen: Exploring the Pruning and Distillation in Large MoE Model Pre-trainingcited by 4
  9. 2026φ-Balancing for Mixture-of-Experts Trainingcited by 4
  10. 2025Point-MoE: Large-Scale Multi-Dataset Training with Mixture-of-Experts for 3D Semantic Segmentationcited by 3
  11. 2026SD-MoE: Spectral Decomposition for Effective Expert Specializationcited by 3
  12. 2026Expert Divergence Learning for MoE-based Language Modelscited by 3
  13. 2026Expert Threshold Routing for Autoregressive Language Modeling with Dynamic Computation Allocation and Load Balancingcited by 3
  14. 2026EMO: Pretraining Mixture of Experts for Emergent Modularitycited by 3
  15. 2026Extracting Small Translation Specialists from LLMs by Aggressively Pruning Expertscited by 3
  16. 2026DAG-MoE: From Simple Mixture to Structural Aggregation in Mixture-of-Expertscited by 3
  17. 2025Sigma-MoE-Tiny Technical Reportcited by 2
  18. 2026EMoE: Eigenbasis-Guided Routing for Mixture-of-Expertscited by 2
  19. 2026Multi-Head LatentMoE and Head Parallel: Communication-Efficient and Deterministic MoE Parallelismcited by 2
  20. 2026The Rise of Sparse Mixture-of-Experts: A Survey from Algorithmic Foundations to Decentralized Architectures and Vertical Domain Applicationscited by 2
  21. 2026MagicAgent: Towards Generalized Agent Planningcited by 2
  22. 2025Fine-grained MoE Load Balancing with Linear Programmingcited by 1
  23. 2026Improving MoE Compute Efficiency by Composing Weight and Data Sparsitycited by 1
  24. 2026MoE Routing Testbed: Studying Expert Specialization and Routing Behavior at Small Scalecited by 1
  25. 2026Laguna M.1/XS.2 Technical Reportcited by 1
  26. 2025Neural network task specialization via domain constrainingcited by 0
  27. 2025Mixture of Lookup Key-Value Expertscited by 0
  28. 2026How to Set the Batch Size for Large-Scale Pre-training?cited by 0
  29. 2026Scalable Heterogeneous Graph Learning via Heterogeneous-aware Orthogonal Prototype Expertscited by 0
  30. 2026A Multi-Agent Framework for Automated Exploit Generation with Constraint-Guided Comprehension and Reflectioncited by 0
  31. 2026UB-SMoE: Universally Balanced Sparse Mixture-of-Experts for Resource-adaptive Federated Fine-tuning of Foundation Modelscited by 0
  32. 2026GEM: GPU-Variability-Aware Expert to GPU Mapping for MoE Systemscited by 0
  33. 2026PithTrain: A Compact and Agent-Native MoE Training Systemcited by 0
  34. 2026Confidence-Adaptive SwiGLU for Mixture-of-Expertscited by 0
  35. 2026STAR: Rethinking MoE Routing as Structure-Aware Subspace Learningcited by 0
  36. 2026Fisher-Routed Mixture of Experts for Federated Class-Incremental Learningcited by 0
  37. 2026Structure Aware Neural Architecture Search for Mixture of Expertscited by 0
  38. 2026A.X K2 Technical Reportcited by 0

50 most cited of 53 All 53 (JSON)

Similar, not cited

Close in wording, no citation either way

  1. 2025Binary-Integer-Programming Based Algorithm for Expert Load Balancing in Mixture-of-Experts Models
  2. 2026LAER-MoE: Load-Adaptive Expert Re-layout for Efficient Mixture-of-Experts Training
  3. 2025Advancing Expert Specialization for Better MoE
  4. 2026CIPHER-MoE: Balancing Efficiency and Routing Fidelity in Trillion-Scale MoE Training
  5. 2026MoE-DisCo:Low Economy Cost Training Mixture-of-Experts Models
  6. 2025Mixture-of-Experts Can Surpass Dense LLMs Under Strictly Equal Resource
  7. 2025SYMI: Efficient Mixture-of-Experts Training via Model and Optimizer State Decoupling
  8. 2026MoEless: Efficient MoE LLM Serving via Serverless Computing