Rethinking Routing Exploration in Modern Mixture-of-Experts

Summary (Overview)

• Identifies a fundamental exploration dilemma in modern Mixture-of-Experts (MoE) architectures that combine fine-grained expert segmentation with shared experts: the combinatorially enlarged routing space demands more exploration, while shared experts simultaneously weaken routing exploration by dominating learned token representations.

• Provides empirical evidence across four model scales (100M to 2.4B parameters) showing that shared-expert MoEs suffer from three concrete inefficiencies: eroded benefits of shared experts as $K$ increases, diminished expressive gains from larger $K$ (even negative gains at some scales), and slower routing specialization.

• Proposes MirrorMoE, a conceptually simple pretraining approach that adaptively shrinks the exploration space through token-adaptive suppression of unnecessary experts via mirror routing (sign negation of routing logits) combined with mirror aggregation using standard routing probabilities.

• Achieves consistent improvements: On the 2.4B model, MirrorMoE improves pretraining perplexity by 0.67 points and downstream performance by 1.64 points on average across six commonsense reasoning tasks, while restoring the performance gains from increasing $K$.

• Theoretically characterizes adaptivity through sufficient conditions for both expert suppression (Proposition 3.2) and expert preservation (Proposition 3.3), establishing MirrorMoE's token-adaptive behavior.

Introduction and Theoretical Foundation

Background

Mixture-of-experts (MoE) models have become instrumental in scaling large language models (LLMs) to trillions of parameters. MoEs employ sparse conditional computation, where each token activates only a subset of experts. During pretraining, MoEs can outperform dense counterparts under matched total parameter counts, compute, and data budgets; at inference, sparse activation allows computational cost to grow sublinearly with total parameters.

Two Key Architectural Trends

Modern MoE architectures increasingly combine two design choices:

1. Fine-grained expert segmentation: Using larger numbers of total and activated experts (e.g., $(E, K) = (64, 8)$ preferred over $(8, 1)$ at the same activation ratio). This yields combinatorially richer routing spaces: the number of expert subsets increases from 8 when $(E, K) = (8, 1)$ to $\binom{64}{8} \approx 4.4 \times 10^9$ with $(64, 8)$.

2. Shared experts: Activated for every token, these capture common knowledge that would otherwise be learned redundantly by routed experts, allowing routed experts to devote more capacity to specialization.

The Exploration Dilemma

The paper identifies a fundamental tension between these two trends:

• Fine-grained segmentation demands more exploration: Realizing the expressive potential requires searching a combinatorially larger routing space within each MoE layer.

• Shared experts weaken exploration: Shared experts dominate learned token representations, leaving routed experts with less influence on layer output. Even activating different routed experts leads to limited variation in learned features, making alternative routes harder to distinguish.

MoE Formulation

For a token representation $\mathbf{h}_l^n \in \mathbb{R}^d$ entering the $l$-th MoE layer, the router produces a probability distribution over routed experts:

$$p_{l,e}(\mathbf{h}_l^n) = \frac{\exp(\mathbf{G}_l[e,:]\mathbf{h}_l^n)}{\sum_{j=1}^{R} \exp(\mathbf{G}_l[j,:]\mathbf{h}_l^n)}, \quad \mathcal{A}_l(\mathbf{h}_l^n) = \text{Top}_{K-S}\big(\{p_{l,e}(\mathbf{h}_l^n)\}, e \in [R]\big)$$

The learned token representation combines shared and routed expert outputs:

$$\mathbf{h}_{l+1}^n = \sum_{s=1}^{S} \mathcal{SE}_{l,s}(\mathbf{h}_l^n) + \sum_{e \in \mathcal{A}_l(\mathbf{h}_l^n)} \frac{p_{l,e}(\mathbf{h}_l^n)}{Z(\mathbf{h}_l^n)} \cdot \mathcal{RE}_{l,e}(\mathbf{h}_l^n)$$

where $Z(\mathbf{h}_l^n) = \sum_{j \in \mathcal{A}_l(\mathbf{h}_l^n)} p_{l,j}(\mathbf{h}_l^n)$ renormalizes weights (or $Z = 1$ as in DeepSeek-MoE and Qwen1.5-MoE).

Methodology

Diagnostic Analysis of the Exploration Dilemma

The paper employs three complementary metrics to diagnose the exploration dilemma:

1. Expert ablation: Zeroing out each expert's output and measuring perplexity change $\Delta \text{PPL} = \text{PPL}_{\text{after}} - \text{PPL}_{\text{base}}$. Shared experts show substantially larger perplexity increases when ablated, demonstrating their disproportionate importance.

2. Routing criticality: Replacing the learned router with uniform random routing and measuring relative perplexity increase $\frac{\text{PPL}_{\text{after}} - \text{PPL}_{\text{base}}}{\text{PPL}_{\text{base}}}$. Shared-expert MoEs show much smaller increases, indicating less dependence on learned routing.

3. Variance metrics:

• Expert variance: Measures how much the layer output varies when different experts are activated:

$$\text{ExpertVar}_l = \frac{1}{R} \sum_{e=1}^{R} \frac{\|\tilde{\mathbf{h}}_{l+1,e} - \mathbf{h}_{l+1}\|^2}{\|\mathbf{h}_{l+1}\|^2}$$

• Router variance: Measures output variability when routing decisions are sampled from the learned distribution:

$$\text{RouterVar}_l = \frac{1}{M} \sum_{m} \frac{\|\hat{\mathbf{h}}_{l+1}(m) - \bar{\mathbf{h}}_{l+1}\|^2}{\|\bar{\mathbf{h}}_{l+1}\|^2}$$

MirrorMoE Algorithm

Mirror routing: Activates experts according to a modified routing probability:

$$p_{l,e} = (1-\rho) \underbrace{\frac{\exp(\mathbf{G}_l[e,:]\mathbf{h}_l)}{\sum_{j=1}^{R} \exp(\mathbf{G}_l[j,:]\mathbf{h}_l)}}_{:=\xi_{l,e}} + \rho \underbrace{\frac{\exp(-\mathbf{G}_l[e,:]\mathbf{h}_l)}{\sum_{j=1}^{R} \exp(-\mathbf{G}_l[j,:]\mathbf{h}_l)}}_{:=\zeta_{l,e}}$$

The second term $\zeta_{l,e}$ is obtained by negating the sign of $\mathbf{G}_l$, favoring experts with smaller logits.

Mirror aggregation: Aggregates selected experts using standard routing probabilities $\xi_{l,e}$ instead of $p_{l,e}$:

$$\mathbf{h}_{l+1} = \sum_{s=1}^{S} \mathcal{SE}_{l,s}(\mathbf{h}_l) + \sum_{e \in \mathcal{A}_l} \frac{\xi_{l,e}}{Z(\mathbf{h}_l)} \mathcal{RE}_{l,e}(\mathbf{h}_l)$$

This ensures that even when low-logit experts are activated, they contribute minimally and are functionally suppressed.

Practical modifications:

1. Cosine annealing of $\rho$ to 0 over training (temporal exploration strategy)
2. Modified load-balancing loss:

$$\mathcal{L}_{\text{bal}} = R \sum_{l=1}^{L} \sum_{e=1}^{R} f_{l,e} \bar{q}_{l,e}, \quad \text{with} \quad f_{l,e} = \frac{c_{l,e}}{N(K-S)}, \quad \bar{q}_{l,e} = \frac{1}{N} \sum_{n=1}^{N} \xi_{l,e}(\mathbf{h}_l^n)$$

Effective Experts Definition

$$K_{\text{eff}}(\mathbf{h}_l^n) = \exp\Big(-\sum_{e \in \mathcal{A}_l(\mathbf{h}_l^n)} \bar{\xi}_{l,e}(\mathbf{h}_l^n) \log \bar{\xi}_{l,e}(\mathbf{h}_l^n)\Big), \quad \text{where} \quad \bar{\xi}_{l,e}(\mathbf{h}_l^n) = \frac{\xi_{l,e}(\mathbf{h}_l^n)}{\sum_{j \in \mathcal{A}_l(\mathbf{h}_l^n)} \xi_{l,j}(\mathbf{h}_l^n)}$$

This measures how many activated experts contribute substantially to the aggregated representation.

Empirical Validation / Results

Pretraining Performance

MirrorMoE consistently achieves lower perplexity across all model scales. Notably, MoE-498M-A121M has a larger routing space (more layers with sparse routing), and MirrorMoE yields a larger performance gain on this model, confirming that shrinking the routing space is particularly beneficial when the exploration dilemma is more pronounced.

Downstream Performance

For the 2.4B model, MirrorMoE improves average downstream performance by 1.64 points.

Key Findings Under Enlarged Routing Space

When varying $K \in \{6, 8, 10\}$:

• Standard shared-expert MoEs are outperformed by non-shared counterparts at large $K$
• Increasing $K$ from 6 to 10 yields perplexity gains of {2.23, 1.32, 1.08, 0.87} for MoEs without shared experts, but only {0.29, -0.03, 0.58, -0.03} with shared experts (negative at two scales)
• MirrorMoE enables shared-expert models to consistently outperform non-shared MoEs and restores performance gains from increasing $K$

Ablation Studies

Non-adaptive schedules (linear/cosine increase of activated experts) fail to match MirrorMoE's performance, demonstrating the importance of token/layer adaptivity.

Theoretical and Practical Implications

Theoretical Results

Proposition 3.2 (Sufficient condition for expert suppression): For a given token, let $\xi_{e_{\min}} = \epsilon$ (the smallest routing probability) and $\xi_j = c/R$ (the smallest probability among standard top-$K$ experts). As $\epsilon \to 0$, $\rho > c/(1+c)$ is sufficient for $e_{\min} \in \mathcal{A}^p \setminus \mathcal{A}^\xi$, meaning MirrorMoE activates this expert while functionally suppressing its contribution during aggregation.

Proposition 3.3 (Sufficient condition for expert preservation): Let $\kappa := \sqrt{(1-\rho)/\rho}$. If $\xi_{\max}/\xi_{\min} \leq \kappa$, then $\mathcal{A}^p = \mathcal{A}^\xi$ (the same experts are activated as standard routing). Moreover, $\bar{\xi}_e \geq 1/((K-S)\kappa)$ for all $e \in \mathcal{A}^p$, meaning no selected expert is functionally suppressed.

These results establish that MirrorMoE's suppression is token-adaptive: experts are suppressed only when the router is confident (large logit ratios), while uncertain routing regimes preserve all experts.

Practical Implications

1. Alternative interpretation: MirrorMoE can be viewed as tokens adaptively determining their effective number of activated experts—if a token can be learned well with fewer experts, the capacity of remaining experts can be devoted to other tokens.

2. Computational efficiency: MirrorMoE slightly reduces per-iteration training time despite adding a softmax operation, suggesting improved load balance.

3. Architecture guidance: The findings suggest that shared-expert MoEs should either limit $K$ or employ exploration-aware training strategies like MirrorMoE to realize the full benefits of fine-grained expert segmentation.

4. Generality: MirrorMoE performs comparably to vanilla training when applied to MoEs without shared experts, confirming that the approach specifically addresses the exploration dilemma induced by shared experts.

Conclusion

This work identifies a critical exploration dilemma in modern MoE architectures that combine fine-grained expert segmentation with shared experts: the enlarged routing space demands greater exploration, while shared experts simultaneously weaken routing exploration by dominating learned representations. This tension prevents the additional expressiveness enabled by fine-grained segmentation from being fully realized and diminishes the benefits of shared experts.

MirrorMoE addresses this dilemma by adaptively shrinking the exploration space to better match the model's exploration capability. Through mirror routing (sign negation of logits) and mirror aggregation (using standard probabilities for weighting), MirrorMoE functionally suppresses unnecessary experts in a token-adaptive manner. The approach consistently improves performance across model scales from 100M to 2.4B parameters, restores the benefits of shared experts, and realizes the expressive gains from larger $K$.

Future direction: The current approach addresses the exploration dilemma by shrinking the exploration space. A promising direction is to improve the routing strategy itself for more efficient exploration, potentially enabling even better utilization of fine-grained expert segmentation in shared-expert MoEs.