Branch-Train-MiX: 専門家LLMをMixture-of-Experts LLMに統合する手法

Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM

March 12, 2024

著者: Sainbayar Sukhbaatar, Olga Golovneva, Vasu Sharma, Hu Xu, Xi Victoria Lin, Baptiste Rozière, Jacob Kahn, Daniel Li, Wen-tau Yih, Jason Weston, Xian Li

cs.AI

要旨

大規模言語モデル（LLM）を複数の専門領域（コーディング、数学的推論、世界知識など）で能力を備えるように効率的に訓練する方法を調査します。我々の手法であるBranch-Train-MiX（BTX）は、シードモデルから始まり、高いスループットと低い通信コストで並列的に専門家を訓練するために分岐させます。個々の専門家が非同期に訓練された後、BTXはそれらのフィードフォワードパラメータをMixture-of-Expert（MoE）層の専門家として統合し、残りのパラメータを平均化します。その後、トークンレベルのルーティングを学習するためのMoEファインチューニング段階を経ます。BTXは、ルーティングを学習するMoEファインチューニング段階を持たないBranch-Train-Merge手法と、専門家を非同期に訓練する段階を省略するスパースアップサイクリングという2つの特殊ケースを一般化します。他のアプローチと比較して、BTXは最高の精度と効率のトレードオフを達成します。

English

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.