BAM! Assim Mesmo: Reciclagem Simples e Eficiente de Parâmetros para Mistura de Especialistas

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

Resumo

O framework Mixture of Experts (MoE) tornou-se uma arquitetura popular para grandes modelos de linguagem devido ao seu desempenho superior em relação aos modelos densos. No entanto, treinar MoEs do zero em um regime em larga escala é proibitivamente caro. Métodos existentes mitigam isso pré-treinando vários modelos especialistas densos de forma independente e utilizando-os para inicializar um MoE. Isso é feito usando a rede feed-forward (FFN) dos especialistas para inicializar os especialistas do MoE enquanto mesclam outros parâmetros. No entanto, esse método limita a reutilização dos parâmetros do modelo denso apenas para as camadas FFN, restringindo assim as vantagens ao "reciclar" esses modelos em MoEs. Propomos o BAM (Branch-Attend-Mix), um método simples, porém eficaz, que aborda essa deficiência. O BAM aproveita ao máximo os modelos densos especializados não apenas usando suas camadas FFN para inicializar as camadas do MoE, mas também alavancando totalmente os parâmetros de atenção dos especialistas, inicializando-os em uma variante suave das camadas Mixture of Attention (MoA). Exploramos dois métodos para reciclar os parâmetros de atenção: 1) inicializando especialistas de atenção separados a partir de modelos densos, incluindo todos os parâmetros de atenção para obter o melhor desempenho do modelo; e 2) compartilhando parâmetros de chave e valor entre todos os especialistas para facilitar uma melhor eficiência de inferência. Para melhorar ainda mais a eficiência, adotamos uma arquitetura de transformador de atenção paralela para MoEs, que permite que os especialistas de atenção e especialistas FFN sejam calculados simultaneamente. Nossos experimentos em modelos iniciais variando de 590 milhões a 2 bilhões de parâmetros demonstram que o BAM supera as linhas de base tanto em perplexidade quanto em desempenho de tarefas subsequentes, dentro das mesmas restrições computacionais e de dados.

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.