[論文レビュー] Branch-Train-Merge: Embarrassingly Parallel Training of Expert Language Models

Margaret Li, Suchin Gururangan+5 more|arXiv (Cornell University)|Aug 5, 2022

Topic Modeling被引用数 25

ひとこと要約

BTMは、領域特化データで複数の独立した専門家言語モデルを並行して訓練し、それらを統合または平均させて、サイズと領域を問わず計算量に適合したベースラインを上回るスケーラブルで効率的なアンサンブル/単一モデルを形成します。

ABSTRACT

We present Branch-Train-Merge (BTM), a communication-efficient algorithm for embarrassingly parallel training of large language models (LLMs). We show it is possible to independently train subparts of a new class of LLMs on different subsets of the data, eliminating the massive multi-node synchronization currently required to train LLMs. BTM learns a set of independent expert LMs (ELMs), each specialized to a different textual domain, such as scientific or legal text. These ELMs can be added and removed to update data coverage, ensembled to generalize to new domains, or averaged to collapse back to a single LM for efficient inference. New ELMs are learned by branching from (mixtures of) ELMs in the current set, further training the parameters on data for the new domain, and then merging the resulting model back into the set for future use. Experiments show that BTM improves in- and out-of-domain perplexities as compared to GPT-style Transformer LMs, when controlling for training cost. Through extensive analysis, we show that these results are robust to different ELM initialization schemes, but require expert domain specialization; LM ensembles with random data splits do not perform well. We also present a study of scaling BTM into a new corpus of 64 domains (192B whitespace-separated tokens in total); the resulting LM (22.4B total parameters) performs as well as a Transformer LM trained with 2.5 times more compute. These gains grow with the number of domains, suggesting more aggressive parallelism could be used to efficiently train larger models in future work.

研究の動機と目的

• 大規模言語モデルの分散間同期を減らすため、領域専門のエキスパートを独立して訓練するスケーラブルな訓練パラダイムの動機付け。
• Branch-Train-Merge (BTM) アルゴリズムを開発し、既存のエキスパートから新しいエキスパートをブランチし、領域特定データで訓練し、森に再統合し、融合を改善するためにプレトレーニング済みの seed LM でプロセスを開始。
• エンサンブルまたはパラメータ平均化による推論の柔軟性を実現し、性能と推論コストのバランスを取る。
• BTMを計算量が一致したTransformer-LMsおよび以前の領域特化ベースラインと比較評価し、複数のモデルスケールとデータ領域に渡って実証的に評価。
• 多くの領域へのスケーラビリティを探り、シード訓練、データ出自、初期化などが性能と効率に与える影響を分析。

提案手法

• ELMforestを、データドメインに特化し、訓練時に完全にデカップリングされた独立したエキスパートLMの集合として定義する。
• Branch-Train-Merge: 既存のエキスパートの加重平均から新しいELMを反復的に分岐し、新しいドメインで訓練し、森に統合する。融合を改善するために事前訓練済みの seed LMでプロセスを開始する。
• Ensemble推論は、ドメイン後方確率とともにELMの出力を組み合わせ、専門家ごとの重み付き和として p(X_t|x_<t) を計算する。希少性により上位-k の選択がコストを削減できる。
• パラメータ平均化による代替推論は、ELMのパラメータを重み付き平均することで単一モデルを作成し、ドメイン後方確率やその他の方式で重み付けする。
• 計算量が一致したベースラインと比較: Transformer-LM および DEMix; 125M–1.3B パラメータスケールと 8 トレーニング + 8 評価ドメインで perplexityを評価; シード予算と初期化効果を分析。

実験結果

リサーチクエスチョン

• RQ1 領域特化エキスパート(ELM)による embarrassingly parallel training は、計算量が一致した Transformer-LM や以前の領域ベースラインより、内域・外域の perplexity の双方で優れているか？
• RQ2 64 ドメインまで拡張する際の Branch-Train-Merge 手続きはスケーラビリティにどのように影響するか、より大きなパラメータ数で？
• RQ3 Seedフェーズ訓練、初期化、データ由来が ELM アンサンブルの有効性およびパラメータ平均化の性能に与える影響は？
• RQ4 ドメイン数が増えるほど、アンサンブルとパラメータ平均化は推論コストと性能の点でどう異なる？
• RQ5 ブランチEDLM訓練の更新頻度・通信コストなど、完全同期トレーニングと比較した効率性の影響は？

主な発見

モデル	125M 訓練	125M 評価	125M 全体	350M 訓練	350M 評価	350M 全体	750M 訓練	750M 評価	750M 全体	1.3B 訓練	1.3B 評価	1.3B 全体
Transformer-LM	19.9	25.2	22.5	16.3	20.8	18.5	14.7	19.3	17.0	14.2	18.4	16.3
DEMix	18.2	23.4	20.8	15.0	19.9	17.5	13.5	17.7	15.6	13.7	17.6	15.6
ELMforest	17.2	22.4	19.8	14.7	18.6	16.7	13.4	16.7	15.0	13.0	16.3	14.6

• BTMで訓練されたELMsは、125M, 350M, 750M, 1.3B パラメータの複数モデルスケールで計算量が一致したTransformer-LMsおよび以前の DEMix ベースラインを上回る。
• BTMは完全に同期したベースラインより更新頻度が高く、特に大型モデルでクロス-GPU通信が減るため効率が向上する。
• ドメイン由来ベースのドメイン専門化の利用が重要。ランダムデータ分割はドメイン中心の分割に比べて劣る。
• ELMフォレストのパラメータ平均化はアンサンブル性能に近づくが推論コストは追加されず、ただしドメインが増えるとアンサンブルの方が依然として強い。
• 64ドメインへ拡張すると、2.5倍の計算量で学習したTransformer-LMと同等の22.4Bパラメータのアンサンブル性能を得られ、ドメインが増えるほど利益は大きくなる。
• シード訓練は有効な平均化と堅牢な性能のために不可欠で、最適なシード予算は通常総計算量の約40–60%程度であり、広い範囲で頑健性を示す。