Function Over Form: Distributional Orthogonalization in Mixture-of-Experts with Replica Expert Mechanism

Jinfan He; Yunzhuo Liu; Kai Zhang; Weidong Han; Key; Rayying

cs.AISep 26, 2026

Function Over Form: Distributional Orthogonalization in Mixture-of-Experts with Replica Expert Mechanism

Authors: Jinfan He, Yunzhuo Liu, Kai Zhang, Weidong Han, Key, Rayying

Organizations: School of EECS, Peking University, China · Tencent Hunyuan, China

Abstract

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 $\lambda$ . Validation perplexity as a function of $\lambda$ (4.8BA0.5B, 100B tokens). $\lambda{=}0$ disables DO-loss, leading to expert collapse via redundancy. Excessively large $\lambda$ causes the high-dimensional pairwise regularization to overwhelm the cross-entropy signal. The optimum $\lambda{=}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 ( $\downarrow$ )
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.

Panel of Figure 4 — 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{\times}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.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Require: Historical rank loads $\{L_{r}\}$ , expert token volumes $\{T_{e}\}$ , per-rank capacity $D$
Ensure: Replica assignment $\mathcal{A}\leftarrow\emptyset$
$\tau\leftarrow\frac{1}{\|\mathcal{R}\|}\sum_{r}L_{r}$
repeat
	$r_{r}\leftarrow\arg\min_{\{r:\|\mathcal{A}_{r}\|<D\}}L_{r}$
	$(e^{},r_{s}^{})\leftarrow\arg\max_{e,\,r_{s}\neq r_{r}}S(e,r_{s},r_{r},\tau)$
	$\Delta\leftarrow S(e^{},r_{s}^{},r_{r},\tau)$

Appendix

Algorithm 1 REM Global-Batch Replica Allocation

Require: Replica assignment $\mathcal{A}$ , real-time expert token counts $\{T_{e}\}$ , per-rank capacity $D$
Ensure: Token-to-rank mapping $\mathcal{M}\leftarrow\emptyset$
$T_{e}\leftarrow\text{AllGather}(\{T_{e}\})$
$Q_{r}\leftarrow\sum_{e\in r}T_{e}$
$\tau\leftarrow\frac{1}{\|\mathcal{R}\|}\sum_{r}Q_{r}$
for $d=1$ to $D$ do
	sort ranks by $Q_{r}$ 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 $\lambda$ at finer resolution. Sweep on a 2B/0.2B model (20B tokens). The optimum $\lambda{=}5{\times}10^{-6}$ is roughly two orders of magnitude smaller than typical load-balancing loss coefficients.

Configuration	Micro-batch	GPUs	grad_accum	Eff. Batch ( $G_{ij}$ )	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

G_{ij}

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.

Explore similar work

Distribution-Consistent Inference for Dynamic Sparse Mixture-of-Experts
Sep 8, 2026Dohyeon Kim, Bedionita Soro, Sung Ju HwangParameter-Free Mixture-of-Experts Routing Sparse Mixture-of-Experts
KAIST
Distributionally Robust Mixture-of-Experts Training
Oct 5, 2026Xin Teng, Muxiao Li, Hongyi WenMixture-of-Experts Language Models Sparse Mixture-of-Experts
New York University Center for Data Science, NYU Shanghai
Scaling Multi-Node Mixture-of-Experts Inference Using Expert Activation Patterns
Apr 25, 2026Abhimanyu Bambhaniya, Geonhwa Jeong, Jason Park +6Expert Load Balancing Mixture-of-Experts Inference
1Georgia Institute of Technology, USA · 2Meta Platforms, Inc.

Sep 8, 2026J/K move · Enter open · S save

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

Dohyeon Kim, Bedionita Soro, Sung Ju Hwang

KAIST

Parameter-Free Mixture-of-Experts Routing Sparse Mixture-of-Experts Mixture-of-Experts Inference LLM Inference Acceleration Mixture-of-Experts Models

Open paper AskPDF arXiv

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.

Require: Historical rank loads $\{L_{r}\}$ , expert token volumes $\{T_{e}\}$ , per-rank capacity $D$
Ensure: Replica assignment $\mathcal{A}\leftarrow\emptyset$
$\tau\leftarrow\frac{1}{\|\mathcal{R}\|}\sum_{r}L_{r}$
repeat
	$r_{r}\leftarrow\arg\min_{\{r:\|\mathcal{A}_{r}\|<D\}}L_{r}$
	$(e^{},r_{s}^{})\leftarrow\arg\max_{e,\,r_{s}\neq r_{r}}S(e,r_{s},r_{r},\tau)$
	$\Delta\leftarrow S(e^{},r_{s}^{},r_{r},\tau)$

Abstract

Explore similar work

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

Distributionally Robust Mixture-of-Experts Training

Scaling Multi-Node Mixture-of-Experts Inference Using Expert Activation Patterns