ReMoE : Mélange de spécialistes entièrement différentiable avec routage ReLU

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

December 19, 2024

Auteurs: Ziteng Wang, Jianfei Chen, Jun Zhu

cs.AI

Résumé

Les modèles Mixture-of-Experts (MoE) activés de manière éparses sont largement adoptés pour augmenter la capacité du modèle sans augmenter le budget de calcul. Cependant, les routeurs TopK classiques sont entraînés de manière discontinue et non différentiable, limitant leurs performances et leur extensibilité. Pour résoudre ce problème, nous proposons ReMoE, une architecture MoE entièrement différentiable qui offre un remplacement simple mais efficace pour le routage conventionnel TopK+Softmax, en utilisant ReLU comme routeur à la place. Nous proposons également des méthodes pour réguler la parcimonie du routeur tout en équilibrant la charge entre les experts. La nature continue de ReMoE permet une allocation dynamique efficace des calculs entre les jetons et les couches, tout en présentant une spécialisation de domaine. Nos expériences démontrent que ReMoE surpasse de manière constante les MoE routés en TopK classiques sur diverses tailles de modèles, nombres d'experts et niveaux de granularité. De plus, ReMoE présente une extensibilité supérieure en ce qui concerne le nombre d'experts, dépassant les architectures MoE traditionnelles. L'implémentation basée sur Megatron-LM est disponible sur 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.