ReMoE: Полностью дифференцируемая смесь экспертов с маршрутизацией ReLU

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

December 19, 2024

Авторы: Ziteng Wang, Jianfei Chen, Jun Zhu

cs.AI

Аннотация

Модели Mixture-of-Experts (MoE) с разреженной активацией широко применяются для увеличения объема модели без увеличения вычислительного бюджета. Однако стандартные маршрутизаторы TopK обучаются дискретно, недифференцируемо, что ограничивает их производительность и масштабируемость. Для решения этой проблемы мы предлагаем ReMoE, полностью дифференцируемую архитектуру MoE, которая представляет собой простую, но эффективную замену стандартной маршрутизации TopK+Softmax, используя в качестве маршрутизатора ReLU. Мы также предлагаем методы регулирования разреженности маршрутизатора при балансировке нагрузки между экспертами. Непрерывная природа ReMoE обеспечивает эффективное динамическое распределение вычислений между токенами и слоями, а также проявляет специализацию на область. Наши эксперименты показывают, что ReMoE последовательно превосходит стандартные MoE с маршрутизацией TopK по различным размерам моделей, количеству экспертов и уровням детализации. Более того, ReMoE обладает превосходной масштабируемостью по сравнению с традиционными архитектурами MoE. Реализация на основе Megatron-LM доступна по ссылке 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.