Distributionally Robust Mixture-of-Experts Training

Xin Teng; Muxiao Li; Hongyi Wen

cs.LGOct 5, 2026

Distributionally Robust Mixture-of-Experts Training

Authors: Xin Teng, Muxiao Li, Hongyi Wen

Organizations: New York University Center for Data Science, NYU Shanghai

Abstract

Mixture-of-Experts (MoE) transformers scale capacity by activating only a few experts per token, but this sparsity creates a hidden reliability problem: when routing is imperfect, load-balanced models may send tokens to experts that are insufficiently trained for the assigned inputs. We propose Distributionally Robust MoE Training (DRMoET), a drop-in objective that treats layer-wise experts as endogenous robustness groups and optimizes high-loss routing outcomes rather than merely equalizing traffic. DRMoET updates a per-layer expert distribution by an entropy-regularized softmax rule on EMA-smoothed, activation-weighted expert losses, strengthening plausible non-top routing paths while preserving standard MoE computation. Under the FLAME-MoE recipe at 746M-total and 10.3B-total scales, DRMoET improves downstream averages over both standard FLAME-MoE and auxiliary-loss-free balancing. At 10.3B total parameters and 67B training tokens, DRMoET improves the seven-task average from 0.6625 to 0.6767, while the auxiliary-loss-free baseline achieves 0.6431. Mechanistic analyses show lower expert-loss variance with nearly unchanged mean loss, 4.3% lower excess loss under forced mid- $k$ misrouting, and improved domain-expert specialization. These results position routing robustness-not only utilization balance-as a practical objective for reliable sparse MoE scaling. Project page and code are available at: https://drmoet.github.io/.

Figures & tables

Figure 1 : Distributionally Robust MoE Training workflow. DRMoET treats experts as layer-wise robustness groups and adaptively upweights high-loss routing outcomes, strengthening plausible non-optimal paths rather than only encouraging uniform utilization.

Scale	Method	ARC-C	ARC-E	HellaSwag	PIQA	WinoGrande	SciQ	ReCoRD*	Average
290M-746M		Training tokens: 33.5B
	FLAME-MoE	0.2329	0.5295	0.3447	0.6697	0.5012	0.8020	0.7023	0.5403
	Aux-free ( $u{=}10^{-3}$ )	0.2261	0.4878	0.3283	0.6670	0.4949	0.7100	0.6445	0.5084
	Aux-free ( $u{=}10^{-2}$ )	0.2184	0.4895	0.3301	0.6627	0.4949	0.6990	0.6440	0.5055
	DRMoET ( $\eta{=}0.1$ )	0.2338	0.5522	0.3472	0.6844	0.5209	0.7950	0.7115	0.5493
	DRMoET ( $\eta{=}0.01$ )	0.2270	0.5400	0.3470	0.6839	0.4941	0.8180	0.7007	0.5444

Table 1 : Main results at two scales. The two blocks use approximately 33.5B and 67B training tokens, respectively. All DRMoET rows use

\beta=0.999

. Averages are over the seven displayed tasks. Best per-task and average within each block are in bold . *ReCoRD is reported as F1.

Metric	Baseline	DRMoET	$\Delta$ (%)
Worst loss	2.6930	2.6379	$-2.05$
Best loss	2.1149	2.0722	$-2.02$
Mean loss	2.3712	2.3699	$-0.05$
Range	0.5781	0.5656	$-2.16$
Std.	0.1516	0.1386	$-8.58$
CV	0.0639	0.0585	$-8.45$

Table 2: Expert-loss statistics at convergence.

Model	Normal	Misrouting	Deg.
Baseline	3.187	7.249	$2.27\times$
DRMoET	3.182	7.071	$2.22\times$

Table 3: Forced mid-

k

misrouting.

Tier	FLAME-MoE	DRMoET	Improvement(%)
Top	3.176	3.188	$-0.4\%$
Mid	3.100	3.077	$+0.7\%$
Bottom	3.072	3.058	$+0.4\%$

Table 4 : Expert tier and OOD analysis. Mid-tier experts improve the most, and under distribution shift DRMoET yields more uniform expert quality while keeping mean loss nearly unchanged.

Metric	Improvement (%)
Avg Max Selectivity	$+1.0\%$
Avg $\Delta S$ (specialization margin)	$+1.2\%$
Avg $\Delta Q$ (competence advantage)	$+30.2\%$
Avg MI $I(E;G)$	$+12.9\%$

Table 5: Domain–expert specialization metrics. DRMoET improves all measures.

Ablation	Setting	ARC-E	HellaSwag	PIQA	SciQ	ReCoRD*	Average
Credit	FLAME-MoE	0.5295	0.3447	0.6697	0.8020	0.7023	0.6096
	Activation weighted, $\eta=10^{-2}$	0.5400	0.3470	0.6839	0.8180	0.7007	0.6179
	Activation weighted, $\eta=10^{-3}$	0.5682	0.3452	0.6888	0.8190	0.7105	0.6263
	Raw probability, $\eta=10^{-2}$	0.5455	0.3384	0.6839	0.7920	0.6938	0.6107
	Raw probability, $\eta=10^{-3}$	0.5421	0.3435	0.6861	0.7910	0.7065	0.6138
Top-12 routing	FLAME-MoE	0.5459	0.3506	0.6991	0.7950	0.7033	0.6188

Table 6 : Ablation studies. We ablate credit assignment, routing sparsity, and EMA smoothing. Averages are over the five displayed tasks within each ablation block. *ReCoRD is reported as F1.

Metric	Baseline	DRMoET	$\Delta$
TFLOP/s/GPU	$64.17{\pm}1.81$	$67.19{\pm}0.93$	$+4.7\%$
Time/iter (ms)	$6460{\pm}185$	$6166{\pm}87$	$-4.6\%$
Throughput std	$1.81$	$0.93$	$2{\times}$

Table 7: Training throughput comparison. Mean

\pm

std across measured iterations. DRMoET introduces no slowdown and exhibits lower variance.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 2 : Per-layer expert utilization variability over training. DRMoET stabilizes routing dynamics in early layers and avoids the transient CV spikes observed in the baseline.

Step	Baseline CV	DRMoET CV	Diff
540	2.75%	3.32%	+0.57%
2160	2.46%	3.04%	+0.58%
3780	1.78%	2.73%	+0.95%
5400	2.00%	2.64%	+0.64%
7020	6.47%	2.21%	-4.26%
8640	1.60%	2.07%	+0.46%

Appendix

Table 8: Expert utilization CV over training (six routed experts plus two shared experts). Higher CV indicates greater relative variation in load across experts. DRMoET shows slightly higher average CV but avoids the baseline’s transient mid-training imbalance spike (step 7k), consistent with controlled exploration early and modest specialization later.

Figure 3 : Expert co-activation heatmaps (layer-wise) comparing baseline vs. DRMoET. DRMoET produces more uniform co-activation in early/mid layers (L2–L5) and slightly more concentrated patterns in later layers (L6–L9).

Layer	Variance			Gini			Top-3 Conc.
Layer	Baseline	DRMoET	$\Delta$	Baseline	DRMoET	$\Delta$	Baseline	DRMoET	$\Delta$
2	0.00518	0.00483	-6.6%	0.248	0.238	-4.2%	17.7%	17.5%	-1.2%
3	0.00800	0.00678	-15.3%	0.301	0.275	-8.4%	19.7%	19.3%	-2.4%
4	0.01534	0.01395	-9.0%	0.413	0.393	-4.9%	24.3%	23.8%	-2.4%
5	0.02487	0.02188	-12.0%	0.500	0.475	-5.0%	30.7%	28.7%	-6.4%
6	0.01477	0.01549	+4.9%	0.422	0.427	+1.1%	24.0%	25.1%	+4.3%

Appendix

Table 9 : Coactivation concentration metrics (DRMoET vs. Baseline). Negative

\Delta

= more uniform partnerships under DRMoET; positive

\Delta

= increased specialization.

Method	ARC-C	ARC-E	HellaSwag	PIQA	WinoGrande	SciQ	ReCoRD*	Average
FLAME-MoE (with balancing)	0.2329	0.5295	0.3447	0.6697	0.5012	0.8020	0.7023	0.5403
DRMoET without balancing	0.2389	0.5539	0.3430	0.6877	0.5067	0.7770	0.7066	0.5448
DRMoET with balancing	0.2295	0.5682	0.3452	0.6888	0.5051	0.8190	0.7105	0.5523

Appendix

Table 10 : Load-balancing interaction at 290M-active/746M-total scale. Averages are over seven tasks. Standard FLAME-MoE includes balancing; a FLAME-MoE-without-balancing control is not reported. *ReCoRD is F1.

Explore similar work

Jun 1, 2026cs.LG

ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

Mixture-of-Experts (MoE) models scale by activating only a small subset of experts per token. However, training such models remains challenging because top-

k

routing is discrete and non-differentiable, requiring gradient estimators for expert selection whose design remains a central open problem. We introduce ProbMoE, a probabilistic routing framework that models expert selection as a distribution over cardinality-constrained expert subsets and formulates routing as probabilistic inference in this discrete subset space. We first propose ProbMoE Exact-

k

routing, which samples

k

-expert subsets in the forward pass, and the backward pass uses gradients through each expert's exact marginal probability as a tractable surrogate for the true gradient. ProbMoE naturally generalizes to a dynamic-

k

routing setting, where both training and inference constrain the routing cardinality to the same predefined range, allowing adaptive expert allocation per token. Across benchmarks and model backbones, ProbMoE Exact-

k

achieves strong performance compared to competitive baselines, with improved expert utilization and routing diversity; ProbMoE Dynamic-

k

achieves comparable performance with fewer activated experts.

Heng Zhao, Zilei Shao, Guy Van den Broeck +1

Department of Computer Science, University of Virginia, Charlottesville, USA · Department of Computer Science, University of California, Los Angeles, Los Angeles, USA

Large Language Model Routing Probabilistic Inference Cardinality

Ask PDF

Sep 10, 2026cs.LG

Distribution-Consistent Inference for Dynamic Sparse Mixture-of-Experts

Mixture-of-Experts (MoE) architectures have emerged as a powerful paradigm for scaling model capacity while preserving efficient inference in large foundation models. However, most MoE models use a fixed top-

k

expert selection policy, assigning the same expert budget to every token even when fewer experts may be sufficient. Inference-time dynamic top-

k

routing can reduce computation without retraining, but existing methods often overlook the distributional shift caused by deviating from the training-time routing configuration. We show that reducing the number of activated experts consistently increases the RMS scale and variance of SMoE outputs, inducing a representation mismatch that contributes to downstream performance degradation in addition to the loss of expert capacity. To address this correctable component, we propose Layer-wise Distribution Alignment (LDA), a lightweight inference-time correction that uses layer-wise calibration statistics to align reduced-routing representations with the default configuration. Across multiple SMoE LLMs, benchmarks, and routing strategies, LDA recovers much of the performance lost induced by the distributional shift under reduced routing while preserving sparse-inference efficiency with negligible overhead.

Dohyeon Kim, Bedionita Soro, Sung Ju Hwang

KAIST

Mixture-Of-Expert Inference Mixture-Of-Experts Distributional Alignment Wireless Foundation Models

Ask PDF

May 8, 2026cs.LG

Hierarchical Mixture-of-Experts with Two-Stage Optimization

Sparse Mixture-of-Experts (MoE) models scale capacity by routing each token to a small subset of experts. However, their routers exhibit a fundamental trade-off: strong load balancing can suppress expert specialization, while aggressive diversity often causes routing collapse. We propose Hi-MoE, a grouped MoE framework that decomposes routing control into two coupled levels: (i) inter-group balancing that enforces fair traffic across expert groups, and (ii) intra-group specialization that promotes complementary expert behaviors while preventing within-group collapse. Our analysis provides a principled explanation of how our hierarchical objectives reshape the router, thereby promoting stable specialization and mitigating collapse. We observe consistent improvements over recent sparse-routing and grouped-MoE baselines across NLP and vision benchmarks, and confirm robustness via scaling studies (model size, expert count) and targeted ablations. In large-scale pre-training on 58B tokens, Hi-MoE-7B achieves a 5.6% perplexity reduction and a 40% improvement in expert balance over OLMoE-7B across diverse evaluation domains.

Gleb Molodtsov, Alexander Miasnikov, Aleksandr Beznosikov

BRAIn Lab, Moscow, Russia

Mixture-Of-Experts Load Balancing Pretraining

Ask PDF