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

📄 arXiv: 2407.21770v3 📥 PDF

作者: Xi Victoria Lin, Akshat Shrivastava, Liang Luo, Srinivasan Iyer, Mike Lewis, Gargi Ghosh, Luke Zettlemoyer, Armen Aghajanyan

分类: cs.AI, cs.LG

发布日期: 2024-07-31 (更新: 2024-08-12)

备注: v2 -> update related work section v3 -> fix spelling

💡 一句话要点

MoMa：通过模态感知专家混合加速多模态早期融合预训练。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 多模态学习 早期融合 专家混合 模型预训练 计算效率 模态感知 Transformer

📋 核心要点

  1. 现有混合模态早期融合模型训练成本高昂，参数量大，难以有效扩展到更大规模的数据集。
  2. MoMa通过引入模态感知的专家混合机制，将专家模块划分为模态特定组，实现参数的有效分配和利用。
  3. 实验表明，MoMa在预训练效率上显著优于传统MoE方法，在相同训练预算下，FLOPs节省高达3.7倍。

📝 摘要（中文）

本文提出了一种新颖的模态感知专家混合（MoE）架构MoMa，专为预训练混合模态、早期融合语言模型而设计。MoMa通过将专家模块划分为模态特定的组来处理任意序列的图像和文本。这些组专门处理指定的token，同时在每个组内采用学习到的路由来保持语义感知的自适应性。实验结果表明，通过这种模态特定的参数分配，可以显著提高预训练效率。在1万亿token的训练预算下，具有4个文本专家和4个图像专家的MoMa 1.4B模型实现了令人印象深刻的FLOPs节省：总体节省3.7倍，其中文本处理节省2.6倍，图像处理节省5.2倍，优于具有8个混合模态专家的标准专家选择MoE（总体节省3倍）。将MoMa与深度混合（MoD）相结合，进一步将预训练FLOPs节省提高到总体4.2倍。这些结果表明，MoMa有潜力显著提高混合模态、早期融合语言模型预训练的效率，为更具资源效率和更强大的多模态AI系统铺平道路。

🔬 方法详解

问题定义：论文旨在解决多模态早期融合语言模型预训练效率低下的问题。现有方法，如密集模型和标准MoE模型，在处理混合模态数据时，参数利用率不高，计算成本巨大，难以扩展到更大规模的数据和模型。

核心思路：论文的核心思路是利用模态信息，设计模态感知的专家混合（MoE）架构。通过将专家划分为模态特定的组，每个组只处理特定模态的token，从而实现参数的有效分配和利用，降低计算成本。

技术框架：MoMa模型的整体架构基于Transformer，但在Transformer层中引入了模态感知的MoE层。该MoE层包含多个专家模块，这些专家模块被划分为模态特定的组（例如，文本专家组和图像专家组）。输入token首先被路由到相应的模态专家组，然后在该组内进行进一步的路由选择，最终由选定的专家处理。

关键创新：MoMa的关键创新在于其模态感知的专家混合机制。与标准MoE模型不同，MoMa的专家是模态特定的，这意味着每个专家只处理特定模态的token。这种设计可以更好地利用模态信息，提高参数利用率，降低计算成本。

关键设计：MoMa的关键设计包括：1) 模态特定专家组的划分；2) 组内路由机制，用于在每个模态专家组内选择合适的专家；3) 损失函数的设计，用于优化路由策略和专家参数。论文还探索了将MoMa与深度混合（MoD）相结合的方法，以进一步提高预训练效率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

MoMa 1.4B模型在1万亿token的训练预算下，实现了3.7倍的FLOPs节省，其中文本处理节省2.6倍，图像处理节省5.2倍，优于标准专家选择MoE（总体节省3倍）。与深度混合（MoD）相结合，进一步将预训练FLOPs节省提高到总体4.2倍。这些结果表明MoMa在预训练效率方面具有显著优势。

🎯 应用场景

MoMa的潜在应用领域包括多模态内容理解、图像描述生成、视觉问答、跨模态检索等。该研究可以降低多模态模型训练的成本，使其更容易部署在资源受限的环境中，并促进更强大的多模态AI系统的发展。未来，MoMa可以应用于更多模态的融合，例如视频、音频等。

📄 摘要（原文）

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.