ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

Technical Summary: ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

Problem Statement
Mixture-of-Experts (MoE) models scale efficiently by activating only a small subset of experts per token. However, training these models remains challenging because standard routing relies on a discrete, non-differentiable top-$k$ operator. Conventional training procedures bypass this by ignoring the dependence of the top-$k$ selection on router logits, propagating gradients only through the softmax probabilities of the selected experts. This approximation limits the learning signal for alternative expert subsets, often leading to concentrated routing distributions, poor expert utilization, and unstable training dynamics. Existing attempts to mitigate this, such as dense backpropagation or heuristic gradient estimators, often fail to explicitly model the distribution over feasible expert subsets or provide limited mechanisms for exploring alternative combinations during training.

Methodology
The authors propose ProbMoE, a framework that reformulates expert routing as a probabilistic inference problem over cardinality-constrained expert subsets. Instead of treating routing as a deterministic operator, ProbMoE models the selection of experts as a distribution over discrete subsets.

1. Probabilistic Formulation:

   • Exact-$k$ Routing: The framework defines a distribution over subsets $S$ where $|S|=k$. It conditions independent Bernoulli selection probabilities (derived from router logits via sigmoid) on the constraint that exactly $k$ experts are selected. The normalizing constant for this distribution, $Z_k$, is computed exactly using dynamic programming, making the formulation tractable.
   • Dynamic-$k$ Routing: The framework generalizes to a range-constrained setting where the subset size $|S|$ lies within $[k_{min}, k_{max}]$. The distribution is conditioned on this range, allowing the model to jointly infer the expert identities and the routing cardinality for each token.

2. Forward Pass:

   • During training, ProbMoE samples an expert subset $S$ from the defined distribution (exact-$k$ or range-constrained).
   • The MoE layer computes the output using only the sampled experts, preserving the sparse execution pattern of standard MoE models.
   • At inference, the Maximum-A-Posteriori (MAP) set is selected, ensuring inference costs match conventional MoE routing.

3. Backward Pass:

   • To enable gradient-based optimization, ProbMoE employs the SIMPLE gradient estimator (Ahmed et al., 2023).
   • Instead of differentiating through the discrete sample, the backward pass uses the marginal probabilities of each expert being selected under the constrained distribution.
   • Routing weights are constructed using a straight-through estimator: $w = (\text{stopgrad}(z - m) + m) \odot \pi$, where $z$ is the sampled mask, $m$ is the vector of exact marginal probabilities, and $\pi$ are the softmax routing probabilities.
   • This allows gradients to flow through the marginal probabilities, providing informative updates to the router that reflect how changes in selection probabilities affect the expected loss, while maintaining sparse expert execution.

Key Contributions

1. Probabilistic Inference Formulation: The paper casts MoE routing as probabilistic inference over cardinality-constrained expert subsets, moving beyond heuristic gradient approximations.
2. Tractable Marginal-Based Optimization: By leveraging the SIMPLE estimator, the authors pair stochastic subset sampling with exact marginal-based router gradients, enabling differentiable training without dense expert-side computation.
3. Unified Dynamic-$k$ Framework: The method naturally extends to dynamic routing, where the model infers both the expert identities and the number of active experts per token within a predefined range, without relying on heuristic gating rules.
4. Empirical Improvements:
   • Exact-$k$: ProbMoE achieves strong performance across benchmarks (GSM8K, Law, Translation, MBPP, MMLU) on OLMoE and Qwen backbones, often matching or outperforming competitive baselines like DenseMixer and Conventional Top-$k$.
   • Dynamic-$k$: ProbMoE Dynamic-$k$ achieves comparable performance to Exact-$k$ while activating fewer experts on average, demonstrating effective adaptive allocation.
5. Routing Diversity: Analysis shows ProbMoE induces broader routing distributions and more balanced expert utilization compared to baselines, mitigating expert collapse.

Results and Significance
The paper demonstrates that ProbMoE improves expert utilization and routing diversity. In ablation studies, the combination of stochastic sampling and marginal-based optimization is shown to be the primary driver of performance gains, outperforming configurations that use only one of these components.

The authors claim that ProbMoE provides a principled method for expert selection that improves training stability and computational efficiency. By modeling routing probabilistically, the framework allows for the exploration of alternative expert combinations during training, leading to better-calibrated routing distributions. The Dynamic-$k$ variant further shows that adaptive cardinality can be learned to match token-level complexity, reducing effective routing costs without sacrificing task performance. The work suggests that future directions include extending ProbMoE to pre-training and system-level optimizations to fully exploit sparsity in dynamic settings.