Branch-Train-MiX: Het combineren van expert-LLM's in een Mixture-of-Experts-LLM

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

March 12, 2024

Auteurs: 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

Samenvatting

We onderzoeken efficiënte methoden voor het trainen van Large Language Models (LLMs) om vaardigheden te ontwikkelen in meerdere gespecialiseerde domeinen, zoals coderen, wiskundig redeneren en wereldkennis. Onze methode, genaamd Branch-Train-MiX (BTX), begint met een basismodel, dat wordt vertakt om experts te trainen op een parallelle manier met hoge doorvoer en gereduceerde communicatiekosten. Nadat individuele experts asynchroon zijn getraind, brengt BTX hun feedforward-parameters samen als experts in Mixture-of-Expert (MoE)-lagen en worden de overige parameters gemiddeld, gevolgd door een MoE-finetuningfase om token-level routing te leren. BTX generaliseert twee speciale gevallen: de Branch-Train-Merge-methode, die niet de MoE-finetuningfase heeft om routing te leren, en sparse upcycling, die de fase van het asynchroon trainen van experts overslaat. In vergelijking met alternatieve benaderingen bereikt BTX de beste balans tussen nauwkeurigheid en efficiëntie.

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.