Acoplamento de Especialistas e Roteadores em Mistura de Especialistas via uma Função de Perda 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
Resumo
Os modelos Mixture-of-Experts (MoE) carecem de restrições explícitas para garantir que as decisões do roteador estejam alinhadas com as capacidades dos especialistas, o que limita o desempenho do modelo. Para resolver isso, propomos a perda de acoplamento roteador-especialista (ERC), uma perda auxiliar leve que acopla firmemente as decisões do roteador às capacidades dos especialistas. Nossa abordagem trata a incorporação do roteador de cada especialista como um token proxy para os tokens atribuídos a esse especialista e alimenta incorporações do roteador perturbadas através dos especialistas para obter ativações internas. A perda ERC impõe duas restrições nessas ativações: (1) Cada especialista deve exibir maior ativação para seu próprio token proxy do que para os tokens proxy de qualquer outro especialista. (2) Cada token proxy deve eliciar uma ativação mais forte de seu especialista correspondente do que de qualquer outro especialista. Essas restrições garantem conjuntamente que cada incorporação do roteador represente fielmente a capacidade de seu especialista correspondente, enquanto cada especialista se especializa em processar os tokens realmente roteados para ele. A perda ERC é computacionalmente eficiente, operando apenas em n² ativações, onde n é o número de especialistas. Isso representa um custo fixo independente do tamanho do lote, diferentemente de métodos de acoplamento anteriores que escalam com o número de tokens (frequentemente milhões por lote). Através do pré-treinamento de MoE-LLMs variando de 3B a 15B de parâmetros e análise extensiva em trilhões de tokens, demonstramos a eficácia da perda ERC. Além disso, a perda ERC oferece controle flexível e rastreamento quantitativo dos níveis de especialização dos especialistas durante o treinamento, fornecendo insights valiosos sobre MoEs.
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.