[論文レビュー] Multi-Head Mixture-of-Experts

Xun Wu, Shaohan Huang+2 more|arXiv (Cornell University)|Apr 23, 2024

Speech and dialogue systems被引用数 5

ひとこと要約

MH-MoEは、複数の専門家へサブトークンをルーティングするマルチヘッドトークン分割機構を導入し、より密な活性化とより細かな理解を実現する。追加コストなしに、言語・マルチモーダルタスク全体でSMoEベースラインより性能を向上させる。

ABSTRACT

Sparse Mixtures of Experts (SMoE) scales model capacity without significant increases in training and inference costs, but exhibits the following two issues: (1) Low expert activation, where only a small subset of experts are activated for optimization. (2) Lacking fine-grained analytical capabilities for multiple semantic concepts within individual tokens. We propose Multi-Head Mixture-of-Experts (MH-MoE), which employs a multi-head mechanism to split each token into multiple sub-tokens. These sub-tokens are then assigned to and processed by a diverse set of experts in parallel, and seamlessly reintegrated into the original token form. The multi-head mechanism enables the model to collectively attend to information from various representation spaces within different experts, while significantly enhances expert activation, thus deepens context understanding and alleviate overfitting. Moreover, our MH-MoE is straightforward to implement and decouples from other SMoE optimization methods, making it easy to integrate with other SMoE models for enhanced performance. Extensive experimental results across three tasks: English-focused language modeling, Multi-lingual language modeling and Masked multi-modality modeling tasks, demonstrate the effectiveness of MH-MoE.

研究の動機と目的

• 少数の専門家が使われると容量を無駄にする Sparse Mixtures of Experts (SMoE) における低い専門家活性化への対処。
• サブトークンを複数の専門家に分配して、トークン内のより細かな意味理解を可能にする。
• 実際に活用される専門家の数を増やしつつ、計算量とパラメータ効率を維持する。
• 英語中心、マルチリンガル、及びマスクドマルチモーダルモデリングタスクでの有効性を示す。

提案手法

• 入力トークンを射影するマルチヘッド層を適用し、各トークンを h 個のサブトークンに分割する。
• ゲーティング機構を用いてサブトークンを上位k個の活性化された専門家へルーティングし、より密な活性化を実現する。
• TSM（Token-Splitting-Merging）演算とマージレイヤーを介してサブトークン出力を統合し、後続層に追加コストをかけずに最終トークン表現を生成する。
• 専門家の使用偏りを緩和するため、ロードバランシング項を用いて訓練損失を調整する。
• 既存のSMoEフレームワークとの互換性を維持し、X-MoEのようなバックボーンに最小限の変更で実装する。

実験結果

リサーチクエスチョン

• RQ1 計算コストを増やすことなく、MH-MoE は標準の SMoE よりも密な専門家活性化を達成できるか。
• RQ2 マルチヘッドトークン分割によって、言語とモダリティを跨ぐより細かな意味理解が可能になるか。
• RQ3 Dense および X-MoE のベースラインと比較して、英語中心、マルチリンガル、マスク付きマルチモーダルタスクで MH-MoE はどのように性能を発揮するか。
• RQ4 ヘッド数とMLP/TSM コンポーネントの数が性能と活性化パターンに与える影響は何か。

主な発見

• MH-MoE は、SMoE と比較して専門家活性化を大幅に増加させる（例として、いくつかの設定で最大 90.71% の活性化）。
• サブトークンを多様な専門家へ分配することで、トークン理解をより細分化し、表現学習を改善する。
• MH-MoE は英語中心、マルチリンガル、マスク付きマルチモーダルタスクで Dense および X-MoE ベースラインを上回り、さまざまな設定で困難度（perplexity）が低く、下流精度が高い。
• ヘッド数を最適な範囲（概ね 4–6）まで増やすと性能が向上するが、それを超えると意味内容が希薄になる可能性がある。
• Token-Splitting-Merging (TSM) と MLP 層の組み合わせが、単純な TS や MLP のみではなく、実質的な利得を得るために必要である。
• MH-MoE はより良いスケーラビリティを示し、密な活性化と有用な専門家活用の上限を高める（下流タスクで最大 256 専門家）。
• ビジョン-言語タスク（VQA、NLVR2、COCO Captioning）全般で、MH-MoE は一貫して X-MoE および Dense のベースラインを上回る。