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

> ProbMoE replaces deterministic top-k expert selection with exact probabilistic routing, improving expert utilization and enabling adaptive computation without sacrificing sparse execution.

- **Source:** [arXiv](https://arxiv.org/abs/2606.01509)
- **Published:** 2026-09-19
- **Permalink:** https://picx.dev/p/LHBUiK
- **Whiteboard:** https://picx.dev/p/LHBUiK/image

## Summary

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

| Dataset | OLMoE | Qwen1.5 |
|---------|-------|---------|
| GSM | -1.82 (80.00%) | -4.29 (75.00%) |
| Law | -0.04 (84.50%) | +2.70 (75.00%) |
| Translation | +0.36 (82.00%) | +3.22 (75.00%) |

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

---

_Markdown view of https://picx.dev/p/LHBUiK, served by PicX — AI-generated visual whiteboard summaries of research papers._