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

arXiv: https://arxiv.org/abs/2208.03306 · HTML: https://mlsanity.com/paper/2208.03306/

## Key facts

- Cited by 191 papers in the map (data through 14 Oct 2026).
- Builds on 53 of its 82 references that are in the map.
- First posted on arXiv 5 Aug 2022.

## Builds on (most cited first, 30 of 53)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Parameter-Efficient Transfer Learning for NLP](https://mlsanity.com/paper/1902.00751/) (2019)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Don't Stop Pretraining: Adapt Language Models to Domains and Tasks](https://mlsanity.com/paper/2004.10964/) (2020)
- [fairseq: A Fast, Extensible Toolkit for Sequence Modeling](https://mlsanity.com/paper/1904.01038/) (2019)
- [Competition-Level Code Generation with AlphaCode](https://mlsanity.com/paper/2203.07814/) (2022)
- [Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time](https://mlsanity.com/paper/2203.05482/) (2022)
- [RealToxicityPrompts: Evaluating Neural Toxic Degeneration in Language Models](https://mlsanity.com/paper/2009.11462/) (2020)
- [BitFit: Simple Parameter-efficient Fine-tuning for Transformer-based Masked Language-models](https://mlsanity.com/paper/2106.10199/) (2021)
- [Averaging Weights Leads to Wider Optima and Better Generalization](https://mlsanity.com/paper/1803.05407/) (2018)
- [Aligning Books and Movies: Towards Story-like Visual Explanations by Watching Movies and Reading Books](https://mlsanity.com/paper/1506.06724/) (2015)
- [Plug and Play Language Models: A Simple Approach to Controlled Text Generation](https://mlsanity.com/paper/1912.02164/) (2019)
- [Robust fine-tuning of zero-shot models](https://mlsanity.com/paper/2109.01903/) (2021)
- [CTRL: A Conditional Transformer Language Model for Controllable Generation](https://mlsanity.com/paper/1909.05858/) (2019)
- [Linear Mode Connectivity and the Lottery Ticket Hypothesis](https://mlsanity.com/paper/1912.05671/) (2019)
- [Merging Models with Fisher-Weighted Averaging](https://mlsanity.com/paper/2111.09832/) (2021)
- [Defending Against Neural Fake News](https://mlsanity.com/paper/1905.12616/) (2019)
- [Neural Text Generation with Unlikelihood Training](https://mlsanity.com/paper/1908.04319/) (2019)
- [Rigging the Lottery: Making All Tickets Winners](https://mlsanity.com/paper/1911.11134/) (2019)
- [A Dataset of Information-Seeking Questions and Answers Anchored in Research Papers](https://mlsanity.com/paper/2105.03011/) (2021)
- [MAD-X: An Adapter-Based Framework for Multi-Task Cross-Lingual Transfer](https://mlsanity.com/paper/2005.00052/) (2020)
- [Analysing Mathematical Reasoning Abilities of Neural Models](https://mlsanity.com/paper/1904.01557/) (2019)
- [BASE Layers: Simplifying Training of Large, Sparse Models](https://mlsanity.com/paper/2103.16716/) (2021)
- [One Billion Word Benchmark for Measuring Progress in Statistical Language Modeling](https://mlsanity.com/paper/1312.3005/) (2013)

## Built on by (most cited first, 30 of 191)

- [Editing Models with Task Arithmetic](https://mlsanity.com/paper/2212.04089/) (2022)
- [TIES-Merging: Resolving Interference When Merging Models](https://mlsanity.com/paper/2306.01708/) (2023)
- [Prometheus 2: An Open Source Language Model Specialized in Evaluating Other Language Models](https://mlsanity.com/paper/2405.01535/) (2024)
- [Dataless Knowledge Fusion by Merging Weights of Language Models](https://mlsanity.com/paper/2212.09849/) (2022)
- [OLMoE: Open Mixture-of-Experts Language Models](https://mlsanity.com/paper/2409.02060/) (2024)
- [Personalized Soups: Personalized Large Language Model Alignment via Post-hoc Parameter Merging](https://mlsanity.com/paper/2310.11564/) (2023)
- [Model Merging in LLMs, MLLMs, and Beyond: Methods, Theories, Applications and Opportunities](https://mlsanity.com/paper/2408.07666/) (2024)
- [Rewarded soups: towards Pareto-optimal alignment by interpolating weights fine-tuned on diverse rewards](https://mlsanity.com/paper/2306.04488/) (2023)
- [Task Arithmetic in the Tangent Space: Improved Editing of Pre-Trained Models](https://mlsanity.com/paper/2305.12827/) (2023)
- [ReFT: Representation Finetuning for Language Models](https://mlsanity.com/paper/2404.03592/) (2024)
- [A Review of Sparse Expert Models in Deep Learning](https://mlsanity.com/paper/2209.01667/) (2022)
- [WARM: On the Benefits of Weight Averaged Reward Models](https://mlsanity.com/paper/2401.12187/) (2024)
- [RAPHAEL: Text-to-Image Generation via Large Mixture of Diffusion Paths](https://mlsanity.com/paper/2305.18295/) (2023)
- [Safe LoRA: the Silver Lining of Reducing Safety Risks when Fine-tuning Large Language Models](https://mlsanity.com/paper/2405.16833/) (2024)
- [DiLoCo: Distributed Low-Communication Training of Language Models](https://mlsanity.com/paper/2311.08105/) (2023)
- [Knowledge Fusion of Large Language Models](https://mlsanity.com/paper/2401.10491/) (2024)
- [ERNIE-ViLG 2.0: Improving Text-to-Image Diffusion Model with Knowledge-Enhanced Mixture-of-Denoising-Experts](https://mlsanity.com/paper/2210.15257/) (2022)
- [Language Models are Homer Simpson! Safety Re-Alignment of Fine-tuned Language Models through Task Arithmetic](https://mlsanity.com/paper/2402.11746/) (2024)
- [Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM](https://mlsanity.com/paper/2403.07816/) (2024)
- [Model Ratatouille: Recycling Diverse Models for Out-of-Distribution Generalization](https://mlsanity.com/paper/2212.10445/) (2022)
- [Deep Model Fusion: A Survey](https://mlsanity.com/paper/2309.15698/) (2023)
- [Soft Merging of Experts with Adaptive Routing](https://mlsanity.com/paper/2306.03745/) (2023)
- [Findings of the BabyLM Challenge: Sample-Efficient Pretraining on Developmentally Plausible Corpora](https://mlsanity.com/paper/2504.08165/) (2025)
- [MatFormer: Nested Transformer for Elastic Inference](https://mlsanity.com/paper/2310.07707/) (2023)
- [AdapterSoup: Weight Averaging to Improve Generalization of Pretrained Language Models](https://mlsanity.com/paper/2302.07027/) (2023)
- [Modular Deep Learning](https://mlsanity.com/paper/2302.11529/) (2023)
- [Exploring the Benefits of Training Expert Language Models over Instruction Tuning](https://mlsanity.com/paper/2302.03202/) (2023)
- [Personalized Pieces: Efficient Personalized Large Language Models through Collaborative Efforts](https://mlsanity.com/paper/2406.10471/) (2024)
- [Towards Modular LLMs by Building and Reusing a Library of LoRAs](https://mlsanity.com/paper/2405.11157/) (2024)
- [SILO Language Models: Isolating Legal Risk In a Nonparametric Datastore](https://mlsanity.com/paper/2308.04430/) (2023)

## Abstract (arXiv)

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.