ReMoE: Misto de Especialistas Totalmente Diferenciável com Roteamento ReLU

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

December 19, 2024

Autores: Ziteng Wang, Jianfei Chen, Jun Zhu

cs.AI

Resumo

Os modelos Mixture-of-Experts (MoE) ativados de forma esparsa são amplamente adotados para aumentar a capacidade do modelo sem aumentar o orçamento computacional. No entanto, os roteadores TopK convencionais são treinados de forma descontínua e não diferenciável, limitando seu desempenho e escalabilidade. Para lidar com esse problema, propomos o ReMoE, uma arquitetura MoE totalmente diferenciável que oferece uma substituição simples e eficaz para o roteamento convencional TopK+Softmax, utilizando ReLU como roteador. Propomos ainda métodos para regular a dispersão do roteador enquanto equilibramos a carga entre os especialistas. A natureza contínua do ReMoE permite uma alocação dinâmica eficiente de computação entre tokens e camadas, além de exibir especialização de domínio. Nossos experimentos demonstram que o ReMoE supera consistentemente o MoE roteado por TopK convencional em vários tamanhos de modelo, contagens de especialistas e níveis de granularidade. Além disso, o ReMoE apresenta escalabilidade superior em relação ao número de especialistas, superando arquiteturas MoE tradicionais. A implementação baseada no Megatron-LM está disponível em 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.