BAM! Proprio così: Riciclo semplice ed efficiente dei parametri per le Miscele di Esperti

Abstract

Il framework Mixture of Experts (MoE) è diventato un'architettura popolare per i grandi modelli linguistici grazie alle sue prestazioni superiori rispetto ai modelli densi. Tuttavia, addestrare MoE da zero su larga scala è proibitivamente costoso. I metodi esistenti mitigano questo problema pre-addestrando in modo indipendente più modelli densi di esperti e utilizzandoli per inizializzare un MoE. Questo viene fatto utilizzando la rete feed-forward (FFN) degli esperti per inizializzare gli esperti del MoE, mentre si fondono altri parametri. Tuttavia, questo metodo limita il riutilizzo dei parametri dei modelli densi solo agli strati FFN, limitando così i vantaggi quando si "riciclano" questi modelli in MoE. Proponiamo BAM (Branch-Attend-Mix), un metodo semplice ma efficace che affronta questa carenza. BAM sfrutta appieno i modelli densi specializzati non solo utilizzando la loro FFN per inizializzare gli strati MoE, ma anche sfruttando completamente i parametri di attenzione degli esperti inizializzandoli in una variante soft degli strati Mixture of Attention (MoA). Esploriamo due metodi per riciclare i parametri di attenzione: 1) inizializzare esperti di attenzione separati dai modelli densi includendo tutti i parametri di attenzione per ottenere le migliori prestazioni del modello; e 2) condividere i parametri di chiave e valore tra tutti gli esperti per favorire una migliore efficienza inferenziale. Per migliorare ulteriormente l'efficienza, adottiamo un'architettura transformer con attenzione parallela per i MoE, che consente di calcolare contemporaneamente gli esperti di attenzione e gli esperti FFN. I nostri esperimenti su modelli di partenza che vanno da 590 milioni a 2 miliardi di parametri dimostrano che BAM supera i baseline sia in termini di perplessità che di prestazioni su task downstream, mantenendo gli stessi vincoli computazionali e di dati.

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.

BAM! Proprio così: Riciclo semplice ed efficiente dei parametri per le Miscele di Esperti

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

Abstract

Support