Mixture-of-Expertsにおける専門家とルータの補助損失による結合

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

要旨

Mixture-of-Experts（MoE）モデルは、エキスパートの能力とルーターの決定を適切に整合させる明示的な制約を欠いており、これが最終的にモデル性能を制限している。この問題に対処するため、我々はエキスパートとルーターを緊密に結合する軽量な補助損失関数であるexpert-router coupling（ERC）損失を提案する。本手法では、各エキスパートのルーター埋め込みを、そのエキスパートに割り当てられたトークンの代理トークンとして扱い、摂動を加えたルーター埋め込みをエキスパートに入力して内部活性化を取得する。ERC損失はこれらの活性化に対して2つの制約を課す：（1）各エキスパートは、他のエキスパートの代理トークンよりも自身の代理トークンに対して高い活性化を示さなければならない。（2）各代理トークンは、他のエキスパートよりも対応するエキスパートから強い活性化を引き出さなければならない。これらの制約は共同して、各ルーター埋め込みが対応するエキスパートの能力を忠実に表現すると同時に、各エキスパートが実際にルーティングされたトークンの処理に特化することを保証する。ERC損失は計算効率が高く、n個のエキスパート数に対してn^2個の活性化のみで動作する。これはバッチサイズに依存しない固定コストであり、トークン数（バッチあたり数百万規模になることが多い）に比例してスケールする従来の結合手法とは異なる。3Bから15Bパラメータ規模のMoE-LLMを用いた事前学習と、数兆トークンにわたる詳細な分析を通じて、ERC損失の有効性を実証する。さらにERC損失は、訓練中のエキスパート特化レベルを柔軟に制御し定量的に追跡する機能を提供し、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.