Branch-Train-Merge: Embarrassingly Parallel Training of Expert Language Models

By Margaret Li, Suchin Gururangan, Tim Dettmers, Michael Lewis, Tim Althoff, Noah A. Smith, Luke S. Zettlemoyer. 2022

213 citations recorded by the source corpus.

We present Branch-Train-Merge (BTM), a communication-efficient algorithm for embarrassingly parallel training of large language models (LLMs). We show it is possible to independently train subparts of a new class of LLMs on different subsets of the data, eliminating the massive multi-node synchronization currently required to train LLMs. BTM learns a set of independent expert LMs (ELMs), each specialized to a different textual domain, such as scientific or legal text. These ELMs can be added and removed to update data coverage, ensembled to generalize to new domains, or averaged to collapse back to a single LM for efficient inference. New ELMs are learned by branching from (mixtures of) ELMs in the current set, further training the parameters on data for the new domain, and then merging the resulting model back into the set for future use. Experiments show that BTM improves in- and out-of-domain perplexities as compared to GPT-style Transformer LMs, when controlling for training cost. Through extensive analysis, we show that these results are robust to different ELM initialization schemes, but require expert domain specialization; LM ensembles with random data splits do not perform well. We also present a study of scaling BTM into a new corpus of 64 domains (192B whitespace-separated tokens in total); the resulting LM (22.4B total parameters) performs as well as a Transformer LM trained with 2.5 times more compute. These gains grow with the number of domains, suggesting more aggressive parallelism could be used to efficiently train larger models in future work.

Subfields: Artificial Intelligence, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Model Architectures

Citations in the corpus

A Review of Sparse Expert Models in Deep Learning (2022, background): “This confers a benefit of modularity and composability; Li et al. (2022) shows custom networks can be constructed by composing them of expert language models trained on specific domains.”
Editing Models with Task Arithmetic (2022, unknown): “[55] when averaging the parameters of language models fine-tuned on various domains.”
Prometheus 2: An Open Source Language Model Specialized in Evaluating Other Language Models (2024, unknown): “Prior works have demonstrated that weight merging can enhance performance across various domains, including language modeling (Li et al., 2022; Matena and Raffel, 2022; Ilharco et al., 2022; Don-Yehiya et al., 2022; Gururangan et al., 2023; Yadav et al., 2024; Sukhbaatar et al., 2024),…”
TIES-Merging: Resolving Interference When Merging Models (2023, background): “…target task [30, 23, 82, 9], improving out-of-domain generalization [31, 29, 7, 4, 60, 59], creating a multitask models from different tasks [31, 29, 38], for federated learning [46, 41], compression [39], multimodal merging models [72], continual learning [86, 85], and other settings [38, 13].”
SILO Language Models: Isolating Legal Risk In a Nonparametric Datastore (2023, background): “Our work is closely related to recent studies on modular LMs, which have specialized parameters (or experts) trained on different domains (Gururangan et al., 2022; Li et al., 2022; Gururangan et al., 2023), languages (Pfeiffer et al.”
Dataless Knowledge Fusion by Merging Weights of Language Models (2022, unknown): “Weight averaging was also used by Li et al. (2022) for building language models with multi-domain capabilities where new domain ‘experts’ are initialized using weight averaging from the existing experts.”