ReMoE: Misto di Esperti Completamente Differenziabile con Instradamento ReLU

ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routing

December 19, 2024

Autori: Ziteng Wang, Jianfei Chen, Jun Zhu

cs.AI

Abstract

I modelli Mixture-of-Experts (MoE) attivati in modo sparso sono ampiamente adottati per aumentare la capacità del modello senza aumentare il budget computazionale. Tuttavia, i router TopK vanilla vengono addestrati in modo discontinuo e non differenziabile, limitando le loro prestazioni e scalabilità. Per affrontare questo problema, proponiamo ReMoE, un'architettura MoE completamente differenziabile che offre un semplice ma efficace sostituto per il routing convenzionale TopK+Softmax, utilizzando ReLU come router al suo posto. Proponiamo inoltre metodi per regolare la sparità del router bilanciando il carico tra gli esperti. La natura continua di ReMoE consente un'allocazione dinamica efficiente della computazione tra token e livelli, mostrando anche una specializzazione di dominio. I nostri esperimenti dimostrano che ReMoE supera costantemente i MoE con routing TopK vanilla su varie dimensioni del modello, conteggi di esperti e livelli di granularità. Inoltre, ReMoE mostra una scalabilità superiore rispetto al numero di esperti, superando le architetture MoE tradizionali. L'implementazione basata su Megatron-LM è disponibile su https://github.com/thu-ml/ReMoE.

English

Sparsely activated Mixture-of-Experts (MoE) models are widely adopted to scale up model capacity without increasing the computation budget. However, vanilla TopK routers are trained in a discontinuous, non-differentiable way, limiting their performance and scalability. To address this issue, we propose ReMoE, a fully differentiable MoE architecture that offers a simple yet effective drop-in replacement for the conventional TopK+Softmax routing, utilizing ReLU as the router instead. We further propose methods to regulate the router's sparsity while balancing the load among experts. ReMoE's continuous nature enables efficient dynamic allocation of computation across tokens and layers, while also exhibiting domain specialization. Our experiments demonstrate that ReMoE consistently outperforms vanilla TopK-routed MoE across various model sizes, expert counts, and levels of granularity. Furthermore, ReMoE exhibits superior scalability with respect to the number of experts, surpassing traditional MoE architectures. The implementation based on Megatron-LM is available at https://github.com/thu-ml/ReMoE.