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

TL;DR

The paper tackles the problem of misalignment between MoE routers and expert capabilities by introducing the Expert-Router Coupling (ERC) loss. ERC treats router rows as cluster centers and uses perturbed proxy tokens to probe expert responses, enforcing the activation-based constraints $\mathbf{M}[i,j] < \alpha \mathbf{M}[i,i]$ and $\mathbf{M}[j,i] < \alpha \mathbf{M}[i,i]$, with a loss $\mathcal{L}_{ERC}$ that scales as $O(n^2)$. Across MoE-LLMs from $3\mathrm{B}$ to $15\mathrm{B}$ parameters, ERC yields performance gains over vanilla MoE, remains computationally cheap relative to alternatives like AoE, and enables controlled study of expert specialization through parameters $\alpha$ and perturbation bound $\epsilon$. The results demonstrate improved downstream task performance, preserved throughput, and a tangible trade-off between specialization and collaboration, offering a practical and insightful tool for MoE analysis and deployment.

Abstract

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 internal 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.

Figures (8)

• Figure 1: Three steps for computing the expert-router coupling loss.
• Figure 2: The overview of MoE and AoE models.
• Figure 3: The 3B-parameter MoE model augmented with ERC loss achieves substantial and stable performance gains, while maintaining comparable load balancing to the baseline. Figure \ref{['fig:3b_tasks_details']} shows task-specific details.
• Figure 4: t-SNE projections of ${\bm{W}}_g$ in MoE experts trained without and with the ERC loss. Our ERC loss provides greater expert specialization.
• Figure 5: (a) Since routers are deeply coupled with experts, the distance between neighboring cluster centers (i.e., the maximum noise level $\epsilon$) quantitatively reflects changes in expert specialization during training, which is controlled by $\alpha$. (b) Downstream performance across different values of $\alpha$.