Routeur Récurrent par Couches pour les Mélanges d'Experts

Layerwise Recurrent Router for Mixture-of-Experts

August 13, 2024

Auteurs: Zihan Qiu, Zeyu Huang, Shuang Cheng, Yizhi Zhou, Zili Wang, Ivan Titov, Jie Fu

cs.AI

Résumé

Le passage à l'échelle des grands modèles de langage (LLMs) a révolutionné leurs capacités dans diverses tâches, mais cette croissance doit s'accompagner de stratégies de calcul efficaces. L'architecture Mixture-of-Experts (MoE) se distingue par sa capacité à augmenter la taille des modèles sans augmenter significativement les coûts d'entraînement. Malgré leurs avantages, les modèles MoE actuels présentent souvent une inefficacité en termes de paramètres. Par exemple, un LLM basé sur MoE pré-entraîné avec 52 milliards de paramètres pourrait offrir des performances comparables à un modèle standard de 6,7 milliards de paramètres. Étant une partie cruciale de MoE, les routeurs actuels dans différentes couches attribuent les tokens de manière indépendante sans exploiter les informations historiques de routage, ce qui peut conduire à des combinaisons token-expert sous-optimales et au problème d'inefficacité des paramètres. Pour atténuer ce problème, nous introduisons le Layerwise Recurrent Router for Mixture-of-Experts (RMoE). RMoE utilise une unité récurrente à porte (GRU) pour établir des dépendances entre les décisions de routage à travers les couches consécutives. Une telle récurrence par couche peut être calculée efficacement en parallèle pour les tokens d'entrée et introduit des coûts négociables. Nos évaluations empiriques approfondies démontrent que les modèles de langage basés sur RMoE surpassent systématiquement un éventail de modèles de référence. De plus, RMoE intègre une nouvelle étape de calcul orthogonale aux méthodes existantes, permettant une compatibilité transparente avec d'autres architectures MoE. Nos analyses attribuent les gains de RMoE à son partage efficace d'informations inter-couches, qui améliore également la sélection et la diversité des experts. Notre code est disponible à l'adresse https://github.com/qiuzh20/RMoE.

English

The scaling of large language models (LLMs) has revolutionized their capabilities in various tasks, yet this growth must be matched with efficient computational strategies. The Mixture-of-Experts (MoE) architecture stands out for its ability to scale model size without significantly increasing training costs. Despite their advantages, current MoE models often display parameter inefficiency. For instance, a pre-trained MoE-based LLM with 52 billion parameters might perform comparably to a standard model with 6.7 billion parameters. Being a crucial part of MoE, current routers in different layers independently assign tokens without leveraging historical routing information, potentially leading to suboptimal token-expert combinations and the parameter inefficiency problem. To alleviate this issue, we introduce the Layerwise Recurrent Router for Mixture-of-Experts (RMoE). RMoE leverages a Gated Recurrent Unit (GRU) to establish dependencies between routing decisions across consecutive layers. Such layerwise recurrence can be efficiently parallelly computed for input tokens and introduces negotiable costs. Our extensive empirical evaluations demonstrate that RMoE-based language models consistently outperform a spectrum of baseline models. Furthermore, RMoE integrates a novel computation stage orthogonal to existing methods, allowing seamless compatibility with other MoE architectures. Our analyses attribute RMoE's gains to its effective cross-layer information sharing, which also improves expert selection and diversity. Our code is at https://github.com/qiuzh20/RMoE