[論文レビュー] Gating Dropout: Communication-efficient Regularization for Sparsely Activated Transformers

Rui Liu, Young Jin Kim+2 more|arXiv (Cornell University)|May 28, 2022

Topic Modeling被引用数 11

ひとこと要約

本稿では、特にMixture-of-Experts (MoE) モデルを対象として、通信効率の良い正則化手法であるGating Dropoutを提案する。トレーニング中にトークンのリモートエキスパートへのルーティングをランダムにスキップすることで、高コストなマシン間全対全通信を低減しつつ、一般化性能を向上させる正則化効果をもたらし、収束が速くなり、多言語翻訳タスクで最大0.6 BLEUスコアの向上を達成する。

ABSTRACT

Sparsely activated transformers, such as Mixture of Experts (MoE), have received great interest due to their outrageous scaling capability which enables dramatical increases in model size without significant increases in computational cost. To achieve this, MoE models replace the feedforward sub-layer with Mixture-of-Experts sub-layer in transformers and use a gating network to route each token to its assigned experts. Since the common practice for efficient training of such models requires distributing experts and tokens across different machines, this routing strategy often incurs huge cross-machine communication cost because tokens and their assigned experts likely reside in different machines. In this paper, we propose \emph{Gating Dropout}, which allows tokens to ignore the gating network and stay at their local machines, thus reducing the cross-machine communication. Similar to traditional dropout, we also show that Gating Dropout has a regularization effect during training, resulting in improved generalization performance. We validate the effectiveness of Gating Dropout on multilingual machine translation tasks. Our results demonstrate that Gating Dropout improves a state-of-the-art MoE model with faster wall-clock time convergence rates and better BLEU scores for a variety of model sizes and datasets.

研究の動機と目的

• Mixture-of-Experts (MoE) モデルの分散学習における高いマシン間通信コストを軽減すること。これは、リモートエキスパートにトークンをルーティングする頻繁な全対全操作に起因する。
• 特に大規模な分散学習環境において、モデル性能を損なわず通信オーバーヘッドを低減すること。
• トレーニング中に局所エキスパートの利用を強制することで、一般化性能を向上させる正則化メカニズムを導入すること。
• MoEベースの多言語機械翻訳モデルにおける学習効率と収束速度を向上させること。

提案手法

• Gating Dropoutはトレーニング中にルーティング機構をランダムに無効化し、トークンが割り当てられたリモートエキスパートに送信されず、代わりにローカルマシンで処理されるようにする。
• この手法は、ゲーティングネットワークのルーティング意思決定にドロップアウト率pを適用する。確率pでトークンはルーティングされず、ローカルで処理される。
• 全対全通信操作のランダムなスキップにより、マシン間データシャッフルが著しく低減され、通信コストが低下する。
• 正則化効果は、最適なルーティングにアクセスできなくなるため、エキスパートがより汎用的な特徴を学習する必要が生じることに起因する。これにより一般化性能が向上する。
• 本手法はk=1およびk>1の両方のルーティング戦略と互換性があり、任意のMoEベースのTransformerアーキテクチャに適用可能である。
• 本手法はトレーニング時のみに適用され、推論時には無効化されるため、モデルの精度が保持される。

実験結果

リサーチクエスチョン

• RQ1 ルーティング操作のランダムなスキップは、モデル性能を劣化させることなく、分散MoE学習における通信コストを低減できるか？
• RQ2 Gating Dropoutによるルーティングのランダマイゼーションは、MoEモデルにおける一般化性能の向上をもたらすか？
• RQ3 ドロップアウト率は、MoE学習における通信効率、収束速度、モデル精度のトレードオフにどのように影響するか？
• RQ4 Gating Dropoutは多言語翻訳ベンチマークにおける学習スループットと収束速度を向上させることができるか？
• RQ5 Gating Dropoutは低リソース言語翻訳タスクにおいて特に利点をもたらすか？

主な発見

• Gating Dropoutは、ベースラインのMoEモデルと比較して、目標BLEUスコアに到達するまでの学習時間を最大58%短縮する。
• 本手法は一般化性能を向上させ、多言語翻訳タスクで最大0.6 BLEUスコアの向上を達成する。
• 全対全通信が削減されたため、スループットが最大16%向上する。特に高いドロップアウト率で顕著である。
• 最適なパフォーマンスは、Gate-Expert-Dropでドロップアウト率0.2、Gate-Dropで0.3で達成され、それ以上の率では性能が低下する。
• 正則化効果は特に低リソース言語翻訳において顕著で、少なくとも0.3 BLEUポイントの向上をもたらす。
• 本手法はさまざまなモデルサイズとデータセットで有効であり、収束性と効率性の両面で一貫した改善を示す。