Source-linked AI summary
Rethinking Routing Exploration in Modern Mixture-of-Experts
Bingcong Li, Hao Ma, Michael Muehlebach, Niao He, Yi Ma
TL;DR
Fine-grained expert segmentation enlarges the routing space, while shared experts can make routing less exploratory, limiting the realized expressiveness of modern MoEs. MirrorMoE adaptively suppresses routed experts to shrink that space, and it consistently improves performance across scales. On the 2.4B model, it improves pretraining and downstream performance by 0.67 and 1.64 points, respectively.
Problem
Fine-grained segmentation requires exploration of a combinatorially larger routing space, but shared experts can reduce variation in token representations and make routes harder to distinguish.
Source: §1 · p. 1, §2.2 · p. 4Method
MirrorMoE adaptively suppresses routed experts on a token basis, shrinking the exploration space to better match the model’s exploration capability.
Source: §1 · p. 2Results
MirrorMoE consistently improves performance across scales; on the 2.4B model, it improves pretraining and downstream performance by 0.67 and 1.64 points, respectively.
Source: §1 · p. 2, §5 · p. 10Takeaways & Limitations
MirrorMoE addresses the exploration dilemma by shrinking the routing space, and its adaptivity is theoretically characterized.
Source: §1 · p. 2, §5 · p. 10The current approach shrinks the exploration space; improving the routing strategy itself for more efficient exploration remains a future direction.
Source: §5 · p. 10
Abstract
from arXiv · showhide
Mixture-of-experts (MoE) models have become an important approach for scaling model parameters. Modern MoE architectures increasingly combine two design choices: fine-grained expert segmentation and shared experts. The former increases expressiveness by enabling more expert combinations at a fixed activation ratio, and the latter introduces experts that are active for every token. We identify an exploration dilemma between these designs: shared experts can make routing less exploratory by reducing the variability of token representations. This conflicts with the greater need for exploration induced by fine-grained expert segmentation due to the combinatorially enlarged routing space. Consequently, the additional expressiveness offered by fine-grained segmentation may not be fully realized in practice. MirrorMoE is proposed to reconcile this exploration dilemma by shrinking the exploration space through token-adaptive suppression of unnecessary experts. MirrorMoE consistently improves performance across model scales.
1 INTRODUCTION
Modern MoEs combine fine-grained expert segmentation, which enlarges the routing space, with shared experts, which can weaken routing exploration. MirrorMoE adaptively shrinks that space and improves performance across scales.
1 INTRODUCTION
- Fine-grained segmentation enlarges the set of expert combinations to explore, while shared experts can dominate token representations and make alternative routes harder to distinguish.For (E, K) = (64, 8), the passage reports about 4.4 × 10^9 expert subsets.Source: §1 · pp. 1–2
- On the 2.4B model, MirrorMoE improves pretraining and downstream performance by 0.67 and 1.64 points, respectively.Figure 1 compares MoEs with and without shared experts and shows MirrorMoE realizing gains from larger K.Source: Figure 1 · p. 2, §1 · p. 2
- With shared experts, increasing K can stop providing benefits and can even degrade performance, unlike the gains observed without shared experts.The authors attribute this pattern to inefficient exploration that hinders realization of the additional expressiveness from larger K.Source: §1 · p. 2
- MirrorMoE adaptively suppresses routed experts on a token basis, shrinking the exploration space to better match the model’s exploration capability.It is derived by mirroring routing logits via sign negation.Source: §1 · p. 2
1.1 Related work
- Prior work suggests that available model expressiveness may not be fully realized in practice; this work finds the same issue for fine-grained segmentation in shared-expert MoEs.The cited related-work discussion describes empirical benefits from explicitly incorporating operations that standard FFN experts could in principle learn.Source: §1.1 · pp. 2–3
2 THE EXPLORATION DILEMMA IN SHARED-EXPERT MOES
Fine-grained segmentation expands the routing space and raises the need for exploration, while shared experts can dominate token representations and weaken it. The section measures this mismatch and shows that shared experts can blunt the benefits of activating more experts.
2.1 MoE recap
- Fine-grained segmentation expands the set of expert compositions, but realizing that expressiveness requires more exploration; shared experts are active for every token alongside token-selected routed experts.The section frames these architectural choices as an exploration dilemma because shared experts can dominate learned representations and make alternative routes harder to distinguish.Source: §2.1 · p. 3
2.2 Shared experts weaken routing exploration
- Ablating shared experts causes substantially larger perplexity increases than ablating routed experts across early, middle, and final checkpoints.The two rightmost columns in Figure 2’s right subplot represent shared experts.Source: §2.2 · p. 4, Figure 2 · p. 4
- Replacing the learned router with uniform random routing produces a smaller perplexity increase in early and middle layers of shared-expert MoEs than in standard MoEs.Those shared experts capture enough of the learned representation for the layers to be more robust to random routing.Source: §2.2 · p. 4
- Expert-induced variance is smaller across all layers when shared experts are present, limiting how much changing routed experts can alter representations.Expert variance also bounds the relative change induced by routing policies from above under the stated normalization condition.Source: §2.2 · p. 4
- Shared experts reduce router-induced variance by approximately an order of magnitude, so sampled routing choices produce less diverse token representations.The paper defines greater RouterVar as more diverse, and thus more exploratory, sampled features.Source: §2.2 · p. 5, Figure 3 · p. 5
2.3 Consequences of the exploration dilemma
- With K increasing from 6 to 10, perplexity gains are {2.23, 1.32, 1.08, 0.87} without shared experts and {0.29, −0.03, 0.58, −0.03} with them, negative at two scales.Shared-expert benefits also diminish or disappear as K increases.Source: §2.3 · p. 5
- The relative perplexity gap under uniform random routing rises across training stages without shared experts, but changes little between the 20% and 50% checkpoints with them.The authors interpret the smaller increase as slower routing specialization in shared-expert MoEs.Source: §2.3 · pp. 5–6
3 MIRRORMOE FOR PRETRAINING SHARED-EXPERT MOES
MirrorMoE addresses the exploration dilemma in shared-expert MoEs by shrinking the routing space through token-adaptive suppression of experts. Its routing and aggregation mechanisms reduce effective exploration early in training, while suppression and preservation conditions allow the method to adapt to token-level routing confidence.
3.1 MirrorMoE shrinks the exploration space
- MirrorMoE selects experts using probabilities shifted toward both high- and low-logit experts, then aggregates them using standard routing probabilities so low-logit selections can contribute negligibly.This can reduce K_eff without simply lowering the number of selected experts.Source: §3.1 · p. 6
3.2 Visualization of routing space reduction in MirrorMoE
- During the first half of training, MirrorMoE has smaller K_eff than the baseline, while minimum normalized aggregation weights approach zero, indicating functional suppression of selected experts.Its smaller end-of-training K_eff is consistent with a more concentrated routing policy and stronger expert specialization.Source: §3.2 · p. 7
3.3 Discussions on adaptivity of MirrorMoE
- When an expert’s routing probability is sufficiently small and ρ is appropriately chosen, MirrorMoE can activate it while functionally suppressing its aggregation contribution.When routing uncertainty is high enough that ξ_max/ξ_min ≤ κ, no selected expert is functionally suppressed.Source: §3.3 · pp. 7–8
- MirrorMoE can be viewed as letting each token adaptively determine its number of activated experts, leaving remaining expert capacity available for other tokens.This interpretation follows from its ability to suppress or preserve experts according to token-level routing.Source: §3.3 · p. 8
4 NUMERICAL RESULTS
MirrorMoE improves pretraining and downstream performance across model scales, while retaining benefits as the activated-expert count grows. Additional experiments favor token/layer adaptivity and show little benefit without shared experts.
4.1 Pretraining performance of MirrorMoE
- The experiments use shared-expert MoEs from 100M to 2.4B parameters, with six routed and two shared experts activated per token.The evaluation uses the DCLM dataset and NVIDIA H100 GPUs.Source: §4.1 · p. 8, §4 · p. 8
- MirrorMoE lowers test perplexity across scales and improves the 2.4B model's six-task average performance by 1.64.Table 1 covers pretraining performance across scales; Table 2 reports downstream performance.Source: §4.1 · pp. 8–9, Table 1 · p. 9, Table 2 · p. 9
- MirrorMoE slightly reduces per-iteration training time despite adding a softmax operation, suggesting improved load balance.Source: §4.1 · p. 9
4.2 Comparison under enlarged routing space
- As K increases, MirrorMoE consistently outperforms non-shared MoE and its perplexity consistently decreases.The comparison tests K ∈ {6, 8, 10}.Source: §4.2 · p. 9
4.3 Additional experiments
- A cosine schedule with ρ0 = 0.4 performs best among the tested schedules and initial values.The comparison covers constant, linear, and cosine schedules, with ρ0 ∈ {0.1, 0.2, 0.4}.Source: §4.3 · p. 9, Figure 5 · p. 9
- Neither linear nor cosine nonadaptive schedules matches MirrorMoE.The alternatives use the same activated-expert count for all tokens and layers, unlike MirrorMoE's adaptivity.Source: §4.3 · pp. 9–10, Table 3 · p. 9
- Without shared experts, MirrorMoE performs comparably to vanilla training, consistent with a less pronounced exploration dilemma.Source: §4.3 · p. 10
5 CONCLUSION
The paper concludes that MirrorMoE addresses the exploration dilemma by adaptively shrinking the routing space, improving performance across scales. It identifies more efficient routing and compatibility with system-level optimizations as relevant context for future work.
5 CONCLUSION
- MirrorMoE adaptively shrinks the exploration space to address the tension between fine-grained segmentation and shared experts, and consistently improves performance across scales.Source: §5 · p. 10
- Improving the routing strategy itself for more efficient exploration is left as a future direction.Source: §5 · p. 10
A ADDITIONAL RELATED WORK
- MirrorMoE is largely orthogonal to MoE system optimizations and can be integrated with them.The cited system approaches target hardware utilization, communication efficiency, and throughput.Source: A ADDITIONAL RELATED WORK · p. 14
- Unlike approaches that encourage broader exploration, MirrorMoE adaptively suppresses experts; its focus is routing during pretraining.The related-work comparison distinguishes this from inference-time study of shared experts.Source: A ADDITIONAL RELATED WORK · p. 14
B.1 Training details
The appendix describes the training setup and examines routing variance, exploration strategies, layer-wise suppression, load balancing, and performance without shared experts. These experiments qualify where MirrorMoE's approach is useful and how it behaves.
B.1 Training details
- Training comparisons preserve approximately the same total training tokens, although GPU counts and global batch sizes can differ across no-shared-expert configurations.The appendix also specifies a Llama-style MoE, DCLM data, AdamW, and a warmup-stable-decay schedule.Source: §B.1 · p. 14, Table 4 · p. 15
B.2 Router Variance under Different Levels of Gumbel Noise
- Across tested Gumbel perturbation levels, shared experts consistently reduce routing-induced variance.Router variance is evaluated by stochastic sampling equivalent to adding Gumbel noise to routing logits.Source: §B.2 · p. 15, Figure 6 · p. 15
B.3 Increasing Exploration during Pretraining Can Be Tricky
- On the 100M model, deterministic routing performs best, followed by partly stochastic routing; fully stochastic routing performs worst.The results suggest added exploration instability can outweigh its potential benefits during pretraining.Source: §B.3 · p. 15, Table 5 · p. 16
B.4 Layer-wise effective experts
- With K = 8, effective-expert counts vary by layer: Layers 1–2 retain roughly three early, Layer 3 approaches one, and Layers 7–8 retain around four.This demonstrates layer-adaptive shrinking of the exploration space.Source: §B.4 · p. 16
B.5 On the Load-Balancing Loss
- Balancing original routing probabilities reduces perplexity but slightly increases iteration time, so the modified load-balancing loss is used.Source: §B.5 · p. 16, Table 6 · p. 16
B.6 MirrorMoE performance degrades for MoEs without shared experts
- Without shared experts, MirrorMoE performs comparably to standard training, with no significant improvement; these configurations are not directly comparable to Table 1.The appendix attributes the limited gain to a less pronounced exploration dilemma without shared experts.Source: §B.6 · p. 16
B.7 Comparison with top-𝑝routing
Top-p routing adapts the number of activated experts, but its changing compute makes direct comparison challenging. Even with more experts activated on average, it underperforms MirrorMoE in the reported training comparison.
B.7 Comparison with top-p routing
- Top-p routing activates the smallest expert set whose cumulative probability exceeds threshold p.The comparison reports activated experts averaged over layers and tokens.Source: §B.7 · p. 16, Table 8 · p. 18
- Top-p's average activated-expert count decreases during training, making direct comparison with fixed-activation methods challenging.Source: §B.7 · p. 16
- Even at p = 0.6, which activates more experts on average, top-p routing underperforms MirrorMoE.The top-p training comparison uses entropy regularization with coefficient 10^-4.Source: §B.7 · p. 16, Table 8 · p. 18
C MISSING PROOFS
The proof considers an arbitrary routing probability and invokes convexity to bound routing-induced changes by ExpertVar_l.
C MISSING PROOFS
- For any routing probability, the proof represents the routed representation as a convex combination of expert representations and invokes convexity of the norm to bound routing-induced changes.The passage ends before stating the resulting inequality.Source: §C.1 · p. 18
C.2 Proofs of Prop. 3.2
The proof uses the MirrorMoE score expression to show that the minimum-ξ expert is activated, then functionally suppressed as its ξ score approaches zero.
C.2 Proofs of Prop. 3.2
- MirrorMoE scores are written as p_e = (1 − ρ)ξ_e + ρζ_e, with ζ_e rewritten using the ξ scores.The proof uses the detailed expressions in equation (3).Source: §C.2 · p. 18
- The minimum-ξ expert is activated in MirrorMoE, as its score is at least the score of the lowest-ξ selected expert and of every non-activated expert.The proof uses the condition on ρ stated in the proposition.Source: §C.2 · p. 19
- Once activated, the minimum-ξ expert is functionally suppressed as its ξ score approaches zero.This follows from the aggregation in equation (5).Source: §C.2 · p. 19
C.3 Proof of Prop. 3.3
The proof bounds C, establishes matching expert rankings under ξ and p, and concludes that selected experts have non-vanishing aggregation weights.
C.3 Proof of Prop. 3.3
- C ≤ 1/R^2 follows from the Cauchy–Schwarz inequality.Source: §C.3 · p. 19
- Experts have the same ordering under ξ and p, so top-(K−S) aggregation selects the same set under both scores.Source: §C.3 · p. 19
- Every selected expert has a non-vanishing aggregation weight, ruling out functional path pruning.Source: §C.3 · p. 20