ChatPaper.aiChatPaper

¡ZAS! Así de simple y eficiente: Reciclaje de parámetros para una mezcla de expertos.

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

August 15, 2024
Autores: Qizhen Zhang, Nikolas Gritsch, Dwaraknath Gnaneshwar, Simon Guo, David Cairuz, Bharat Venkitesh, Jakob Foerster, Phil Blunsom, Sebastian Ruder, Ahmet Ustun, Acyr Locatelli
cs.AI

Resumen

El marco de Mezcla de Expertos (MoE) se ha convertido en una arquitectura popular para modelos de lenguaje grandes debido a su rendimiento superior sobre modelos densos. Sin embargo, entrenar MoEs desde cero en un régimen a gran escala es prohibitivamente costoso. Los métodos existentes mitigan esto pre-entrenando múltiples modelos expertos densos de forma independiente y utilizando su inicialización para un MoE. Esto se logra utilizando la red feed-forward (FFN) de los expertos para inicializar los expertos del MoE mientras se fusionan otros parámetros. Sin embargo, este método limita la reutilización de los parámetros del modelo denso solo a las capas FFN, lo que restringe las ventajas al "reciclar" estos modelos en MoEs. Proponemos BAM (Rama-Atender-Mezclar), un método simple pero efectivo que aborda esta deficiencia. BAM aprovecha al máximo los modelos densos especializados no solo utilizando sus capas FFN para inicializar las capas del MoE, sino también aprovechando plenamente los parámetros de atención de los expertos al inicializarlos en una variante suave de capas de Mezcla de Atención (MoA). Exploramos dos métodos para reciclar los parámetros de atención: 1) inicializando expertos de atención separados a partir de modelos densos que incluyen todos los parámetros de atención para obtener el mejor rendimiento del modelo; y 2) compartiendo parámetros clave y de valor entre todos los expertos para facilitar una mejor eficiencia de inferencia. Para mejorar aún más la eficiencia, adoptamos una arquitectura de transformador de atención paralela a MoEs, que permite que los expertos de atención y los expertos de FFN se calculen de forma concurrente. Nuestros experimentos en modelos base que van desde 590 millones hasta 2 mil millones de parámetros demuestran que BAM supera a las líneas de base tanto en perplejidad como en el rendimiento de tareas secundarias, dentro de las mismas limitaciones computacionales y de datos.
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.
PDF133November 26, 2024