Связывание экспертов и маршрутизаторов в смесях экспертов с помощью вспомогательной функции потерь

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

December 29, 2025

Авторы: Ang Lv, Jin Ma, Yiyuan Ma, Siyuan Qiao

cs.AI

Аннотация

Модели со смесью экспертов (MoE) не имеют явных ограничений, гарантирующих соответствие решений маршрутизатора возможностям экспертов, что в конечном счете ограничивает производительность модели. Для решения этой проблемы мы предлагаем функцию потерь связи «эксперт-маршрутизатор» (ERC loss) — легковесную вспомогательную функцию потерь, которая тесно связывает решения маршрутизатора с возможностями экспертов. Наш подход рассматривает эмбеддинг маршрутизатора каждого эксперта как прокси-токен для токенов, назначенных этому эксперту, и пропускает возмущенные эмбеддинги маршрутизатора через экспертов для получения внутренних активаций. Функция потерь ERC накладывает два ограничения на эти активации: (1) Каждый эксперт должен демонстрировать более высокую активацию для своего собственного прокси-токена, чем для прокси-токенов любого другого эксперта. (2) Каждый прокси-токен должен вызывать более сильную активацию у своего соответствующего эксперта, чем у любого другого эксперта. Эти ограничения в совокупности гарантируют, что каждый эмбеддинг маршрутизатора достоверно отражает возможности соответствующего эксперта, в то время как каждый эксперт специализируется на обработке фактически направленных к нему токенов. Функция потерь ERC вычислительно эффективна, оперируя лишь n² активациями, где n — количество экспертов. Это представляет собой фиксированную стоимость, не зависящую от размера пакета, в отличие от предыдущих методов связи, которые масштабируются с количеством токенов (часто миллионы на пакет). Благодаря предварительному обучению MoE-LLM с диапазоном параметров от 3B до 15B и обширному анализу на триллионах токенов мы демонстрируем эффективность функции потерь ERC. Более того, функция потерь ERC обеспечивает гибкий контроль и количественное отслеживание уровней специализации экспертов во время обучения, предоставляя ценные insights о 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.