Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM
Authors: name: Sainbayar Sukhbaatar; name: Olga Golovneva; name: Vasu Sharma; name: Hu Xu; name: Xi Victoria Lin; name: Baptiste Rozière; name: Jacob Kahn; name: Daniel Li; name: Wen-tau Yih; name: Jason Weston; name: Xian Li
Published: 2024-03-12 16:54:58

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.