mlsanityData to 7 Oct 2026

Unbiased Gradient Estimation with Balanced Assignments for Mixtures of Experts

arXiv 2109.11817 ↗First posted 24 Sep 2021cs.LG stat.MLcited by 10

On the map: Mixture-of-experts → MoE theory and gating functions

Abstract (arXiv)

Training large-scale mixture of experts models efficiently on modern hardware requires assigning datapoints in a batch to different experts, each with a limited capacity. Recently proposed assignment procedures lack a probabilistic interpretation and use biased estimators for training. As an alternative, we propose two unbiased estimators based on principled stochastic assignment procedures: one that skips datapoints which exceed expert capacity, and one that samples perfectly balanced assignments using an extension of the Gumbel-Matching distribution [29]. Both estimators are unbiased, as they correct for the used sampling procedure. On a toy experiment, we find the `skip'-estimator is more effective than the balanced sampling one, and both are more robust in solving the task than biased alternatives.

Key facts. Cited by 10 papers in the map (data through 7 Oct 2026). Placed in MoE theory and gating functions, one of 12 directions in Mixture-of-experts. Builds on 21 of its 46 references that are in the map. First posted on arXiv 24 Sep 2021.

Builds on

21 of 46 references.

References · 21

  1. 2014Adam: A Method for Stochastic Optimizationcited by 36,960
  2. 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,413
  3. 2021Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsitycited by 3,014
  4. 2013Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distancescited by 2,277
  5. 2013Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computationcited by 2,031
  6. 2020GShard: Scaling Giant Models with Conditional Computation and Automatic Shardingcited by 1,601
  7. 2021Scaling Vision with Sparse Mixture of Expertscited by 740
  8. 2016Adaptive Computation Time for Recurrent Neural Networkscited by 479
  9. 2021BASE Layers: Simplifying Training of Large, Sparse Modelscited by 315
  10. 2013Learning Factored Representations in a Deep Mixture of Expertscited by 275
  11. 2021Hash Layers For Large Sparse Modelscited by 236
  12. 2018Learning Latent Permutations with Gumbel-Sinkhorn Networkscited by 172
9 more
  1. 2015Conditional Computation in Neural Networks for faster modelscited by 166
  2. 2014A* Samplingcited by 150
  3. 2019Stochastic Optimization of Sorting Networks via Continuous Relaxationscited by 111
  4. 2018Modular Networks: Learning to Decompose Neural Computationcited by 62
  5. 2019Routing Networks and the Challenges of Modular and Compositional Computationcited by 49
  6. 2021M6-T: Exploring Sparse Expert Models and Beyondcited by 22
  7. 2014Techniques for Learning Binary Stochastic Feedforward Neural Networkscited by 14
  8. 2013On Sampling from the Gibbs Distribution with Random Maximum A-Posteriori Perturbationscited by 7
  9. 2019Approximating the Permanent by Sampling from Adaptive Partitionscited by 2

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2022: 4; 2023: 3; 2024: 1; 2025: 1; 2026: 1.

10 papers.

  1. 2022Unified Scaling Laws for Routed Language Modelscited by 249
  2. 2023Soft Merging of Experts with Adaptive Routingcited by 88
  3. 2022Sparsity-Constrained Optimal Transportcited by 32
  4. 2023Approximating Two-Layer Feedforward Networks for Efficient Transformerscited by 27
  5. 2024Routers in Vision Mixture of Experts: An Empirical Studycited by 22
  6. 2023Sparse Backpropagation for MoE Trainingcited by 16
  7. 2022Automatic Expert Selection for Multi-Scenario and Multi-Task Searchcited by 9
  8. 2025Selective Sinkhorn Routing for Improved Sparse Mixture of Expertscited by 1
  9. 2026PR2: Predictive Routing Replay for MoE-Based LLM Reinforcement Learningcited by 1
  10. 2022Stochastic gradient descent with gradient estimator for categorical featurescited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2026φ-Balancing for Mixture-of-Experts Training
  2. 2020Biased Mixtures Of Experts: Enabling Computer Vision Inference Under Data Transfer Limitations
  3. 2024Generalization Error Analysis for Sparse Mixture-of-Experts: A Preliminary Study
  4. 2024Learning Mixtures of Experts with EM: A Mirror Descent Perspective
  5. 2024Sigmoid Gating is More Sample Efficient than Softmax Gating in Mixture of Experts
  6. 2024Beyond Parameter Count: Implicit Bias in Soft Mixture of Experts
  7. 2026Toward Calibrated Mixture-of-Experts Under Distribution Shift
  8. 2024Understanding Expert Structures on Minimax Parameter Estimation in Contaminated Mixture of Experts