Branch-Train-MiX: Integrazione di LLM esperti in un modello Mixture-of-Experts LLM

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

March 12, 2024

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

Abstract

Investighiamo metodi efficienti per addestrare Modelli Linguistici di Grande Scala (LLMs) affinché acquisiscano competenze in molteplici domini specializzati, come la programmazione, il ragionamento matematico e la conoscenza generale. Il nostro metodo, denominato Branch-Train-MiX (BTX), parte da un modello iniziale, che viene suddiviso per addestrare esperti in modo parallelizzabile con elevata produttività e ridotti costi di comunicazione. Dopo che i singoli esperti sono stati addestrati in modo asincrono, BTX combina i loro parametri feedforward come esperti in strati Mixture-of-Expert (MoE) e media i parametri rimanenti, seguito da una fase di fine-tuning MoE per apprendere il routing a livello di token. BTX generalizza due casi particolari: il metodo Branch-Train-Merge, che non prevede la fase di fine-tuning MoE per apprendere il routing, e lo sparse upcycling, che omette la fase di addestramento asincrono degli esperti. Rispetto ad approcci alternativi, BTX raggiunge il miglior compromesso tra accuratezza ed efficienza.

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.