MoMa: Pré-treinamento Eficiente de Fusão Antecipada com Mistura de Especialistas Conscientes de Modalidade

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

July 31, 2024

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

cs.AI

Resumo

Apresentamos o MoMa, uma arquitetura inovadora de mistura de especialistas (MoE) consciente de modalidade, projetada para pré-treinar modelos de linguagem de fusão precoce mistos-modais. O MoMa processa imagens e texto em sequências arbitrárias dividindo os módulos especialistas em grupos específicos de modalidade. Esses grupos processam exclusivamente tokens designados, empregando roteamento aprendido dentro de cada grupo para manter adaptabilidade semanticamente informada. Nossos resultados empíricos revelam ganhos substanciais de eficiência de pré-treinamento por meio dessa alocação de parâmetros específicos de modalidade. Sob um orçamento de treinamento de 1 trilhão de tokens, o modelo MoMa 1.4B, com 4 especialistas em texto e 4 especialistas em imagem, alcança economias impressionantes de FLOPs: 3,7x no geral, com 2,6x para texto e 5,2x para processamento de imagem em comparação com uma linha de base densa equivalente em computação, medida pela perda de pré-treinamento. Isso supera o MoE padrão de escolha de especialistas com 8 especialistas mistos-modais, que alcança economias de FLOPs de 3x no geral (3x para texto, 2,8x para imagem). Combinar o MoMa com mistura de profundidades (MoD) melhora ainda mais as economias de FLOPs de pré-treinamento para 4,2x no geral (texto: 3,4x, imagem: 5,3x), embora essa combinação prejudique o desempenho na inferência causal devido à sensibilidade aumentada à precisão do roteador. Esses resultados demonstram o potencial do MoMa para avançar significativamente a eficiência do pré-treinamento de modelos de linguagem de fusão precoce mistos-modais, abrindo caminho para sistemas de IA multimodais mais eficientes em recursos e capazes.

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.