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

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Lin, Xi Victoria, Shrivastava, Akshat, Luo, Liang, Iyer, Srinivasan, Lewis, Mike, Ghosh, Gargi, Zettlemoyer, Luke, Aghajanyan, Armen
Formato: Preprint
Publicado: 2024
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866917746800852992
author Lin, Xi Victoria
Shrivastava, Akshat
Luo, Liang
Iyer, Srinivasan
Lewis, Mike
Ghosh, Gargi
Zettlemoyer, Luke
Aghajanyan, Armen
author_facet Lin, Xi Victoria
Shrivastava, Akshat
Luo, Liang
Iyer, Srinivasan
Lewis, Mike
Ghosh, Gargi
Zettlemoyer, Luke
Aghajanyan, Armen
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2407_21770
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts
Lin, Xi Victoria
Shrivastava, Akshat
Luo, Liang
Iyer, Srinivasan
Lewis, Mike
Ghosh, Gargi
Zettlemoyer, Luke
Aghajanyan, Armen
Artificial Intelligence
Machine Learning
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.
title MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts
topic Artificial Intelligence
Machine Learning
url https://arxiv.org/abs/2407.21770