The scaling of LLMs increasingly relies on MoE architectures to decouple active computation from total parameter count. However, the efficacy of MoE is often constrained by expert collapse and representation redundancy, both leading to underutilization of model capacity. To address these challenges, this paper proposes Distributional Orthogonalization Loss (DO-loss), an auxiliary regularization that shifts the focus from static weight diversity to dynamic routing behavior. By representing each expert's token assignment history as a high-dimensional binary load signature, DO-loss penalizes signature overlap to prevent expert collapse while encouraging functional specialization. To align this algorithmic design with system efficiency, we further introduce the Replica Expert Mechanism (REM), which improves load balancing through a two-tiered strategy: adjusting replica expert placement at the global-batch level and performing real-time token dispatching at the micro-batch level. Empirical evaluations demonstrate that our method outperforms the evaluated routing algorithms on downstream tasks for both 4.8BA0.5B and 30BA3B MoE models, while maintaining comparable training efficiency.
Figures & tables
Figure 1: Token assignment histories encoded as binary load signatures. While the total sum of the pairwise Hamming distances remains constant under top- K routing, the DO-loss drives the distances between all pairs of experts toward equidistance to penalize high routing overlap.
Figure 2: The architecture of REM. Left: The global-batch controller (top) allocates replicas once per global batch; a collector asynchronously retrieves metrics from the router. The micro-batch dispatcher (bottom) synchronizes router information and routes tokens in real time. Right: Load statistics are gathered hierarchically—the coordinator (R0) of each EP group aggregates intra-group rank loads locally, and R0 nodes subsequently exchange summaries across EP groups to construct a global view. This two-level aggregation avoids all-to-all collection across all ranks while enabling decentralized load-balancing decisions at every rank in real time.
Figure 3: Ablation of DO-loss coefficient λ . Validation perplexity as a function of λ (4.8BA0.5B, 100B tokens). λ=0 disables DO-loss, leading to expert collapse via redundancy. Excessively large λ causes the high-dimensional pairwise regularization to overwhelm the cross-entropy signal. The optimum λ=10−5 is approximately two orders of magnitude smaller than typical load-balancing loss coefficients.
Domain
Description
GitHub code
Multi-language code repositories
STEM data
Academic papers and textbooks
Chinese/English web
General web crawl
Academic books
Multi-discipline textbooks
Encyclopedias
Encyclopedia-style knowledge
Q&A banks
Exam and quiz datasets
Table 1: Domain composition of the pretraining corpus.
Model
Method
AGIEval
MMLU
MMLU-PRO
CEval
GSM8K
TriviaQA
PPL ( ↓ )
4.8BA0.5B
Local-LBL
22.24
36.01
10.95
36.63
12.71
23.56
8.4650
Global-LBL
23.02
36.80
11.96
37.24
13.19
23.75
8.4604
Orth-loss
22.60
34.44
11.42
34.63
12.25
23.47
8.4664
Auxfree
23.55
36.85
12.02
38.63
13.57
23.06
8.4620
DO-loss
24.05
36.33
13.50
37.10
13.51
23.89
8.4518
Global-DO-loss
24.13
37.18
13.88
39.64
14.55
24.03
8.4433
Table 2: Validation perplexity and benchmark accuracy at two scales (100B and 200B tokens). At 30BA3B, we additionally evaluate Global-LBL for completeness; Global-LBL underperforms Auxfree on most benchmarks, particularly on knowledge-intensive tasks such as TriviaQA, indicating that load-balancing objectives can conflict with knowledge-intensive tasks at larger scales.
Figure 4: Expert behavior in the trained 4.8BA0.5B model. (a) Top- K routing score sum per token: higher values indicate the router assigns tokens with greater confidence, reflecting stronger expert specialization. (b) CKA similarity matrix (Appendix D ) at layer 4, sampled every 8th expert ( 16×16 grids). CKA measures functional similarity between expert outputs on the same inputs; high off-diagonal values indicate representational redundancy, while a near-identity matrix indicates functional diversity. (c) Distribution of token counts per expert at layer 4.
Figure 5: Ablation of REM replica count. GPU throughput and MaxVio over 1000 steps on H20 GPUs (EP =8 , 4.8BA0.5B) with 0, 1, and 2 replica experts. A single replica recovers most of the throughput gap; two replicas drive MaxVio to near zero. MaxVio (Appendix E ) measures the fractional load excess of the busiest rank relative to the ideal average.
Require: Replica assignment A , real-time expert token counts {Te} , per-rank capacity D
Ensure: Token-to-rank mapping M←∅
Te←AllGather({Te})
Qr←∑e∈rTe
τ←∣R∣1∑rQr
for d=1 to D do
sort ranks by Qr ascending
Appendix
Algorithm 2 REM Micro-Batch Token Dispatching
Method
Layer 4 CKA mean (std)
Layer 8 CKA mean (std)
Auxfree
0.263 (0.031)
0.115 (0.023)
Global-LBL
0.089 (0.016)
0.077 (0.008)
Orth-loss
0.084 (0.010)
0.084 (0.011)
Global-DO-loss
0.086 (0.010)
0.071 (0.007)
Appendix
Table 3: Mean and standard deviation of CKA cosine similarity across all 128 expert pairs. Lower values indicate less representational redundancy. Best results in bold .
Figure 6: Ablation of load imbalance without REM. MaxVio over the course of 100B-token training for DO-loss alone (no replica scheduling). MaxVio stabilizes below 1.0, indicating that even without REM the busiest rank processes at most twice the average load, though it remains substantially higher than with replica scheduling enabled.
Figure 7: Ablation of DO-loss coefficient λ at finer resolution. Sweep on a 2B/0.2B model (20B tokens). The optimum λ=5×10−6 is roughly two orders of magnitude smaller than typical load-balancing loss coefficients.
Configuration
Micro-batch
GPUs
grad_accum
Eff. Batch ( Gij )
Val. PPL
DO-loss (local)
4
–
–
4 (no sync)
8.4518
DO-32
4
8
1
32
8.4475
DO-64
4
8
2
64
8.4458
DO-128
4
8
4
128
8.4446
Global-DO-loss
4
8
16
512
8.4433
Appendix
Table 4: Effect of the effective batch size for Gij statistics on validation perplexity (4.8BA0.5B, 100B tokens).
Figure 8: Ablation of expert granularity. Validation perplexity vs. expert count (2B/0.2B, 20B tokens). Global-DO-loss trails Auxfree at 32 experts but surpasses it as granularity increases to 64 and 128, demonstrating that DO-loss benefits more from finer expert specialization.
H20
H800
Method
DP2 EP8
DP1 EP16
DP2 EP8
DP1 EP16
Auxfree
81.2
70.1
161.3
101.8
DO-only
75.3
62.8
152.7
93.1
DO+REM
79.5
69.7
159.6
101.3
DO+EfficientMoE
77.8
67.5
155.2
95.7
DO+Prophet
76.8
65.9
152.3
92.2
Appendix
Table 5: Ablation of expert scheduling method. Steady-state throughput (TFLOPs) averaged over 1,000 steps. All methods use two replica experts per rank on the DO-loss-trained 4.8BA0.5B model.
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.
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/.
Xin Teng, Muxiao Li, Hongyi Wen
New York University Center for Data Science, NYU Shanghai
Most recent state-of-the-art (SOTA) large language models (LLMs) use Mixture-of-Experts (MoE) architectures to scale model capacity without proportional per-token compute, enabling higher-quality outputs at manageable serving costs. However, MoE inference at scale is fundamentally bottlenecked by expert load imbalance and inefficient token routing, especially in multi-node deployments where tokens are not guaranteed to be routed to local experts, resulting in significant inter-node all-to-all communication overhead. To systematically characterize these challenges, we profile SOTA open-source MoE models, including Llama 4 Maverick, DeepSeek V3-671B, and Qwen3-230B-A22B, on various datasets and collected over 100k real expert activation traces. Upon studying the expert activation patterns, we uncover various persistent properties across all the frontier MoE models: variable expert load imbalance, domain-specific expert activation where expert popularity shifts across task families (code, math, chat, general), and a strong correlation between prefill and decode expert activations. Motivated by these findings, we propose workload-aware micro-batch grouping and an expert placement strategy to maximize token locality to the destination expert, thereby reducing inter-node communication. Across models and datasets, these optimizations help reduce all2all communication data up to 20, resulting in lower MoE decode latency and better accelerator utilization.
Abhimanyu Bambhaniya, Geonhwa Jeong, Jason Park +6
1Georgia Institute of Technology, USA · 2Meta Platforms, Inc.