[論文レビュー] Scaling Expert Language Models with Unsupervised Domain Discovery

Suchin Gururangan, Margaret Li+5 more|arXiv (Cornell University)|Mar 24, 2023

Topic Modeling被引用数 7

ひとこと要約

この論文では、文書コーパスを教師なしクラスタリングによってドメイン固有のグループに分類し、通信効率的でありながらも、ありとあらゆる並列処理が可能なスパースエキスパート言語モデルを訓練するためのc-BTMを提案する。エキスパートを学習済みクラスタに特化させることで、FLOPsを著しく削減しつつ、密度型モデルを上回る性能を達成する。クラスタ数が増えるほど、データ量が大きくなるほど、その恩恵は大きくなる。

ABSTRACT

Large language models are typically trained densely: all parameters are updated with respect to all inputs. This requires synchronization of billions of parameters across thousands of GPUs. We introduce a simple but effective method to asynchronously train large, sparse language models on arbitrary text corpora. Our method clusters a corpus into sets of related documents, trains a separate expert language model on each cluster, and combines them in a sparse ensemble for inference. This approach generalizes embarrassingly parallel training by automatically discovering the domains for each expert, and eliminates nearly all the communication overhead of existing sparse language models. Our technique outperforms dense baselines on multiple corpora and few-shot tasks, and our analysis shows that specializing experts to meaningful clusters is key to these gains. Performance also improves with the number of experts and size of training data, suggesting this is a highly efficient and accessible approach to training large language models.

研究の動機と目的

• 数千のGPUを用いた大規模な密度型言語モデルを訓練する際の、高い通信コストと同期コストを低減すること。
• メタデータを用いないで、スケーラブルで通信効率の良いスパースエキスパート言語モデルの訓練を可能にすること。
• 教師なしクラスタリングが、密度型ベースラインと比較して、モデル性能を向上させる意味のあるドメインを発見できるかどうかを調査すること。
• 自動的に発見されたクラスタに訓練されたエキスパートモデルのスケーリング特性、特にエキスパート数とデータサイズとの関係を調査すること。

提案手法

• メタデータを一切使用せずに、テキストコーパスに対してk-meansクラスタリングを適用し、ドメインに類似したグループを発見する。
• 事前学習済みベースモデル（例：OPT-1.3B）からエキスパート言語モデル（ELM）を初期化し、それぞれのELMを割り当てられたクラスタ上で独立して微調整する。
• 入力コンテキストの埋め込みと各エキスパートのクラスタ中心とのコサイン類似度に基づき、上位k個のエキスパートを選択するスパースアンサンブル推論機構を用いる。
• GPU間の同期的パラメータ更新を避けることで、相互通信を最小限に抑えた、ありとあらゆる並列処理が可能なエキスパートの訓練を行う。
• クラスタのバランスをとるために、オフラインクラスタリングを用いてトレーニング中にエキスパートの利用を均等に保ち、ルーティングのボトルネックを回避する。
• 推論時に上位k個のエキスパートの出力を重み付き平均で統合することで、効率的なスパース計算を実現する。

実験結果

リサーチクエスチョン

• RQ1 教師なしクラスタリングが、密度型訓練と比較して、言語モデル性能を向上させる意味のあるドメインを効果的に発見できるか？
• RQ2 クラスタ数（＝エキスパート数）を増やすと、モデル性能と訓練効率にどのような影響を与えるか？
• RQ3 より大きなデータセットにスケーリングした際、c-BTMは困惑度とFLOP効率の両面で密度型ベースラインを上回るか？
• RQ4 学習ルーティング機構を備えた既存のスパースMoEモデルと比較して、c-BTMの性能はどの程度か？
• RQ5 ランダムまたは非バランスクラスタリングと比較して、バランスの取れたクラスタリングがモデル性能に与える影響はどの程度か？

主な発見

• c-BTMは、計算量に匹敵する密度型言語モデルをバリデーションの困惑度で上回り、クラスタ数が増えるほどその差が大きくなる。
• 168Bトークンのデータを用いて、128個のエキスパートLM（合計168Bパラメータ）を、同時に8つのGPUで訓練することで、効果的な並列化とスケーラビリティを実証した。
• 128個のエキスパートを用い、上位4つのエキスパートを選択する推論を行うことで、すべてのエキスパートを使用する場合と同等の性能を達成し、上位1つの推論でも密度型モデルを上回った。
• 1.3Bパラメータのエキスパートとスパース推論を用いたc-BTMは、6.7Bパラメータの密度型モデルと同等の困惑度を達成したが、トレーニングにおけるFLOPsはわずか29%にまで削減された。
• バランスの取れたクラスタリングは、ゴールメタデータを用いたエキスパート割り当てと同等の性能を示し、ランダムまたは非バランスクラスタリングを著しく上回った。
• 最適なクラスタ数はトレーニングデータサイズに応じて増加し、クラスタ数とデータサイズが増えるほど性能が単調に向上した。