[論文レビュー] Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM

Sainbayar Sukhbaatar, Olga Golovneva+9 more|arXiv (Cornell University)|Mar 12, 2024

Topic Modeling被引用数 5

ひとこと要約

BTX は seed モデルからドメイン特化の専門 LLMs を並列に訓練し、それらをトークンレベルのルーティングとファインチューニングを備えた統一 MoE モデルに統合することで、ベースラインより高い効率で強力な精度を達成します。

ABSTRACT

We investigate efficient methods for training Large Language Models (LLMs) to possess capabilities in multiple specialized domains, such as coding, math reasoning and world knowledge. Our method, named Branch-Train-MiX (BTX), starts from a seed model, which is branched to train experts in embarrassingly parallel fashion with high throughput and reduced communication cost. After individual experts are asynchronously trained, BTX brings together their feedforward parameters as experts in Mixture-of-Expert (MoE) layers and averages the remaining parameters, followed by an MoE-finetuning stage to learn token-level routing. BTX generalizes two special cases, the Branch-Train-Merge method, which does not have the MoE finetuning stage to learn routing, and sparse upcycling, which omits the stage of training experts asynchronously. Compared to alternative approaches, BTX achieves the best accuracy-efficiency tradeoff.

研究の動機と目的

• 複数のドメイン（数学、コード、世界知識）に対する seed LLM の性能を阻害的な同期コストなしに改善する。
• 完全に独立して並列に実行できる専門家トレーニングと MoE ファインチューニングを活用して、統一モデル内でドメイン特化を可能にする。
• 精度と計算効率の観点から、Dense Baselines、Branch-Train-Merge、Sparse upcycling と比較して BTX を評価する。
• MoE ファインチューニングがドメイン全体の性能にどう影響するかを理解するために、ルーティング戦略とアブレーションを調査する。

提案手法

• 事前学習済み seed LLM からブランチを作成し、ドメインデータセット上で N 個のドメイン特化型エキスパート LLM を並行してトレーニングする。
• エキスパートから構成されるフィードフォワード部分を用いて単一の MoE 層を形成し、他のモジュールを平均化することでドメインエキスパートを統合する。
• ルーターネットワークを介してトークンレベルのルーティングを学習するために、統一 MoE モデルをファインチューニングする。
• Top-1/Top-2、ソフト対ハード、ロードバランシングの有無を含むルーティングスキームと、分割/ブレンディングなどの追加的なバリアントを探索して、性能と効率を評価する。
• 数学、コード、知識に跨る多ドメイン評価ベンチマークを用いて、BTX をベースライン（Dense DM、Sparse upcycling、BTM）と比較する。
• GSM8K、MATH、HumanEval、MBPP、Natural Questions、TriviaQA、ARC、MMLU などのタスクで BTX を評価し、横断的なドメイン性能を把握する。

実験結果

リサーチクエスチョン

• RQ1 ドメイン専用エキスパートへ分岐し、その後の MoE ファインチューニングが、ベースラインと比較して全体的な性能にどう影響するか？
• RQ2 非同期で恥ずかしくなるほど並列なエキスパート訓練と MoE ルーティングの組み合わせが、精度と計算効率の有利なトレードオフを生み出すか？
• RQ3 複数のドメインにわたる BTX にとって、どのルーティング戦略が性能と計算量の最良のバランスを提供するか？

主な発見

MoE モデル	数学	コード	知識	推論	MMLU	平均
Llama-2 7B	14.7	2.5	12.8	20.8	16.4	58.5?	46.1?
CodeLlama 7B	0	8.1	36.3	22.2	56.6	38.6	37.9
Llemma 7B	28.0	33.5	17.2	38.8	33.5	32.1
Llama-2 7B	0	8.6	16.8	37.4	63.3	46.1	40.7
Llama-2 13B	16.3	24.5	40.0	66.1	52.8	45.4
Dense (DM)	18.3	25.8	39.6	63.3	49.8	44.5
Sparse upcycling (DM), Top-2	28.1	34.7	34.0	62.9	52.1	46.3
BTX, Top-1	21.3	36.4	26.5	61.0	44.3	43.1
BTX, Top-2	27.4	34.0	41.0	63.5	52.5	47.9
BTX, Sample Top-1	26.4	31.5	40.1	63.7	53.2	47.3

• 各ドメインエキスパートは個別に評価した場合、それぞれのドメインで卓越しており（Math と Code で大きな利得を示す）。
• BTX は seed モデルと比べてすべてのタスクを改善し、数学とコードの特化モデルに近づきつつ、非ドメインタスクを保持する。
• BTX は Branch-Train-Merge よりも優れており、同一データ予算下で Dense や Sparse upcycling のベースラインより計算効率が高い。
• MoE ファインチューニングを用いてトークンレベルのルーティングを学習すると、単にエキスパート出力を平均した場合よりも顕著な利得が得られる（BTX 対 BTM）。
• ロードバランシングを伴う Top-2 ルーティングはエキスパートの利用をより均一にし、全体的な性能も高い。いくつかのルーティングバリアントは数学 vs コードの性能間でトレードオフを提供する。
• BTX はより多くのデータで訓練でき、同程度の計算予算内で純粋な MoE や単一の大規模モデルより高い平均性能を達成できる。