Koppeling van Experts en Routers in Mixture-of-Experts via een Hulpverlies

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

December 29, 2025

Auteurs: Ang Lv, Jin Ma, Yiyuan Ma, Siyuan Qiao

cs.AI

Samenvatting

Mixture-of-Experts (MoE)-modellen missen expliciete beperkingen om ervoor te zorgen dat de beslissingen van de router goed aansluiten bij de capaciteiten van de experts, wat uiteindelijk de modelprestaties beperkt. Om dit aan te pakken, stellen we expert-router coupling (ERC)-verlies voor, een lichtgewicht hulpverlies dat de beslissingen van de router nauw koppelt aan de expertcapaciteiten. Onze aanpak behandelt de router-embedding van elke expert als een proxy-token voor de tokens die aan die expert zijn toegewezen, en voert verstoorde router-embeddings door de experts om interne activaties te verkrijgen. Het ERC-verlies legt twee beperkingen op aan deze activaties: (1) Elke expert moet een hogere activatie vertonen voor zijn eigen proxy-token dan voor de proxy-tokens van enige andere expert. (2) Elk proxy-token moet een sterkere activatie oproepen van zijn corresponderende expert dan van enige andere expert. Deze beperkingen zorgen er gezamenlijk voor dat elke router-embedding de capaciteit van zijn corresponderende expert getrouw vertegenwoordigt, terwijl elke expert gespecialiseerd is in het verwerken van de tokens die er daadwerkelijk naartoe worden gerouteerd. Het ERC-verlies is computationeel efficiënt, omdat het alleen werkt op n² activaties, waarbij n het aantal experts is. Dit vertegenwoordigt een vaste kost onafhankelijk van de batchgrootte, in tegenstelling tot eerdere koppelingsmethoden die schalen met het aantal tokens (vaak miljoenen per batch). Door pre-training van MoE-LLMs variërend van 3B tot 15B parameters en uitgebreide analyse op biljoenen tokens, demonstreren we de effectiviteit van het ERC-verlies. Bovendien biedt het ERC-verlies flexibele controle en kwantitatieve tracking van expertspecialisatieniveaus tijdens de training, wat waardevolle inzichten biedt in MoE's.

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.