Branch-Train-MiX: Смешивание экспертов LLM в смесь экспертов 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 обобщает два особых случая: метод Branch-Train-Merge, который не имеет этапа донастройки MoE для изучения маршрутизации, и разреженное повышение, которое исключает этап асинхронного обучения экспертов. По сравнению с альтернативными подходами, 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.