ReMoE: Mezcla de Expertos Totalmente Diferenciable con Enrutamiento ReLU

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

December 19, 2024

Autores: Ziteng Wang, Jianfei Chen, Jun Zhu

cs.AI

Resumen

Los modelos de Mezcla de Expertos (MoE) activados de forma dispersa son ampliamente adoptados para aumentar la capacidad del modelo sin incrementar el presupuesto computacional. Sin embargo, los enrutadores TopK convencionales se entrenan de manera discontinua y no diferenciable, lo que limita su rendimiento y escalabilidad. Para abordar este problema, proponemos ReMoE, una arquitectura MoE completamente diferenciable que ofrece un reemplazo simple pero efectivo para el enrutamiento convencional TopK+Softmax, utilizando ReLU como enrutador en su lugar. Además, proponemos métodos para regular la dispersión del enrutador mientras se equilibra la carga entre los expertos. La naturaleza continua de ReMoE permite una asignación dinámica eficiente de la computación entre tokens y capas, al mismo tiempo que muestra especialización de dominio. Nuestros experimentos demuestran que ReMoE supera consistentemente a los modelos MoE convencionales enrutados con TopK en distintos tamaños de modelo, cantidades de expertos y niveles de granularidad. Además, ReMoE muestra una escalabilidad superior en relación con el número de expertos, superando a las arquitecturas MoE tradicionales. La implementación basada en Megatron-LM está disponible en 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.