Branch-Train-Merge: Embarrassingly Parallel Training of Expert Language Models
Authors: Margaret Li
Suchin Gururangan
Tim Dettmers
Mike Lewis
Tim Althoff
Noah A. Smith
Luke Zettlemoyer
Venue: Submitted to ICLR 2023 (2023)
Keywords: sparsity
language model
efficient

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 on new domains, and then merging the resulting models 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; ensembles with random data splits do not perform well. Our results suggest that aggressive parallelism could be used to efficiently scale larger LMs in future work.

Under review as a conference paper at ICLR 2023
EMBARRASSINGLY PARAL- LEL TRAINING OF EXPERT LANGUAGE MODELS
Anonymous authors Paper under double-bling review
ABSTRACT
We present Branch-Train-Mes (BTM), a communication--fficieenn algorithm for embarrassingly parallel training language models (LMs). We show it possible to independently train subparts of new class of LMs on different subsets of the data, eliminating the massive multi-node synchronization currently required to train LMs. BTM learns 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 generalize 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 on new domains, and then merging the resulting models back into the set for future use. Experiments show that BTM improves in- and -fomain 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; ensembles with random data splits do not perform well. Our results suggest that extreme parallelism could be used to efficiently scale LMs in future work.
INTRODUCTION
Training and inference in language models (LMs) typically require access to supercomputers that can achieve the massive multi-node synchronization required to compute model activations and gradients (Brown et al., 2020; Fedus et al., 2022; Zhang et al., 2022). We develop instead embarrassingly parallel: different parts of the model are independently trained on different subsets of the data, with no need for multi-node training or inference (Figure 1).
class of LMs that is
Our new MS model consists of a set of EXPERT LMS (ELMs), each specialized to a distinct domain in the training corpus, e.g., scientific or legal text. ELMS are independently functional LMs with no shared parameters, unlike recent mixture- contturr-offee-ppee ef-eets models that only specialize the transformer feedforward layers domains (Gururangan al., 2022). ELMSs can be added removed at any time to update data coverage, ensembled to generalize to new domains, parameter averaged to collapse back to a single LM for more efficient inference.
We also present the Branch- Techh--rrin--aa Train-Merge Merge (B (BTM) algorithm for learning ELMs. BTM repeatedly expands the MFOORTT by adding one or more new ELMs in parallel. Each new ELM is first brunched by initializing new LM with an average of the parameters of the most relevant LMS in the current set, then further trained on new domains with a standard cross-entropy loss, and finally merged into the ELMFOREST (Figure 2). The ELMFOREST is initalized with single LM, trained on heterogeneous data to establish strong shared representations for future domain specialization.
When evaluated in- and out-of-domain, ELMFORESTS trained with BTM outperform GPT-style transformer LMs, a domain- -peeciai--peiillieeed mixture- -feeeers baseline (Gururangan et al. 2022), and non- domain-based ensemble baselines across range of computational budgets - up 1.3B parameters per ELM trained for 7000 GPU-hours in aggregate ($4.2). These gains are biggest for E.M.FORSTT ensembles, which activate all experts during inference, but also hold when combine experts with parameter averaging. We also do extensive analysis of these results; expert specialization to domains crucial ($5.1), while the compute budget allocation ($5.2) and the choice of training data for the initial ELM ($5.3) are much less so. We release our code and models publicly.
Experrt Language Models For Efficient Sparse Training URL anonymized for review.

---

Under review as a conference paper at ICLR 2023
Each GPU has
copy the LM
Each GPU has chadinct LM
(a) Fully Synchronized Training
(b) Embarrassingly Parallel Training
LM paramotors
data milihais
Train k independent LMs in parallel one data domain each, without synchronizing across LMSs
Train one LM on mono-corpus, synchronizing weights across all GPUs
mandomll drewn
drtan assigned
unified data corpus
k data domains
Figure I: Fully Synchronized Embarrassingly Parallel Training ($3). (a) In fully synchronized data-parallel training of a TRANSFORMER-IE all parameters are synchronized across all GPUs. This synchronization incurs hefty cross-node communication costs. (b) embarrassingly parallel training (our work), individual models are trained on each domain, eliminating expensive cross-node parameter synchronization between those models.
2
ELMFORESTS
We define an E.MFOREST to be set of EXPERT LMs (ELMs), each independently trained to specialize to a different subset a corpus. ELMs are inspired by the experts in earlier MoE models (Jacobs et al., 1991), but we define ours to be domain specialists and specialize the full LM instead of components. We follow Gururangan et al. 2022 and define domains by provenance, or the source of the document (e.g., legal document, computer science research paper), which yields simple and interpretable corpus segmentations useful for identifying ELMs our experiments." Potential extensions multi-lingual -modal, -task, other types of data splits are left for future work. ELMs remain independent throughout training and inference, enabling the functions below.
2.1 ADDING AND REMOVING ELMS
Existing control techniques to steer LMs towards (Keskar et al., 2019; Gururangan et al., 2020; Dathathri et al., 2020) away (Welleck al., 2019) from certain behaviors tend to be expensive, require retraining the model, or do not provide strong guarantees on test-time behavior (Gehman al., 2020). In contrast, ELMFORESTS allow for explicit inference-time application of constraints on the provenance of training data. We modify the domain coverage of an ELMFOREST at any time by incorporating new ELMs specialized different domains or removing existing ELMs in the set.
2.2 INFERENCE FROM ELMFORESTS
ELMFORESTS support two inference modes, which trade off efficiency for performance.
Output Ensembling In the first method, ensemble the output probabilities of multiple ELMs. This allows us to generalize to texts of unknown domain. We use the cached prior method proposed in Gururangan et al. (2022).
Consider the probabilistic view of language modeling, where we estimate p( Xx (<)). We introduce domain variable D, alongside each sequence. Then the next-step conditional distribution the history & Eet is:
p(X, x xZ
Ex(x, aece D = p(D = x<x))
We estimate domain posterior, probability of a sequence belonging to each of the k domains using Bayes' rule:
p(D j [2()= p<a<1 D = jD
pl D= =
(2)
p(x2()
E D == =') - pDDD=)))) j]
See $A.1 for discussion the possible limitations of this domain definition.
2

---

Under review as a conference paper at ICLR 2023
Branch-Train-Merge
Step initialization
corpus
Step branch from existing experts, or seed LM
repeat with another batch domains
randomly drawn data minibatches
Step 2: branched training on domains in parallel
Step 3: merge k domain expert LMs into ELMforest
seed corpus
branch copies of seed parameters
marged ELMforest
domains
Figure 2: BTM training process overview ($3). In the seed phase (Step 0), an LM trained on one data corpus. We branch, copy, the parameters & times (Step 1), and continue to train each copy on a unique data domain, resulting in & ELMs (Step 2), which are merged into the ELMFOREST (Step 3). After the seed phase, ELMs are fully disconnected, with no communication between them.
ELMs are used to compute the likelihood over contexts given domain label. To compute the cached prior, we maintain an exponential moving average posterior probabilities over domains, updated only at the end of each sequence block: p(D = j) = = X p(D = j x x<r) Following Gururangan al. 2022, we use N 100 sequences (of length T 1024 each) of development data, and set EMA decay À = 0.3. We fix this prior at test time for each domain.
Output ensembling naively requires forward pass through all ELMs in the ELMFOREST, but observe in practice that the domain posterior is sparse, which suggests that top-k selection of EXPERT LMS can reduce inference time costs with negligible effects on performance.
Parameter Averaging Alternatively, we also use weighted parameter averaging (Izmailov et al., 2018; Wortsman et al., 2022; Matena & Raffel, 2021) to collapse the ELMFOREST into single LM, using the cached prior from above as averaging coefficients. This operation keeps inference cost constant regardless of how many ELMs are added to the set.
BRANCH-TRAIN-MERGE (BTM)
BRANCH-TRAIN-MER training of ELMFOREST models is incremental and embarrassingly parallel (Figure 1). EXPERT LMs are trained fully independently, starting from a seed LM (Figure 2).
3.1 THE BRANCH-
BRANCH-TRAIN-MERS TRAIN-MERGE ITERATION
Each BRANCH- BRANCH-TRAIN-IE TRAIN-MERGE iteration begins with an existing ELMORETST F= {0,) + Each
ELM O represents corresponding domain d; in the dataset of & domains DE = } modeled
by E. We first describe the inductive case of then describe how to train the initial model fg.
Step (Branch): Sprouting a new ELM The new ELM parameters are function of the current ex- pert set E. Given some vector of weights w = fwi, W2 .... wg over the existing experts 01, f Gb. initialize the expert with the weighted parameter average 0 Σx W:O.,
Step 2 (Train): Growing the ELM We train the new ELM 0x+1 data domain de+1 with the log likelihood objective. None of the existing ELMs E are involved in the training of the new ELM. We also refer this step branched training distinguish it from other training regimens.

---

Under review as a conference paper at ICLR 2023
Step 3 (Merge); Transplanting the ELM We merge the new ELM Bk-1, which represents the domain dk+1, into E to create an updated set: E' E {fx-1} and DE Dg (dx+1)-
These operations can be parallelized to add multiple ELMs in batch iterated for many batches.
3.2 STEP O (INITAAIIZATON): SEEDING THE ELMFOREST
In the first iteration of BTM, E 0; have ELMs in the set to branch from. Instead of initializing the first ELMs randomly, we hypothesize that ELM performance boosted by branching from pretrained LM parameters, since multi-phase adaptive pretraining is effective way develop domain-specific language models (Gururangan et al., 2020), and model merging techniques work best with models that have shared initialization (lzmailov et al., 2018; Frankle al., 2020) To this end, we perform a seed phase, training a seed LM Owed on some data corpus dwnnt which can be used initialize the first batch of ELMs in the set.
CORE EXPERIMENTS AND RESULTS
We first compare BTM training to compute-matched matched baselines, to carefully measure the efficiency gains. We use the simplest form of BTM, training set of & = 8 domains.
4.1 EXPERIMENTAL. SETUP
Data We use data from Gururangan al. (2022), which consists of 8 diverse training and 8 evaluation (primarily English-language domains, from web text and U.S court opinions to GitHub and COVID-19 research papers. Details in Appendix Table 7.
Model hyperparameters The model architecture is randomly -nittalliid LM with the GPT-3 (Brown et al., 2020) architecture implemented in Fairseq (Ott et al., 2019). We use 125M (small), 350M (medium), 750M (large), 1.3B (xl) parameter models. Following Brown al. 2020, we the GPT-2 (Radford et al., 2019) vocabulary of 50,264 BPE types, and train with 1,024-token sequences, across document boundaries. We prepend beginning- cieiiiing-of-docuuee of-document token each document.
Compared Models
TRANSFORMER-LAN LM The first baseline is standard transformer LM, implemented with distributed data parallelism (Li, 2021). This is identical the DENSE model from Guru- rangan et al. (2022), in which data from each domain is balanced. We find, in line with Gururangan et al. (2022), that balancing data domains achieves better performance than without data balancing (Appendix Table 9).
We follow the training procedure of Gururangan et al. (2022), where feedforward layers in the transformer are trained to specialize as domain experts, resulting in better domain specialization and generalization than other sparsely activated (c.g., MoE) models.
DEMIX:
ELMFOREST: We first conduct seed phase initialize the ensemble with LM parameters ($3.2), then conduct branched training on the ELMs ($3.1), all of which are initialized with the seed L.M. Between the seed and branched phases, continue training from the saved optimizer state.
These models are compute-maten matched, since computation typically the limiting factor in model training. Like other sparse models (Fedus et al., 2022; Lepikhin et al., 2021; Gururangan et al., 2022), ELMFORESTS decouple compute and parameters; we can train many more parameters at the same computational cost as the equivalent TRANSFORMER-L.M. Total parameter counts are in Table 1.
Training settings To disentangle variations in GPU speed, we number of updates computational budget: 80k, 32k, 24k, and 12k updates on 16, 32, 64, and 128 GPUs in parallel for the 125M, 350M, 750M, 1.3B parameter TRANSFORMER- -LM and DEMIX baselines, respectively. We use the same GPU counts for the seed phase in the ELMRORSST. For branched training, we divide these GPU counts equally among the ELMM;; for example, the 1.3B parameter per GPU MFORETTT
4

---

Under review as a conference paper at ICLR 2023
125M
350M
T-LM 125M
DEMIX 512M
ELMFOREST 1B
T-LM 350M
DEMIX 1.8B
ELMPOREST
2.8B
Train
19.90.2
18.20.62
17.2p.02
Train
16.3
15.0
14.7
Eval
25.2o.18
23.40.54
22.40.12
Eval
19.9
18.6
All
22.50.14
20.80.63
19.8p.05
All
18.5
17.5
16.7
750M
1.3B
T-LM
DEMIX 3.8B
ELMFOREST 6B
T-LM 1.3B
DEMIX
ELMFOREST
750M
7B
10.4B
Train
14.7
13.5
13.4
Train
14.2
13.7
13.0
Eval
19.3
17.7
16.7
Eval
18.4
17.6
16.3
All
17.0
15.6
15.0
All 16.3
15.6
14.6
Table ELMFORESTS trained with BTM outperform baselines across multiple model scales ($4.2). Average test-set perplexity (4) for each model scale across the 8 training, 8 evaluation, and 16 data domains. Total parameters are shown for each model and scale. At 125M parameters per GPU, we show the mean and standard deviation over 8 random seeds. For BTM, we show results with 50% of compute dedicated to the seed phase.
Average updates per second, normalized (t) partially synchronized BTM: embarrassingly parallel (branched EL.Ms)
fully synchronized (TRANSFORMER--MM)
(DEMix)
125M
1.00
1.01
1.05
350M
1.11
1.23
750M
1.00
1.01
1,27
1.3B
1.00
0.97
1.33
Table BTM is more efficient ($4.3), Average updates per second (f) for each setup and model size, normalized by the average updates per second during fully synchronized training of the TRANSFORMER-LAN L.. The efficiency gains from embarrassingly parallel training (the branched phase of B BTM) become more substantial with larger model size and more nodes used in parallel.
uses 16 GPUs for each of the ELMs. For all models, we fix the learning rate at 0.0005 with a polynomial (linear) decay learning rate schedule and 8% warmup, which found to be optimal for most settings after a large grid search. We use batch size of 16 for each GPU, with gradient accumulation of 32 steps, and train with fpl 6. We train NVIDIA V100 32GB GPUs.
4.2 PERFORMANCE COMPARISONS
Our results are shown in Table 1. At these model scales, ELMFORESTS trained with the BTM procedure outperform both the sparsely trained DEMIX LM and the densely trained TRANSFORMER- LM baselines. The improvements in performance we observe over DEMIX layers suggest that isolation of all LM parameters results in better specialization of domain experts.
4.3 EFFICIENCY COMPARISONS
Training ELMFORESTS requires less inter-GPU communication than TRANSFORMER--LM or DEMIX models, since no synchronization occurs between GPUs assigned to different ELMs. This results in higher updates per second and therefore shorter train times (Table 2). Additionally, the embarrassingly parallel branched training provides flexibility resource consumption; GPUs dedicated to different ELMs may be online at different times, and ELMs may be trained serially the same GPUs. Specifically, none of our branched training required more than 16 GPUs simultancously, while TRANSFORMER-IE training experiments consumed 128 GPUs simultaneously. Empirically,

---

Under review as a conference paper at ICLR 2023
Eval Domains PPL (J)
125M
350M
760M
1.3B
TRANSFORMER-LIN
25.2
20.8
19.3
18.4
ELMFOREST parameter average (uniform weights)
31.0
22.4
20.8
19.5
Argmax ELM (one-hot weights)
28.3
22.3
22.3
20.3
ELMFOREST parameter average (posterior weights)
28.5
20.3
18.0
17.0
ELMFOREST ensemble
22.4
18.6
16.7
16.3
Table 3: ELMs can be combined through parameter averaging ($4.4). Average test-set perplexity across the 8 evaluation domains, from the models in Table 1, comparing techniques to collapse ELMFOREST into single LM. The relatively poor performance of ELMFOREST parameter averaging for the 125M parameter model is investigated (and improved) in $5.2.
ELMFOREST training jobs were scheduled and run more quickly, and with less preemption, than the TRANSFORMER-LM and DEMIX training jobs at the same overall budget.
4.4 ELMFOREST PARAMETER AVERAGE
While ELMFOREST substantially improves performance at lower training cost relative to the TRANSFORMER-LM. it comes at the price of larger model size and higher associated inference costs when ensembling. Here, we explore an alternative way to combine experts to improve generalization with additional inference costs relative the TRANSFORMER- LM bascline: parameter averaging. Given some weight vector w over ELMs {0,,... ..., 0 }. we define a single model such that all of its parameters are weighted average of the ELM parameters, according to w: 0 Σ w.O;, For w, we consider a number of options:
We set to be a uniform; i.e., T. This setting disregards the relevance of each ELM to the target domain, assuming all ELMS should contribute equally to the average.
Uniform:
Argmax: We set w to be an indicator vector corresponding to the maximum probability in the domain posterior ($2.2), thus activating only the estimated best-performing ELM.
Posterior: We set to be the domain posterior ($2.2), computed on the validation set.
Results on the evaluation domains are Table 3.4 Using uniform weights underperforms all baselines, even TR TRANSFORMER-LE -.Ms, highlighting the importance of domain relevance in output ensembling and parameter averaging ELMs. Using the argmax ELM outperforms uniform averaging for small models, but not larger models. Weighting the average with the domain posterior outperforms all other techniques, and consistently improves performance over TRANSFORMER- LMM at no additional inference cost. Though parameter averaging achieves lower performance than output ensembling, the lower inference costs and simplicity of deployment may make averaging the preferred inference technique for resource-constrained applications. Due its superior performance, report results for ELMFOREST with output ensembling, unless otherwise noted.
Surprisingly, we observe poor performance model averaging at the 125M scale. We see later that the amount of compute allocated to the seed phase critically affects the viability of ELMFOR- EST parameter averaging ($5.2). With sufficient seed training, parameter averaging outperforms LM at all scales.
ANALYSIS
In $4, we fix the training setup to conduct controlled comparison of BTM to baseline methods. We now analyze the importance of various training and inference decisions on language modeling performance.
We display similar findings with training domains in Appendix Table

---

Under review as a conference paper at ICLR 2023
750M
1.3B
Random Ensemble (seed init)
ELM FOREST (random init)
ELM FOREST (seed init)
Random Ensemble (seed init)
ELM FOREST (random init)
ELM
FOREST
(seed init)
Train
17.4
14.4
13.4
Train
17.4
13.3
13.0
Eval
20.9
19.3
16.7
Eval
20.4
17.8
16.3
All
19.2
16.9
15.0
All
18.9
15.6
14.6
Table 4: Domain expert ensemble outperforms random split ensemble ($5.1), Average test-set perplexity (1) for our largest model scales across the training, 8 evaluation, and all 16 domains. We show similar results for the 125M and 350M parameter scale models in Appendix Figure 10.
5.1 ELMFOREST OUTPERFORMS PARAMETER- MATCHED ENSEMBLES
We compare our method to other LM ensembles to study the effect of increased parameter counts:
Random Ensemble (seed init) A set LMs trained on random data splits, to assess the importance of ELM domain specialization. We pool the training and development sets of the train domains, divide into 8 random splits, then conduct BTM on those splits, with 50% seed training.
ELMFOREST (random init) An ELMFOREST trained with BTM where all ELMs randomly initialized, to assess the effect of seed training. This is equivalent to setting the seed training compute budget to zero updates. We fix the random initialization across models.
ELMFOREST (seed init) The ELMFOREST setting of $4. We conduct BTM on the train domains, and dedicate 50% of the updates in the budget seed and branched ELM training.
Results with the largest models are in Table 4.5 ELMFOREST (random init) nearly matches ELM- FOREST training domains but performs poorly on evaluation domains. The random ensemble is consistently worse than both variants of ELLOFOSSST, showing that the performance improvement is only due ensembling or increased total model size."
5.2 ELMFOREST PERFORMANCE ROBUST SEED LM TRAINING COMPUTE ALLOCATION
The ELMFOREST (random init), which has no seed training. underperforms ELMFOREST (LM init) in 1,.1, indicating that seed training is essential. On the other hand, TRANSFORMER-LMN equivalent to 100% seed training. also underperforms ELOFOOEST (LM init) in $4, which suggests the importance of branched ELM training. We now study the changes to performance when we vary the portion of the compute budget dedicated to seed training. We control for the total compute budget (across seed and branched training).
Our results, in Figure 3, show that the optimal amount of seed training is about 40 60% of the total budget. At both ends of the full range, performance deteriorates, approaching the ELMFOREST (random init) and TRANSFORMER-L.M. performance (at 0% and 100% seed training, respectively).
However, as little as 10% of seed training can be performed to result in strong gains over the ELMFOREST (random init) and TRANSFORMER-LM. This suggests that the majority of BTM training may focus on branched training to dramatically reduced computational costs ($4.3). The optimal share of compute to use towards each training phase likely depends many factors, including the total compute budget. We leave more thorough study of this trend to future work.
E.MFFOEST averaging performance strongest with about 60-70% seed LM training Separate experiments, shown in Appendix Figure 4, suggest strong effect of the seed phase on the viability of ELMFOREST averaging. We find that ELMFOREST averaging does not work with ELMs trained
'ee display similar findings with smaller models in Appendix Table 10.
'We speculate that the random ensemble is poor because its constituent models make correlated errors during evaluation (Gontjj--Lopes al., 2022).
7

---

Under review as a conference paper at ICLR 2023
Training Domains
8 Evaluation Domains
20
26
Params/ELM
Params/ELM
125M
125M
350M
350M
750M
750M
1.38
1.38
22
Perplexity
Perplexity
14
18
16
12
20
40
60
80
100
20 40
60
80
100
% Seed Training
% Seed Training
Figure 3: ELMFOREST ensembling performance is robust to most seed training compute allocations ($5.2). Test perplexity averaged across the training (left) or evaluation (right) domains (from $4.1) when fixing total compute budget but varying the portion allocated to seed training.
Average Test PPL (.)
Train
Evaluation
Overall
TRANSFORMER-IN
19.8
25.5
22.7
seed corpus
train domains
17.2
22.7
20.0
Wikipedia
17.7
23.2
20.5
C4
17.9
23.5
20.7
StackOverflow
18.4
24.6
21.5
JavaScript
19.2
24.9
22.0
Table 5: ELMFOREST ensembling performance robust to seed training corpus ($5.3), Test set perplexity averages the 8 training, 8 evaluation, and all 16 data domains, using different training corpora used in seed LM training. All models are of the 125M parameters per GPU scale.
from random initialization (i.e., with seed phase), resulting in perplexities in the thousands. Since these ELMs still share the same random initialization, we conclude that there is importance to seeding the ELMFOREST with a shared set of partially trained LM parameters.
We observe that parameter averaging performance on training domains is relatively robust to seed training. On evaluation domains, however, the smallest scale ELMFOREST does not achieve optimal performance until about 60% or more updates are dedicated to seed training. This explains the poor performance of the 125M parameter scale ELMFOREST average on evaluation domains in Table 3. ELMFOREST parameter averaging at 50% seed training outperforms TRANS TRANSFORMER- baselines on evaluation domains 350M, 750M, and 1.3B parameter scales (Table 3).
5.3 ELMFOREST PERFORMANCE ROBUST THE CHOICE SEED TRAINING CORPUS
We compare the effects of using different training corpora for seed training in BTM. Here, we fix the compute budget allocations studied in $5.2 so that 50% updates are allocated to seed training and 50% to branched training. As seen in Table 5, our experiments using the most diverse corpora for seed training resulted the best performance, but even seed training on only JavaScript code yielded better results than the compute- matched TRANSFORMER-LE -LM baseline, and far better than the EFMRRTS (random init) models in Table 1, which use identical random initialization. This suggests that initializing EL.Ms with parameters of any model checkpoint is critical.
8

---

Under review as a conference paper at ICLR 2023
Domain forgetting through ELM removal is mostly robust to seed training corpus We evaluate the empirical effect of removing ELMs to reduce the influence of specific training domains at inference time (c.g., those that contain stale or harmful text). In Appendix Table 11, show the performance of MFFORSTT ensembles on the training domains when using all ELMs (from Table 5), and compare performance when removing the relevant ELM;; e.g., evaluate on REDDIT, we keep active all ELMS except the REDDIT ELM. Removal of an ELM from an ELMFOREST guarantees that the domain will be forgotten, if the seed training corpus did include that domain's data ($2.1). 7 We find that performance does indeed degrade appreciably on ELMFORESTS when removing the ELM, indicating that ELMFORESTS are capable of effectively forgetting data domain without any gradient updates to the parameters.
6 RELATED WORK
Sparsely activated language models have been considered in few forms (Evci et al., 2020; Mostafa & Wang, 2019; Dettmers & Zeitlemoyer, 2019), but most related this work is the Eexture-of-xpeeee (MoE) model (Jacobs et al., 1991). Recent MoE models (Shazeer et al., 201 1) have been proposed with boken-based routing Lepikhin et al. (2021), I); Fedus et al. (2022); Lewis et al. (2021); Roller et al. (2021). Of this line of work, ours is most closely related to DEMix layers Gururangan et al. (2022), which replace transformer feedforward layers as domain experts. Similarly, Pfeiffer et al. (2022) develop multilingual expert model with language-specinn routing, and Kudugunta et al. (2021) develop a multi-task expert model with task-specific routing.
Previous work explored extending the capacity of model with additional specialized parameters (e.g., adapters; Houlsby et al., 2019; Pfeiffer et al., 2020; Ben Zaken et al., 2022). Our approach is simpler, as our ELMs each consist of an entire model which requires no additional and no shared parameters. Future work may explore combining ELMs with adapters.
Ensemble methods are classic techniques in machine learning (Breiman, 1996; Freund, 1995; Wolpert, 1992). Similar to our work, recent work has considered growing ensembles, in which new models are trained sequentially on streaming data Caccia et (2021).
Our averaging mechanism is inspired by the model merging techniques in the vision and NLP literature (Wortsman et al., 2022; Izmailov et al., 2018; Matena & Raffel, 2021). Our posterior weighted average is highly related to Bayesian model averaging techniques used in classic ensembling methods (Fragoso al., 2018). Model averaging has also been explored for federated learning (MeMahan et al., 2017), where different models are trained to fit privacy-sensitivy data on different devices and merged. Future work may explore variations of ELM weighted averages.
The importance of the seed training as a critical warm-up phase for BTM aligns with findings that model merging only works when models share part of their optimization trajectory (Frankle et al., gating,
2020; Entezari et al., 2022). Similar to seed training, Nie et al. (2021) propose where mixture-of-e-eee -eeeeessss routing mechanisms are gradually sparsified during the course of training.
7 CONCLUSION
We introduce BTM training, a new algorithm to train an ELMFOREST, which contains many EXPERT LMs that can be added and removed, ensembled, parameter averaged any time for efficient scaling and rapid customization. Our extensive experiments show that ELMFOREST ensembles trained with BTM outperform baselines averaged MO closely approach ELMFFORES ensemble performance while enabling substantially cheaper inference. These results provide compelling evidence for the promise of scaling language models with many smaller, independently trained ELMs.
additional training cost. Additionally, parameter
Thhe actual effect of E.MM removal on model performance may depend on the overlap hetween the removed domain and other training domains.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Roee Aharoni and Yoav Goldberg. Unsupervised domain clusters in pretrained language models. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 7747-7763, Online, July 2020. Association for Computational Linguistics. doi: 10. 10.18653/1//0220
acl-main.692. URL https:tpp//. /ttp://wwllccccccco web...e..... srrg/anthologyy wwtp://22000000 acl.0-.cll-m main. 692.
Mikel Artetxe, Shruti Bhosale, Naman Goyal, Todor Mihaylov, Myle Ott, Sam Shleifer, Xi Victoria Lin, Jingfei Du, Srinivasan Iyer, Ramakanth Pasunuru, Giri Anantharaman, Xian Li, Shuohui Chen, Halil Akin, Mandeep Baines, Louis Martin, Xing Zhou, Punit Singh Koura, Brian O' Horo, Jeff Wang, Luke Zettlemoyer, Mona Diab, Zornitsa Kozareva, and Ves Stoyanov. Efficient large scale language modeling with mixtures of experts, 2021. URL https: /arxivivvvv..///v///aa////aaavva org/ /ab///////// 2212. 10684.
Jason Baumgartner, Savvas Zannettou, Brian Keegan, Megan Squire, and Jeremy Blackburn. The pushshift reddit dataset, 2020.
Elad Ben Zaken, Yoav Goldberg, and Shauli Ravfogel. BitFit: Simple parameter-efficient fine-tuning for transformer-basee masked aogguage-mdeel. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Pupers), pp. 1-9, Dublin, Ireland, URL.
May 2022. Association for Computational Linguistics. doi: 10.18665//1/20222aclsshho https: a antho log wtgg//0220 www.cssssstttt..
Douglas Biber. Variation across Speech and Writing. Cambridge University Press, 1988. doi: 10.0077/CH09980511222222
John Blitzer, Ryan McDonald, and Fernando Pereira. Domain adaptation with structural correspon- dence learning. In Proceedings the 2006 Conference on Empirical Methods in Natural Language Processing, pp. 120-128, Sydney, Australia, July 2006. Association for Computational Linguisties. URL https::pp//.......tt. www acl web org wwwpg//gggogggg/g/oooooooooooo/t wwwhhhooogyy/o-0o0 /100m1/006-111100
Su Lin Blodgett, Lisa Green, and Brendan O'Connor: Demographic dialectal variation in social media: A case study of African-American English. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pp. 1119-1130, Austin, Texas, November 2016. Association for Computational Linguistics. doi: 10. 10.18653//1/DI--11 1120. URL https:
//www. wwwcleebbb.. wwwp:aawww.....gaaanaaaee httpb//.../aattoooooo////l--111-
Leo Breiman. Bagging predictors. Machine learning, 24(2):123-140, 1996.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss. Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020. URL https: //arxiv. / abs 2005 14165.
Lucas Caccia, Jing Xu, Myle Ott, Marc' Aurelio Ranzato, and Ludovic Denoyer. On anytime learning at macroscale. CoRR, abs/2106.09563, 2021. URL https: http:///axvvvvvo wwgpaiiw.oooa/aaaooooa//aa////ae attps:///axdvv..o//aabb/ 2106. 09563.
Caselaw Access Project. Caselaw access project. URL http::/s: law/
Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robinson. One billion word benchmark for measuring progress in statistical language modeling, 2014.
Alexandra Chronopoulou, Matthew Peters, and Jesse Dodge. Efficient hierarchical domain adaptation for pretrained language models. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 1356-1351, Seattle, United States, July 2022. Association for Computational Linguistics. URL https: / ////lalam lanthol l ogy. wwt./rg/22022 2022 naacl -ma i. 96.
10

---

Under review as a conference paper at ICLR 2023
Pradeep Dasigi, Kyle Lo, Iz Beltagy, Arman Cohan, Noah A. Smith, and Matt Gardner. A dataset of information-seeking questions and answers anchored in research papers. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, pp. 4599-4610, Online, June 2021. Association for Computational Linguistics. doi: 10. (0.18653/v1/2011 maacl-main.366 URL https: //acl anthology org/ 2021. maaacl-main...
Sumanth Dathathri, Andrea Madotto, Janice Lan, Jane Hung, Eric Frank, Piero Molino, Jason Yosinski, and Rosanne Liu. Plug and play language models: A simple approach to controlled text generation. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL https: httpss//openreviee. openrevie.. net/forum? iiddiid--i hieeenidddlliiiee HldEEeE....DS.
Tim Dettmers and Luke Zettlemoyer. Sparse networks from scratch: Faster training without losing performance. CoRR, abs/1907.04840, 2019. URL
org/ abs / 1907. 04840.
Rahim Entezari, Hanie Sedghi, Olga Saukh, and Behnam Neyshabur. The role of permutation invariance in linear mode connectivity of neural networks. In International Conference Learning Representations, 2022. URL https://///.. / oooooeeeeee wewirevieww... eee net/ forum? ? d=dNi wyt.:diiigggtiieeo.
Utku Evci, Trevor Gale, Jacob Menick, Pablo Samuel Castro, and Erich Elsen. Rigging the lottery: Making all tickets winners. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 19943-2952 PMLR. 13-18 Jul 2020. URL https: /ttp:://proee http://pooeeeeegeeeeeeeeee mlr. press/vll 9/ evci 20a html.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Reseurch, 23(120): 2022. URL http:// / fmlr...... org/ http//apes/////ee 12233/21-099 html.
Tiago Fragoso, Wesley Bertoli, and Francisco Louzada. Bayesian model averaging: A systematic review and conceptual classification. International Statistical Review, 86:1-28, 04 2018. doi: 10.111 1/insr. 12243.
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode connectivity and the lottery ticket hypothesis. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning. volume 119 of Proceedings of Machine Learning Research, pp. 3259 3269. PMLR, 13-18 Jul 2020. URL https: BtLtttt::////ooeoeo ppppeeeieee..s... mlr. wwwpppressvvvv... wwt.svvvll//// (tt.1///faaeeee01 html.
Yoav Freund. Boosting weak learning algorithm by majority. Information and computation, 121(2): 256-285, 1995.
Samuel Gehman, Suchin Gururangan, Maarten Sap, Yejin Choi, and Noah A. Smith. RealTox- icityPrompts: Evaluating neural toxic degeneration in language models. In Findings of the
Association for Computational Linguisties: EMNLP 2020, pp. 3356 3369,,,,,,,- Online, November
2020. Association for Computational Linguistics. doi: 10. 10.18653/v1/2020, 11166/////222..mdminnnmmmmm/ findings-emmlp.301...
URL hwtpt:tt//// /atps://aalaaaten /awaaaaaaannooooocooaaaacaooooee .rg. 2020/220..0 findinngnnnne emnlp.300
Github Archive Project.
Github archive project.
URL https:
Aaron Gokaslan and Vanya Cohen. Openwebtext corpus, 2019. URL http:://skyllnnn////
Raphael Gontijo-Lopes, Yann Dauphin, and Ekin Dogus Cubuk. No one representation to rule them all: Overlapping features of training methods. In International Conference on Learning Representations, 2022. URL https: /pwpreveeeeeeeeeee/eeeeeee net/ wtw.fforummmiideee M00-BB-444GGGIB3 4qbGgIB3.
Suchin Gururangan, Ana Marasovié, Swabha Swayamdipta, Kyle Lo, Iz Beltagy, Doug Downey, and Noah A. Smith. Don't t stop pretraining: Adapt language models to domains and tasks. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp-
11

---

Under review as a conference paper at ICLR 2023
1342-8360, Online, July 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020.
acl-main.740) 740. URL httpt::///// www. wewweb........ wwtpg//tttggggggotootooooooooott acl.::///..22.........mm........
Suchin Gururangan, Mike Lewis, Ari Holtzman, Noah A. Smith, and Luke Zettlemoyer. DEMix layers: Disentangling domains for modular language modeling. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 5557-5576, Seattle, United States, July 2022. Association for Computational Linguistics. URL http:: http:////l///aalcolooo.ooo/////
naacl -maain.
407.
Dan Hendrycks, Collin Burns, Anya Chen, and Spencer Ball. Cuad: An expert-annotated nlp dataset for legal contract review, 2021.
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficien transfer learning for NLP. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 2790-2799 PMLR, 09-15 Jun 2019. URL https: http://occeeddeeee mlr.
press/v9 97 hhooohoooooo 1by1 9a html.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization, 2018. htth::/::// //arxiv..
org/aba /a////////// 1803. 05407.
Robert A Jacobs, Michael Jordan, Steven J Nowlan, and Geoffrey Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79-87, 1991.
Nitish Shirish Keskar, Bryan McCann, Lav R. Varshney, Caiming Xiong, and Richard Socher. Ctrl: conditional transformer language model for controllable generation, 2019.
Sneha Kudugunta, Yanping Huang, Ankur Bapna, Maxim Krikun, Dmitry Lepikhin, Minh- Thah-Thang Luong, and Orhan Firat. Beyond distillation: Task-level mixture-of-experts for efficient inference. In Findings of the Association for Computational Linguistics: EMNLP 2021, pp. 3577-3599, Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10.18653/v 1/20211111111111// 1...02...inniige. aandings-emnlp.... find Andings-emn 1 p 304.
URL https: https///aaaaaatthoooooooo org 2021.
Angeliki Lazaridou, Adhiguna Kuncoro, Elena Gribovskaya, Devang Agrawal, Adam Liska, Tayfun Terzi, Mai Gimenez, Cyprien de Masson d' A'ttume Tomáš Kočiský, Sebastian Ruder, Dani Yogatama, Kris Cao, Susannah Young, and Phil Blunsom. Mind the gap: Assessing temporal generalization in neural language models. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https:// openreview net forum? ? id 30mmrrmmyms
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. (GS) hhard: Scaling giant models with condi- tional computation and automatic sharding. In International Conference on Learning Representa- tions, 2021. URL https: wpw.eeeeeveeeeeeeeeeeeeeeeeeeeee vivviiese iew, net/ forum ? (ecrummiidie id==qree7xHTnYb 7xHTmYb.
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large, sparse models. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 6205-66666666 6274. PMLR, 18-24 Jul 2021. URL https: htttts/://ppoceeeeeee /ttp:////cocccces....... mar.ppres//oe vv..svv1139 pttess/vv19//ee 1tt1//////wwwaa.aa.e.. html.
Shen Li. Getting started with distributed data parallel, 2021. URL http::/ts// wattorraas/intereddiii
html.
Kyle Lo, Lucy Lu Wang, Mark Neumann, Rodney Kinney, and Daniel Weld. S2ORC: The semantic scholar open research corpus. In Proceedings of the 58th Anmual Meeting of the Association for Computational Linguistics. pp. 4969-4983 4983, Online. July 2020. Association for Computa- tional Linguistics. doi: 10. 101118661//2200000 m47n.447. URL https: / /www. aclweb. org/ anthology/222 oomolllyy22000 acl main 447.
12

---

Under review as a conference paper at ICLR 2023
Li Lucy and David Bamman. Characterizing English variation across social media communities with BERT. Transactions of the Association for Computational Linguistics, 9:538-555 556, 2021. doi: 10.1162/tacl 00383. URL https: httss///aaaaaatttaaaaeeeaaa wrg. www......c........
Kelvin Luu, Daniel Khashabi, Suchin Gururangan, Karishma Mandyam, and Noah A. Smith. Time waits for no one! analysis and challenges of temporal misalignment. 2021. doi: 10.48550/ARXIV. 2111.07408. URL https:/ /aaaiiivvvv.....aa org/ abs//bs//211/ /2111, 07408.
Michael Matena and Colin Raffel. Merging models with fisher-weighted averaging. arXiv preprint arXiv:2111.09832, 2021.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. learning of deep networks from decentralized data. In Artificial intelli- gence and statistics. pp. 1273-1282. PMLR, 2017.
Communication-
Hesham Mostafa and Xin Wang. Parameter efficient training of deep convolutional neural networks by dynamic sparse reparameterization. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. proceedings mwr.mreeeeee..............e.eeee / v97 /tt////////////// /two//ttaaa9999 9a. html.
PMLR, 09-15 Jun 2019. URL https:/
Jianmo Ni, Jiacheng Li, and Julian McAuley. Justifying recommendations using distantly-labeled reviews and tme-grained aspects. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Nutural Language Processing (EMNLP-IJCNLP), pp. 188-197, Hong Kong, China, November 2019. Association for Computational I Linguistics. doi: 10. 10.18////11D19-1 1018. URL https: www. a web. org/ anthology DD19- 1018.
Xiaonan Nie, Shijie Cao, Xupeng Miao, Lingxiao Ma, Jilong Xue, Youshan Miao, Zichao Yang, Zhi Yang, and Bin Cui. Dense-to-sparse gate for mixture-of-experts, 2021. URL https: //arxiv. www./rr//aaaao abs / 221/222111../ 14397.
Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: fast, extensible toolkit for sequence modeling. Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics (Demonstrations)) pp. 48-53, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10. 10111//////11190009 URL https:/ /ww. aclweb. org/ anthology/ N1 4009.
Ou- Yang Lucas. Newspaper3k. latest/.
URL https: reattppep wwtpiidocccsoon/ieeeeeeeee in/en/
Jonas Pfeiffer, Ivan Vulić, Iryna Gurevych, and Sebastian Ruder. MAD-X: An Adapter-Based Framework for Multi-Tast -aak Cross- Lingull Transfer. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 7654-7673, Online, November 2020. Association for Computational Linguisties. doi: 10. (0.18653/v1/2020 https: htwwww/ac... web. org/ anthooo ogy/ 2020. emn lomnlp-mal iep.miii..... 6.7.
www.:200mmmmmmmmm....mm URL
Jonas Pfeiffer, Naman Goyal, Xi Victoria Lin, Xian Li, James Cross, Sebastian Riedel, and Mikel Artetxe. Lifting the curse of multilinguality by pre-training modular transformers, 2022. URL
https: /ttps///arae /twpvva.vvvv///////////../...av. org/ alss 2205. 0 6666.
Project Gutenberg. Project gutenberg. URL www www.gueeebbrggggom
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019. URL https: // //d/mucfpksyww....mmmmmmoopkk/// hts.mmccpcsyy...ooofffototo net http:bbtteer--nngageg mode 1s /slelllsllan_... models are aww_uunuuerrvieeeeeeeeeeeeeeeeee wwwprprrrreeeeeeteea.eea..e.
multitask Leareereeeerrrreeeeeeeeeeeeeeeee pdf.
John R. Rickford, Ethnicity sociolinguistic boundary. American Speech, 60:99, 1985.
13

---

Under review as a conference paper at ICLR 2023
Stephen Roller, Sainbayar Sukhbaatar, Arthur Szlam, and Jason E Weston. Hash layers for large sparse models. A. Beygelzimer, Y. Dauphin, Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL httpp:ss//peeeerroii.ettttttttte idd [d==MMDDWbl U.M.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc V. Le, Geoffrey E. Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-ofeexperts layer. In 5th International Conference Learning Representations, ICLR 2017, Toulon, France, April 24- 26, 2017, Conference Track Proceedings. OpenReview.net, 2017. URL httph:://ppeeeeeeeeeee net / Eorum www.mmnddd1cckdaee2--
Lucy Lu Wang, Kyle Lo, Yoganand Chandrasekhar, Russell Reas, Jiangjiang Yang. Doug Burdick, Darrin Eide, Kathryn Funk, Yannis Katsis, Rodney Kinney, Yunyao Li, Ziyang Liu, William Merrill, Paul Mooney, Dewey Murdick, Devvret Rishi, Jerry Sheehan, Zhihong Shen, Brandon Stilson, Alex Wade, Kuansan Wang, Nancy Xin Ru Wang, Chris Wilhelm, Boya Xie, Douglas Raymond, Daniel Weld, Oren Etzioni, and Sebastian Kohlmeier. Cord-199 19: The covid-19 open research dataset, 2020.
Sean Welleck, Ilia Kulikov, Stephen Roller, Emily Dinan, Kyunghyun Cho, and Jason Weston. Neural text generation with unlikelihood training. In International Conference Learning Representations, 2019.
David H Wolpert. Stacked generalization. Neural networks, 5(2):241-259) 1992.
Mitchell Wortsman, Gabriel Ilharco, Samir Yitzhak Gadre, Rebecca Roclofs, Raphael Gontijo-Lppes, Ari S. Morcos, Hongscok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, and Ludwig Schmidt. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time, 2022. URL https: http///xxvvvv. org. 220g2333b/222...../// 05482.
Jing Xu, Da Ju, Margaret Li, Y-Lan Boureau, Jason Weston, and Emily Dinan. Bot-adversarial dia- logue for safe conversational agents. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 2950-2968, Online, June 2021. Association for Computational Linguistics. doi: 10.18653/v1/ 2021 maiacclmmii.2253 235. URL ht / org/2021 www.clacllae... -main.ainm.mmmm 235.
maacc-main....
Yelp Reviews. Yelp reviews. URL https://:
Rowan Zellers, Ari Holtzman, Hannah Rashkin, Yonatan Bisk, Ali Farhadi, Franziska Roesner, and Yejin Choi. Defending against neural fake news. In NewrlPS, 2019. URL http::ttpp//// /ttttprr//pppr..... ni
de in neuraanneualll
.dff.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Pumit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettlemoyer. Opt: Open pre- trained transformer language models, 2022. URL https : / arxiv. org/ abs/ 2205. 01068.
14

---

Under review as a conference paper at ICLR 2023
Category
Link to Regex
Dummy Token
Email
http:://rggee//..ccm///////////
<EMA <DART>
DART
FB User ID
httpss//aagg2200.co//////22112
<FB USERRID>>>>> <PHONE NNMBER>>>>>>
Phone Number
https////ege/10..co//////////
Credit Card Number
Social Security Number
<SSN>
User handles
<USER>
Table 6: De-identification schema. We de-identify text using the regexes provided in the above links for the categories listed.
A
APPENDIX
A.1 LIMITATIONS
The definition of a domain The nature of domains in NLP is a matter of active research. Textual domains reflect language variation that stems from factors such vocabulary differences (Blitzer et al., 2006), sociolinguistic (Biber, 1988) or demographic (Rickford, 1985; Blodgett et al., 2016) variables, community membership (Lucy & Bamman, 2021), end-tasks (Gururangan et al., 2020), temporal shifts (Lazaridou et al., 2021; Luu et al., 2021). this work, we follow Gururangan et al. (2022) and define domains by provenance, or the source of the document. Provenance labels yield simple and interpretable segmentations of a corpus, which are useful for identifying ELMs in our experiments. However, other methods for discovering domains, including unsupervised techniques (Aharoni & Goldberg, 2020; Chronopoulou et al., 2022), may yield better expert assignments. We leave experimentation with other definitions of domain for future work.
Domain posterior data requirement To calculate the domain posteriors used for our ensembling and parameter averaging weights, we assume access to a small additional sample of data to train the vector w. While is easy to imagine that extra data may be available for most applications to estimate the posterior, future work may explore the possibility of eliminating this requirement.
Other distributed training baselines Our TRANSFORMER-LYY baseline is implemented with distributed data- aata-parallel. Model-part Morellparalle,, fully sharded data parallel, and other distributed training strategies (Artetxe al., 2021) confer different scaling patterns that may change the conclusions that we report in this work. However, we expect that BTM will provide strong efficiency gains against these alternatives.
Harms of language models BTM results in an LM whose test time behaviors can be controlled with much stronger guarantees after training due to the isolation of domains in ELMs. However, ELMFORESTS exposed to large datasets scraped from the Internet may contain toxic language (e.g.. hatespeech) that are difficult to identify with coarse provenance domain labels, and nevertheless result harmful output from the ELMs (Gchman et al., 2020). Future work may explore recipes for training and deploying ELMFORESTS to better support user safety.
15

---

Under review as a conference paper at ICLR 2023
Domain
Corpus
# Train (Eval.) Tokens
IB
30M WewsWire sentences (Chelba et al., 2014)
700M (10M)
CS
1.89M full-text CS papers from S2ORC (Lo et al., 2020)
4.5B (10M)
TRAINING
LEGAL
2.22M U.S. court opinions (Caselaw Access Project)
10.5B (10M)
MED
3.2M full-text medical papers from S2ORC (I.o et al., 2020)
9.5B (10M)
WEBTEXT
8M Web documents (Gokaslan & Cohen, 2019)
6.5B (10M)
REALNEWS!
35M articles from REALNEWS Zellers al. (2019)
15B (10M)
REDDIT
Reddit comments from pushshift.io (Baumgartner et al., 2020)
25B (10M)
REVIEWS
30M Amazon product reviews (Ni al., 2019)
2.1B (10M)
Total
73.8B (80M)
Domain
Corpus
# Train (Eval.) Tokens
ACL PAPERS
1.5K NLP papers from ACL (Dasigi et al., 2021)
1M (IM)
BREAKING NEWS
20K English news articles, scraped using YOngY Lucas)
11M (IM)
EVALUATION
CONTRACTS!
500 commercial legal contracts (Hendrycks al., 2021)
1.5M (1M)
CORD-19
400K COVID-19 research papers (Wang et al., 2020)
60M (10M)
GITHUB
230K public Github code (Github Archive Project)
200M (10M)
GUTENBERG
3.2M copyright-expire books (Project Gutenberg)
3B (10M)
TWBETS
IM English tweets from 2013-2018
8M (1M)
YELP REVIEWS!
6M Yelp restaurant reviews (Yelp Reviews)
600M (10M)
Table Multi-domain data corpus used in $4 and $5. Details of this corpus, both training and evaluation domains, including the size of our training and evaluation (i.e. validation and test) data in whitespace -eiispaace-eparatee tokens. We borrow these datasets from Gururangan et (2022). f indicates datasets we de-identify with regexes in Table 6. REDDIT was de-identified by Xu et al. (2021); we use their version. Meta researchers did not collect any the Reddit or Twitter data and the data was not collected behalf of Meta.
Train Domains PPL. (L)
125M
350M
760M
1.3B
TRANSFORMER-LN
19.9
16.3
14.7
14.2
ELMFOREST parameter average (uniform weights)
47.4
19.9
19.0
18.0
Argmax ELM (one-hot posterior)
18.0
15.3
14.1
13.8
ELMFOREST parameter average (posterior weights)
18.0
15.1
13.9
13.4
ELMFOREST ensemble
17.2
14.7
13.4
13.0
Table 8: Performance of ELM parameter averaging training domains ($4.4). Average test-set perplexity across the training domains, from the models in Table 1, comparing techniques to collapse ELMFOREST into single LM. As with evaluation domain results in the main paper, parameter averaging (with posterior weights) generally yields better average perplexities than TRANSFORMER- LM no additional inference cost, but underperforms ELMFOREST ensembling, which uses more effective parameters and is included for comparison as a lower bound.
16

---

Under review as a conference paper at ICLR 2023
125M - 16 GPUs 80k updates
T-LM
T-LM
DEMIX
ELMFOREST (random init)
RANDOM ENSEMBLE
ELMFOREST
UNBALANCED
(seed init)
125M
125M
512M
1B
IB
1B
Train
19.8
20.7
17.7
18.0
23.0
17.2
Novel
25.6
26.4
23.1
24.1
26.0
22.4
All
22.7
23.5
20.4
21.0
24.7
19.8
350M 32 GPUs
32k updates
T-LM
T-LM
DEMIX
ELMFOREST (random init) 2.8B
RANDOM ENSEMBLE
ELMFOREST
UNBALANCED 350M
(seed init)
350M
1.8B
2.8B
2.8B
Train
16.3
16.7
15.0
15.3
19.9
14.7
Novel
20.8
21.2
19.9
21.3
23.1
18.6
All
18.5
19.0
17.5
18.3
21.5
16.7
750M
64 GPUs
24k updates
T-LM
T-LM
DEMIX
ELMFOREST (random init) 6B
RANDOM ENSEMBLE
ELMFOREST
UNBALANCED 750M
(seed init)
750M
3.8B
6B
6B
Train
14.7
14.9
13.5
14.4
17.4
13.4
Novel
19.3
19.8
17.7
19.3
20.9
16.7
All
17.0
17.4
15.6
16.9
19.2
15.0
1.3B
128 GPUs
12k updates
T-LM
T-LM UNBALANCED
DEMIX
ELMFOREST (random init) 10.4B
RANDOM ENSEMBLE
ELMFOREST
(seed init)
1.3B
1.3B
10.4B
10.4B
Train
14.2
15.0
13.7
13.3
17.4
13.0
Novel
18.4
19.5
17.6
20.4
16.3
All
16.3
17.3
15.6
15.6
18.9
14.6
Table ELMFOOESTS trained with BTM outperform all baselines and ensemble variations across multiple model scales. Average rest-set perplexity (L) for each model scale (125M, 350M, 750M, 1.3B parameters) across the training, novel, and all domains described in $4.1. Total compute budget (in update numbers) and GPU usage are shown for each model size, and total parameters are shown for each model type at each size. TRANSFORMER-LMM. (here, abbreviated to T-LM) trained without balancing between data domains performs worse than T-LM trained with data balancing: hence, we only compare against the balanced T-LM setting in $4. For ELMFOREST, show results with 50% dense training.
125M
350M
Random Ensemble (seed init)
ELM FOREST (random init)
ELM FOREST (seed init)
Random Ensemble (seed init)
ELM FOREST (random init)
ELM FOREST (seed init)
Train
23.0
18.2
17.2
Train
19,9
15.3
14.7
Eval
26.0
23.4
22.4
23.1
21.3
18.6
All
24.7
20.8
19.8
All
21.5
18.3
16.7
Table 10: Domain expert ensemble outperforms random split ensemble ($5.1). Average test-set perplexity (L) for our smallest model scales across the training, evaluation, and all 16 domains.
17

---

Under review as a conference paper at ICLR 2023
Training Domains
Evaluation Domains
20
40
seed training (>1000 ppl)
19
0% seed training 1>1000 ppl)
35
18
Perplexity
perplexity
30
17
25
15
Params
Params
125M
125M
350M
350M
20
% Seed Training
% Seed Training
Figure The seed phase is vital to our ability to parameter average ELMs ($5.2). Test perplexity averaged across the 8 training (left) and 8 evaluation (right) domains when averaging ELMFOREST with different seed training compute allocations for the 125M and 350M parameter LMs. Seed training is critical to reasonable averaging performance.
w/ ELM
-ELM
(Average Test PPL 4)
(A PPL)
8 train domains
17.2
(+9.4)
Wikipedia
17.7
(+11.9)
seed corpus
C4
17.9
(+11.8)
StackOverflow JavaScript
18.4
(+12.7)
19.2
(+13.6)
Table 11:
The ability to reduce the influence of domains through ELM removal is (mostly) robust to seed training corpus ($5.3). We present the average test perplexity for the train domains in ELMFORESTS where all ELMs active. We vary the seed training corpora. In parentheses, show the increase in perplexity when the ELM trained to specialize on each domain is removed at inference time. Large increases desired and suggest the ease of removing (e.g., stale harmful) data from the ELMFOREST'S distribution after training.
18


Reviews:
[{"review_id": "y7V_XAO6L-", "reviewer": "Reviewer_Rw95", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper introduces BTM training, a new algorithm to train an ELMForest, which contains many expert LMs that can be ensembled or parameter averaged at any time for efficient training and rapid customization. \n\nIn the training process, it first trains a seed model and then make k copies of the model, and train the k models independently, which makes the training process highly parallelizable. ", "summary_of_the_review": "The paper proposes a new way to train ensemble of neural language models, in order to reduce the model inference cost, the paper proposed using domain posterior to activate only related trained experts. The assumption is not used in other models, so it seems unfair. ", "strength_and_weaknesses": "Strength\n1. The experiments are very extensive, the ELMforest ensembles outperform baselines at no additional training cost.\n\nWeekness\n1. The evaluation for the parameter averaged ELMforests seems not fair, given that during evaluation the parameter averaged model can activate domain related experts on the test set. The hidden assumption is that the consecutive sentences in the test set are from the same domain. \n2. If the parameter average approach doesn't give reasonable performance, the inference cost can be very high if we want to gain good performance. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "PgL-1jlJip", "reviewer": "Reviewer_kH7N", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper propose a new LM architecture that trains branches of LMs on different dataset domains and ensemble the results by taking a weighted average of the output logits or the model parameters. Adding a new LM branch can be done by first taking a weighted average over the existing weights, and then training on the new domain.", "summary_of_the_review": "I have concerns about the novelty of the work and the experiments to evaluate the proposed idea.", "strength_and_weaknesses": "The paper proposes a promising modification to the current LLM training scheme. Specifically, I found the way the authors combine the outputs with domain posterior interesting. However, I have the following concerns regarding the proposed method and its evaluation:\n\n- The technique of ensemble of many deep learning models is a well-known idea in deep learning. The paper uses a standard ensemble method with limited algorithmic innovations. Please highlight the difference between your work and previous works on\n- The results in Table 1 are not fair because different models have different numbers of parameters. A more meaningful experiment is to compare the accuracy of the ELM model with a Transformer-LM with the same number of parameters and check whether the ELM model can match the performance with less training cost, or can outperform Transformer-LM under the same training budget.\n- In addition, one can also ensemble the baseline Transformer-LM across different checkpoints and this can also boost the final inference performance. How does this baseline compare with the proposed method?\n- The study on the speed/efficiency of the proposed method is not complete. Table 2 only provides the per-step speed during training. Please provide a more detailed study on the convergence/inference speed of the expert-based LM.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "McjcFSRDhn", "reviewer": "Reviewer_B2er", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper is on build language models via embarrassingly parallel training mechanism. More specifically, the author propose a branch train and merge method, which could build the large-scale language models by independently training subparts of a new class of LLMs on different subsets of the data. Such a way is able to eliminate the massive multi-node synchronization currently required to train LLMs. The results show that this training paradigm could be used to efficiently scale LLMs. \n", "summary_of_the_review": "The authors design an elegant and impactful pre-training design, which could potential change the current pre-training mechanism and reduce the training costs. The findings are valuable and will be more impactful if more insights and hypothesis could be provided. The performance on the downstream tasks are missing. The paper quality will be further improved if the downstream task experiments could be added. ", "strength_and_weaknesses": "Strength:\nThe studied topic is at the core of community's interest. How to scale the large scale language models  with reduced costs is important for many tasks. The branch, train and merge framework has a clear design and the empirical studies supported the design idea and indicate the problems that the authors has faced. The proposed method is validated on the large scale datasets and suggest the effectiveness of the paradigm.\n\nWeakness:\n1. Some of issues still remain unknown and insights behind this design is also lacking. For example, why the model need to branch from some LMs on the similar domains? Why does the merging of models trained on different domains lead to better performance rather than the models trained on random splits?\n2. The trained LLMs are not evaluated on the downstream tasks. The evaluation of downstream tasks are useful to understand the effectiveness of this training paradigm. GPT style models usually have strong few-shot abilities and strong generation abilities. Look forward to seeing the corresponding results to support the effectiveness claims. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "K04aOfMXyx", "reviewer": "Reviewer_uT5X", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper introduces BTM training to train a collection of Expert LMs that built an ELMFOREST. Each component in the forest can be dynamically added/removed/ensembled (merged) at any time. The idea of the method is to train scalable and parallel individual models that each specialize in a particular domain. The paper presents comprehensive experiments and analyses. The study improves in- and out-of-domain perplexities as compared to GPT-style transformer LMs.\n\nMain contributions:\n- Propose a new framework to efficiently train language models\n- Evaluate the proposed method on various datasets (in- and out-domains)", "summary_of_the_review": "The paper has new ideas for learning efficient language models in a parallel and non-synchronized manner. And the proposed method can be effectively used to build language models via ensemble.\n\nI would gladly see the paper at the conference and give a \"marginally above the acceptance threshold\" score. However, there are some questions and concerns, and I will adjust the score if the authors address my concerns.", "strength_and_weaknesses": "**Strengths:**\n-  Interesting approach of learning models in parallel and running on multi-node. The paper introduces a new learning framework to train multiple models on different domains and merge them.\n\n**Weaknesses:**\n- The evaluation is limited to language model perplexity. The work would be more complete to have the results on other downstream tasks. Did you have any chance to evaluate the model on downstream NLP tasks? And does the model have the same capability as standard GPT models (i.e., in-context learning or generation tasks)?\n- The motivation is unclear on why the authors proposed to branch-train-and merge. Why do you choose to branch the LMs? How can they be effective?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject