BAM ! En un clin d'œil : un recyclage de paramètres simple et efficace pour les mélanges d'experts

BAM! Just Like That: Simple and Efficient Parameter Upcycling for Mixture of Experts

August 15, 2024

Auteurs: Qizhen Zhang, Nikolas Gritsch, Dwaraknath Gnaneshwar, Simon Guo, David Cairuz, Bharat Venkitesh, Jakob Foerster, Phil Blunsom, Sebastian Ruder, Ahmet Ustun, Acyr Locatelli

cs.AI

Résumé

Le cadre du Mixture of Experts (MoE) est devenu une architecture populaire pour les grands modèles de langage en raison de ses performances supérieures par rapport aux modèles denses. Cependant, l'entraînement des MoE à partir de zéro à grande échelle est prohibitivement coûteux. Les méthodes existantes atténuent ce problème en pré-entraînant plusieurs modèles denses experts de manière indépendante et en les utilisant pour initialiser un MoE. Cela se fait en utilisant le réseau feed-forward (FFN) des experts pour initialiser les experts du MoE tout en fusionnant les autres paramètres. Cependant, cette méthode limite la réutilisation des paramètres des modèles denses aux seules couches FFN, restreignant ainsi les avantages lors de la "revalorisation" de ces modèles en MoE. Nous proposons BAM (Branch-Attend-Mix), une méthode simple mais efficace qui remédie à cette lacune. BAM tire pleinement parti des modèles denses spécialisés en utilisant non seulement leur FFN pour initialiser les couches du MoE, mais en exploitant également pleinement les paramètres d'attention des experts en les initialisant dans une variante douce des couches Mixture of Attention (MoA). Nous explorons deux méthodes pour revaloriser les paramètres d'attention : 1) initialiser des experts d'attention séparés à partir de modèles denses en incluant tous les paramètres d'attention pour une performance optimale du modèle ; et 2) partager les paramètres de clé et de valeur entre tous les experts pour favoriser une meilleure efficacité lors de l'inférence. Pour améliorer encore l'efficacité, nous adoptons une architecture de transformateur à attention parallèle pour les MoE, ce qui permet aux experts d'attention et aux experts FFN d'être calculés simultanément. Nos expériences sur des modèles de départ allant de 590 millions à 2 milliards de paramètres démontrent que BAM surpasse les méthodes de référence à la fois en termes de perplexité et de performance sur les tâches en aval, dans les mêmes contraintes de calcul et de données.

English

The Mixture of Experts (MoE) framework has become a popular architecture for large language models due to its superior performance over dense models. However, training MoEs from scratch in a large-scale regime is prohibitively expensive. Existing methods mitigate this by pre-training multiple dense expert models independently and using them to initialize an MoE. This is done by using experts' feed-forward network (FFN) to initialize the MoE's experts while merging other parameters. However, this method limits the reuse of dense model parameters to only the FFN layers, thereby constraining the advantages when "upcycling" these models into MoEs. We propose BAM (Branch-Attend-Mix), a simple yet effective method that addresses this shortcoming. BAM makes full use of specialized dense models by not only using their FFN to initialize the MoE layers but also leveraging experts' attention parameters fully by initializing them into a soft-variant of Mixture of Attention (MoA) layers. We explore two methods for upcycling attention parameters: 1) initializing separate attention experts from dense models including all attention parameters for the best model performance; and 2) sharing key and value parameters across all experts to facilitate for better inference efficiency. To further improve efficiency, we adopt a parallel attention transformer architecture to MoEs, which allows the attention experts and FFN experts to be computed concurrently. Our experiments on seed models ranging from 590 million to 2 billion parameters demonstrate that BAM surpasses baselines in both perplexity and downstream task performance, within the same computational and data constraints.