전문가 혼합 모델에서 보조 손실을 통한 전문가와 라우터의 결합

초록

전문가 혼합(MoE) 모델은 라우터의 결정이 전문가의 역량과 잘 일치하도록 보장하는 명시적 제약이 부족하여 궁극적으로 모델 성능을 제한합니다. 이를 해결하기 위해 우리는 라우터의 결정과 전문가 역량을 긴밀하게 결합하는 경량 보조 손실 함수인 전문가-라우터 결합(ERC) 손실을 제안합니다. 우리의 접근 방식은 각 전문가의 라우터 임베딩을 해당 전문가에 할당된 토큰의 프록시 토큰으로 간주하고, 교란된 라우터 임베딩을 전문가에 통과시켜 내부 활성화를 얻습니다. ERC 손실은 이러한 활성화에 두 가지 제약을 적용합니다: (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.

전문가 혼합 모델에서 보조 손실을 통한 전문가와 라우터의 결합

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

초록

Support