ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

Summary (Overview)

• ProbMoE is a novel probabilistic routing framework for Mixture-of-Experts (MoE) models that models expert selection as a distribution over cardinality-constrained expert subsets, treating routing as probabilistic inference rather than a deterministic top-k operation.

• The framework introduces ProbMoE Exact-k routing, which samples k-expert subsets in the forward pass and uses exact marginal probabilities as a tractable gradient surrogate in the backward pass, enabling informative router gradients while preserving sparse expert execution.

• ProbMoE naturally generalizes to Dynamic-k routing, where both training and inference constrain routing cardinality to a predefined range $[k_{\min}, k_{\max}]$, allowing adaptive expert allocation per token based on token-level complexity.

• Empirical results show ProbMoE Exact-k achieves state-of-the-art performance across multiple benchmarks (GSM8K, Law, Translation, Summary) on both OLMoE-1B-7B and Qwen1.5-MoE-A2.7B backbones, with improved expert utilization and routing diversity.

• ProbMoE Dynamic-k achieves comparable performance to Exact-k while activating fewer experts on average (75–84.5% of the maximum), demonstrating effective adaptive computation allocation.

Introduction and Theoretical Foundation

Background and Motivation

Mixture-of-Experts (MoE) architectures scale large language models while keeping computational costs manageable by activating only a small subset of experts per token. However, training remains challenging because top-k routing is discrete and non-differentiable. Standard approaches ignore the dependence of the top-k operator on router logits during backpropagation, propagating gradients only through softmax probabilities. This leads to:

• Highly concentrated routing distributions
• Poor expert utilization
• Unstable training dynamics

Theoretical Foundation

The paper formalizes why top-k selection poses fundamental challenges. For a MoE layer with $N$ experts, the router produces logits $r = \text{Router}_\theta(x) \in \mathbb{R}^N$, converted to softmax weights:

$$\pi_i = \frac{\exp(r_i)}{\sum_{j=1}^{N}\exp(r_j)}, \quad \forall i \in [N]$$

The MoE output for a selected subset $S$ is:

$$y_S(x; r) \triangleq \sum_{j \in S} \pi_j f_j(x) \tag{1}$$

Differentiating the routed output with respect to router logits reveals two paths:

$$\frac{\partial y(x;r)}{\partial r_i} = \underbrace{\sum_{S\subseteq[N], |S|=k}\mathbb{I}[S = S_{\text{top-}k}(r)]\frac{\partial y_S(x;r)}{\partial r_i}}_{\text{softmax-weight path}} + \underbrace{\sum_{S\subseteq[N], |S|=k}y_S(x;r)\frac{\partial\mathbb{I}[S = S_{\text{top-}k}(r)]}{\partial r_i}}_{\text{discrete-selection path}} \tag{2}$$

The second term is non-differentiable because the indicator function $\mathbb{I}[S = S_{\text{top-}k}(r)]$ is piecewise-constant. Conventional MoE training drops this term, limiting the learning signal.

Probabilistic Formulation

ProbMoE optimizes the expected loss under a router-induced distribution over feasible subsets:

$$\mathcal{J}(\theta) = \mathbb{E}_{S \sim \mathbb{P}_r(\cdot \mid |S| = k)}\left[\mathcal{L}\big(y_S(x; r)\big)\right] \tag{3}$$

For each expert $i$, let $p_i = \sigma(r_i)$ be the Bernoulli selection probability. The exact-k distribution is:

$$\mathbb{P}_r(S \mid |S| = k) = \frac{1}{Z_k} \prod_{j \in S} p_j \prod_{j \notin S} (1 - p_j)$$

with normalizing constant:

$$Z_k = \sum_{S \subseteq [N], |S| = k} \prod_{j \in S} p_j \prod_{j \notin S} (1 - p_j) \tag{4}$$

Theorem 4.1 (Ahmed et al., 2023): The normalizing constant $Z_k$ can be computed exactly in time $\mathcal{O}(Nk)$, with vectorized complexity $\mathcal{O}(\log N \log k)$ under perfect parallelization.

Methodology

Forward Pass

ProbMoE samples an expert subset from the exact-k distribution:

$$S \sim \mathbb{P}_r(S \mid |S| = k)$$

This preserves sparse execution (only $k$ experts evaluated per token) while enabling stochastic exploration of alternative expert combinations. At inference, the Maximum-A-Posteriori (MAP) set is selected.

Backward Pass

The key insight is using conditional marginals as a differentiable surrogate. The marginal probability of expert $j$ is:

$$m_j \triangleq \mathbb{P}_r(j \in S \mid |S| = k) = \frac{\partial \log Z_k}{\partial \log p_j} \tag{6}$$

Since $Z_k$ is computed exactly via dynamic programming, these marginals are exact.

Marginal-Integrated Routing Weights

ProbMoE combines the sampled forward mask, marginals, and soft routing weights via a straight-through estimator:

$$w = \left(\text{stopgrad}(z - m) + m\right) \odot \pi \tag{7}$$

where $z$ is the k-hot sampled mask, $m$ are the marginals, and $\pi$ are softmax routing weights. In the forward pass, $w = z$ (preserving the sampled mask); in the backward pass, gradients flow through both $m$ and $\pi$.

The resulting router gradient decomposes as:

$$\frac{\partial \mathcal{L}}{\partial r_i} = \sum_{j=1}^{N}\left\langle\frac{\partial \mathcal{L}}{\partial y}, f_j(x)\right\rangle\left(m_j\frac{\partial \pi_j}{\partial r_i} + \pi_j\frac{\partial m_j}{\partial r_i}\right) \tag{8}$$

Dynamic-k Extension

ProbMoE Dynamic-k conditions on a range constraint $k \in [k_{\min}, k_{\max}]$:

$$\mathbb{P}_r(S \mid k_{\min} \leq |S| \leq k_{\max}) = \frac{1}{Z^*}\prod_{j \in S} p_j \prod_{j \notin S}(1 - p_j)$$

Theorem 5.1: The range-constrained normalizing constant $Z^* = \sum_{k=k_{\min}}^{k_{\max}} Z_k$ can be computed in $\mathcal{O}(Nk_{\max})$ time.

The conditional distribution over subset sizes enables adaptive cardinality inference:

$$\mathbb{P}_r(|S| = k \mid k_{\min} \leq |S| \leq k_{\max}) = \frac{\mathbb{P}_r(|S| = k)}{\sum_{k'=k_{\min}}^{k_{\max}}\mathbb{P}_r(|S| = k')} \tag{9}$$

Proposition 5.2: The range-constrained marginal probability satisfies:

$$m_j^* \triangleq \mathbb{P}_r(z_j = 1 \mid k_{\min} \leq |S| \leq k_{\max}) = \frac{\partial \log Z^*}{\partial \log p_j}$$

Empirical Validation / Results

Experimental Setup

• Models: OLMoE-1B-7B (64 experts/layer, top-8 activated) and Qwen1.5-MoE-A2.7B (60 routed + 4 shared experts/layer, top-4 activated)
• Baselines: Conventional, DenseMixer, Frozen Router, DefaultMoE, SparseMixer, ReMoE
• Tasks: GSM8K (math reasoning), Law (legal understanding), Translation, Summarization, MBPP (code generation), MMLU

Key Results (Table 1)

OLMoE Backbone (Top-k = 8):

• ProbMoE achieves best performance on GSM (50.19), Law (29.00), Translation (31.63), and Summary (39.29)
• Outperforms DenseMixer on most tasks while preserving sparse expert execution

Qwen Backbone (Top-k = 4):

• ProbMoE achieves best performance on Law (34.40), Translation (39.23), Summary (44.40), and MMLU Stem (53.82)
• Competitive on GSM (53.29) and MBPP (35.00)

Ablation Study

Four routing configurations were evaluated:

1. ProbMoE (Sample + Marginal): 50.24% EM, σ ≈ 0.09
2. Sample + Dense STE: 46.6% EM, σ ≈ 0.37
3. Top-k + Marginal and Top-k + Dense STE: intermediate performance

This confirms gains stem from aligning probabilistic selection with marginal-based optimization.

Routing Diversity Analysis

• ProbMoE requires more experts to reach 99% cumulative routing probability mass, indicating broader routing support
• On Qwen, ProbMoE shows lower Top-4 mass and higher normalized entropy, indicating more even expert assignment
• This broader distribution mitigates expert collapse and improves specialization

Dynamic-k Results (Table 2)

Dynamic-k achieves comparable performance with fewer activated experts, with Law consistently activating more experts than Translation, and GSM the fewest.

Token-Level Analysis

• Rare tokens receive more experts; frequent tokens receive fewer
• Higher-cardinality tokens include punctuation, morphological fragments, and context-sensitive symbols
• Lower-cardinality tokens are typically numerals or semantically concrete words

Training–Inference Mismatch (Table 3)

When evaluated under dynamic-k MAP inference, ProbMoE Dynamic-k achieves 44.50% EM with avg. k = 5.018, while DenseMixer (38.97%, k = 5.292) and Conventional (38.59%, k = 5.039) underperform despite similar cardinality, revealing that exact-k training produces peaked distributions that don't support flexible allocation.

Theoretical and Practical Implications

Theoretical Contributions

1. Principled routing formulation: Casting MoE routing as probabilistic inference over discrete subsets provides a mathematically grounded alternative to heuristic gradient approximations.

2. Tractable inference: The use of SIMPLE enables exact normalization and marginals for cardinality-constrained distributions, making probabilistic routing computationally feasible.

3. Unified framework: ProbMoE unifies exact-k and dynamic-k routing within a single probabilistic formulation, treating cardinality as a modeling choice rather than a fixed architectural constraint.

Practical Implications

1. Improved expert utilization: ProbMoE distributes routing probability more broadly, mitigating expert collapse and improving model capacity utilization.

2. Adaptive computation: Dynamic-k enables token-level compute allocation, reducing inference cost without substantial performance loss.

3. Architecture-agnostic: ProbMoE integrates with existing MoE architectures without modifying expert parameterization.

4. Preserved sparsity: Sparse expert execution is maintained at both training and inference, unlike DenseMixer which requires dense expert-side computation.

Conclusion

ProbMoE demonstrates that modeling MoE routing probabilistically provides a principled method for expert selection. Key takeaways:

• Exact-k routing improves expert utilization and routing diversity through probabilistic subset sampling and marginal-based gradients
• Dynamic-k routing achieves competitive performance with fewer activated experts, adapting naturally to token and task complexity
• The framework is tractable, architecture-agnostic, and preserves sparse execution

Future Directions

• Extending ProbMoE to pre-training scenarios
• System-level optimizations to fully exploit sparsity in the dynamic setting
• Further exploration of probabilistic routing for other discrete decision-making in large-scale models