Zack! Ganz einfach: Einfaches und effizientes Parameter-Upcycling für ein Expertengemisch

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

August 15, 2024

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

cs.AI

Zusammenfassung

Das Mixture of Experts (MoE) Framework hat sich aufgrund seiner überlegenen Leistung gegenüber dichten Modellen zu einer beliebten Architektur für große Sprachmodelle entwickelt. Das Training von MoEs von Grund auf in einem groß angelegten Regime ist jedoch prohibitiv teuer. Bestehende Methoden mildern dies, indem sie mehrere dichte Expertenmodelle unabhängig voneinander vorab trainieren und sie zur Initialisierung eines MoE verwenden. Dies wird durch Verwendung des Feedforward-Netzwerks (FFN) der Experten zur Initialisierung der Experten des MoE durchgeführt, während andere Parameter fusioniert werden. Diese Methode beschränkt jedoch die Wiederverwendung von Parametern des dichten Modells nur auf die FFN-Schichten, was die Vorteile bei der "Aufwertung" dieser Modelle in MoEs einschränkt. Wir schlagen BAM (Branch-Attend-Mix) vor, eine einfache, aber effektive Methode, die diese Schwäche angeht. BAM nutzt spezialisierte dichte Modelle voll aus, indem nicht nur ihr FFN zur Initialisierung der MoE-Schichten verwendet wird, sondern auch die Aufmerksamkeitsparameter der Experten vollständig genutzt werden, indem sie in eine weiche Variante der Mixture of Attention (MoA) Schichten initialisiert werden. Wir untersuchen zwei Methoden zur Aufwertung von Aufmerksamkeitsparametern: 1) Initialisierung separater Aufmerksamkeitsexperten aus dichten Modellen, einschließlich aller Aufmerksamkeitsparameter für die beste Modellleistung; und 2) gemeinsame Initialisierung von Schlüssel- und Werteparametern über alle Experten hinweg, um eine bessere Inferenzeffizienz zu ermöglichen. Um die Effizienz weiter zu verbessern, übernehmen wir eine parallele Aufmerksamkeits-Transformer-Architektur für MoEs, die es ermöglicht, dass die Aufmerksamkeitsexperten und FFN-Experten gleichzeitig berechnet werden. Unsere Experimente mit Ausgangsmodellen von 590 Millionen bis 2 Milliarden Parametern zeigen, dass BAM die Baselines sowohl in Bezug auf Perplexität als auch auf die Leistung bei nachgelagerten Aufgaben innerhalb derselben Rechen- und Datenbeschränkungen übertrifft.

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.