MoMa : Pré-entraînement efficace par fusion précoce avec un mélange d'experts conscients des modalités

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

July 31, 2024

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

cs.AI

Résumé

Nous présentons MoMa, une architecture novatrice de mélange d'experts (MoE) sensible aux modalités, conçue pour le pré-entraînement de modèles de langage à fusion précoce et multimodaux. MoMa traite les images et le texte dans des séquences arbitraires en divisant les modules experts en groupes spécifiques à chaque modalité. Ces groupes traitent exclusivement les tokens désignés tout en utilisant un routage appris au sein de chaque groupe pour maintenir une adaptabilité sémantiquement informée. Nos résultats empiriques révèlent des gains substantiels d'efficacité lors du pré-entraînement grâce à cette allocation de paramètres spécifique aux modalités. Avec un budget d'entraînement de 1 trillion de tokens, le modèle MoMa 1.4B, doté de 4 experts pour le texte et 4 experts pour les images, réalise des économies impressionnantes en termes de FLOPs : 3,7x globalement, avec 2,6x pour le texte et 5,2x pour le traitement des images par rapport à une base dense de calcul équivalent, mesuré par la perte lors du pré-entraînement. Cela surpasse le MoE standard avec choix d'experts utilisant 8 experts multimodaux, qui atteint des économies globales de FLOPs de 3x (3x pour le texte, 2,8x pour les images). La combinaison de MoMa avec le mélange de profondeurs (MoD) améliore encore les économies de FLOPs lors du pré-entraînement à 4,2x globalement (texte : 3,4x, images : 5,3x), bien que cette combinaison nuise aux performances en inférence causale en raison d'une sensibilité accrue à la précision du routeur. Ces résultats démontrent le potentiel de MoMa à faire progresser significativement l'efficacité du pré-entraînement des modèles de langage à fusion précoce et multimodaux, ouvrant la voie à des systèmes d'IA multimodaux plus efficaces en termes de ressources et plus performants.

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.