ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

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. Code is available at: https://github.com/HengHugoZhao/ProbMoE.git

Mixture-of-Experts models scale large language models efficiently by routing each token to a small set of experts. However, choosing which experts to use is difficult: most routers rely on non-differentiable top-$k$ selection, often overusing a few experts while providing almost no learning signals to the rest. We propose ProbMoE, a probabilistic routing framework for fine-tuning Mixture-of-Experts models. Instead of treating routing as a fixed top-$k$ ranking problem, ProbMoE casts expert selection as probabilistic inference over possible expert subsets. The key idea is that the router is no longer trained as if there is only one correct expert subset, but learns from a distribution of valid expert subset choices. By using exact expert-selection probabilities, ProbMoE gives the router more informative training signals while still activating only a sparse set of experts. This probabilistic perspective also makes dynamic-$k$ routing natural. Rather than requiring every token to use the same number of experts, ProbMoE can infer how many experts a token needs. In our experiments, ProbMoE improves expert utilization and routing diversity, while ProbMoE Dynamic-$k$ achieves comparable performance with fewer activated experts on average.