MoMa: Pre-addestramento Efficiente a Fusione Precoce con Miscela di Esperti Consapevoli della Modalità

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

July 31, 2024

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

cs.AI

Abstract

Presentiamo MoMa, una innovativa architettura modality-aware mixture-of-experts (MoE) progettata per il pre-training di modelli linguistici a fusione precoce e multimodali. MoMa elabora immagini e testo in sequenze arbitrarie suddividendo i moduli expert in gruppi specifici per modalità. Questi gruppi elaborano esclusivamente i token designati, impiegando un routing appreso all'interno di ciascun gruppo per mantenere un'adattività semanticamente informata. I nostri risultati empirici rivelano sostanziali guadagni di efficienza nel pre-training grazie a questa allocazione di parametri specifica per modalità. Con un budget di training di 1 trilione di token, il modello MoMa 1.4B, dotato di 4 expert per il testo e 4 per le immagini, ottiene impressionanti risparmi di FLOPs: 3.7x complessivi, con 2.6x per il testo e 5.2x per l'elaborazione delle immagini rispetto a una baseline densa equivalente in termini di calcolo, misurata dalla perdita di pre-training. Questo supera lo standard expert-choice MoE con 8 expert multimodali, che raggiunge risparmi complessivi di FLOPs di 3x (3x per il testo, 2.8x per le immagini). Combinare MoMa con mixture-of-depths (MoD) migliora ulteriormente i risparmi di FLOPs nel pre-training a 4.2x complessivi (testo: 3.4x, immagini: 5.3x), sebbene questa combinazione comprometta le prestazioni nell'inferenza causale a causa di una maggiore sensibilità all'accuratezza del router. Questi risultati dimostrano il potenziale di MoMa di far progredire significativamente l'efficienza del pre-training di modelli linguistici multimodali a fusione precoce, aprendo la strada a sistemi AI multimodali più efficienti in termini di risorse e capaci.

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.