PaperlayerView exact revision on arXiv

Source-linked AI summary

Coupling Experts and Routers in Mixture-of-Experts via an Auxiliary Loss

Ang Lv, Jin Ma, Yiyuan Ma, Siyuan Qiao

arXiv:2512.23447v2cs.CLcs.LG

TL;DR

MoE models lack explicit constraints aligning routing decisions with expert capabilities, and denser activation guidance can be costly. ERC uses perturbed router embeddings as proxies and expert activations to couple routers with experts. The authors report improved downstream performance with negligible training overhead, alongside a way to control and quantify specialization.

  • Problem

    MoE routers lack explicit constraints ensuring their decisions align with expert capabilities, while denser activation guidance can incur substantial computational and memory costs.

    Source: §1 · p. 1

  • Method

    ERC treats perturbed router rows as proxies for routed-token clusters and uses expert activation norms to couple router representations with expert capabilities.

    Source: §1 · p. 2

  • Results

    ERC improves MoE-based LLM downstream performance with negligible training overhead, significantly enhancing results and narrowing the gap to more costly MoE variant [25].

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

  • Takeaways & Limitations

    α controls specialization while δ_i supplies a quantitative measure, revealing a specialization–performance trade-off during training.

    Source: §1 · p. 2

    Even when routers and experts are orthogonalized, there is no guarantee that each router vector aligns with its corresponding expert parameters.

    Source: §4.5 · p. 9

Abstract

from arXiv · show

Mixture-of-Experts (MoE) models lack explicit constraints to ensure the router's decisions align well with the experts' capabilities, which ultimately limits model performance. To address this, we propose expert-router coupling (ERC) loss, a lightweight auxiliary loss that tightly couples the router's decisions with expert capabilities. Our approach treats each expert's router embedding as a proxy token for the tokens assigned to that expert, and feeds perturbed router embeddings through the experts to obtain intermediate activations. The ERC loss enforces two constraints on these activations: (1) Each expert must exhibit higher activation for its own proxy token than for the proxy tokens of any other expert. (2) Each proxy token must elicit stronger activation from its corresponding expert than from any other expert. These constraints jointly ensure that each router embedding faithfully represents its corresponding expert's capability, while each expert specializes in processing the tokens actually routed to it. The ERC loss is computationally efficient, operating only on $n^2$ activations, where $n$ is the number of experts. This represents a fixed cost independent of batch size, unlike prior coupling methods that scale with the number of tokens (often millions per batch). Through pre-training MoE-LLMs ranging from 3B to 15B parameters and extensive analysis on trillions of tokens, we demonstrate the effectiveness of the ERC loss. Moreover, the ERC loss offers flexible control and quantitative tracking of expert specialization levels during training, providing valuable insights into MoEs.

1 Introduction

MoE routers select experts for tokens, but the introduction identifies no explicit guarantee that routing aligns with expert capabilities. It proposes ERC to couple them and reports large-scale pre-training gains.

1 Introduction

  • MoE routers lack explicit constraints ensuring decisions match expert capabilities, while denser activation guidance can incur substantial computational and memory costs.Without direct access to expert parameters, routers learn routing strategies by trial and error; misrouted-token gradients can interfere with expert specialization.Source: §1 · p. 1
  • ERC treats each router row as the center of its routed-token set and adds bounded noise to create a proxy token.The noise is bounded by half the minimum distance between adjacent cluster centers, simulating token variation without crossing cluster boundaries.Source: §1 · p. 2
  • ERC feeds proxies through all experts and uses intermediate activation norms as capability-alignment indicators, producing M[i,j] for proxy i and expert j.M[i,j] is the activation norm from expert j given proxy input i.Source: §1 · p. 2
  • The loss penalizes cross-activations above α-scaled self-activation, encouraging each expert to specialize for its routed cluster and each router vector to match its expert.The two effects are described as expert specialization and precise token routing.Source: §1 · p. 2
  • Across 3B–15B pre-training, ERC significantly enhances performance, narrows the gap with more computationally expensive MoE variant [25], and retains vanilla-MoE efficiency.The experiments use several trillion tokens.Source: §1 · p. 2
  • ERC enables specialization analysis: α controls specialization, δ_i provides a quantitative measure, and analysis reveals a trade-off with model performance.The authors say this analysis challenges beliefs about specialization drawn from small-scale experiments.Source: §1 · p. 2

2 Background

The background describes standard MoE routing and output combination, then contrasts it with AoE's expert-encoded routing and token-dependent computation costs.

2 Background

  • An MoE router assigns each token to the top-K experts, whose outputs are combined as a weighted sum.The router is represented by weight matrix R.Source: §2 · p. 3
  • AoE encodes routing in expert parameters and uses activation norms to select experts, but its computational overhead scales with token count.The top-K experts with the highest activation norms continue their forward computation; other experts stop early.Source: §2 · p. 3
  • Pham et al. supervise router logits with experts’ final-output norms, avoiding inference overhead but activating the model densely during training.The authors include this method only for background discussion, not as a baseline.Source: §2 · p. 3

3 Method

ERC adds an auxiliary loss to couple routers and experts while retaining the router and vanilla-MoE activation density. It uses perturbed router centers as proxies for routed tokens, with α controlling coupling strength.

3 Method

  • ERC is designed to retain routers, couple routers and experts through an auxiliary loss, and add neither token-scaled complexity nor activation density beyond vanilla MoE.Source: §3 · p. 3

3.1 Expert-router coupling loss

  • ERC treats each router center as a proxy for its routed token cluster, perturbs proxies for loss computation, and processes each through all experts to form n^2 activation norms.Source: §3.1 · p. 4
  • The loss requires each proxy to activate its corresponding expert most strongly and each expert to respond most to its own proxy; lowering α tightens coupling and increases specialization.Source: §3.1 · p. 4

3.2 Bounded random noise for generating proxy tokens

  • Proxy noise δ_i is sampled uniformly within a bound set by nearest-center distance to keep perturbations within the original cluster; its maximum permitted level is recomputed each layer and step.Source: §3.2 · p. 5

3.3 Efficiency analysis

  • ERC adds 2n^2Dd FLOPs, versus AoE's 2T(n −K)dr FLOPs, which the authors say becomes prohibitive when n > 4K.Source: §3.3 · p. 5
  • 0.2–0.8% overhead during LLM pretraining, and no additional inference overhead when the auxiliary loss is not applied.Source: §3.3 · p. 5

4 Experiments

ERC improves MoE performance while retaining vanilla MoE efficiency and comparable load balancing. Experiments also show how ERC can track and control specialization, with results spanning 3B and 15B models and ablations identifying useful design choices.

4.1 Experimental settings

  • Experiments compare ERC-augmented MoE with vanilla MoE and AoE using 3B models trained on 500B tokens, with 64 experts and 8 selected per token.The models use 500M activated parameters and share the same load-balancing loss; ERC loss weight is fixed at 1, with α = 1 by default.Source: §4.1 · p. 5

4.2 Performance, efficiency, and load balancing

  • ERC yields stable gains over vanilla MoE and narrows the performance gap with AoE, while keeping load-balancing loss close to vanilla MoE.Figure 3(a) reports average accuracy across tasks; Figure 3(b) shows the load-balancing-loss difference is on the order of 10^-5.Source: §4.2 · pp. 5–6, Figure 3 · p. 6

4.3 Validating ERC Loss in 15B-Parameter MoEs

  • At 15B, ERC improves performance across more challenging public benchmarks, with no loss spikes or abnormal gradients during training.The benchmarks include MMLU, C-Eval, MMLU-Pro, AGI-Eval, BBH, MATH, GSM8K, and TriviaQA.Source: Table 1 · p. 6, §4.3 · p. 6

4.4 The ERC loss is an effective tool for exploring expert specialization

  • t-SNE projections show distinct expert-parameter clusters with ERC, unlike vanilla MoE, whose parameter features lack meaningful clusters.The projections use rows of Wg from experts where i mod 8 = 0 in layer 6.Source: §4.4 · p. 6, Figure 4 · p. 7
  • α controls coupling and ϵ tracks specialization; tested α values beat vanilla MoE, but overly strict α degrades performance, exposing a specialization–performance trade-off.Increasing α reduces specialization and corresponds to a decrease in ϵ.Source: §4.4 · p. 7
  • Optimal specialization settings remain empirical because reliable quantitative metrics for n and K are lacking, and other architectural choices also shape the trade-off.The reported optimal α differs between the 3B model with n = 64 and the 15B model with n = 256.Source: §4.4 · pp. 7–8

4.5 Ablation studies

  • Among five activation choices, ˜RWg is most effective; final-output activations are comparable but costlier, so ˜RWg is the default.The five candidates include norms of ˜RWg, ˜RWp, SiLU(˜RWg), post-SwiGLU activations, and final expert outputs.Source: §4.5 · p. 8
  • Removing noise δ greatly degrades performance, while router orthogonalization yields limited gains on an already-near-orthogonal baseline, underscoring the importance of router–expert coupling.The baseline routers have average absolute cosine similarity 0.15; the authors attribute ERC’s gains to addressing weak router–expert coupling.Source: §4.5 · pp. 8–9

5 Related works

Prior MoE auxiliary losses largely focus on training stability, while specialization methods leave gaps in quantitative control, trade-off analysis, or expert-router coupling. ERC is presented as an efficient approach addressing those gaps.

5 Related works

  • Most existing MoE auxiliary losses target training stability, whereas ERC is presented as the first tailored to strengthen expert-router coupling.Prior examples include load-balancing losses and z-loss.Source: §5 · p. 9
  • Specialization approaches include shared experts [6], selected-expert overlap loss, and millions of tiny experts [12] [29] [47], but have limitations in control, trade-off analysis, or coupling.The selected-expert loss incurs K^2 cosine-similarity calculations per token, while tiny experts are memory-bounded.Source: §5 · pp. 9–10
  • Orthogonal-expert methods [13] are evaluated through fine-tuning or reinforcement learning, while the optimal specialization degree remains a complex, open question.The paper argues orthogonalizing expert outputs is not equivalent to extreme specialization.Source: §5 · p. 10
  • Applying contrastive learning only to routers or only to experts leaves their weak coupling unaddressed.The paper notes related MoE work applying contrastive learning to expert outputs [24].Source: §5 · p. 10

6 Conclusions

The conclusion identifies weak router-expert coupling as a limitation of MoE models and presents ERC as a low-overhead way to strengthen that coupling. It reports downstream-task improvements and describes ERC as a tool for studying MoE behavior and expert specialization.

6 Conclusions

  • Weak coupling between router decisions and expert capabilities limits MoE models in multiple important aspects.Source: §6 · p. 10
  • ERC tightly couples router parameters with their corresponding experts and improves MoE-based LLMs on downstream tasks with negligible training overhead.Source: §6 · p. 10
  • ERC's properties offer deeper insight into MoE behavior and a promising tool for future research on expert specialization.Source: §6 · p. 10

A.1 What happens if α > 1?

Post-hoc evaluation finds that the vanilla MoE baseline requires α = 5 for zero ERC loss across all layers. Experiments with α > 1 show diminishing gains, and the authors advise against those settings.

A.1 What happens if α > 1?

  • α = 5 is required to achieve zero ERC loss across all layers in the pre-trained vanilla MoE baseline.Table 2 reports deterministic post-hoc loss values computed using the original R across α values.Source: §A.1 · p. 14, Table 2 · p. 14
  • α = 2 yields only limited improvement, while α = 3 shows almost no improvement over vanilla MoE; the authors do not recommend α > 1.These experiments demonstrate the effects of loosening ERC constraints.Source: §A.1 · p. 14

A.2 Do models reduce ERC loss through manipulating parameter norms?

The appendix argues that ERC loss cannot be minimized simply by inflating parameter norms: reducing one loss term through norm manipulation increases other terms. It also bounds multiplicative noise so perturbed router embeddings remain closer to their own cluster centers than to others.

A.2 Do models reduce ERC loss through manipulating parameter norms?

  • Increasing router-embedding norms decreases one ERC-loss term but increases another, and manipulating expert-weight norms likewise trades off loss terms.The stated result is that minimizing total ERC loss requires comparable router-embedding norms and meaningful router-expert coupling.Source: §A.2 · p. 15
  • Models trained with ERC have comparable average parameter norms but lower norm standard deviations, while the baseline has substantially higher ERC loss.Table 3 reports parameter norms and layer-wise ERC loss at the last checkpoint.Source: §A.2 · p. 15, Table 3 · p. 15

B Determining the maximum multiplicative noise level

  • The multiplicative perturbation uses independently sampled components, and its noise bound keeps each perturbed router embedding closer to its own center than to any other.The paper adopts a tighter upper bound because computing the denominator of the more complex expression is relatively costly.Source: B Determining the maximum… · pp. 15, 17

C Efficiency analysis

Appendix C analyzes the computational cost of vanilla MoE, AoE, and ERC, including real-world parallelism, and concludes that ERC is practical.

C Efficiency analysis

  • The ideal FLOPs breakdown and analysis of real-world MoE pre-training parallelism both demonstrate the practicality of ERC.The analyses cover AoE and ERC overhead as well as vanilla MoE cost.Source: C Efficiency analysis · p. 17

C.1 FLOPs cost breakdown of three methods

The appendix compares FLOPs for vanilla MoE, AoE, and ERC, reporting the vanilla MoE cost, AoE's parameter-matched low-rank factorization, and ERC's additional cost.

C.1 FLOPs cost breakdown of three methods

  • For K experts processing T tokens, vanilla MoE costs 6KTdD FLOPs.Each expert contributes 6TdD FLOPs.Source: §C.1 · p. 17
  • AoE factorizes the expert matrix Wg into two low-rank matrices while maintaining the original matrix's parameter count.The passage specifies both factors have rank r.Source: §C.1 · p. 17
  • ERC adds 2n^2Dd FLOPs through n^2 matrix multiplications on tensors shaped 1 × d and d × D.This cost is reported as the computational overhead of ERC loss.Source: §C.1 · p. 17

C.2 Throughputs under multiple parallelism strategies

Under large-scale parallel pre-training, ERC's estimated and measured overheads are small, and its distributed cost is expressed relative to a base MoE forward pass.

C.2 Throughputs under multiple parallelism strategies

  • With data and expert parallelism, ERC's computational cost is equivalent to a forward pass on n^2/3 tokens.The passage gives base MoE forward cost as K · T / dp_size and ERC overhead as n · (n / ep_size) / 3.Source: §C.2 · p. 18
  • For the 15B model, ERC overhead is estimated at 0.72%; throughput is 62.03B tokens/day for baseline versus 61.52B tokens/day for our model, a 0.82% reduction.With n = 64 in the 3B models, the reported overhead ratio is 0.18%.Source: §C.2 · p. 18
Loading 2512.23447v2…