title: Demons in the Detail: On Implementing Load Balancing Loss for Training
    Specialized Mixture-of-Expert Models
summary: This paper revisits the implementation of
  Load-balancing Loss (LBL) when training
  Mixture-of-Experts (MoEs) models. Specifically, LBL for MoEs is defined as N_E
  sum_{i=1}^{N_E} f_i p_i, where N_E is the total number of experts, f_i
  represents the frequency of expert i being selected, and p_i denotes the
  average gating score of the expert i. Existing MoE training frameworks
  usually employ the parallel training strategy so that f_i 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 (e.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 f_i 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.
paper_authors._id: 679071da11a3f67d8f498649
paper_authors.name: Zihan Qiu
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2025-01-22T15:32:39.315Z
paper_authors._id: 647ccbd6e07cf9bb2d485244
paper_authors.avatarUrl: /avatars/e8915abaff04f6762247e196b7cf84df.svg
paper_authors.fullname: Zihan Qiu
paper_authors.type: user
paper_authors.user: QwQZh
paper_authors._id: 679071da11a3f67d8f49864a
paper_authors.name: Zeyu Huang
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2025-01-22T15:35:10.842Z
paper_authors._id: 64a5520dea3c5861303064b0
paper_authors.avatarUrl: /avatars/7e618bcfc4d8d3291cf9bc8f2c4c1b15.svg
paper_authors.fullname: Zeyu hUANG
paper_authors.type: user
paper_authors.user: zeroyhuang
paper_authors._id: 679071da11a3f67d8f49864b
paper_authors.name: Bo Zheng
paper_authors._id: 679071da11a3f67d8f49864c
paper_authors.name: Kaiyue Wen
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2025-01-22T20:49:51.765Z
paper_authors._id: 6690a44e76c0fa097f9f65c9
paper_authors.avatarUrl: /avatars/0c6efa10889d1b6dd9a3f9985f4d2d97.svg
paper_authors.fullname: Kaiyue Wen
paper_authors.type: user
paper_authors.user: KaiyueWen
paper_authors._id: 679071da11a3f67d8f49864d
paper_authors.name: Zekun Wang
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2025-01-22T15:31:36.863Z
paper_authors._id: 656832dfbd65fd41ee7aa8cd
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/656832dfbd65fd41ee7aa8cd/HHkyetTqNq1wIBPipzjQA.jpeg
paper_authors.fullname: Zekun Wang
paper_authors.type: user
paper_authors.user: kugwzk
paper_authors._id: 679071da11a3f67d8f49864e
paper_authors.name: Rui Men
paper_authors._id: 679071da11a3f67d8f49864f
paper_authors.name: Ivan Titov
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2025-01-22T15:30:57.878Z
paper_authors._id: 666f6375623133f1ce79c021
paper_authors.avatarUrl: /avatars/b57a457cff81d685fdd9eadee53445fc.svg
paper_authors.fullname: Ivan Titov
paper_authors.type: user
paper_authors.user: Ivanchoo
paper_authors._id: 679071da11a3f67d8f498650
paper_authors.name: Dayiheng Liu
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2025-01-22T10:08:15.815Z
paper_authors._id: 6434d4989bd5a84b5dd0b0f5
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/6434d4989bd5a84b5dd0b0f5/0Elf9qbfG9Hkgypm9pTGm.jpeg
paper_authors.fullname: Dayiheng Liu
paper_authors.type: user
paper_authors.user: Losin94
paper_authors._id: 679071da11a3f67d8f498651
paper_authors.name: Jingren Zhou
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2025-01-22T15:30:48.272Z
paper_authors._id: 602f88f5e8149a962412a667
paper_authors.avatarUrl: /avatars/b78f0e583df8e5d5e3365934fe5f4900.svg
paper_authors.fullname: Zhou
paper_authors.type: user
paper_authors.user: Jingren
paper_authors._id: 679071da11a3f67d8f498652
paper_authors.name: Junyang Lin
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2025-01-22T15:30:37.427Z
paper_authors._id: 620760a26e3b7210c2ff1943
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/620760a26e3b7210c2ff1943/VC-rKqimF6yxGESNVlPoR.jpeg
paper_authors.fullname: Junyang Lin
paper_authors.type: user
paper_authors.user: JustinLin610
paper_title: Demons in the Detail: On Implementing Load Balancing Loss for Training
    Specialized Mixture-of-Expert Models
paper_summary: This paper revisits the implementation of
  Load-balancing Loss (LBL) when training
  Mixture-of-Experts (MoEs) models. Specifically, LBL for MoEs is defined as N_E
  sum_{i=1}^{N_E} f_i p_i, where N_E is the total number of experts, f_i
  represents the frequency of expert i being selected, and p_i denotes the
  average gating score of the expert i. Existing MoE training frameworks
  usually employ the parallel training strategy so that f_i 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 (e.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 f_i 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.
paper_ai_keywords: Load-balancing Loss (LBL) Mixture-of-Experts (MoEs) micro-batch experts frequency gating score parallel training strategy token sequence router global-batch corpus level communication step domain specialization