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

2022/08/05 by Margaret Li, Suchin Gururangan, Li, Margaret +11 · 2 voices · 10 citations
Computer Science · #Computation and Language (cs.CL) #FOS: Computer and information sciences #Natural Language Processing Techniques #Speech Recognition and Synthesis #Topic Modeling #cs.CL

paper · pdf · doi:10.48550/arxiv.2208.03306

openalex publication_date 2022/08/05 · arxiv published 2022/08/05 · arxiv updated 2022/08/05 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

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.

Cited by

• VersatileFFN: Achieving Parameter Efficiency in LLMs via Adaptive Wide-and-Deep Reuse
• MixtureKit: A General Framework for Composing, Training, and Visualizing Mixture-of-Experts Models
• Grow Up and Merge: Scaling Strategies for Efficient Language Adaptation
• MLPMoE: Zero-Shot Architectural Metamorphosis of Dense LLM MLPs into Static Mixture-of-Experts
• CREST: Improving Interpretability and Effectiveness of Troubleshooting at Ericsson through Criterion-Specific Trouble Report Retrieval
• Research in Collaborative Learning Does Not Serve Cross-Silo Federated Learning in Practice
• Pretraining with hierarchical memories: separating long-tail and common knowledge
• Personality Vector: Modulating Personality of Large Language Models by Model Merging
• UNIFORM: Unifying Knowledge from Large-scale and Diverse Pre-trained Models
• Learn Faster and Remember More: Balancing Exploration and Exploitation for Continual Test-time Adaptation

Discussions

• No need to apologise! Sorry if I came across as judging or something like that. MoEs take a bit to wrap our heads around and there's a lot of bad explanations and terminology around. For something mor [bsky, 3 points, 1 comments]
• BTM (and follow ups) splice on the topic axis, models can be trained on coherent but different corpuses then optionally merged for instance. This is just an example, but all in all MoE feels like it' [bsky, 0 points, 1 comments]

Related

• Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM
• BTS: Harmonizing Specialized Experts into a Generalist LLM
• Scaling Expert Language Models with Unsupervised Domain Discovery
• Merge to Learn: Efficiently Adding Skills to Language Models with Model Merging
• Expert Merging: Model Merging with Unsupervised Expert Alignment and Importance-Guided Layer Chunking