# Unbiased Gradient Estimation with Balanced Assignments for Mixtures of Experts

arXiv: https://arxiv.org/abs/2109.11817 · HTML: https://mlsanity.com/paper/2109.11817/

## Key facts

- Cited by 10 papers in the map (data through 14 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.

## Placed in

- Mixture-of-experts / [MoE theory and gating functions](https://mlsanity.com/moe/moe-theory-and-gating-functions/)

## Builds on (most cited first, 21 of 21)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Scaling Vision with Sparse Mixture of Experts](https://mlsanity.com/paper/2106.05974/) (2021)
- [Adaptive Computation Time for Recurrent Neural Networks](https://mlsanity.com/paper/1603.08983/) (2016)
- [BASE Layers: Simplifying Training of Large, Sparse Models](https://mlsanity.com/paper/2103.16716/) (2021)
- [Learning Factored Representations in a Deep Mixture of Experts](https://mlsanity.com/paper/1312.4314/) (2013)
- [Hash Layers For Large Sparse Models](https://mlsanity.com/paper/2106.04426/) (2021)
- [Learning Latent Permutations with Gumbel-Sinkhorn Networks](https://mlsanity.com/paper/1802.08665/) (2018)
- [Conditional Computation in Neural Networks for faster models](https://mlsanity.com/paper/1511.06297/) (2015)
- [A* Sampling](https://mlsanity.com/paper/1411.0030/) (2014)
- [Stochastic Optimization of Sorting Networks via Continuous Relaxations](https://mlsanity.com/paper/1903.08850/) (2019)
- [Modular Networks: Learning to Decompose Neural Computation](https://mlsanity.com/paper/1811.05249/) (2018)
- [Routing Networks and the Challenges of Modular and Compositional Computation](https://mlsanity.com/paper/1904.12774/) (2019)
- [M6-T: Exploring Sparse Expert Models and Beyond](https://mlsanity.com/paper/2105.15082/) (2021)
- [Techniques for Learning Binary Stochastic Feedforward Neural Networks](https://mlsanity.com/paper/1406.2989/) (2014)
- [On Sampling from the Gibbs Distribution with Random Maximum A-Posteriori Perturbations](https://mlsanity.com/paper/1309.7598/) (2013)
- [Approximating the Permanent by Sampling from Adaptive Partitions](https://mlsanity.com/paper/1911.11856/) (2019)

## Built on by (most cited first, 10 of 10)

- [Unified Scaling Laws for Routed Language Models](https://mlsanity.com/paper/2202.01169/) (2022)
- [Soft Merging of Experts with Adaptive Routing](https://mlsanity.com/paper/2306.03745/) (2023)
- [Sparsity-Constrained Optimal Transport](https://mlsanity.com/paper/2209.15466/) (2022)
- [Approximating Two-Layer Feedforward Networks for Efficient Transformers](https://mlsanity.com/paper/2310.10837/) (2023)
- [Routers in Vision Mixture of Experts: An Empirical Study](https://mlsanity.com/paper/2401.15969/) (2024)
- [Sparse Backpropagation for MoE Training](https://mlsanity.com/paper/2310.00811/) (2023)
- [Automatic Expert Selection for Multi-Scenario and Multi-Task Search](https://mlsanity.com/paper/2205.14321/) (2022)
- [Selective Sinkhorn Routing for Improved Sparse Mixture of Experts](https://mlsanity.com/paper/2511.08972/) (2025)
- [PR2: Predictive Routing Replay for MoE-Based LLM Reinforcement Learning](https://mlsanity.com/paper/2606.00395/) (2026)
- [Stochastic gradient descent with gradient estimator for categorical features](https://mlsanity.com/paper/2209.03771/) (2022)

## 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.