[论文解读] MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts

Xi Lin, Akshat Shrivastava+6 more|arXiv (Cornell University)|Jul 31, 2024

Robotics and Automated Systems被引用 4

一句话总结

MoMa 提出了一种模态感知的专家混合（MoE）架构，通过按模态（文本或图像）分组专家，实现模态特定的参数分配和组内学习路由，从而高效预训练早期融合的多模态语言模型。在 1 万亿 token 的预算下，MoMa 相较于密集基线模型实现了 3.7 倍的整体 FLOPs 节省，图像处理部分节省达 5.2 倍，优于使用混合模态专家的标准 MoE 架构。

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.

研究动机与目标

• 通过利用模态特定的稀疏性，提升早期融合多模态语言模型预训练效率。
• 解决当多模态基础模型在参数量和数据量上扩展时，密集模型存在的计算低效问题。
• 探索模态感知稀疏性（MaS），在保持跨模态融合的同时，优化模态特定的特征学习。
• 评估将模态感知 MoE 与深度混合（MoD）结合对 FLOPs 节省和推理性能的影响。
• 证明模态感知稀疏性能够保持经验缩放律，从而实现可扩展且高效的多模态预训练。

提出的方法

• 模型采用基于 Chameleon 的 Transformer 架构，使用统一的文本与图像 token 序列。
• 专家被划分为模态特定的组别：4 个文本专家和 4 个图像专家，每组仅处理其指定模态的 token。
• 在每组模态内，采用学习路由机制为每个 token 选择活跃专家，实现语义感知的自适应，同时保持模态特定的参数分配。
• 路由机制采用专家选择路由（expert-choice routing），以平衡负载并保留静态计算图，从而实现高训练吞吐量。
• 该方法将模态感知 MoE 与深度混合（MoD）结合，使 token 可选择性跳过某些层，以减少计算量。
• 模型在受 FLOPs 控制的训练设置下进行训练，采用 1 万亿 token 的预算，将 MoMa 与密集基线及使用混合模态专家的标准 MoE 进行对比。

Figure 1 : Overview of our proposed multimodal early-fusion architecture. (a) The early-fusion architecture processes interleaved text and image data as a unified token sequence using a transformer model. (b) Detailed architecture of the mixture of modality-aware experts (MoMa) transformer block, wh

实验结果

研究问题

• RQ1 模态感知专家分组是否能在不损失性能的前提下提升早期融合多模态模型的预训练效率？
• RQ2 与混合模态专家路由相比，模态特定的参数分配在 FLOPs 节省和训练稳定性方面表现如何？
• RQ3 将模态感知 MoE 与深度混合（MoD）结合，对整体 FLOPs 减少和推理准确率有何影响？
• RQ4 模态感知稀疏性是否保持了密集模型中观察到的经验缩放律？
• RQ5 MoD 增强模型的性能对因果推理场景中路由准确性的敏感度如何？

主要发现

• 在 1 万亿 token 的训练预算下，MoMa 1.4B 模型相较计算量匹配的密集基线，实现了 3.7 倍的整体 FLOPs 节省，其中文本处理节省 2.6 倍，图像处理节省 5.2 倍。
• MoMa 模型优于使用 8 个混合模态专家的标准专家选择 MoE，后者仅实现 3 倍的整体 FLOPs 节省（文本 3 倍，图像 2.8 倍）。
• 将 MoMa 与深度混合（MoD）结合后，FLOPs 节省提升至 4.2 倍整体（文本 3.4 倍，图像 5.3 倍），显示出进一步的效率增益。
• 尽管 FLOPs 节省更高，MoMaD 模型在因果推理任务中性能下降，归因于对路由准确性的更高敏感度。
• MoMa 架构保持了稳定的实证缩放律，表明其具备面向未来大规模多模态预训练的可扩展性。
• 一种无需模态绑定的再利用技术（modality-untied upcycling）进一步提升了性能，表明架构优化具有潜在改进空间。

Figure 2 : Architecture of transformer layer consisting of MoMa combined with mixture-of-depths (MoD).