$φ$-Balancing for Mixture-of-Experts Training

• URL: http://arxiv.org/abs/2605.15403v1
• Date: Thu, 14 May 2026 20:39:28 GMT
• Title: $φ$-Balancing for Mixture-of-Experts Training
• Abstract summary: Mixture-of-Experts (MoE) models rely on balanced expert utilization to fully realize their scalability.<n>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.
• Score: 20.061328126956028
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• Distribution-Consistent Inference for Dynamic Sparse Mixture-of-Experts [48.294824672909876]
  Inference-time dynamic top-$k$ routing can reduce computation without retraining.<n>We show that reducing the number of activated experts consistently increases the RMS scale and variance of SMoE outputs.<n>We propose Layer-wise Distribution Alignment (LDA), a lightweight inference-time correction that uses layer-wise calibration statistics to align reduced-routing representations with the default configuration.
  arXiv  Detail & Related papers  (2026-09-08T06:23:07Z)
• Estimating Mixture Distributions via Stochastic Mirror Descent [15.58350224981605]
  We revisit the problem of estimating an unknown distribution from its samples by fitting a mixture model that minimizes cross-entropy loss.<n>We propose a family of estimators derived from the mirror descent (SMD) algorithm.<n>We demonstrate that under mild conditions, the proposed estimators achieve near-optimal convergence rates.
  arXiv  Detail & Related papers  (2026-05-24T08:19:42Z)
• UB-SMoE: Universally Balanced Sparse Mixture-of-Experts for Resource-adaptive Federated Fine-tuning of Foundation Models [2.3099144596725565]
  We propose Universally Balanced Sparse Mixture-of-Experts (UB-SMoE) to address expert utilization imbalance and Top-K routing problems.<n>UB-SMoE achieves up to $45.0%$ computational reduction on low-resource clients while improving their performance by $8.7 times$ compared to existing heterogeneous LoRA-rank methods.
  arXiv  Detail & Related papers  (2026-05-15T23:06:59Z)
• Self-Distilled Trajectory-Aware Boltzmann Modeling: Bridging the Training-Inference Discrepancy in Diffusion Language Models [65.89572755202245]
  Diffusion Language Models (DLMs) offer stronger global awareness and highly parallel generation.<n>Post-training DLMs with standard Negative Evidence Lower Bound (NELBO)-based supervised fine-tuning remains inefficient.<n>We propose a self-distilled trajectory-based post-training framework that aligns training with the easy-to-hard structure of inference.
  arXiv  Detail & Related papers  (2026-05-12T09:39:06Z)
• SparseBalance: Load-Balanced Long Context Training with Dynamic Sparse Attention [19.076502375242352]
  We propose SparseBalance, which exploits the sparsity and sequence heterogeneity to optimize model accuracy and system efficiency jointly.<n>First, we propose workload-aware dynamic sparsity tuning, which employs a bidirectional sparsity adjustment to eliminate stragglers inherent to free accuracy.<n>Second, we propose a sparityaware strategy to achieve coarse-grained sparsity tuning, which exploits sparse up to 1.33$times$ end-to-end speedup while still improving long-context capability by 0.46% on the LongBench benchmark.
  arXiv  Detail & Related papers  (2026-04-15T13:18:07Z)
• Stability and Generalization of Push-Sum Based Decentralized Optimization over Directed Graphs [55.77845440440496]
  Push-based decentralized communication enables optimization over communication networks, where information exchange may be asymmetric.<n>We develop a unified uniform-stability framework for the Gradient Push (SGP) algorithm.<n>A key technical ingredient is an imbalance-aware generalization bound through two quantities.
  arXiv  Detail & Related papers  (2026-02-24T05:32:03Z)
• 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)
• Benefits of Online Tilted Empirical Risk Minimization: A Case Study of Outlier Detection and Robust Regression [7.238889207632064]
  Empirical Risk Minimization (ERM) is a foundational framework for supervised learning.<n>In online or streaming settings where data arrive one sample at a time, the classical TERM objective degenerates to standard ERM.<n>We propose an online TERM that removes the logarithm from the classical objective, preserving tilt effects without additional computational or memory overhead.
  arXiv  Detail & Related papers  (2025-09-18T16:51:23Z)
• Generalized Linear Bandits: Almost Optimal Regret with One-Pass Update [70.38810219913593]
  We study the generalized linear bandit (GLB) problem, a contextual multi-armed bandit framework that extends the classical linear model by incorporating a non-linear link function.<n>GLBs are widely applicable to real-world scenarios, but their non-linear nature introduces significant challenges in achieving both computational and statistical efficiency.<n>We propose a jointly efficient algorithm that attains a nearly optimal regret bound with $mathcalO(1)$ time and space complexities per round.
  arXiv  Detail & Related papers  (2025-07-16T02:24:21Z)
• Latent Prototype Routing: Achieving Near-Perfect Load Balancing in Mixture-of-Experts [0.0]
  Latent Prototype Routing (LPR) is a novel routing framework that promotes balanced expert utilization without compromising downstream performance.<n>LPR reduces the Gini coefficient of expert load from 0.70 to 0.035 on average, improves the min-max expert load ratio from 1e-6 to 0.70, achieving near-perfect load balancing.
  arXiv  Detail & Related papers  (2025-06-26T14:41:18Z)
• Distributionally Robust Models with Parametric Likelihood Ratios [123.05074253513935]
  Three simple ideas allow us to train models with DRO using a broader class of parametric likelihood ratios. We find that models trained with the resulting parametric adversaries are consistently more robust to subpopulation shifts when compared to other DRO approaches.
  arXiv  Detail & Related papers  (2022-04-13T12:43:12Z)

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.