[論文レビュー] EvoMoE: An Evolutional Mixture-of-Experts Training Framework via Dense-To-Sparse Gate

Xiaonan Nie, Xupeng Miao+8 more|arXiv (Cornell University)|Dec 29, 2021

Expert finding and Q&A systems被引用数 7

ひとこと要約

EvoMoEは、エキスパート多様化段階とゲートスパース化段階を経て、単一エキスパートを段階的にスパースなMixture-of-Experts（MoE）モデルに進化させる二段階訓練フレームワークを提案する。Dense-to-Sparseゲート（DTS-Gate）を導入し、初期は密な状態から出発し、徐々にスパース化することで、訓練の安定性と性能を向上させ、GLUEベンチマークでベースライン比最高で0.562および0.403の向上を達成した。

ABSTRACT

Mixture-of-experts (MoE) is becoming popular due to its success in improving the model quality, especially in Transformers. By routing tokens with a sparse gate to a few experts (i.e., a small pieces of the full model), MoE can easily increase the model parameters to a very large scale while keeping the computation cost in a constant level. Most existing works just initialize some random experts, set a fixed gating strategy (e.g., Top-k), and train the model from scratch in an ad-hoc way. We identify that these MoE models are suffering from the immature experts and unstable sparse gate, which are harmful to the convergence performance. In this paper, we propose an efficient end-to-end MoE training framework called EvoMoE. EvoMoE starts from training one single expert and gradually evolves into a large and sparse MoE structure. EvoMoE mainly contains two phases: the expert-diversify phase to train the base expert for a while and spawn multiple diverse experts from it, and the gate-sparsify phase to learn an adaptive sparse gate and activate a dynamic number of experts. EvoMoE naturally decouples the joint learning of both the experts and the sparse gate and focuses on learning the basic knowledge with a single expert at the early training stage. Then it diversifies the experts and continues to train the MoE with a novel Dense-to-Sparse gate (DTS-Gate). Specifically, instead of using a permanent sparse gate, DTS-Gate begins as a dense gate that routes tokens to all experts, then gradually and adaptively becomes sparser while routes to fewer experts. Evaluations are conducted on three popular models and tasks, including RoBERTa for masked language modeling task, GPT for language modeling task and Transformer for machine translation task. The results show that EvoMoE outperforms existing baselines, including Switch, BASE Layer, Hash Layer and StableMoE.

研究の動機と目的

• ランダムに初期化されたエキスパートとスパースゲートによる既存のMoE訓練における不安定さと収束不良の問題を解消すること。
• エキスパートとゲートの学習を分離することで、訓練効率とモデル品質を向上させること。
• 適応的スパース化により、密ルーティングからスパースルーティングへの滑らかな移行を可能にすること。
• 既存のMoE手法と比較して、より優れた性能とFLOPs効率を達成すること。
• 大規模MoEモデル向けに、より安定的かつスケーラブルな訓練パラダイムを確立すること。

提案手法

• エキスパート多様化段階では、まず単一のベースエキスパートを訓練し、その後ランダムマスキングを用いて構造的多様性を導入することで複数のエキスパートに多様化させる。
• ゲートスパース化段階では、Dense-to-Sparseゲート（DTS-Gate）を導入し、初期段階ではすべてのエキスパートに対してソフトで密なルーティングを実施し、Gumbel-Softmaxにおける温度の減少に伴い徐々にスパース化を促進する。
• DTS-Gateはコンテンツベースのルーティング機構を採用し、ゲート重みが学習されたしきい値を超えるエキスパートのみを活性化することで、適応的かつ動的選択を可能にする。
• フレームワークはエキスパートとゲートの訓練を分離する：スパースルーティングを導入する前に、エキスパートを共有された形で事前学習する。
• Gumbel-Softmaxにおける温度の減少スケジュールを適用し、ゲート出力のスパース性を段階的に促進する。
• 本手法は、RoBERTa、GPT、Transformerなど複数のモデルにエンドツーエンドで適用され、多様なNLPタスクで検証された。

実験結果

リサーチクエスチョン

• RQ1 単一エキスパートからスパースなMoEアーキテクチャへの段階的進化は、訓練の安定性とモデル性能の向上に寄与するか？
• RQ2 Dense-to-Sparseゲート機構は、固定スパースゲート（例：Top-k）と比較して、より良好な収束性とルーティング精度を実現するか？
• RQ3 共有ベースエキスパートからのエキスパート多様化は、知識の一般化を向上させるとともに、訓練の不安定性を低減するか？
• RQ4 EvoMoEは、Switch、BASELayer、StableMoEといった既存のMoE手法と比較して、性能とFLOPs効率の両面で優れているか？
• RQ5 エキスパートとゲートの訓練を分離することで、MoEモデルの学習ダイナミクスはどの程度向上するか？

主な発見

• EvoMoEはGLUEベンチマークにおいて、ベースライン比平均スコアで最高で0.562および0.403の向上を達成し、優れたモデル品質を示した。
• DTS-Gate機構により、より滑らかなルーティング適応が可能となり、固定Top-kゲートと比較して初期化のランダム性への感受性が低減した。
• エキスパート多様化段階により、より安定的かつ効果的なエキスパート特化が実現され、全体のモデルの頑健性が向上した。
• EvoMoEは、密なMoEよりも優れたFLOPs効率を達成し、スパースなMoEベースラインと比較しても競争力のある性能を示した。
• マスク言語モデル、言語モデル、機械翻訳など、複数のタスクで一貫した性能向上が確認された。
• アブレーションスタディの結果、エキスパート多様化段階とゲートスパース化段階の両方が最適性能を達成するために不可欠であることが確認された。