# ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Expert routing strategies, one of 12 directions in Mixture-of-experts.
- Builds on 30 of its 37 references that are in the map.
- First posted on arXiv 1 Jun 2026.

## Placed in

- Mixture-of-experts / [Expert routing strategies](https://mlsanity.com/moe/expert-routing-strategies/)

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

- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [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)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Mixtral of Experts](https://mlsanity.com/paper/2401.04088/) (2024)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [GLaM: Efficient Scaling of Language Models with Mixture-of-Experts](https://mlsanity.com/paper/2112.06905/) (2021)
- [DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models](https://mlsanity.com/paper/2401.06066/) (2024)
- [DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale](https://mlsanity.com/paper/2201.05596/) (2022)
- [BASE Layers: Simplifying Training of Large, Sparse Models](https://mlsanity.com/paper/2103.16716/) (2021)
- [OLMoE: Open Mixture-of-Experts Language Models](https://mlsanity.com/paper/2409.02060/) (2024)
- [Unified Scaling Laws for Routed Language Models](https://mlsanity.com/paper/2202.01169/) (2022)
- [Taming Sparsely Activated Transformer with Stochastic Experts](https://mlsanity.com/paper/2110.04260/) (2021)
- [Dynamic Mixture of Experts: An Auto-Tuning Approach for Efficient Transformer Models](https://mlsanity.com/paper/2405.14297/) (2024)
- [MoE++: Accelerating Mixture-of-Experts Methods with Zero-Computation Experts](https://mlsanity.com/paper/2410.07348/) (2024)
- [ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routing](https://mlsanity.com/paper/2412.14711/) (2024)
- [AdaMoE: Token-Adaptive Routing with Null Experts for Mixture-of-Experts Language Models](https://mlsanity.com/paper/2406.13233/) (2024)
- [SIMPLE: A Gradient Estimator for k-Subset Sampling](https://mlsanity.com/paper/2210.01941/) (2022)
- [Multi-Head Mixture-of-Experts](https://mlsanity.com/paper/2404.15045/) (2024)
- [Probabilistically Rewired Message-Passing Neural Networks](https://mlsanity.com/paper/2310.02156/) (2023)
- [Let the Expert Stick to His Last: Expert-Specialized Fine-Tuning for Sparse Architectural Large Language Models](https://mlsanity.com/paper/2407.01906/) (2024)
- [Sparse Backpropagation for MoE Training](https://mlsanity.com/paper/2310.00811/) (2023)
- [Ada-K Routing: Boosting the Efficiency of MoE-based LLMs](https://mlsanity.com/paper/2410.10456/) (2024)
- [DA-MoE: Towards Dynamic Expert Allocation for Mixture-of-Experts Models](https://mlsanity.com/paper/2409.06669/) (2024)
- [A Unified Approach to Count-Based Weakly-Supervised Learning](https://mlsanity.com/paper/2311.13718/) (2023)
- [MoEs Are Stronger than You Think: Hyper-Parallel Inference Scaling with RoE](https://mlsanity.com/paper/2509.17238/) (2025)

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

- [How Sparse Probability Maps Shape Mixture-of-Experts Routing](https://mlsanity.com/paper/2610.06677/) (2026)

## Abstract (arXiv)

Mixture-of-Experts (MoE) models scale by activating only a small subset of experts per token. However, training such models remains challenging because top-k routing is discrete and non-differentiable, requiring gradient estimators for expert selection whose design remains a central open problem. We introduce ProbMoE, a probabilistic routing framework that models expert selection as a distribution over cardinality-constrained expert subsets and formulates routing as probabilistic inference in this discrete subset space. We first propose ProbMoE Exact-k routing, which samples k-expert subsets in the forward pass, and the backward pass uses gradients through each expert's exact marginal probability as a tractable surrogate for the true gradient. ProbMoE naturally generalizes to a dynamic-k routing setting, where both training and inference constrain the routing cardinality to the same predefined range, allowing adaptive expert allocation per token. Across benchmarks and model backbones, ProbMoE Exact-k achieves strong performance compared to competitive baselines, with improved expert utilization and routing diversity; ProbMoE Dynamic-k achieves comparable performance with fewer activated experts.