ReMoE: Vollständig differenzierbare Expertenmischung mit ReLU-Routing

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

December 19, 2024

Autoren: Ziteng Wang, Jianfei Chen, Jun Zhu

cs.AI

Zusammenfassung

Spärlich aktiviert Mixture-of-Experts (MoE) Modelle werden weit verbreitet eingesetzt, um die Modellkapazität zu erhöhen, ohne das Berechnungsbudget zu erhöhen. Allerdings werden herkömmliche TopK-Router auf eine diskontinuierliche, nicht differenzierbare Weise trainiert, was ihre Leistungsfähigkeit und Skalierbarkeit einschränkt. Um dieses Problem zu lösen, schlagen wir ReMoE vor, eine vollständig differenzierbare MoE-Architektur, die einen einfachen, aber effektiven Ersatz für die herkömmliche TopK+Softmax-Routing bietet und dabei ReLU als Router verwendet. Wir schlagen außerdem Methoden vor, um die Sparsamkeit des Routers zu regulieren und die Last zwischen den Experten auszugleichen. Die kontinuierliche Natur von ReMoE ermöglicht eine effiziente dynamische Zuweisung von Berechnungen über Token und Schichten hinweg und zeigt auch Domänenspezialisierung. Unsere Experimente zeigen, dass ReMoE konsequent bessere Leistungen erbringt als herkömmliche TopK-geroutete MoE-Modelle bei verschiedenen Modellgrößen, Expertenzahlen und Granularitätsstufen. Darüber hinaus zeigt ReMoE eine überlegene Skalierbarkeit in Bezug auf die Anzahl der Experten und übertrifft traditionelle MoE-Architekturen. Die Implementierung basierend auf Megatron-LM ist verfügbar unter 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.