Couplage des experts et des routeurs dans les modèles Mixture-of-Experts via une fonction de coût auxiliaire

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

December 29, 2025

papers.authors: Ang Lv, Jin Ma, Yiyuan Ma, Siyuan Qiao

cs.AI

papers.abstract

Les modèles à mélange d'experts (MoE) manquent de contraintes explicites pour garantir que les décisions du routeur s'alignent correctement sur les capacités des experts, ce qui limite in fine les performances du modèle. Pour remédier à cela, nous proposons la perte de couplage expert-routeur (ERC), une fonction de perte auxiliaire légère qui couple étroitement les décisions du routeur aux capacités des experts. Notre approche traite l'embedding du routeur de chaque expert comme un token proxy pour les tokens assignés à cet expert, et fait passer ces embeddings perturbés à travers les experts pour obtenir leurs activations internes. La perte ERC impose deux contraintes sur ces activations : (1) Chaque expert doit présenter une activation plus élevée pour son propre token proxy que pour les tokens proxy des autres experts. (2) Chaque token proxy doit provoquer une activation plus forte de son expert correspondant que de tout autre expert. Ces contraintes garantissent conjointement que chaque embedding du routeur représente fidèlement la capacité de son expert correspondant, tandis que chaque expert se spécialise dans le traitement des tokens qui lui sont effectivement routés. La perte ERC est efficace sur le plan computationnel, opérant uniquement sur n² activations, où n est le nombre d'experts. Cela représente un coût fixe indépendant de la taille du lot, contrairement aux méthodes de couplage antérieures qui évoluent avec le nombre de tokens (souvent des millions par lot). Grâce au pré-entraînement de MoE-LLMs allant de 3B à 15B de paramètres et à une analyse approfondie sur des milliers de milliards de tokens, nous démontrons l'efficacité de la perte ERC. De plus, la perte ERC offre un contrôle flexible et un suivi quantitatif des niveaux de spécialisation des experts pendant l'entraînement, fournissant des insights précieux sur les modèles MoE.

English

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.