MoMa: Entrenamiento previo de fusión temprana eficiente con mezcla de expertos conscientes de la modalidad

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

Resumen

Presentamos MoMa, una arquitectura novedosa de mezcla de expertos (MoE) consciente de la modalidad, diseñada para el preentrenamiento de modelos de lenguaje de fusión temprana mixtos. MoMa procesa imágenes y texto en secuencias arbitrarias dividiendo los módulos expertos en grupos específicos de modalidad. Estos grupos procesan exclusivamente tokens designados, empleando enrutamiento aprendido dentro de cada grupo para mantener una adaptabilidad semánticamente informada. Nuestros resultados empíricos revelan ganancias sustanciales en eficiencia de preentrenamiento a través de esta asignación de parámetros específica de modalidad. Bajo un presupuesto de entrenamiento de 1 billón de tokens, el modelo MoMa 1.4B, con 4 expertos en texto y 4 expertos en imagen, logra impresionantes ahorros de FLOPs: en general, 3.7 veces, con 2.6 veces para texto y 5.2 veces para procesamiento de imagen en comparación con una línea base densa equivalente en cómputo, medido por la pérdida de preentrenamiento. Esto supera al estándar MoE de elección de expertos con 8 expertos mixtos modales, que logra ahorros de FLOPs en general de 3 veces (3 veces para texto, 2.8 veces para imagen). La combinación de MoMa con mezcla de profundidades (MoD) mejora aún más los ahorros de FLOPs en preentrenamiento a 4.2 veces en general (texto: 3.4 veces, imagen: 5.3 veces), aunque esta combinación afecta el rendimiento en inferencia causal debido a una mayor sensibilidad a la precisión del enrutador. Estos resultados demuestran el potencial de MoMa para avanzar significativamente en la eficiencia del preentrenamiento de modelos de lenguaje de fusión temprana mixtos, allanando el camino para sistemas de IA multimodales más eficientes en recursos y capaces.

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.