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

要約

まばらに活性化された混合混合物（MOE）モデルは、計算予算を増やすことなくモデル容量を拡大するために広く採用されています。
ただし、バニラTOPKルーターは、不連続で変化する可能性のない方法で訓練されており、パフォーマンスとスケーラビリティを制限しています。
この問題に対処するために、レミーを提案します。これは、リリーターをルーターとして使用して、従来のTOPK+ソフトマックスルーティングにシンプルで効果的なドロップイン交換を提供する完全に微分可能なMOEアーキテクチャを提案します。
さらに、専門家の負荷のバランスをとりながら、ルーターのスパースを調節する方法を提案します。
Remoeの連続性は、トークンとレイヤー間の計算の効率的な動的割り当てを可能にし、ドメインの専門化も示します。
私たちの実験は、Remoeが、さまざまなモデルサイズ、専門家数、および粒度のレベルにわたってバニラTopkルーティングのMOEを常に上回ることを示しています。
さらに、Remoeは、従来のMOEアーキテクチャを上回る専門家の数に関して優れたスケーラビリティを示します。
Megatron-LMに基づく実装は、https：//github.com/thu-ml/remoeで入手できます。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google