MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware
  Experts
Authors: name: Xi Victoria Lin; name: Akshat Shrivastava; name: Liang Luo; name: Srinivasan Iyer; name: Mike Lewis; name: Gargi Gosh; name: Luke Zettlemoyer; name: Armen Aghajanyan
Published: 2024-07-31 17:46:51

Abstract: 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.