Branch-Train-Merge: Embarrassingly Parallel Training of Expert Language Models (2208.03306v1)

Abstract: We present Branch-Train-Merge (BTM), a communication-efficient algorithm for embarrassingly parallel training of 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.

• The paper introduces the Branch-Train-Merge algorithm that enables embarrassingly parallel training of expert language models without multi-node synchronization.
• It details independent training of domain-specific sub-models, achieving efficiency improvements and lower perplexity across varied domains.
• The approach scales language model capacity while reducing resource requirements, paving the way for more accessible and specialized LLM training.

Overview of Branch-Train-Merge: Embarrassingly Parallel Training of Expert LLMs

The paper "Branch-Train-Merge: Embarrassingly Parallel Training of Expert LLMs" introduces Branch-Train-Merge (BTM), a novel training algorithm aimed at enhancing the efficiency of constructing LLMs. This approach is particularly distinguished by its ability to train LLMs through a communication-efficient and parallelizable framework. The key to this innovation is the Branch-Train-Merge algorithm, which eliminates the need for multi-node synchronization—a requirement that traditionally incurs significant computational overhead.

Key Contributions

The authors propose a new paradigm for training models called Expert LLMs (ELMs), each dedicated to a specific textual niche or domain. This results in several advantages, including the ease of scaling model capacity and tailoring inference to specific data domains. BTM capitalizes on this by training sub-models independently on different data subsets, which can subsequently be combined or further specialized.

Methodology

1. ELMforest Composition:

• An ELMforest comprises multiple ELMs, each specialized for distinct domains such as legal, scientific, or general prose.
• These ELMs are built to function as standalone models without shared parameters, diverging from mixture-of-expert models that typically involve parameter sharing.

2. Branch-Train-Merge Algorithm:

• Branch: Initiate new ELMs by copying parameters from existing ELMs. The new ELMs receive averaged parameters relevant to the new domain.
• Train: Independently train each branched ELM on domain-specific data, maintaining complete disconnection from others.
• Merge: Incorporate the trained ELM into the set, enriching the domain coverage portfolio of the ELMforest.

3. Inference Techniques:

• ELMs can be combined during inference either through ensembling or parameter averaging. The former involves integrating output probabilities, while the latter consolidates parameters into a singular model offering reduced inference overhead.

Results and Analysis

The BTM method outperforms traditional transformer-based LLMs in several metrics, particularly in computational efficiency and textual domain specialization. Through rigorous experiments across various domains and computational setups, the research demonstrated that the BTM-trained ELMforests require substantially less computational resource compared to compute-intensive monolithic transformer models, such as those akin to GPT-3.

Quantitatively, the BTM achieves perplexity (a measure of predictive certainty) improvements across both in-domain and out-of-domain datasets. Notably, the approach remains resilient under diverse initialization conditions and data domain compositions, emphasizing its robustness and flexibility. The ELMforest constructs trained with BTM even achieve performance parity with larger transformer models, doing so with a compute budget significantly lower—demonstrating both efficiency and scalability.

Implications and Future Directions

The BTM framework's implications extend across the computational landscape of natural language processing. It articulates a path towards democratizing LLM training, potentially enabling contributions to model development from diverse researchers and entities with varying computational capabilities. This collaborative potential could pave the way for collective model-building endeavors, promoting inclusivity and resource efficiency.

Further explorations could delve into dynamic domain allocation, parameter sharing strategies that minimize initialization complexity, and adaptive pruning to refine inference performance. Additionally, exploring synergies with federated learning, where privacy considerations demand decentralized model training, could uncover novel applications for BTM.

In conclusion, the "Branch-Train-Merge" paper contributes a significant methodological advancement to the domain of LLMs. By restructuring the training process to facilitate domain specificity and parallel training, it sets the stage for more efficient and scalable LLMs compatible with the evolving landscape of textual data diversity.