PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2609.32398v1cs.AI

TL;DR

MoE models face expert collapse and redundant representations, while existing approaches address load balance and diversity separately with trade-offs. The paper proposes DO-loss on routing histories and REM for system-level load balancing. Across two model scales, the design outperforms evaluated routing algorithms on downstream tasks while maintaining comparable training throughput.

  • Problem

    Expert collapse and redundant representations constrain MoE effectiveness, while load-balancing and orthogonality approaches address these challenges separately and carry trade-offs.

    Source: §1 · pp. 1–2

  • Method

    DO-loss penalizes overlap in experts’ binary token-assignment histories, while REM rebalances workloads through replica placement and real-time token dispatching.

    Source: §1 · p. 2, §6 · p. 10

  • Results

    On 4.8BA0.5B and 30BA3B models, the design outperforms evaluated routing algorithms on downstream tasks while maintaining comparable training throughput.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    The paper separates functional data partitioning from system load balancing, using distinct mechanisms for routing objectives and workload imbalance.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

MoE specialization is limited by expert collapse and redundant representations, while existing balancing and orthogonality methods address these problems separately and introduce trade-offs. The paper proposes DO-loss for routing-history diversity and REM for system-level load balancing; evaluations report improved downstream results across two model scales.

1 Introduction

  • Load-balancing losses can suppress specialization, while parameter-space orthogonality does not ensure experts receive tokens; combining approaches can destabilize training, and output-space methods add system overhead.The introduction frames DO-loss as addressing collapse and specialization through routing statistics rather than model weights.Source: §1 · pp. 1–2
  • DO-loss penalizes overlap between binary token-assignment histories, using pairwise Hamming distances to encourage specialization and make expert collapse self-correcting.It operates on routing statistics rather than model weights and is described as having negligible computational overhead.Source: §1 · p. 2, Figure 1 · p. 2
  • REM places replicas of hot experts on underloaded ranks, then routes tokens to the shortest available queue to handle residual load fluctuations without constraining specialization.This addresses stragglers and memory-footprint disparities associated with overloaded ranks.Source: §1 · p. 2
  • Across six downstream benchmarks, Global-DO-loss outperforms all baselines at both 4.8BA0.5B and 30BA3B, while REM recovers specialization-related throughput loss without degrading model quality.The reported results also say REM reduces load imbalance below traditional load-balancing methods.Source: §1 · p. 2

2 Related Work

Related work covers token routing and auxiliary balancing, orthogonality-based diversity methods, and system-level expert scheduling. It contrasts their mechanisms and trade-offs with the paper's focus on routing and scheduling.

2.1 Routing and Auxiliary Load Balancing Loss

  • Token-choice routing lets each token select Top-K experts but can collapse; balancing losses and bias updates address imbalance, while expert-choice routing guarantees balance but can introduce causal leakage.Per-device auxiliary losses can hinder specialization; global synchronization is proposed to preserve semantic clustering (Qiu et al. (2025)).Source: §2.1 · p. 3

2.2 Orthogonality in MoE

  • Orthogonality methods reduce representation redundancy through parameter- or output-similarity penalties, but their overhead can restrict use to post-training or fine-tuning.MOORE uses Gram-Schmidt regularization (Hendawy et al. (2024)); AES-MoE regularizes output similarity (Wang et al. (2025)).Source: §2.2 · p. 3

2.3 System Efficiency and Expert Scheduling

  • Expert-scheduling methods differ in granularity: EfficientMoE places replicas per global batch, whereas Pro-Prophet re-dispatches at micro-batch scale with time and overlap trade-offs.Pro-Prophet's overlap of scheduling with computation depends on the compute-to-communication ratio (Wang et al. (2026)).Source: §2.3 · p. 3

3 Methods

DO-loss uses expert token-assignment histories to regularize routing diversity, while REM redistributes hot-expert load through global-batch replica placement and micro-batch dispatch.

3.1.1 Binary Load and Zero-Sum Constraint

  • With Top-K routing, the batch-wide sum of pairwise Hamming distances is fixed, and minimizing their variance drives pairs toward equal distances while penalizing idle experts.The resulting gradient pushes the router to reactivate idle experts.Source: §3.1.1 · p. 3

3.1.2 Equalization Loss and Gradient Derivation

  • Straight-Through Estimation replaces discrete load signatures with routing probabilities plus a stop-gradient correction to enable gradient flow.The gradient derivation uses binary assignments to simplify the sign term.Source: §3.1.2 · p. 4

3.1.3 Surrogate Loss Implementation

  • The surrogate loss separates a detached magnitude from a gradient-direction term and applies a normalization scalar to the final loss.The normalization scalar is 1/m^2k, as stated in the passage.Source: §3.1.3 · p. 4

3.1.4 Global-Level Regularization via Cross-DP Synchronization

  • Global-DO-loss synchronizes the distance-matrix magnitude across data-parallel ranks while keeping the gradient-direction indicator local, using cross-entropy plus λ-weighted DO-loss.The stated computation complexity is O(mn^2), and no additional load-balancing auxiliary term is used.Source: §3.1.4 · p. 4

3.2 Replica Expert Mechanism (REM)

  • REM exploits stable hot-expert phases through two decoupled loops: global-batch replica placement and independent micro-batch token dispatch.A collector aggregates load statistics asynchronously; the controller places replicas, while the dispatcher routes tokens in real time.Source: §3.2 · p. 5, Figure 2 · p. 5

3.2.1 Global-Batch Replica Allocation

  • At each global-batch boundary, REM greedily places replicas on underloaded ranks using a Benefit Score for the receiving rank and expert.The ideal token threshold is the average load across ranks; allocation continues until available replica slots are occupied.Source: §3.2 · p. 5, §3.2.1 · p. 5

3.2.2 Dynamic Routing via Micro-step Scheduling

  • At each micro-batch, the dispatcher uses synchronized live queue depths to route tokens greedily, rather than relying on historical load averages.It performs D iterations without re-sorting ranks after each allocation; replica gradients return to the original rank after each global step.Source: §3.2.2 · p. 6

4 Experimental Setup

The experiments use two Qwen3MoE model scales, compare five routing strategies, train on a proprietary corpus, and evaluate across six reasoning, knowledge, and mathematics benchmarks.

4.1 Model Architecture

  • The 4.8BA0.5B and 30BA3B Qwen3MoE models use Sigmoid-gated Top-8 routing over 128 experts, with 15 and 48 layers, respectively.Both use a 120,832-token vocabulary and 32-head GQA; their hidden sizes are 1024 and 2048, respectively.Source: §4.1 · p. 6

4.2 Baseline Methods

  • Five routing strategies span local and global load-balancing losses, bias-based balancing, routing-weight orthogonality, and sequence-level balancing.Seq-LBL uses coefficient 10−3 and is included only in MaxVio and throughput evaluations.Source: §4.2 · pp. 6–7

4.3 Training Settings

  • Training uses Megatron-Core on NVIDIA H20 GPUs, with sequence length 4096 and global batch size 512, on a 2T-token proprietary corpus.The 4.8BA0.5B and 30BA3B models train for 100B and 200B tokens, respectively; Table 1 summarizes corpus domain composition.Source: §4.3 · p. 7, Table 1 · p. 7

4.4 Evaluation Benchmarks

  • Six benchmarks cover reasoning, broad language understanding, Chinese academic knowledge, arithmetic reasoning, and open-domain factual recall.The evaluations include AGIEval, MMLU, MMLU-Pro, CEval, GSM8K, and TriviaQA.Source: §4.4 · p. 7

5 Results and Analysis

The results assess DO-loss regularization, benchmark performance, expert behavior, and REM’s load-balancing and throughput effects. Across these analyses, the reported comparisons include scale-dependent benchmark outcomes and a trade-off between replica-based balancing and throughput.

5.1 Regularization Coefficient Sensitivity

  • λ = 10−5 is optimal in the reported coefficient sweep; a finer-grained 2B/0.2B sweep refines the optimum to λ = 5×10−6.Too-high coefficients overwhelm the cross-entropy signal, while too-low coefficients fail to encourage specialization and lead to collapse.Source: §5.1 · pp. 7–8

5.2 Downstream Benchmark Performance

  • Global-LBL underperforms Auxfree on most 30BA3B benchmarks, particularly knowledge-intensive TriviaQA.Table 2 covers validation perplexity and benchmark accuracy at two scales and training tokens of 100B and 200B.Source: Table 2 · p. 8

5.3 Ablation Study

  • Global statistics benefit both load-balancing loss and DO-loss, but the DO objective dominates Local-LBL at both levels and Global-DO-loss performs best across seven metrics.The authors attribute the improvements primarily to pairwise orthogonalization rather than global statistics alone.Source: §5.3 · p. 8

5.4 Expert Behavior: A Comparative Analysis

  • Auxfree balances token counts but has higher CKA redundancy and lower routing confidence; DO-loss matches Orth-loss on CKA sparsity and routing confidence with more controlled loads.Across 128 expert pairs, Global-DO-loss has the best separation in deeper layers with the smallest variance; the figure compares routing confidence, CKA similarity, and expert token counts.Source: §5.4 · p. 9, Figure 4 · p. 9

5.5 REM: Load Balancing and System Efficiency

  • On H20 GPUs, one REM replica recovers most of the throughput gap and lowers MaxVio below Auxfree; two replicas bring MaxVio near zero, with slightly lower throughput than Auxfree.REM communication overhead is negligible, and the method consistently outperforms EfficientMoE and Pro-Prophet on H20 and H800 clusters.Source: §5.5 · pp. 9–10, Figure 5 · p. 10

6 Conclusion

The conclusion frames expert collapse and representational redundancy as distinct MoE challenges and separates routing specialization from system-level load balancing. The accompanying material describes DO-loss normalization and REM’s allocation and dispatch loops.

6 Conclusion

  • DO-loss addresses collapse by penalizing routing-trajectory overlap, while REM restores workload balance through adaptive replica placement and real-time dispatch, decoupled from the routing objective.The conclusion reports downstream gains at both model scales with comparable training throughput and says REM can be optional on wafer-scale architectures.Source: §6 · p. 10

A DO-loss Normalization Coefficient

  • The DO-loss normalization coefficient 1/m²k makes the loss scale-free and comparable to cross-entropy; its derivation evaluates the unnormalized loss at balanced equilibrium.The derivation uses normalized routing probabilities, with the mean Top-K routing-score sum approaching 1 as specialization intensifies.Source: A DO-loss Normalization Coefficient · p. 14

B Replica Expert Mechanism

B.1 Global-Batch Replica Allocation

  • At each global-batch boundary, REM greedily places replicas on underloaded ranks using a Benefit Score and updated load estimates.Placement stops when rank capacity is filled or no beneficial placement remains.Source: §B.1 · p. 15

B.2 Micro-Batch Token Dispatching

  • Within each micro-batch, a local dispatcher sorts ranks by current load and assigns replica tokens from lightest to heaviest without waiting for the global controller.Queue-depth synchronization sends only |R| integers per EP group.Source: §B.2 · p. 15

B.3 Training Trajectory Invariance

REM’s training-trajectory equivalence depends on synchronized replica parameters, summed replica gradients, and optimizer updates on the original expert rank. The paper reports a precision test for gradient alignment.

B.3 Training Trajectory Invariance

  • Optimizer states remain on the original expert rank, where the standard DDP all-reduce is applied.Source: §B.3 · p. 15
  • A 10,000-step precision test confirmed gradient alignment to three decimal places.Source: §B.3 · p. 16
  • Trajectory equivalence requires synchronized replica parameters before each forward pass, summed gradients, and optimizer updates only on the original expert rank.Replica parameters are overwritten from the original expert, while gradients computed across replicas are accumulated back to that rank.Source: §B.3 · pp. 15–16

C Gradient Analysis of DO-loss

DO-loss minimizes variation in pairwise expert-signature distances; its gradient can revive idle experts, while equidistance implies uniform token counts and minimum pairwise co-occurrence.

C.1 Zero-Sum Constraint and Variance Minimization

  • Because Top-K routing fixes the total pairwise Hamming distance, minimizing its variance favors equal distances across expert pairs.The total number of set bits across signatures is mK.Source: §C.1 · p. 17

C.2 Anti-Collapse Gradient

  • For an idle expert, the DO-loss gradient raises routing probabilities for tokens assigned to active experts, creating a self-correcting anti-collapse feedback loop.As the idle expert receives tokens, its signature distances from active experts decrease, lowering overall variance.Source: §C.2 · p. 17

C.3 Equidistance and Optimal Specialization

  • Equal pairwise distances imply uniform per-expert token counts and the minimum possible co-occurrence for every expert pair under Top-K routing.The paper connects this minimum, uniform co-occurrence to distinct semantic patterns being processed by maximally disjoint expert combinations.Source: §C.3 · p. 17

D Centered Kernel Alignment (CKA)

The CKA analysis compares expert feature representations on a fixed token probe set. Global-DO-loss shows the lowest mean similarity and variance at Layer 8, while Orth-loss yields the sparsest Layer 4 representations.

Feature Extraction

  • Each expert’s feature vector combines gating magnitude with the value branch to represent its functional selectivity.A fixed probe set of tokens is passed independently through every expert, bypassing the router.Source: D Centered Kernel Alignment (CKA) · p. 18

CKA Interpretation

  • The square-root CKA metric ranges from 0 for linearly independent representations to 1 for representations identical up to an invertible linear transformation.High off-diagonal values indicate correlated features across distinct experts and serve as a symptom of representation redundancy.Source: D Centered Kernel Alignment (CKA) · p. 18

Quantitative CKA Results

  • Global-DO-loss attains the lowest mean and smallest variance in Layer 8 CKA similarity across 128 expert pairs.Orth-loss achieves the sparsest representations in Layer 4; lower CKA values indicate less representational redundancy.Source: Table 3 · p. 18, D Centered Kernel Alignment (CKA) · p. 18

E MaxVio Metric

The supplementary analyses characterize rank-level load imbalance and test DO-loss across coefficient, batch-size, and expert-count settings. They report improved validation perplexity with wider statistical scope and stronger relative performance at finer expert granularity.

E MaxVio Metric

  • Without REM, MaxVio stabilizes below 1.0, meaning the busiest rank processes at most twice the average load, but imbalance remains substantially higher than with replica scheduling.MaxVio measures busiest-rank load deviation from the ideal uniform load and captures the straggler effect.Source: Figure 6 · p. 19, E MaxVio Metric · p. 19

F.1 Fine-grained λ Sweep

  • The fine-grained sweep identifies λ=5×10−6 as optimal, roughly two orders of magnitude below typical load-balancing loss coefficients.The sweep uses a 2B/0.2B model trained for 20B tokens.Source: Figure 7 · p. 20, §F.1 · p. 20

F.2 DO-loss Batch Size Ablation

  • Validation perplexity improves monotonically as the effective batch scope for computing G_ij widens, with diminishing marginal returns.The ablation uses the 4.8BA0.5B model, fixed micro-batch size 4, and synchronization across 8 GPUs.Source: Table 4 · p. 20, §F.2 · p. 20

F.3 Expert Count Scaling

  • As expert count rises from 32 to 64 and 128, Global-DO-loss moves from trailing Auxfree to surpassing it, with the gap widening.The benefit of distributional orthogonalization grows with expert granularity.Source: §F.3 · pp. 20–21, Figure 8 · p. 21

F.4 REM Memory and Throughput Overhead

REM outperforms EfficientMoE and Pro-Prophet across the reported cluster configurations and recovers most of DO-only’s throughput gap to Auxfree. Its two-tiered scheduling is paired with modest replica memory overhead.

Throughput Recovery

  • DO-only runs roughly 10% below Auxfree throughput under EP8, while REM recovers most of the gap and nearly matches Auxfree under EP16.EP8 has higher throughput than EP16 across methods; REM brings throughput within a few percent of Auxfree across configurations.Source: §F.4 · p. 21

Throughput Comparison

  • REM consistently outperforms EfficientMoE and Pro-Prophet across H20 and H800 configurations and approaches Auxfree throughput.The steady-state comparison averages 1,000 steps, with two replica experts per rank on the DO-loss-trained 4.8BA0.5B model.Source: Table 5 · p. 21, §F.4 · p. 21

Scheduling Comparison

  • REM combines global-batch replica placement with lightweight micro-batch dispatching, avoiding unresolved transient imbalance and communication overhead that can offset scheduling benefits.The comparison attributes EfficientMoE’s lag to global-batch-only scheduling and notes Pro-Prophet’s micro-batch replication overhead under communication-bound conditions.Source: §F.4 · p. 21

Memory and Communication Overhead

  • One or two replicas increase peak GPU memory by a single-digit percentage, while queue-depth AllGather sends only |R| integers per EP group.The queue-depth payload is dwarfed by the AlltoAll dispatch of full hidden states.Source: §F.4 · p. 22

Related papers

Loading 2609.32398v1…