Acoplamiento de Expertos y Enrutadores en Mezcla de Expertos mediante una Pérdida Auxiliar
Coupling Experts and Routers in Mixture-of-Experts via an Auxiliary Loss
December 29, 2025
Autores: Ang Lv, Jin Ma, Yiyuan Ma, Siyuan Qiao
cs.AI
Resumen
Los modelos Mixture-of-Experts (MoE) carecen de restricciones explícitas para garantizar que las decisiones del enrutador se alineen adecuadamente con las capacidades de los expertos, lo que finalmente limita el rendimiento del modelo. Para abordar esto, proponemos la pérdida de acoplamiento enrutador-experto (ERC, por sus siglas en inglés), una pérdida auxiliar ligera que acopla estrechamente las decisiones del enrutador con las capacidades de los expertos. Nuestro enfoque trata la incrustación del enrutador de cada experto como un token proxy para los tokens asignados a ese experto, y alimenta las incrustaciones del enrutador perturbadas a través de los expertos para obtener activaciones internas. La pérdida ERC impone dos restricciones sobre estas activaciones: (1) Cada experto debe exhibir una activación más alta para su propio token proxy que para los tokens proxy de cualquier otro experto. (2) Cada token proxy debe provocar una activación más fuerte de su experto correspondiente que la de cualquier otro experto. Estas restricciones garantizan conjuntamente que cada incrustación del enrutador represente fielmente la capacidad de su experto correspondiente, mientras que cada experto se especializa en procesar los tokens que realmente se le enrutan. La pérdida ERC es computacionalmente eficiente, ya que opera solo en n² activaciones, donde n es el número de expertos. Esto representa un costo fijo independiente del tamaño del lote, a diferencia de los métodos de acoplamiento anteriores que escalan con el número de tokens (a menudo millones por lote). Mediante el pre-entrenamiento de MoE-LLMs que van desde 3B hasta 15B de parámetros y un análisis extenso en billones de tokens, demostramos la efectividad de la pérdida ERC. Además, la pérdida ERC ofrece un control flexible y un seguimiento cuantitativo de los niveles de especialización de los expertos durante el entrenamiento, proporcionando información valiosa sobre los modelos 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.