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

ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts
Jun 1, 2026Heng Zhao, Zilei Shao, Guy Van den Broeck +1Large Language Model Routing Probabilistic Inference
Department of Computer Science, University of Virginia, Charlottesville, USA · Department of Computer Science, University of California, Los Angeles, Los Angeles, USA
Distribution-Consistent Inference for Dynamic Sparse Mixture-of-Experts
Sep 10, 2026Dohyeon Kim, Bedionita Soro, Sung Ju HwangMixture-Of-Expert Inference Mixture-Of-Experts
KAIST
Hierarchical Mixture-of-Experts with Two-Stage Optimization
May 8, 2026Gleb Molodtsov, Alexander Miasnikov, Aleksandr BeznosikovMixture-Of-Experts Load Balancing
BRAIn Lab, Moscow, Russia

Jun 1, 2026, cs.LGJ/K move · Enter open · S save

ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

Heng Zhao, Zilei Shao, Guy Van den Broeck, Zhe Zeng

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

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

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

Hierarchical Mixture-of-Experts with Two-Stage Optimization