PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Margaret Li, Suchin Gururangan, Tim Dettmers, Mike Lewis, Tim Althoff, Noah A. Smith, Luke Zettlemoyer

arXiv:2208.03306v1cs.CL

TL;DR

Large-LM training typically requires costly multi-node synchronization, motivating a more communication-efficient approach. BTM independently trains domain-specialized expert LMs, then ensembles or averages them for inference; these models outperform matched baselines and scale to 64 domains with comparable performance at 2.5 times less compute.

  • Problem

    Large language models typically require massive multi-node synchronization for training and inference, creating a need for embarrassingly parallel alternatives.

  • Method

    BTM branches new expert LMs from existing models, trains them independently on new domains, and merges them into an ELMFOREST for ensembling or parameter averaging.

  • Results

    ELMFORESTs outperform GPT-style transformer LMs and DEMIX across computational budgets, while a 64-expert model performs comparably with a TRANSFORMER-LM trained with 2.5 times the GPU hours.

  • Takeaways & Limitations

    Independent domain-specialized ELMs can scale large-language-model training while parameter averaging reduces ensemble inference costs to those of an equivalent TRANSFORMER-LM.

  • Takeaways & Limitations

    The approach assumes provenance-based domain definitions, a small additional sample to estimate domain posteriors, and training-budget choices that remain open questions.

Abstract

from arXiv · show

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.

1 Introduction

Branch-Train-Merge trains independently specialized expert language models on separate domains, avoiding the multi-node synchronization required by conventional large-LM training. The resulting ELMFOREST models outperform matched baselines, with benefits extending to large-scale training and collaborative model development.

  • ELMFOREST models: ELMFORESTs contain independently functional LMs specialized to distinct domains such as scientific or legal text.Experts can be ensembled for generalization or parameter-averaged into a single LM.
  • Empirical results: ELMFORESTs trained with BTM outperform GPT-style transformer LMs and DEMIX in- and out-of-domain across computational budgets.The comparison extends to models with up to 1.3B parameters per ELM and 7000 aggregate GPU-hours.
  • Empirical results: Domain-specialized ELMFOREST ensembles outperform ensembles formed from random data splits, while performance remains robust to varied initialization choices.These analyses identify domain specialization as a critical component of the approach.
  • Scaling study: 64 experts trained on 192B tokens perform comparably with a 1.3B-parameter TRANSFORMER-LM trained with 2.5 times the total GPU hours.The reported benefits increase with the number of domains.
  • Implications: The authors position independently trained ELMs as a basis for open-sourced, community-authored models built under differing resource constraints.The proposed model repository is intended to evolve as groups contribute experts.
  • Branch-Train-Merge: BTM trains many expert LMs in parallel, each specializing to one data domain after an initial shared training phase.This design removes parameter synchronization between independently trained experts.

2 ELMFORESTs

ELMFORESTs are collections of independently trained, domain-specialized language models that can be added, removed, ensembled, or collapsed by parameter averaging. Their inference and training procedures preserve domain separation while trading inference efficiency against ensemble performance.

  • 2.1 Model definition: An ELMFOREST is a set of expert LMs independently trained to specialize to different corpus subsets.The paper defines domains by document provenance, such as legal documents or computer science papers.
  • 2.2 Adding and removing ELMs: ELMFORESTs can add or remove domain experts to modify coverage, keeping each domain’s influence isolated in its associated LM.Removing an expert after training ablates its associated data from future predictions.
  • 2.3 Ensembling the ELMFOREST: Ensembling combines output probabilities from multiple ELMs to generalize to text with unknown domain provenance.The procedure estimates domain posteriors from ELM likelihoods and uses them to weight expert outputs.
  • 2.3 Ensembling the ELMFOREST: The ensemble procedure may require forward passes through all experts, although sparse domain posteriors suggest top-k selection can reduce inference time with negligible performance effects.The paper evaluates posterior sparsity and top-k expert use in scaled experiments.
  • 2.4 Averaging ELM parameters: Parameter averaging collapses an ELMFOREST into one LM, keeping inference cost constant as experts are added.Weighted averages using cached domain priors are reported as strongest and approach ensemble performance while outperforming TRANSFORMER-LMs.

3 BRANCH-TRAIN-MERGE (BTM)

BTM incrementally builds an ELMFOREST through shared seeding, independent domain-specialized training, and merging, without communication among ELMs after seeding. New experts can be added in parallel or asynchronously, using existing experts or a pretrained seed model for initialization.

  • Step 0 (Initialization): BTM begins with a shared seed LM, then branches parameters into independently trained experts for distinct data domains.After seeding, ELMs are fully disconnected and communicate no longer during training.
  • Step 1 (Branch): Each new ELM is initialized from a weighted parameter average of existing experts, with domain-posterior weighting reported as the best-performing approach.The weighted average can instead use only the nearest ELM.
  • Step 2 (Train): A branched ELM is trained solely on its new domain with a log-likelihood objective, while existing ELMs remain uninvolved.This independent training is the core communication-saving step.
  • Step 3 (Merge): After training, the new expert is merged by adding it to the existing ELMFOREST and its represented-domain set.The updated forest can then be used for subsequent iterations.
  • Incremental training: BTM operations can add multiple ELMs in parallel batches or asynchronously, enabling incremental expansion across new domains.Subsequent experts may leverage parameter averages from existing LMs during branching.
  • Step 0 (Initialization): The initial seed phase uses a pretrained LM trained on a corpus to initialize the first ELM batch rather than branching from an empty expert set.The paper reports that seeding is important for enabling parameter averaging.

4 Core Experiments and Results

BTM trains domain-specialized expert LMs with reduced synchronization, improving compute-matched performance while enabling flexible ensemble or averaged inference. Across scales, ELMFORESTs outperform Transformer and DEMIX baselines, though averaging trades some performance for lower inference cost.

  • 4.2 Performance Comparisons: ELMFORESTs trained with BTM consistently achieve the lowest average test perplexity across 125M, 350M, 750M, and 1.3B parameter scales.Table 1 evaluates average test-set perplexity across eight training, eight evaluation, and all sixteen domains.
  • 4.2 Performance Comparisons: BTM improves performance over both densely trained TRANSFORMER-LM and sparsely trained DEMIX baselines across the evaluated model scales.The results suggest that isolating all LM parameters improves expert domain specialization relative to isolating only feedforward layers.
  • 4.3 Efficiency Comparisons: BTM reduces cross-GPU communication because GPUs assigned to different ELMs do not synchronize during branched training.This disconnected setup also permits resources to operate at different times or train ELMs serially; branched jobs used no more than two eight-GPU nodes simultaneously, versus sixteen such nodes for Transformer-LM experiments.
  • 4.3 Efficiency Comparisons: 33% speedup at 1.3B parameters per GPU: BTM branched training achieves higher updates per second than fully synchronized Transformer-LM training.The efficiency advantage grows with model size and the number of nodes used in parallel, although estimates vary with hardware and environment.
  • 4.4 ELMFOREST Parameter Average: Posterior-weighted parameter averaging consistently improves over TRANSFORMER-LM at no additional inference cost, but remains below ELMFOREST ensembling.Uniform averaging underperforms all baselines, indicating that domain relevance matters when combining ELMs.
  • 4.4 ELMFOREST Parameter Average: Sufficient seed training enables parameter averaging to outperform TRANSFORMER-LM at all model scales.The smallest model performs poorly when too little compute is allocated to the seed phase, especially on evaluation domains.

5 Analysis

The analysis finds that BTM gains depend on domain-specialized experts and a shared seed phase, while remaining robust to several initialization choices. ELMFOREST parameter averaging additionally requires adequate seed training, and ELM removal can support domain forgetting under specific conditions.

  • 5.1 ELMFOREST performance is not simply the result of ensembling more parameters: Domain-specialized ELMFOREST experts outperform ensembles trained on random data splits, showing that increased parameter count alone is insufficient.The random ensemble is consistently worse than both ELMFOREST variants across the evaluated domains.
  • 5.2 ELMFOREST performance robust to wide range of seed LM training compute allocations: Seed training and branched ELM training are both important: no seed training underperforms, while 100% seed training also trails BTM.The analysis varies seed compute while holding total compute fixed.
  • 5.2 ELMFOREST performance robust to wide range of seed LM training compute allocations: 40–60% of total compute is approximately optimal for seed training, with performance deteriorating toward 0% and 100%.As little as 10% seed training still produces strong gains over random initialization and the Transformer baseline.
  • 5.2 ELMFOREST performance robust to wide range of seed LM training compute allocations: Parameter averaging fails with randomly initialized ELMs, whereas evaluation-domain performance at smaller scale requires about 60% or more seed updates.At larger scales, the optimal seed share shifts lower, to about 40%, and 50% seed training outperforms Transformer baselines at several scales.
  • 5.3 ELMFOREST performance is robust to the choice of seed training corpus: Diverse seed-training corpora perform best overall, but even JavaScript-only seed training outperforms the compute-matched Transformer baseline.These experiments allocate 50% of updates to seed training and 50% to branched training.
  • 5.3 ELMFOREST performance is robust to the choice of seed training corpus: Removing an ELM can forget its associated domain, but the effect on other performance depends on overlap between the removed domain and other training domains.The authors note that stronger forgetting guarantees may require carefully curated seed and domain data.

6 Incrementally Training an ELMFOREST on 64 Domains

The 64-domain scaling study incrementally branches and merges experts across domain batches, achieving Transformer-comparable perplexity at substantially lower training compute. Sparse domain posteriors also permit reduced inference-time expert usage or parameter averaging.

  • 6.1 Experimental Setup: BTM incrementally trains ELMFOREST experts across 64 training domains organized into four batches, branching from weighted averages of existing experts.The corpus contains 64 training domains and 16 evaluation domains.
  • 6.3 Results: ELMFOREST achieves perplexities comparable to a large Transformer LM on training and evaluation domains using only 40% of the total compute.The comparison is intentionally not compute-matched to demonstrate efficiency gains.
  • 6.3 Results: The order and composition of domain batches likely affect ELMFOREST's performance-efficiency tradeoff, motivating further analysis.The authors leave optimal scaling settings and batch effects to future work.
  • 6.4 Inference-Time Sparsity: Domain posteriors are sparse on both training and evaluation domains, and their most likely ELMs are usually relevant to the evaluation domain.This motivates activating only a small subset of experts at inference time.
  • 6.4 Inference-Time Sparsity: Using the top-8 ELMs causes negligible performance loss relative to all 64, while even the top-1 ELM beats the Transformer baseline for the training cost.Top-1 inference adds no inference cost over the baseline.
  • 6.5 Parameter Averaging: Weighted parameter averaging provides the best performance at the original Transformer inference cost, although ensembles of at least two ELMs generally perform better.Averaging keeps inference cost constant regardless of the number of ELMs.
  • 6.6 Summary: Small independently trained ELMs can perform comparably to a large Transformer trained with at least 2.5× the compute.Parameter averaging can reduce the inference cost of using the full ensemble with minimal performance degradation.

7 Related Work

Related work spans sparsely activated language models, adapters, ensembles, parameter averaging, and seed training. BTM differs by using independently functional full-model experts and posterior-weighted averaging, while its seed phase aligns with prior findings on shared optimization trajectories.

  • Sparse Language Models: Prior sparse language models include mixture-of-experts systems with independent experts and token- or domain-based routing.The paper identifies DEMix and multilingual expert models as especially close related work.
  • Adapters: Unlike adapters and feedforward-only domain experts, each ELM is an entire independently functional model with no shared parameters.The paper presents this as a simplification relative to those approaches.
  • Ensembles: Earlier growing ensembles trained on randomly created batches underperform non-incremental methods, contrasting with BTM's domain-specialized training.The comparison highlights the role of structured specialization rather than random partitioning.
  • Parameter Averaging: BTM uses posterior-weighted parameter averaging because uniform averaging does not work for combining domain-specific models.The weights are estimated from model relevance to the target domain.
  • Seed Training: The importance of BTM's seed phase is consistent with findings that parameter averaging works when models share part of their optimization trajectory.The paper leaves open what the seed phase learns that supports specialization across corpora.

8 Limitations

The paper identifies limitations involving domain definitions, baseline training strategies, posterior-data requirements, and potential harms from toxic web data. These constraints motivate future work on domain discovery, broader baselines, posterior estimation, and safety.

  • Domain definition: Provenance labels provide interpretable domain segmentations, but other domain-discovery methods may yield better expert assignments.The paper leaves experimentation with alternative domain definitions for future work.
  • Posterior estimation: Ensembling and parameter averaging require a small additional data sample to estimate domain posteriors.The paper notes that future work could investigate eliminating this requirement.
  • Training baselines: BTM’s comparison with alternative distributed training strategies is limited because the Transformer baseline uses distributed data-parallel training.Model-parallel and fully sharded strategies may exhibit different scaling patterns and could change the reported conclusions.
  • Safety: ELMFORESTs trained on Internet-scraped data may contain toxic language that coarse provenance labels fail to identify.The paper calls for future recipes for safer ELMFOREST training and deployment.

9 Conclusion

The conclusion presents BTM as a way to train collections of independently trained, specialized expert LMs that can be customized and combined efficiently. Experiments support its performance advantages and the promise of scaling through many smaller experts.

  • 9 Conclusion: BTM trains ELMFORESTs whose expert LMs can be added, removed, ensembled, or parameter averaged for efficient scaling and rapid customization.These operations support updating data coverage and choosing between ensemble quality and cheaper inference.
  • 9 Conclusion: ELMFOREST ensembles trained with BTM outperform baselines without additional training cost.The conclusion reports this as a result of the extensive experiments.
  • 9 Conclusion: Parameter-averaged ELMFORESTs closely approach ensemble performance while enabling substantially cheaper inference.Averaging collapses the expert collection into a single language model.
  • 9 Conclusion: The results provide evidence that many smaller, independently trained ELMs can support scaling large language models at inclusive compute budgets.The authors envision repositories of continually evolving, community-authored expert LMs.

A Appendix

The appendix documents the multi-domain corpora, de-identification procedures, baseline evaluations, and scaling-study comparisons. Across increasing domain counts, ELMFORESTs retain or improve performance while compute-matched Transformer LMs degrade.

  • Corpus: The multi-domain corpus includes training and evaluation domains with dataset sizes reported in whitespace-separated tokens.The corpus is used for the paper’s experiments and analyses.
  • Corpus: The scaling corpus contains 64 training domains organized into batches for the multi-domain study.Table 10 identifies the batch associated with each domain and notes de-identification for selected datasets.
  • Corpus: The novel-domain corpus reports training and evaluation sizes for 32 domains in whitespace-separated tokens.These datasets are borrowed from Gururangan et al. (2022).
  • Data processing: The appendix specifies regex-based de-identification schemas for listed data categories.The schema references the associated dataset links and categories.
  • Evaluation: ELMFORESTs trained with BTM outperform baselines and ensemble variations across 125M, 350M, 750M, and 1.3B parameter scales.Table 13 reports average test-set perplexity across 8 training, 8 novel, and all 16 domains, alongside compute and parameter totals.
  • Scaling study: As training domains increase, ELMFORESTs retain or improve performance while compute-matched Transformer LMs degrade.The comparison averages test-set perplexity over batches containing 8, 16, 32, and 64 domains, with 6144 GPU hours for each Transformer-LM experiment.
Loading 2208.03306v1…