ReMoE: Volledig Differentieerbare Mixture-of-Experts met ReLU Routing

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

December 19, 2024

Auteurs: Ziteng Wang, Jianfei Chen, Jun Zhu

cs.AI

Samenvatting

Spaarzaam geactiveerde Mixture-of-Experts (MoE) modellen worden veelvuldig gebruikt om de modelcapaciteit te vergroten zonder het rekenbudget te verhogen. Echter, standaard TopK-routers worden op een discontinu, niet-differentieerbare manier getraind, wat hun prestaties en schaalbaarheid beperkt. Om dit probleem aan te pakken, stellen we ReMoE voor, een volledig differentieerbare MoE-architectuur die een eenvoudige maar effectieve vervanging biedt voor de conventionele TopK+Softmax-routing, waarbij ReLU als router wordt gebruikt. We stellen ook methoden voor om de spaarzaamheid van de router te reguleren terwijl de belasting onder de experts wordt gebalanceerd. De continue aard van ReMoE maakt efficiënte dynamische toewijzing van berekeningen over tokens en lagen mogelijk, terwijl ook domeinspecialisatie wordt getoond. Onze experimenten tonen aan dat ReMoE consequent beter presteert dan standaard TopK-gerouteerde MoE-modellen bij verschillende modelgroottes, expertaantallen en granulariteitsniveaus. Bovendien vertoont ReMoE superieure schaalbaarheid wat betreft het aantal experts, waarbij traditionele MoE-architecturen worden overtroffen. De implementatie gebaseerd op Megatron-LM is beschikbaar op 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.