MoMa: Effizientes Frühfusion-Vortraining mit einer Mischung von modalitätssensiblen Experten

MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts

July 31, 2024

Autoren: Xi Victoria Lin, Akshat Shrivastava, Liang Luo, Srinivasan Iyer, Mike Lewis, Gargi Gosh, Luke Zettlemoyer, Armen Aghajanyan

cs.AI

Zusammenfassung

Wir stellen MoMa vor, eine neuartige Modality-aware Mixture-of-Experts (MoE)-Architektur, die für das Pre-Training von gemischt-modalen, frühfusionierten Sprachmodellen entwickelt wurde. MoMa verarbeitet Bilder und Text in beliebigen Sequenzen, indem Expertenmodule in modalitätsspezifische Gruppen unterteilt werden. Diese Gruppen verarbeiten ausschließlich zugewiesene Tokens und nutzen erlerntes Routing innerhalb jeder Gruppe, um semantisch informierte Anpassungsfähigkeit zu gewährleisten. Unsere empirischen Ergebnisse zeigen signifikante Effizienzgewinne beim Pre-Training durch diese modalitätsspezifische Parameterzuweisung. Unter einem Trainingsbudget von 1 Billion Tokens erreicht das MoMa 1.4B-Modell mit 4 Textexperten und 4 Bildexperten beeindruckende FLOPs-Einsparungen: Insgesamt 3,7-fach, mit 2,6-fach für Text und 5,2-fach für die Bildverarbeitung im Vergleich zu einem rechenäquivalenten dichten Baseline, gemessen am Pre-Training-Verlust. Dies übertrifft das Standard-Expertenauswahl-MoE mit 8 gemischt-modalen Experten, das 3-fache FLOPs-Einsparungen insgesamt erreicht (3-fach für Text, 2,8-fach für Bild). Die Kombination von MoMa mit Mixture-of-Depths (MoD) verbessert die FLOPs-Einsparungen beim Pre-Training auf insgesamt 4,2-fach (Text: 3,4-fach, Bild: 5,3-fach), obwohl diese Kombination die Leistung bei kausaler Inferenz aufgrund erhöhter Empfindlichkeit gegenüber der Routergenauigkeit beeinträchtigt. Diese Ergebnisse zeigen das Potenzial von MoMa auf, die Effizienz des Pre-Trainings von gemischt-modalen, frühfusionierten Sprachmodellen signifikant voranzutreiben und den Weg für ressourceneffizientere und leistungsfähigere multimodale KI-Systeme zu ebnen.

English

We introduce MoMa, a novel modality-aware mixture-of-experts (MoE) architecture designed for pre-training mixed-modal, early-fusion language models. MoMa processes images and text in arbitrary sequences by dividing expert modules into modality-specific groups. These groups exclusively process designated tokens while employing learned routing within each group to maintain semantically informed adaptivity. Our empirical results reveal substantial pre-training efficiency gains through this modality-specific parameter allocation. Under a 1-trillion-token training budget, the MoMa 1.4B model, featuring 4 text experts and 4 image experts, achieves impressive FLOPs savings: 3.7x overall, with 2.6x for text and 5.2x for image processing compared to a compute-equivalent dense baseline, measured by pre-training loss. This outperforms the standard expert-choice MoE with 8 mixed-modal experts, which achieves 3x overall FLOPs savings (3x for text, 2.8x for image). Combining MoMa with mixture-of-depths (MoD) further improves pre-training FLOPs savings to 4.2x overall (text: 3.4x, image: 5.3x), although this combination hurts performance in causal inference due to increased sensitivity to router accuracy. These results demonstrate MoMa's potential to significantly advance the efficiency of mixed-modal, early-fusion language model pre-training, paving the way for more resource-efficient and capable multimodal AI systems.