[論文レビュー] 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兆トークンの予算下で、密度型ベースラインと比較して全体で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.

研究の動機と目的

• モダリティ固有のスパarsityを活用することで、エアリー・フェュージョンマルチモーダル言語モデルの事前学習の効率性を向上させること。
• 大規模なパラメータとデータにスケーリングされるマルチモーダル基盤モデルにおける密度型モデルの計算非効率性を是正すること。
• クロスモダリティ統合を維持しつつ、モダリティ固有の特徴学習を最適化するためのモダリティに配慮したスパarsity（MaS）を検討すること。
• モダリティに配慮したMoEと混合深さ（MoD）を組み合わせた際のFLOPs削減および推論性能への影響を評価すること。
• モダリティに配慮したスパarsityが、スケーラブルで効率的なマルチモーダル事前学習を可能にする実験的スケーリング法則を保持していることを示すこと。

提案手法

• テキストと画像トークンの統一シーケンスを扱う、Chameleonベースのトランスフォーマー・アーキテクチャを採用する。
• エキスパートをモダリティ固有のグループに分類する：4つのテキストエキスパートと4つの画像エキスパートで、それぞれが指定されたモダリティのトークンのみを処理する。
• 各モダリティグループ内で、トークンごとにアクティブなエキスパートを選択する学習済みルーティング機構を採用し、意味論に配慮した適応性を実現しながらも、モダリティ固有のパラメータ割り当てを維持する。
• 負荷バランスと高いトレーニングスループットを確保するため、エキスパート選択ルーティングを用い、静的計算グラフを保持する。
• モダリティに配慮したMoEと混合深さ（MoD）を統合し、特定のレイヤーをトークンが選択的にスキップできるようにする。
• 1兆トークンのFLOPs制御下でトレーニングを行い、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 モダリティに配慮したスパarsityは、密度型モデルで観察される実験的スケーリング法則を保持しているか？
• RQ5 因果推論シナリオにおけるルーターの正確性に、MoDを強化したモデルの性能はどれほど敏感か？

主な発見

• 1兆トークンのトレーニング予算下で、MoMa 1.4Bモデルは、計算量が同等の密度型ベースラインと比較して、全体で3.7倍のFLOPs削減を達成。テキスト処理では2.6倍、画像処理では5.2倍の削減を実現。
• 標準的なエキスパート選択MoE（8つの混合モダリティエキスパートを有する）は、全体で3倍のFLOPs削減（テキスト：3倍、画像：2.8倍）にとどまり、MoMaはこれを上回る。
• MoMaと混合深さ（MoD）を組み合わせることで、FLOPs削減は全体で4.2倍（テキスト：3.4倍、画像：5.3倍）に増加し、さらなる効率性向上が確認された。
• FLOPs削減率が高まっても、MoMaDモデルは因果推論タスクで性能が低下しており、ルーターの正確性に高い感受性を示している。
• MoMaアーキテクチャは安定した実験的スケーリング法則を維持しており、将来の大規模マルチモーダル事前学習へのスケーラビリティを示している。
• モダリティに束縛されないアップサイクリング技術により、さらに性能向上が達成された。これは、アーキテクチャ的改善の可能性を示唆している。

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