Load Balancing Mixture of Experts with Similarity Preserving Routers

• URL: http://arxiv.org/abs/2506.14038v2
• Date: Sat, 11 Oct 2025 22:22:55 GMT
• Title: Load Balancing Mixture of Experts with Similarity Preserving Routers
• Authors: Nabil Omi, Siddhartha Sen, Ali Farhadi,
• Abstract summary: Sparse Mixture of Experts (MoE) models offer a scalable and efficient architecture for training large neural networks.<n>We introduce a novel load balancing loss that preserves token-wise relational structure.<n>Our results show that applying our loss to the router results in 36% faster convergence and lower redundancy.
• Score: 30.279616888339543
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Sparse Mixture of Experts (MoE) models offer a scalable and efficient architecture for training large neural networks by activating only a subset of parameters ("experts") for each input. A learned router computes a distribution over these experts, and assigns input tokens to a small subset. However, without auxiliary balancing mechanisms, routers often converge to using only a few experts, severely limiting model capacity and degrading performance. Most current load balancing mechanisms encourage a distribution over experts that resembles a roughly uniform distribution of experts per token. During training, this can result in inconsistent routing behavior, resulting in the model spending its capacity to learn redundant knowledge. We address this by introducing a novel load balancing loss that preserves token-wise relational structure, encouraging consistent expert choices for similar inputs during training. Our experimental results show that applying our loss to the router results in 36% faster convergence and lower redundancy compared to a popular load balancing loss.

Related papers

• A Replicate-and-Quantize Strategy for Plug-and-Play Load Balancing of Sparse Mixture-of-Experts LLMs [64.8510381475827]
  Sparse Mixture-of-Experts (SMoE) architectures are increasingly used to scale large language models efficiently.<n>SMoE models often suffer from severe load imbalance across experts, where a small subset of experts receives most tokens while others are underutilized.<n>We present a systematic analysis of expert routing during inference and identify three findings: (i) load imbalance persists and worsens with larger batch sizes, (ii) selection frequency does not reliably reflect expert importance, and (iii) overall expert workload and importance can be estimated using a small calibration set.
  arXiv  Detail & Related papers  (2026-02-23T15:11:16Z)
• Least-Loaded Expert Parallelism: Load Balancing An Imbalanced Mixture-of-Experts [74.40169987564724]
  Expert parallelism (EP) is designed to scale MoE models by distributing experts across multiple devices.<n>Under extreme imbalance, EP can funnel a disproportionate number of tokens to a small number of experts, leading to compute- and memory-bound failures.<n>We propose Least-Loaded Expert Parallelism (LLEP), a novel EP algorithm that dynamically reroutes excess tokens and associated expert parameters from overloaded devices to underutilized ones.
  arXiv  Detail & Related papers  (2026-01-23T18:19:15Z)
• RepetitionCurse: Measuring and Understanding Router Imbalance in Mixture-of-Experts LLMs under DoS Stress [16.010076395422264]
  We show that out-of-distribution prompts can manipulate the routing strategy, which creates computational bottlenecks on certain devices while forcing others to idle.<n>We propose RepetitionCurse, a low-cost black-box strategy to exploit this vulnerability.
  arXiv  Detail & Related papers  (2025-12-30T05:24:26Z)
• ERMoE: Eigen-Reparameterized Mixture-of-Experts for Stable Routing and Interpretable Specialization [13.182475975397251]
  ERMoE is a sparse MoE transformer that replaces learned gating logits with an "Eigenbasis Score"<n>We show that ERMoE achieves state-of-the-art accuracy on ImageNet classification and cross-modal image-text retrieval benchmarks.<n>A 3D MRI variant (ERMoE-ba) improves brain age prediction accuracy by more than 7% and yields interpretable expert specializations.
  arXiv  Detail & Related papers  (2025-11-14T05:31:37Z)
• Mixture-of-Transformers Learn Faster: A Theoretical Study on Classification Problems [59.94955550958074]
  We study a tractable theoretical framework in which each transformer block acts as an expert governed by a continuously trained gating network.<n>We show that expert specialization reduces gradient conflicts and makes each subtask strongly convex.<n>We prove that the training drives the expected prediction loss to near zero in $O(log(epsilon-1)$ steps, significantly improving over the $O(epsilon-1)$ rate for a single transformer.
  arXiv  Detail & Related papers  (2025-10-30T21:07:36Z)
• From Score Distributions to Balance: Plug-and-Play Mixture-of-Experts Routing [52.01745035243826]
  Mixture-of-Experts (MoE) models can scale parameter capacity by routing each token to a subset of experts.<n> conditional routing shifts the burden on inference memory, limiting the number of experts per device.<n>We present LASER, a plug-and-play, inference-time routing algorithm that balances load while preserving accuracy.
  arXiv  Detail & Related papers  (2025-09-29T16:29:17Z)
• On the Role of Discrete Representation in Sparse Mixture of Experts [33.809432499123275]
  We propose a new architecture dubbed Vector-Quantized Mixture of Experts (VQMoE)<n>VQMoE is an effective solution for scaling up model capacity without increasing the computational costs.<n>We show that VQMoE achieves a 28% improvement in routers compared to other SMoE routing methods.
  arXiv  Detail & Related papers  (2024-11-28T22:32:01Z)
• LocMoE: A Low-Overhead MoE for Large Language Model Training [13.153904674287546]
  We propose a novel routing strategy that combines load balance and locality by converting partial inter-node communication to that of intra-node. The proposed LocMoE reduces training time per epoch by 12.68% to 22.24% compared to classical routers.
  arXiv  Detail & Related papers  (2024-01-25T03:36:39Z)
• Soft Merging of Experts with Adaptive Routing [38.962451264172856]
  We introduce Soft Merging of Experts with Adaptive Routing (SMEAR) SMEAR avoids discrete routing by using a single "merged" expert constructed via a weighted average of all of the experts' parameters. We empirically validate that models using SMEAR outperform models that route based on metadata or learn sparse routing through gradient estimation.
  arXiv  Detail & Related papers  (2023-06-06T15:04:31Z)
• Sparse MoE as the New Dropout: Scaling Dense and Self-Slimmable Transformers [107.3726071306935]
  We propose a new plug-and-play training framework, SMoE-Dropout, to enable scaling transformers to better accuracy in their full capacity without collapse. SMoE-Dropout consists of a randomly and fixed router network to activate experts and gradually increases the activated expert number as training progresses over time. Our experiments demonstrate the superior performance and substantial computation savings of SMoE-Dropout, compared to dense training baselines with equivalent parameter counts.
  arXiv  Detail & Related papers  (2023-03-02T22:12:51Z)
• Gating Dropout: Communication-efficient Regularization for Sparsely Activated Transformers [78.77361169167149]
  We propose emphGating Dropout, which allows tokens to ignore the gating network and stay at their local machines. Similar to traditional dropout, we also show that Gating Dropout has a regularization effect during training, resulting in improved generalization performance.
  arXiv  Detail & Related papers  (2022-05-28T05:12:43Z)
• On the Representation Collapse of Sparse Mixture of Experts [102.83396489230375]
  Sparse mixture of experts provides larger model capacity while requiring a constant computational overhead. It employs the routing mechanism to distribute input tokens to the best-matched experts according to their hidden representations. However, learning such a routing mechanism encourages token clustering around expert centroids, implying a trend toward representation collapse.
  arXiv  Detail & Related papers  (2022-04-20T01:40:19Z)
• Federated Learning with Unreliable Clients: Performance Analysis and Mechanism Design [76.29738151117583]
  Federated Learning (FL) has become a promising tool for training effective machine learning models among distributed clients. However, low quality models could be uploaded to the aggregator server by unreliable clients, leading to a degradation or even a collapse of training. We model these unreliable behaviors of clients and propose a defensive mechanism to mitigate such a security risk.
  arXiv  Detail & Related papers  (2021-05-10T08:02:27Z)
• BASE Layers: Simplifying Training of Large, Sparse Models [53.98145464002843]
  We introduce a new balanced assignment of experts (BASE) layer for large language models. Sparse layers can dramatically improve the efficiency of training and inference by routing each token to specialized expert modules. We formulate token-to-expert allocation as a linear assignment problem, allowing an optimal assignment in which each expert receives an equal number of tokens.
  arXiv  Detail & Related papers  (2021-03-30T23:08:32Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.