Venue: NAACL 2022

# DEMix Layers: Disentangling Domains for Modular Language Modeling

Suchin Gururangan\({}^{\dagger\diamond}\) Mike Lewis\({}^{\diamondsuit}\) Ari Holtzman\({}^{\dagger}\)Noah A. Smith\({}^{\spadesuit}\) Luke Zettlemoyer\({}^{1\diamondsuit}\)

\({}^{\dagger}\)Paul G. Allen School of Computer Science & Engineering, University of Washington

\({}^{\spadesuit}\)Allen Institute for AI

\({}^{\diamondsuit}\)Meta AI

Seattle, WA, USA

sg01@cs.washington.edu

###### Abstract

We introduce a new domain expert mixture (DEMix) layer that enables conditioning a language model (LM) on the domain of the input text. A DEMix layer includes a collection of expert feedforward networks, each specialized to a domain, that makes the LM _modular_: experts can be mixed, added, or removed after initial training. Extensive experiments with autoregressive transformer LMs (up to 1.3B parameters) show that DEMix layers reduce test-time perplexity (especially for out-of-domain data), increase training efficiency, and enable rapid adaptation. Mixing experts during inference, using a parameter-free weighted ensemble, enables better generalization to heterogeneous or unseen domains. We also show it is possible to add experts to adapt to new domains without forgetting older ones, and remove experts to restrict access to unwanted domains. Overall, these results demonstrate benefits of domain modularity in language models.

+
Footnote †: July 10-15, 2022 ©2022 Association for Computational Linguistics

## 1 Introduction

Most language models (LMs) are trained with data homogeneity: all parameters are updated to minimize the loss on all of the data. We refer to this as _dense training_. Dense training leaves variation in the data, or _domains_, to be implicitly discovered Aharoni and Goldberg (2020), assuming that models will be able to fit all domains equally well.

While dense training is convenient, and densely trained LMs achieve impressive results Brown et al. (2020), the approach has drawbacks with respect to generalization, efficiency, and flexibility. Even if training data is sourced from many domains, dense training can in practice emphasize subsets of the data in proportion to their ease of access Oren et al. (2019); Fan et al. (2020), limiting generalization to less prevalent domains. Updating all parameters of the network gets substantially more expensive as model size grows Strubell et al. (2019), making fine-tuning or domain adaptation hard to perform with smaller computational budgets. It is also difficult to adapt to new domains without forgetting the original data McCloskey and Cohen (1989); Aghajanyan et al. (2021) or to restrict access to certain domains the LM has been exposed to during training (e.g., those that contain hatespeech; Bender et al. 2021), leading to risks of unwanted behavior Gehman et al. (2020).

To address these limitations of dense training, we argue that LMs should be designed with _modularity_. We propose a modular LM that has components specialized to distinct domains in the training data, and can be customized at inference-time by mixing,

Figure 1: Illustration of a DEMix layer in a single transformer block. During training, expert feedforward networks are conditionally activated based on the domain (here, document provenance) of the input sequence (i.e., scientific papers or court opinions). At inference time, the language model has new modular functions: domain experts can be mixed to handle heterogeneous domains (e.g., COVID-19 papers), added to adapt to novel domains (e.g., Github code), or removed to reduce the influence of unwanted domains (e.g., social media). Image attribution in §A.1.



adding, or removing these separated components as needed. This design principle emphasizes the ability to rapidly adapt the LM after training, a need that has been broadly advocated for language systems (Dinan et al., 2021; Lazaridou et al., 2021).

We introduce modularity into an LM with a new domain expert (DEMix) layer that explicitly conditions the LM on the domain of the input text (when it is known), or estimates the input domain during inference (when it is not known). A DEMix layer is a drop-in substitute for a feedforward layer in a transformer LM (e.g., GPT-3), creating a specialized version of the layer (or _expert_) per domain (see Figure 1; SS3).

This is an example of conditional computation (Fedus et al., 2021; Lepikhin et al., 2020; Lewis et al., 2021; Roller et al., 2021). Unlike dense training, conditional computation activates different parameters for different inputs. Instead of learning how to route data to experts, the DEMix layer routing mechanism follows from a natural, observable segmentation of the data.1

Footnote 1: We perform a detailed comparison of learned and DEMix routing in §A.5.

We identify domains using coarse provenance categories (e.g., whether a document is a medical research paper or a Reddit post; SS2). Training on data from eight different domains, we find that replacing every feedforward layer in the transformer with a DEMix layer consistently improves in-domain performance (SS4). To improve performance in settings in which the target data does not clearly align with a single domain, we introduce a parameter-free probabilistic approach to dynamically estimate a _weighted mixture_ of domains during inference (SS5). We observe that expert mixing provides especially strong performance gains on _novel_ test-time domains, as well as consistent performance improvements on test data from the _training_ domains, which may themselves be heterogeneous.

Our results suggest that DEMix consistently improves model generalization, especially out-of-domain, while enabling many new modular capabilities. Because DEMix forces experts to specialize to domains, the overall model can be (partially) disentangled after training. Beyond mixing, we can add (SS6) or remove (SS7) domain experts, predictably changing model behavior at inference time. Adding experts allows for model adaptation without updating all parameters (hence avoiding forgetting), and removing experts allows for simulating the removal of training domains without additional training. These results, in aggregate, demonstrate the considerable benefits of moving away from treating data homogeneously during language modeling. Our code is publicly available.2

Footnote 2: github.com/kernelmachine/demix

## 2 Multi-Domain Corpus

To better measure domain modularity, we introduce a new multi-domain corpus constructed with _domain provenance_ that records the original dataset each document appeared in (Table 1). Defining domains in this way is intuitive and conveys a great deal about the type of language that can be expected in each document. Other accounts of domains (e.g., Lucy and Bamman, 2021; Gururangan et al., 2020) may be studied in future work. While other multi-domain corpora (Koh et al., 2021; Gao et al., 2020) cover many more domains, our corpus is restricted to datasets with more permissive licensing to support reproducibility.

We divide our data into training and test domains. The **training** domains text from eight English corpora (top of Table 1), each of which varies in complexity and coverage, totaling 73.8B whitespace-separated tokens. Our test (or **novel**) domains include eight collections of English text (bottom of Table 1), which may or may not align with the training domains. The novel domains allow us to measure how models generalize to a more challenging data distribution shift, where domain boundaries may be less clear.

§A.2 has more details on how these data were collected. For larger domains, we use an additional 10M tokens for the validation and test sets each. Smaller domains have 1M tokens in each. To support future work with the data, we also release an API to download and preprocess it into a format compatible with Fairseq (Ott et al., 2019).3

Footnote 3: github.com/kernelmachine/demix–data

## 3 DEMiX Layer

### Background: Mixture-of-Experts Transformers

The transformer architecture interleaves multi-head self-attention, layer-norms, and feedforward networks (Vaswani et al., 2017). Our focus is on the feedforward component:

\[{\bf h}_{t,\ell}={\rm FFN}({\bf h}_{t,\ell-1}),\] 

where \({\bf h}_{t,\ell}\) is the vector for the \(t\)th token produced by layer \(\ell\).

Shazeer et al. (2017) propose to replace dense feedforward layers with an ensemble of \(n\) experts \({\rm FFN}_{1},\ldots,{\rm FFN}_{n}\), assigned weights respectively by functions \(g_{1},\dots,g_{n}\):

\[{\rm FFN}({\bf h}_{t,\ell-1})=\sum_{j=1}^{n}g_{j}({\bf h}_{t,\ell-1})\cdot{\rm F FN }_{j}({\bf h}_{t,\ell-1}) \tag{2}\]

The \(g\) function routes tokens to different experts, usually each a separate dense feedforward network. If \(g\) routes to a single expert, then the computational cost (in floating-point operations; FLOPs) will be same as a corresponding dense network, even though it has more than \(n\) times as many parameters.

### DEMix Routing

Previous approaches _learn_ the weighting functions \(g\) at a token-level, and either assign at most one Fedus et al. (2021) or two Lepikhin et al. (2020) experts per token. This requires careful load balancing to encourage the model to use all experts, motivating work on explicit balancing mechanisms Lewis et al. (2021).

Instead of learning \(g\), we use domain metadata to route data to experts at the _document_ (i.e., sequence) level. During training, every token in an input text is assigned to the same expert based on the domain label.

Let \(\mathcal{D}\) denote the set of domain labels (i.e., the eight labels in Table 1). If we index the experts by \(\mathcal{D}\) and \(d\in\mathcal{D}\) is the domain label for the current training instance, then

\[g_{j}({\bf h}_{t,\ell})=\left\{\begin{array}{ll}1&\text{if $j=d$}\\ 0&\text{otherwise}\end{array}\right. \tag{3}\]

We assume that each _training_ document is associated with a single domain label. However, we relax this requirement at inference time (SS5), to model unseen or heterogeneous domains.

We perform a detailed comparison of DEMix routing with GShardLepikhin et al. (2020), a mixture-of-experts transformer LM with learned token-level routing, in SSA.5. Our results suggest that learned token-level routing does not enable modularity, underperforms DEMix at similar computational budgets (especially on novel domains), and is much less efficient to train and evaluate.

### DEMix Architecture

Our design results in one expert in a DEMix layer per domain (i.e., eight experts for eight training domains in our multi-domain corpus).

We replace _every_ feedforward layer in the transformer with a separate DEMix layer, in contrast to previous work Fedus et al. (2021); Lepikhin et al. (2020) that interleaves shared and expert layers.

\begin{table}
\begin{tabular}{l l l r} \hline \hline  & **Domain** & **Corpus** & **\# Train (Eval.) Tokens** \\ \cline{2-4}  & 1B & 30M NewsWire sentences Chelba et al. (2014) & 700M (10M) \\  & CS & 1.89M full-text CS papers from S2ORC Lo et al. (2020) & 4.5B (10M) \\  & Legal & 2.22M U.S. court opinions, 1658 to 2018 Caselaw Access Project) & 10.5B (10M) \\  & Med & 3.2M full-text medical papers from S2ORC Lo et al. (2020) & 9.5B (10M) \\  & WebText\({}^{\dagger}\) & 8M Web documents Goksalan and Cohen (2019) & 6.5B (10M) \\  & RealNews\({}^{\dagger}\) & 35M articles from RealNews Zellers et al. (2019) & 15B (10M) \\  & REddit & Reddit comments from pushshiftio Baumgartner et al. (2020) & 25B (10M) \\  & Reviews\({}^{\dagger}\) & 30M Amazon product reviews Ni et al. (2019) & 2.1B (10M) \\ \cline{2-4}  & \multicolumn{2}{c}{**Total**} & \multicolumn{1}{r}{**73.8B (80M)**} \\ \cline{2-4}  & **Domain** & **Corpus** & **\# Train (Eval.) Tokens** \\ \hline  & ACL Papers & 1.5K NLP papers from ACL Dasigi et al. (2021) & 1M (1M) \\  & Breaking News\({}^{\dagger}\) & 20K latest articles from 400 English news sites Baly et al. (2018) & 11M (1M) \\  & Contracts\({}^{\dagger}\) & 500 commercial legal contracts Hendrycks et al. (2021) & 1.5M (1M) \\  & ORD-19 & 400K excerpts from COVID-19 research Wang et al. (2020) & 60M (10M) \\  & Github & 230K public Github repository contents (Github Archive Project) & 200M (10M) \\  & Gutenberg & 3.2M copyright-expired books Project Gutenberg & 3B (10M) \\  & Tweets\({}^{\dagger}\) & 1M English tweets from 2013-2018 & 8M (1M) \\  & Yelp Reviews\({}^{\dagger}\) & 6M Yelp restaurant reviews Yelp Reviews & 600M (10M) \\ \hline \hline \end{tabular}
\end{table}
Table 1: Domains that make up our multi-domain training corpus, including the size of our training and evaluation (i.e. validation and test) data, in whitespace-separated tokens. \(\dagger\) indicates datasets that we (partially) anonymize §2). See Appendix §A.2 for more details on how these data were collected.

 

Preliminary experiments showed that interleaving led to worse in-domain performance (see SSA.6 for more details). Future work may comprehensively compare different architectural choices.

Each expert \({\rm FFN}_{j}\) is a two-layer MLP with the same dimensions as the original \({\rm FFN}\) layer of the transformer. This means that the effective number of parameters in the overall DEM1x LM increases (Table 2), without increasing inference runtime.

### DEMiX Training

To train an LM with DEM1x layers, we partition the GPUs among the domains, so that each GPU is assigned a single domain (along with its corresponding expert). Each mini-batch contains \(k\) sequences from a particular domain, and we send each mini-batch to its dedicated expert. We use larger batch sizes with distributed data parallel between expert parameters on GPUs assigned to the same domain; we assign \(n\)/8 GPUs to each domain (Table 2).

Compared to dense LMs, DEM1x layers achieve the same or slightly higher throughput (measured in TFLOPs/GPU) for the same total FLOPs per update, despite adding significantly more parameters (Table 2). DEM1x achieves higher throughput because we while we sync shared parameters across all GPUs, we only sync expert parameters allocated to the same domain. Dense models sync all parameters across all GPUs. As we increase model size, DEMix reduces latency costs between GPUs, and hence, leads to faster training.

### Comparison with Adapters

DEM1x experts are related to adapters Bapna and Firat (2019), which add a small feedforward network into a frozen pretrained LM to enable parameter efficient finetuning. In contrast, our focus is on efficiently training all of the parameters of a modular LM from scratch, and as such is not directly comparable to existing adapter schemes. Adapters could enable more fine-grained control over which parts of the LM are domain-specific, and may circumvent the need to train domain-aware LMs from scratch. However, the shared parameters in the frozen pretrained LM may limit modularity. We leave exploring such architectural variants and their tradeoffs to future work.

## 4 In-Domain Performance

Our first set of experiments measure in-domain performance when replacing the feedforward layers in a transformer LM with DEM1x layers.

### Experimental Setup

Architecture, Input and HyperparametersWe follow the GPT-3 Brown et al. (2020) architecture (small, medium, large, and XL) implemented in Fairseq Ott et al. (2019). We use the GPT-2 Radford et al. (2019) vocabulary of 50,264 BPE types, and train with 1,024-token sequences. See SSA.7 for training hyperparameters.

EvaluationWe follow previous work by reporting performance for a given computational budget Lewis et al. (2021). For each model, we report test perplexity after a single run of 48 hours of training on differing numbers of NVIDIA V100 32GB GPUs (Table 2).

### Models

DenseThe first baseline is a dense LM, implemented with distributed data parallel Li (2021). There is no explicit conditioning on domain.

Dense (balanced)We train dense models on an equal amount of data from each domain. While there is still no explicit conditioning on domain, the gradient updates during training average across all domains represented in a batch.

+Domain-TokenThis model is trained identically to Dense (balanced), but we prepend a token to every sequence indicating its domain (during training and test time). We ignore the domain token when computing complexity during evaluation.

\begin{table}
\begin{tabular}{c c r r r r} \hline \hline  & & \multicolumn{4}{c}{**Parameters per GPU**} \\  & & 125M & 350M & 760M & 1.3B \\ \cline{3-5}  & **GPUs** & 32 & 64 & 128 & 128 \\  & **Total Experts** & 0 & 0 & 0 & 0 \\  & **GPUs/expert** & 0 & 0 & 0 & 0 \\  & **Total params** & 125M & 350M & 760M & 1.3B \\  & **TFLOPs/update** & 556 & 3279 & 13,637 & 23,250 \\  & **TFLOPs/GPU** & **31** & **37** & 45 & 51 \\ \cline{3-5}  & **GPUs** & 32 & 64 & 128 & 128 \\  & **Total Experts** & 8 & 8 & 8 & 8 \\  & **GPUs/expert** & 4 & 8 & 16 & 16 \\  & **Total params** & 512M & 1.8B & 3.8B & 7.0B \\  & **TTLOPs/optu** & 556 & 3279 & 13,637 & 23,250 \\  & **TFLOPs/GPU** & **31** & **37** & **48** & **55** \\ \hline \hline \end{tabular}
\end{table}
Table 2: Our specifications for training Dense and DEMMix LMs. All models are trained for about 48 hours on V100 GPUs. DEM1x layers increase the total parameters of the LM while maintaining (or increasing) throughput, measured in TFLOPs/GPU. We use the formula described in Narayanan et al. (2021) to calculate these metrics. See §A.4 for more details.

 

DEMix (naive)We replace every feedforward layer in the transformer with a DEMix layer, and use DEMix training (SS3). Under this _naive_ setting, the test data domain is known (e.g., the CS expert is used for CS test data). We relax this assumption in the next section.

### Results

Table 3 shows test perplexities, averaged across the eight training domains. Domain balancing is consistently helpful for dense training. Additional domain information always helps (i.e., domain tokens or DEMix layers), but the effects are largest for the smaller models. Overall, domain information enables the model to better specialize to different training domains. However, as the model size grows, the dense baseline improves, catching up to the DEMix (naive) model, at least when considering the average perplexity across domains.

### Domain Hetereogeneity

However, a more complete view of the experiments with the largest model is shown in Table 4. We see that even at scale, most training domains benefit from DEMix layers in a naive setting (where the domain label is revealed at test time), but some do not; WebText, RealNews, and Reddit fare worse than the dense baseline. We hypothesize that dense training is advantageous for hetereogenous domains. Heterogeneous domains have a higher overlap with other training domains, and therefore, benefit from parameter sharing.

Indeed, we observe that experts perform best on their assigned domain, and the experts assigned to domains that benefit from dense training perform relatively well on many training domains (SA.8). These findings suggest overall that a discrete notion of domain is too rigid. In the next section, we soften Equation 3 into a mixture of experts.

## 5 Mixing Experts at Inference Time

The previous section established that incorporating DEMix layers improves LM performance on test data from _known_ training domains. In practice, however, text may not come with a domain label, may straddle multiple domains, or may not belong to any of the domains constructed at training time.

In these cases, rather than a hard choice among experts (Equation 3), we propose to treat \(g_{1},\ldots,g_{n}\) as mixture coefficients, transforming the domain membership of an input text into a matter of probabilistic belief. Unlike previously proposed mixture-of-experts formulations Shazeer et al. (2017); Lepikhin et al. (2020), this approach is parameter-free and computed only at test time.

### Dynamic Domain Mixtures

Consider the probabilistic view of language modeling, where we estimate \(p(\mathbf{x}_{t}~{}|~{}\mathbf{x}_{<t})\). We introduce a domain variable, \(D_{t}\), alongside each word. We assume that this hidden variable depends on the history, \(\mathbf{x}_{<t}\), so that:

\[p(\mathbf{x}_{t}~{}|~{}\mathbf{x}_{<t}) =\sum_{j=1}^{n}p(\mathbf{x}_{t}~{}|~{}\mathbf{x}_{<t},D_{t}=j)\cdot \underbrace{p(D_{t}=j~{}|~{}\mathbf{x}_{<t})}_{g_{j}} \tag{4}\]

This model is reminiscent of class-based \(n\)-gram LMs Brown et al. (1992); Saul and Pereira (1997).

We have already designed the DEMix LM to condition on a domain label, giving a form for \(p(X_{t}~{}|~{}\mathbf{x}_{<t},D_{t}=j)\). We now further treat \(g_{1},\ldots,g_{n}\) as a posterior probability over domains, calculated either globally or at each timestep.

\begin{table}
\begin{tabular}{r r c c} \hline \hline  & \multicolumn{3}{c}{**1.3B parameters per GPU**} \\ \cline{2-4}
**Domain** & **Dense** & **DEMix** & **DEMix** \\  & & **(naïve)** & **(cached prior; §5.4)** \\ \hline
1B & 11.8 & 11.5 & **11.3** \\ CS & 13.5 & 12.2 & **12.1** \\ Legal & 6.8 & **6.7** & **6.7** \\ Med & 9.5 & 9.2 & **9.1** \\ WebText & **13.8** & 14.6 & 14.3 \\ RealNews & **12.5** & 13.3 & 13.1 \\ Reddit & 28.4 & 30.6 & **28.1** \\ Reviews & 14.0 & 12.6 & **12.5** \\ \hline
**Average** & 13.8 & 13.8 & **13.4** \\ \hline \hline \end{tabular}
\end{table}
Table 4: Test perplexity by domain for the largest models. We discuss the last column in §5.4.

\begin{table}
\begin{tabular}{r r r r r} \hline \hline  & \multicolumn{4}{c}{**Parameters per GPU**} \\  & 125M & 350M & 760M & 1.3B \\ \cline{2-5} Dense & 20.6 & 16.5 & 14.5 & 13.8 \\ Dense (balanced) & 19.9 & 15.8 & 14.3 & 13.6 \\ +Domain-Token & 19.2 & 15.9 & 14.3 & **13.4** \\ DEMix (naive) & 18.4 & 15.5 & 14.2 & 13.8 \\ DEMix (scaled; §5.4) & **17.8** & **14.7** & **13.9** & **13.4** \\ \hline \hline \end{tabular}
\end{table}
Table 3: Average in-domain test-set perplexity across the 8 domains in the training data. We discuss the last row in §5.4. See §A.9 for per-domain results.

 

To do this, we apply Bayes' rule:

\[p(D_{t}=j\mid\mathbf{x}_{<l}) =\frac{p(\mathbf{x}_{<t}\mid D_{t}=j)\cdot p(D_{t}=j)}{p(\mathbf{x}_{<t})} \tag{5}\] \[=\frac{p(\mathbf{x}_{<t}\mid D_{t}=j)\cdot p(D_{t}=j)}{\sum_{j^{\prime }=1}^{n}p(\mathbf{x}_{<t}\mid D_{t}=j^{\prime})\cdot p(D_{t}=j^{\prime})} \tag{6}\]

The conditional probabilities of word sequences given a domain label, as noted above, are already defined by the DEMix LM. For the prior over domain labels, we consider three alternatives:

UniformSet a uniform prior across domains.

UpdatingSet the prior at timestep \(t\) to be an exponentially weighted moving average of the posteriors from previous timesteps:

\[p(D_{t}=j)\propto\sum_{t^{\prime}=1}^{t-1}\lambda^{t-t^{\prime}}\cdot p(D_{t^ {\prime}}=j\mid\mathbf{x}_{<t^{\prime}}) \tag{7}\]

During evaluation, this moving average is calculated over the posterior at the end of each sequence. The decay factor avoids putting too much weight on calculations made early in the dataset, when posterior calculations are noisier (SA.10). We performed a small grid search to set the value \(\lambda\), and found that \(\lambda=0.3\) worked well.

CachedWe calculate the posterior over domain labels from additional data from the test distribution, and fix the prior to that estimate. We use 100 sequences from the validation set to estimate the prior, which we found to result in stable posterior probabilities. See SSA.10 for more details, and Figure 2 for an illustration of expert mixing.

### Visualizing Domain Membership

In Figure 3, we plot domain posteriors calculated using the largest DEMix LM from SS4 and the updating prior, after 100 sequences of validation data. For training domains, the associated domain label has the highest probability, but some of the domains are more hetereogeneous than we assumed. More variation is observed for the novel domains. However, generally we find the domain posterior distribution to be sparse; suggesting that after estimating the domain posterior, not all experts need to be active for test evaluation.

### Experimental Setup

Here, we experiment with the corpus of novel domains (Table 1). We evaluate the three mixture treatments of DEMix layers (SS5.1) against five

Figure 3: Estimates of posteriors \(p(D_{t}\mid\mathbf{x}_{<t})\) with a DEMix LM (1.3B parameters per GPU), after 100 sequences (i.e., 102,400 tokens) of data in training (top heatmap) and novel domains (bottom heatmap).

Figure 2: Illustration of inference with domain expert mixing. For a given input text \(\mathbf{x}_{<t}\) from CORD-19, we estimate a posterior domain probabilities \(p(D_{t}\mid\mathbf{x}_{<t})\), informed by a prior that is either iteratively updated during inference, or is precomputed and cached on heldout data. In this example, the model assigns highest domain probabilities to the medical and news domains. We use these probabilities in a weighted mixture of expert outputs to compute the output \(\mathbf{x}_{t}\).

 

baselines. No new models are trained for this experiment beyond those used in SS4.

Dense and Dense (balanced)These are the basic baselines from SS4.

+Domain-TokenHere test data is evaluated using each domain label token, and we choose the lowest among these perplexity values per test set.

DEMix (naive)Similar to +Domain-Token, we evaluate the data separately with each of the eight experts, and report the lowest among these perplexity values per test set.

DEMix (average)At every timestep, we take a simple average of the eight experts' predictions.

### Results

Novel Domain PerformanceEnsembling DEMix experts outperforms dense baselines and using experts individually (i.e., the "naive" baseline), and caching a prior before evaluation results in the best average performance (Table 5). Ensembling DEMix experts with a cached prior allows smaller models to match or outperform much larger dense models. Weighted ensembling outperforms simple averaging and mixing with a uniform prior, confirming the importance of sparsity in the expert mixture. These results demonstrate that modularity need not come at a cost to generalization to new domains.4

Footnote 4: We have separately observed that with expert mixing, our largest DEMix LM closely approaches the performance of GPT-3 _Da-Vinci_Brown et al. (2020) on another novel domain, the LM benchmark PTB Marcus et al. (1993). See §A.11 for more details.

In-Domain PerformanceWe can also apply the expert mixture variant of inference (using a cached prior) to the training domains; see the last line of Table 3. We see performance improvements across all training domains for every scale, though the largest gains come from hetereogeneous domains (Table 4 and SSA.9; across all model sizes, Reddit improves on average 10.7%, WebText 2.4%, RealNews 1.9%), confirming that domain labels may not align with the most effective boundaries.

### Summary

As opposed to other token-level routing mechanisms (e.g., Lepikhin et al.2020), expert mixing in DEMix is introduced at test-time and is parameter-free; it instead makes use of Bayesian inference with specialized experts to improve generalization. Expert mixing dynamically increases model capacity at test-time, while avoiding the need to learn token-level routing patterns during training, which is expensive and breaks modularity (SSA.5).

## 6 Domain Adaptation with New Experts

Domain adaptation is an important technique to improve LM performance in new domains that are rare or unseen during training. A popular technique for adapting LMs is domain-adaptive pretraining (DAPT; Gururangan et al.2020), which involves continued dense training of the LM on the target domain. However, DAPT with dense training (or Dense-DAPT) is expensive Strubell et al. (2019) and may entail forgetting domains learned during earlier training phases Aghajanyan et al. (2021), since it updates all parameters of the LM towards the target domain. These issues make adapting large LMs less feasible, especially in domains that change frequently over time Luu et al. (2021).

DEMix layers allow for _cheap_ adaptation _without_ forgetting through a technique we call DEMiX-DAPT (Figure 4). To adapt to a new domain, we initialize a new expert in each DEMix layer using the parameters of the nearest pretrained expert, which we identify using domain posteriors from SS5. We then train the added expert on target data, _updating only the new expert parameters_. For inference, we mix experts with a cached prior (SS5).

### Experimental Setup

We compare DEMix-DAPT to Dense-DAPT on the novel domains. We report test perplexity after adapting to each domain for 1 hour with 8 NVIDIA V100 32GB GPUs, tracking validation perplexity every 10 minutes for early stopping. We adapt to

\begin{table}
\begin{tabular}{l c c c c} \hline \hline  & \multicolumn{4}{c}{**Parameters per GPU**} \\  & 12SEM & 350M & 760M & 1.3B \\ \cline{2-5}
**Dense** & 25.9 & 21.4 & 18.4 & 17.8 \\
**Dense (balanced)** & 25.3 & 19.6 & 18.3 & 17.1 \\
**+Domain-Token** & 24.8 & 20.4 & 18.4 & 18.0 \\
**DEMix (naïve)** & 28.8 & 23.8 & 21.8 & 21.1 \\
**DEMix (average)** & 27.2 & 22.4 & 21.5 & 20.1 \\
**DEMix (uniform)** & 24.5 & 20.5 & 19.6 & 18.7 \\
**DEMix (padding)** & 21.9 & 18.7 & 17.6 & 17.1 \\
**DEMix (catened)** & **21.4** & **18.3** & **17.4** & **17.0** \\ \hline \hline \end{tabular}
\end{table}
Table 5: Average perplexity on novel domains. Mixing domain experts with a prior estimated using a small amount of data in the target domain outperforms all other baselines. See §A.9 for per-domain results.

 

each novel domain with the same hyperparameters as SS4, except with a 10x smaller learning rate. DEMix-DAPT updates about 10% of the total parameters in the DEMix LM, while Dense-DAPT updates all parameters of the dense LM.

### Results

Adding One ExpertWe display examples of DEMix-DAPT and Dense-DAPT on a single domain in Figure 5. As Dense-DAPT proceeds, its performance on the training domains progressively worsens (see SSA.12 for results with larger LMs). In contrast, DEMix-DAPT reduces perplexity on the novel domain _without_ forgetting.

Adding Eight ExpertsWe find that adding _all_ eight experts adapted to novel domains to the DEMix model from SS4 significantly reduces perplexity on novel _and_ previously seen domains (Table 6) while also helping in-domain for smaller models (perhaps surprisingly, given the fact that their domain experts are frozen). For example, across all model sizes, on average, we see a 2.4% reduction on Med, 1.8% reduction on RealNews, and 2% reduction on Reddit (see SSA.9 for details).

## 7 Removing Experts

Dense LMs are also prone to unexpected behavior when deployed. For example, they may generate hatespeech Gehman et al. (2020), which is undesirable for user-facing tasks Xu et al. (2020).

We argue that dense training contributes to unexpected model behavior, as domains are learned diffusely over the parameter space, and it is difficult to restrict the model's access to certain training domains during inference. Some mechanisms have been introduced to steer a dense model _towards_ Keskar et al. (2019); Dathathri et al. (2020) and _away_Welleck et al. (2019) from certain behaviors, but they tend to be expensive or require retraining the model with a different objective, which

\begin{table}
\begin{tabular}{l c c c c c} \hline \hline  & & \multicolumn{3}{c}{**Parameters per GPU**} \\ \multirow{2}{*}{**Domains**} & \multirow{2}{*}{**\# Experts**} & \multirow{2}{*}{125M} & \multirow{2}{*}{350M} & \multirow{2}{*}{760M} & \multirow{2}{*}{1.3B} \\ \cline{3-3} \cline{5-7}  & & & & & \\ \cline{3-3} \cline{5-7} Training & 8 & 17.8 & 14.7 & 13.9 & **13.4** \\  & 16 & **17.7** & **14.6** & **13.7** & **13.4** \\ \cline{3-3} \cline{5-7} Novel & 8 & 21.4 & 18.3 & 17.4 & 17.0 \\  & 16 & **16.0** & **14.0** & **13.5** & **12.5** \\ \hline \hline \end{tabular}
\end{table}
Table 6: Average perplexity in training and novel domains before and after adding 8 experts adapted to the novel domains (via DEMix-DAPT). Adding experts reduces perplexity on novel and training domains.

Figure 4: Illustration of DEMix-DAPT. First, we estimate domain posteriors on a held out sample of the target domain (e.g., CORD-19). We then initialize a new expert with the parameters of the most likely expert under the domain posterior distribution. Finally, we adapt the parameters of the new expert to the target domain, keeping all other parameters in the LM frozen.

Figure 5: Adapting an LM (125M parameters per GPU) to CORD-19 or Gutenberg. Top row: with Denseq-DAPT, average perplexity on all training domains degrades. Bottom row: DEMix-DAPT avoids forgetting while achieving close (for Gutenberg) or better (for CORD-19) performance on the target domain.

 

becomes less feasible as the LM grows in size.

DEMix layers offer a simple mechanism for cheap, lightweight control of large LMs: since domain experts specialize (SSA.8), experts assigned to unwanted domains can be disabled at test-time.5

Footnote 5: Removing an expert offers no guarantee of having fully forgotten content from the removed domain, since there are shared parameters in the model.

### Experimental Setup

Does disabling an expert simulate a model that has not been exposed to a particular training domain? To answer this question, we compare three settings: **+Expert**, a DEMix LM with all experts active, **-Expert**, a DEMix LM with a domain expert deactivated, and **-Domain**, a DEMix LM trained from scratch without a particular domain.6

Footnote 6: We replace the removed domain with Gutenberg, since our cluster allocates training jobs via 8-GPU nodes.

For all settings, we use a DEMix LM (125M parameters per GPU) from SS4 and expert mixing with a cached prior (SS5) for inference.

### Results

Removing a domain expert harms model performance on the associated domain, in most cases approaching the performance of a model that has not been exposed to data from that domain (Table 7). In some cases (e.g., WebText and RealNews), **-Expert** even underperforms **-Domain**. This leads us to conjecture that most domain-specific learning happens within the DEMix layer.

Our preliminary analysis here suggests that DEMix enables LMs with _removable parts_, for quick adaptation to situations in which a particular training domain is unwanted for inference. We leave further exploration of this mechanism and its potential for LM control to future work.

## 8 Related Work

Document metadata has been used to improve topic models Mimno and McCallum (2012), adapt RNN-based LMs Jaech and Ostendorf (2018), learn document representations Card et al. (2018), and improve text generation control Zellers et al. (2019); Keskar et al. (2019). Other inference-time methods Dathathri et al. (2020); Liu et al. (2021) may be used to steer text generation with DEMix experts.

Future work may explore applying DEMix to multilingual settings, as multilingual models benefit from language-specific parameters Fan et al. (2020); Pfeiffer et al. (2020); Chau et al. (2020).

DEMix-DAPT is related to model expansion techniques in reinforcement learning or vision Rusu et al. (2016); Draelos et al. (2017) and adapters for pretrained LMs Houlsby et al. (2019); Pfeiffer et al. (2020).

Multi-domain models have been studied in machine translation Pham et al. (2021) and supervised settings Wright and Augenstein (2020), and with smaller dense LMs Maronikolakis and Schutze (2021). Previous studies have shown the importance of considering domains when adapting LMs Ramponi and Plank (2020); Gururangan et al. (2020). Our study establishes the importance of considering domains when training LMs from scratch.

## 9 Conclusion

We introduce DEMix layers, which provide modularity to an LM at inference time, addressing limitations of dense training by providing a rapidly adaptable system. DEMix layers experts can be mixed to handle heterogeneous or unseen domains, added to iteratively incorporate new domains, and removed to restrict unwanted domains. Future work may combine domain and token-level routing, discover domains automatically with unsupervised learning, or scale the number of training domains.

#### Acknowledgments

The authors thank Shruti Bhosalie, Tim Dettmers, Emily Dinan, Doug Downey, Margaret Li, Myle Ott, Ofir Press, and Swabha Swayamdipta, for helpful comments. At UW, this work was partially supported by NSF grant 1562364, the Office of Naval Research under MURI grant N00014-18-1-2670, and an Amazon research award.

\begin{table}
\begin{tabular}{l c c c} \hline \hline  & \multicolumn{3}{c}{**125M Parameters per GPU**} \\ \cline{2-4} Domain & **+Expert** & **–Expert** & **–Domain** \\ \hline
1B & 13.7 & 25.5 & **30.4** \\ CS & 15.7 & 22.4 & **25.4** \\ Legal & 8.9 & 20.9 & **22.7** \\ Med & 12.4 & 18.6 & **21.9** \\ WebText & 20.9 & **27.3** & 25.4 \\ RealNews & 18.9 & **26.7** & 25.0 \\ Reddit & 34.4 & 47.8 & **51.3** \\ Reviews & 20.5 & 39.0 & **43.0** \\ \hline
**Average** & 18.2 & 28.5 & **30.6** \\ \hline \hline \end{tabular}
\end{table}
Table 7: Removing a domain expert (“Expert”) degrades perplexity on the corresponding domain, approaching the performance of an LM that has not been exposed to that domain (“Domain”). Here we bold the _worst_ performing model for each domain.

 



## Ethical Considerations

While DEMiix offers new opportunities to reduce the influence of unwanted training domains (e.g., those that contain hatespeech) at inference time, shared parameters in the LM may prevent the model from fully forgetting the unwanted domain after expert removal. Therefore, DEMix LMs may still be prone to producing harmful generations when deployed, and further research is required to understand the bounds on the probability of toxic degeneration after expert removal.

While we partially anonymize our corpus with simple regexes, it is difficult to guarantee that sensitive information is not exposed in large datasets. To protect data authors and subjects, we do not publicly release our models or data, although we provide instructions and code to replicate them to support reproducibility.

## References

* Agahanyan et al. (2021) Armen Aghajanyan, Akshat Shrivastava, Anchit Gupta, Naman Goyal, Luke Zettlemoyer, and Sonal Gupta. 2021. Better fine-tuning by reducing representational collapse. In _9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021_. OpenReview.net.
* Aharoni and Goldberg (2020) Roee Aharoni and Yoav Goldberg. 2020. Unsupervised domain clusters in pretrained language models. In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_, pages 7747-7763, Online. Association for Computational Linguistics.
* Baly et al. (2018) Ramy Baly, Georgi Karadzhov, Dimitar Alexandrov, James Glass, and Preslav Nakov. 2018. Predicting factuality of reporting and bias of news media sources. In _Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing_, pages 3528-3539, Brussels, Belgium. Association for Computational Linguistics.
* Bapna and Firat (2019) Ankur Bapna and Orhan Firat. 2019. Non-parametric adaptation for neural machine translation. In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_, pages 1921-1931, Minneapolis, Minnesota. Association for Computational Linguistics.
* Baumgartner et al. (2020) Jason Baumgartner, Savvas Zannettou, Brian Keegan, Megan Suire, and Jeremy Blackburn. 2020. The pushshift reddit dataset.
* Bender et al. (2021) Emily M. Bender, Timnit Gebru, Angelina McMillanMajor, and Shmargaret Shmitchell. 2021. On the dangers of stochastic parrots: Can language models be too big? In _Proceedings of the 2021 ACM Conference on Fairness, Accountability, and Transparency_, FAccT '21, page 610-623, New York, NY, USA. Association for Computing Machinery.
* Brown et al. (1992) Peter F. Brown, Vincent J. Della Pietra, Peter V. deSouza, Jenifer C. Lai, and Robert L. Mercer. 1992. Class-based \(n\)-gram models of natural language. _Computational Linguistics_, 18(4):467-480.
* Brown et al. (2020) 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 Ramesch, 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. 2020. Language models are few-shot learners.
* Card et al. (2018) Dallas Card, Chenhao Tan, and Noah A. Smith. 2018. Neural models for documents with metadata. _Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_.
* Creaselaw Access Project (2020) Caselaw Access Project. Caselaw access project.
* Chau et al. (2020) Ethan C. Chau, Lucy H. Lin, and Noah A. Smith. 2020. Parsing with multilingual BERT, a small corpus, and a small treebank. In _Findings of the Association for Computational Linguistics: EMNLP 2020_, pages 1324-1334, Online. Association for Computational Linguistics.
* Chelba et al. (2014) Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robinson. 2014. One billion word benchmark for measuring progress in statistical language modeling.
* Dasigi et al. (2021) Pradeep Dasigi, Kyle Lo, Iz Beltagy, Arman Cohan, Noah A. Smith, and Matt Gardner. 2021. 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 Linguistics: Human Language Technologies_, pages 4599-4610, Online. Association for Computational Linguistics.
* Dathathri et al. (2020) Sumanth Dathathri, Andrea Madotto, Janice Lan, Jane Hung, Eric Frank, Piero Molino, Jason Yosinski, and Rosanne Liu. 2020. 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.
* Dinan et al. (2021) Emily Dinan, Gavin Abercrombie, A. Stevie Bergman, Shannon Spruit, Dirk Hovy, Y-Lan Boureau, and Verena Rieser. 2021. Anticipating safety issues in e2e conversational ai: Framework and tooling

T. Draelos, N. Miner, Christopher C. Lamb, Jonathan A. Cox, Craig M. Vineyard, Kristofor D. Carlson, William M. Severa, C. James, and J. Aimone (2017)Neurogenesis deep learning: extending deep networks to accommodate new classes. pp. 526-533. External Links: Link Cited by: SS4.
* A. Fan, S. Bhosalie, H. Schwenk, Z. Ma, A. El-Kishky, S. Goyal, M. Baines, O. Celebi, G. Wenzek, V. Chaudhary, N. Goyal, T. Birch, V. Liptchinsky, S. Edunov, E. Grave, M. Auli, and A. Joulin (2020)Beyond english-centric multilingual machine translation. External Links: Link Cited by: SS4.
* W. Fedus, B. Zoph, and N. Shazeer (2021)Switch transformers: scaling to trillion parameter models with simple and efficient sparsity. External Links: Link Cited by: SS4.
* L. Gao, Stella Biderman, S. Black, L. Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy (2020)The pile: an 800gb dataset of diverse text for language modeling. External Links: Link Cited by: SS4.
* S. Gehman, S. Gururangan, M. Sap, Y. Choi, and N. A. Smith (2020)RealToxicityPrompts: evaluating neural toxic degeneration in language models. In Findings of the Association for Computational Linguistics: EMNLP 2020, Online, pp. 3356-3369. External Links: Link Cited by: SS4.
* G. Archive Project (2019)Github archive project.. External Links: Link Cited by: SS4.
* S. Gururangan, A. Marasovic, S. Swayamdipta, K. Lo, I. Beltagy, D. Downey, and N. A. Smith (2020)Don't stop pretraining: adapt language models to domains and tasks. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, Online, pp. 8342-8360. External Links: Link Cited by: SS4.
* D. Hendrycks, C. Burns, A. Chen, and P. Ball (2021)Cuad: an expert-annotated nlp dataset for legal contract review. External Links: Link Cited by: SS4.
* N. Houlsby, A. Giurgiu, R. Jusxtzegbski, B. Morrone, Q. de Laroussilhe, A. Gesmundo, M. Attariyan, and S. Gelly (2019)Parameter-efficient transfer learning for nlp. External Links: Link Cited by: SS4.
* A. Jaech and M. Ostendorf (2018)Low-rank rnn adaptation for context-aware language modeling. External Links: Link Cited by: SS4.
* N. Shirish Keskar, B. McCann, L. R. Varshney, C. Xiong, and R. Socher (2019)Ctrl: a conditional transformer language model for controllable generation. External Links: Link Cited by: SS4.
* D. P. Kingma and J. Ba (2015)Adam: a method for stochastic optimization. In 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, External Links: Link Cited by: SS4.
* P. F. Koh, S. Sagawa, H. Marklund, S. Michel Xie, M. Zhang, A. Balsubramani, W. Hu, M. Yasunaga, R. Lanas Phillips, I. Gao, T. Lee, E. David, I. Stavness, W. Guo, B. Earnshaw, I. Haque, S. M. Beery, J. Leskovec, A. Kundaje, E. Pierson, S. Levine, C. Finn, and P. Liang (2021)WILDS: a benchmark of in-the-wild distribution shifts. In Proceedings of the 38th International Conference on Machine Learning, Vol. 139 of _Proceedings of Machine Learning Research_, pp. 5637-5664. External Links: Link, Document Cited by: SS4.
* A. Lazaridou, A. Koncoro, E. Gribovskaya, D. Agrawal, A. Liska, T. Terzi, M. Gimenez, C. de Masson d'Autume, S. Ruder, D. Yogatama, K. Cao, T. Kocisky, S. Young, and P. Blunsom (2021)Pitfalls of static language modelling. External Links: Link Cited by: SS4.
* D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen (2020)Gshard: scaling giant models with conditional computation and automatic sharding. External Links: Link Cited by: SS4.
* M. Lewis, S. Bhosalie, T. Dettmers, N. Goyal, and L. Zettlemoyer (2021)Base layers: simplifying training of large, sparse models. In Proceedings of the 38th International Conference on Machine Learning, Vol. 139 of _Proceedings of Machine Learning Research_, pages 6265-6274. External Links: Link, Document Cited by: SS4.
* S. Li (2021)Getting started with distributed data parallel. External Links: Link Cited by: SS4.
* A. Liu, M. Sap, X. Lu, S. Suwayamdipta, C. Bhagavatula, N. A. Smith, and Y. Choi (2021)Dexperts: decoding-time controlled text generation with experts and anti-experts. External Links: Link Cited by: SS4.
* K. Lo, L. L. Wang, M. Neumann, R. Kinney, and D. Weld (2020)S2ORC: The semantic scholar open research corpus. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, Online, pp. 4969-4983. External Links: Link, Document Cited by: SS4.
* L. Luo and D. Bamman (2021)Characterizing English Variation across Social Media Communities with BERT. Transactions of the Association for Computational Linguistics9, pp. 538-556. External Links: Link, Document Cited by: SS4.
* K. Luu, D. Khashabi, S. Gururangan, K. Mandyam, and N. A. Smith (2021)Time waits for no one! analysis and challenges of temporal misalignment. In Proc. of NAACL 

Mitchell P. Marcus, Beatrice Santorini, and Mary Ann Marcinkiewicz. 1993. Building a large annotated corpus of English: The Penn Treebank. _Computational Linguistics_, 19(2):313-330.
* Maronikolakis and Schutze [2021] Antonios Maronikolakis and Hinrich Schutze. 2021. Multidomain pretrained language models for green NLP. In _Proceedings of the Second Workshop on Domain Adaptation for NLP_, pages 1-8, Kyiv, Ukraine. Association for Computational Linguistics.
* McCloskey and Cohen [1989] M. McCloskey and N. Cohen. 1989. Catastrophic interference in connectionist networks: The sequential learning problem. _Psychology of Learning and Motivation_, 24:109-165.
* Mimno and McCallum [2012] David Mimno and Andrew McCallum. 2012. Topic models conditioned on arbitrary features with dirichlet-multinomial regression.
* Narayanan et al. [2021] Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Anand Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phanishayee, and Matei Zaharia. 2021. Efficient large-scale language model training on gpu clusters using megatron-lm.
* Ni et al. [2019] Jianmo Ni, Jiacheng Li, and Julian McAuley. 2019. Justifying recommendations using distantly-labeled reviews and fine-grained aspects. In _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP)_, pages 188-197, Hong Kong, China. Association for Computational Linguistics.
* Oren et al. [2019] Yonatan Oren, Shiori Sagawa, Tatsunori Hashimoto, and Percy Liang. 2019. Distributionally robust language modeling. In _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP)_, pages 4227-4237, Hong Kong, China. Association for Computational Linguistics.
* Ott et al. [2019] Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. 2019. fairseq: A fast, extensible toolkit for sequence modeling. In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics (Demonstrations)_, pages 48-53, Minneapolis, Minnesota. Association for Computational Linguistics.
* Pfeiffer et al. [2020] Jonas Pfeiffer, Ivan Vulic, Iryna Gurevych, and Sebastian Ruder. 2020. MAD-X: An Adapter-Based Framework for Multi-Task Cross-Lingual Transfer. In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_, pages 7654-7673, Online. Association for Computational Linguistics.
* Pham et al. [2021] MinhQuang Pham, Josep Maria Crego, and Francois Yvon. 2021. Revisiting multi-domain machine translation. _Transactions of the Association for Computational Linguistics_, 9:17-35.
* Gutenberg [2019] Project Gutenberg. Project guttenberg.
* Radford et al. [2019] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. 2019. Language models are unsupervised multitask learners.
* Ramponi and Plank [2020] Alan Ramponi and Barbara Plank. 2020. Neural unsupervised domain adaptation in NLP--A survey. In _Proceedings of the 28th International Conference on Computational Linguistics_, pages 6838-6855, Barcelona, Spain (Online). International Committee on Computational Linguistics.
* Roller et al. [2021] Stephen Roller, Sainbayar Sukhbbaatar, Arthur Szlam, and Jason Weston. 2021. Hash layers for large sparse models.
* Rusu et al. [2016] Andrei A. Rusu, Neil C. Rabinowitz, Guillaume Desjardins, Hubert Soyer, James Kirkpatrick, Koray Kavukcuoglu, Razvan Pascanu, and Raia Hadsell. 2016. Progressive neural networks.
* Saul and Pereira [1997] Lawrence Saul and Fernando Pereira. 1997. Aggregate and mixed-order Markov models for statistical language processing. In _Second Conference on Empirical Methods in Natural Language Processing_.
* Shazeer et al. [2017] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziazz, Andy Davis, Quoc V. Le, Geoffrey E. Hinton, and Jeff Dean. 2017. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In _5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings_. OpenReview.net.
* Strubell et al. [2019] Emma Strubell, Ananya Ganesh, and Andrew McCallum. 2019. Energy and policy considerations for deep learning in NLP. In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_, pages 3645-3650, Florence, Italy. Association for Computational Linguistics.
* Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, L ukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In _Advances in Neural Information Processing Systems_, volume 30. Curran Associates, Inc.
* Wang et al. [2020] Lucy Lu Wang, Kyle Lo, Yoganand Chandrasekhar, Russell Reas, Jiangjiang Yang, Doug Burdick, Darrin Eide, Kathryn Funk, Yannis Kats, Rodney Kinney, Yunyao Li, Ziyang Liu, William Merrill, Paul Mooney, Dewey Murdick, Devvret Rishi, Jerry Shehan, Zhihong Shen, Brandon Stilson, Alex Wade, Kuansan Wang, Nancy Xin Ru Wang, Chris Wilhelm, Boya Xie, Douglas Raymond, Daniel S. Weld, Oren Etzioni, and Sebastian Kohlmeier. 2020. Cord-19: The covid-19 open research dataset.
* Welleck et al. [2019] Sean Welleck, Ilia Kulikov, Stephen Roller, Emily Dinan, Kyunghyun Cho, and Jason Weston. 2019. Neural text generation with unlikelihood training 

Dustin Wright and Isabelle Augenstein. 2020. Transformer based multi-source domain adaptation. In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_, pages 7963-7974, Online. Association for Computational Linguistics.
* Xu et al. [2020] Jing Xu, Da Ju, Margaret Li, Y-Lan Boureau, Jason Weston, and Emily Dinan. 2020. Recipes for safety in open-domain chatbots.
* Reviews [2019] Yelp Reviews. Yelp reviews.
* Zellers et al. [2019] Rowan Zellers, Ari Holtzman, Hannah Rashkin, Yonatan Bisk, Ali Farhadi, Franziska Roesner, and Yejin Choi. 2019. Defending against neural fake news. In _NeurIPS_.

 

Appendix

### Image Attribution

Images retrieved from emojipedia.org or istockphoto.com.

### Collecting Domains

All datasets are fair use for research purposes according to their original licenses. For most domains, we use the associated sources, listed in Table 1, without modification. Reddit was extracted and obtained by a third party and made available on pushshift.io, and was anonymized by Xu et al. (2020); we use their version. For Gutenberg, we use the scraping tool provided in [https://github.com/aparrish/gutenberg-dammit](https://github.com/aparrish/gutenberg-dammit). For Breaking News, we identify a list of factually reliable English news sources, using the list curated by Baly et al. (2018). Specifically, we filter on "high" factuality in the data provided in this repository: [https://github.com/ramybaly/News-Media-Reliability](https://github.com/ramybaly/News-Media-Reliability). We then use Newspaper3K ([https://newspaper.readthedocs.io/en/latest/](https://newspaper.readthedocs.io/en/latest/)) to scrape the latest 1000 articles from each site. After dropping duplicates, we arrive at about 20K articles from 400 news sources. We provide downloading links and general instructions at github.com/kernelmachine/demix-data.

### Dataset Anonymization

To anonymize certain datasets, we apply a suite of regexes that aim to identify common patterns of user-identifiable data and substitute them with dummy tokens. We display anonymization regexes and associated dummy tokens in Table 8.

### Calculating TFLOPs/GPU

We use the formula presented in Narayanan et al. (2021) to calculate TFLOPs/GPU and TFLOPs/update. The spreadsheet that contains the calculations and formula can be accessed here: [https://docs.google.com/spreadsheets/d/1NO-lz_VqZGPZFpJTFxtXyjhmaoYi6qnz50XrBwW8hgGw/edit?usp=sharing](https://docs.google.com/spreadsheets/d/1NO-lz_VqZGPZFpJTFxtXyjhmaoYi6qnz50XrBwW8hgGw/edit?usp=sharing).

### Gshard Comparison

Here we describe empirical comparisons between DEMix and GShard, the token-level mixture of experts transformer proposed by Lepikhin et al. (2020). As opposed to DEMix, which uses domain labels to route data to experts, GShard learns a token-level routing mechanism during training. Each token in every other layer is sent to two of \(k\) experts, and this routing is updated via backpropagation.

As GShard is emblematic of an learned routing procedure, we are generally interested if GShard naturally learns to specialize experts to domains, whether its experts are modular, and how GShard LM generally performs compared to DEMix and Dense models on our multi-domain corpus.

Experimental SetupWe aim to make minimal changes to the overall architecture of the model, to focus on the differences afforded by token-level routing (vs. DEMix routing). As such, we keep all architecture and computational budgets the same as our DEMix and Dense baselines (we generally display results for the 125M, 350M, and 760M parameter LMs). We only add the GShard routing procedure to every other layer, which involves routing each token to the top-2 experts of that layer. This additionally necessitates a load balancing loss to prevent only a minority of experts from being used (Lepikhin et al., 2020). All GShard experts are of the same size as our DEMix experts, i.e., each expert is a two layer MLP with the same dimensions as the original feedforward layer of the transformer. We display hyperparameters used to train GShard in SSA.7.

Model ScaleIn DEMix, we always add the same number of experts as the number of training domains (in our case -- eight experts), and use extra computation to increase the batch size for each expert. Our GShard implementation, on the other hand, allocates one expert per GPU. This means that GShard adds many more experts to the system, which results in a substantially larger increase in model size (Table 9). Unlike DEMix, GShard results in an increase in FLOP count relative to the Dense model, due to a variety of additional computation during training, like load balancing and routing to two experts for every token, which DEMix does not need.

Training efficiencyHowever, unlike DEMix, which increases model size while maintaining or improving GPU throughput, GShard in fact _reduces_ GPU throughput during training (Table 9). This is due to the necessity of expensive _all-to-all_ operations in GShard which mediate communication between experts on different GPUs that are ac 

tivated for different tokens of the same document.7 These _all-to-all_ operations are bottlenecked by the quality of GPU communication channels on the cluster. We also found that additional inefficiencies are introduced via GShard's load balancing, since some experts are not used at test time. DEMix has no load balancing or all-to-all communication. It uses all experts to maximum efficiency, because we simply assign GPUs to domains for our routing protocol.

Footnote 6: [https://github.com/hugging-and-play/](https://github.com/hugging-and-play/)

Evaluation efficiencyAnother benefit to DEMix is that its experts specialize to their domain, and only a sparse subset of them are activated at test time. Does token-level routing via GShard also result in a modular model? We explore this question by computing the average gating probabilities in the GShard router across all experts for all test data in each domain. We generally find that gating probabilities in GShard have high entropy across experts regardless of domain, suggesting that the token-level routing procedure does not in fact result in modularity out-of-the-box and all experts are needed for all input texts (Figure 6). As we increase computational budget, this issue is exacerbated; we need 128 GPUs to evaluate on the test data for the final model. Whereas with DEMix, we only need 8 GPUs to compute the domain posterior on a subset of the validation data. Moreover, because the domain posterior is usually sparse, one can use an even smaller number of GPUs for evaluating on test data, loading only those experts with non-zero probabilities.

\begin{table}
\begin{tabular}{c c c c c} \hline \hline  & \multicolumn{4}{c}{**Parameters per GPU**} \\  & 125M & 350M & 760M & 1.3B \\ \cline{2-5}
**Dense (balanced)** & 19.9 & 15.8 & 14.3 & 13.6 \\
**DEMix** & **17.8** & **14.7** & **13.9** & 13.4 \\
**GShard** & 17.2 & 14.3 & 14.2 & **12.7** \\ \hline \hline \end{tabular}
\end{table}
Table 10: Average in-domain test-set perplexity across the 8 domains in the training data. We discuss the last row in §5.4. See SSA.9 for per-domain results.

\begin{table}
\begin{tabular}{l l r} \hline \hline
**Category** & **Link to Regex** & Dummy Token \\ \hline Email & [https://regex101.com/r/ZqSF9x/1](https://regex101.com/r/ZqSF9x/1) & <EMAIL> \\ DART & [https://regex101.com/r/QtQGEN/1](https://regex101.com/r/QtQGEN/1) & <DART> \\ FB User ID & [https://regex101.com/r/GZ15EZ/1](https://regex101.com/r/GZ15EZ/1) & <FB_USERID> \\ Phone Number & [https://regex101.com/r/YlrDPPD/1](https://regex101.com/r/YlrDPPD/1) & <PHONE_NUMBER> \\ Credit Card Number & [https://regex101.com/r/SPMTO6W/1](https://regex101.com/r/SPMTO6W/1) & <RECDIT_CARD_NUMBER> \\ Social Security Number & [https://regex101.com/r/VGSDN/1](https://regex101.com/r/VGSDN/1) & <SSND> \\ User handles & [https://regex101.com/r/Vpey04/1](https://regex101.com/r/Vpey04/1) & <USER> \\ \hline \hline \end{tabular}
\end{table}
Table 8: Anonymization schema. We anonymize text using the regexes provided in the above links for the categories listed.

\begin{table}
\begin{tabular}{c c c c c c} \hline \hline  & & \multicolumn{4}{c}{**Parameters per GPU**} \\  & & 125M & 350M & 760M & 1.3B \\ \cline{2-5}
**Dense (balanced)** & 19.9 & 15.8 & 14.3 & 13.6 \\
**DEmix** & **17.8** & **14.7** & **13.9** & 13.4 \\
**GShard** & 17.2 & 14.3 & 14.2 & **12.7** \\ \hline \hline \end{tabular}
\end{table}
Table 11: Average perplexity on novel domains. Mixing domain experts with a prior estimated using a small amount of data in the target domain outperforms all other baselines. See SSA.9 for per-domain results.

\begin{table}
\begin{tabular}{c c c c c c} \hline \hline  & & \multicolumn{4}{c}{**Parameters per GPU**} \\  & & 125M & 350M & 760M & 1.3B \\ \cline{2-5}
**Dense (balanced)** & 125M & 350M & 760M & 1.3B \\ \cline{2-5}
**Dense (balanced)** & 19.9 & 15.8 & 14.3 & 13.6 \\
**DEMix** & **17.8** & **14.7** & **13.9** & 13.4 \\
**GShard** & 17.2 & 14.3 & 14.2 & **12.7** \\ \hline \hline \end{tabular}
\end{table}
Table 10: Average in-domain test-set perplexity across the 8 domains in the training data. We discuss the last row in §5.4. See SSA.9 for per-domain results.

\begin{table}
\begin{tabular}{c c c c c} \hline \hline  & & \multicolumn{4}{c}{**Parameters per GPU**} \\  & & 125M & 350M & 760M & 1.3B \\ \cline{2-5}
**Dense (balanced)** & 125M & 350M & 760M & 1.3B \\ \cline{2-5}
**Dense (balanced)** & 19.9 & 15.8 & 14.3 & 13.6 \\
**DEMix** & **17.8** & **14.7** & **13.9** & 13.4 \\
**GShard** & 17.2 & 14.3 & 14.2 & **12.7** \\ \hline \hline \end{tabular}
\end{table}
Table 10: Average in-domain test-set perplexity across the 8 domains in the training data. We discuss the last row in §5.4. See SSA.9 for per-domain results.

 

Model performanceAs noted earlier, our GShard implementation substantially increases the effective parameter count of the model relative to DEMix (Table 9). While this expansion of model size by GShard translates to better in-domain performance than DEMix for the 32 and 64 GPU settings, we observe the DEMix LMs consistently outperform GShard on the novel domains regardless of computational budget (Table 11). Surprisingly, GShard underperforms DEMix even in-domain for the 760M parameter model (Table 10), despite being 4x larger in effective parameter count (Table 9). This suggests that domain-modularity is an important mechanism to improve model generalization, in addition to model size. We believe there is a rich area of future work to investigate how to combine token- and domain-level routing, to realize the benefits of increasing parameter count while maintaining domain modularity at scale.

SummaryOur results suggest that while GShard is an effective method for substantially increasing model size under a fixed budget, it comes with large costs to training and evaluation efficiency, does not result in a modular LM. The lack of modularity also implies that GShard suffers from similar downstream issues as Dense models, e.g., forgetting after adaptation and lack of lightweight controllability, though we leave a close exploration of those phenomena to future work. Overall, DEMix LMs are substantially simpler and more efficient for training and evaluation, and even outperform GShard (especially out of domain) despite being substantially smaller, suggesting the importance of domain modularity as an alternative mechanism to model scaling for improving generalization in LMs.

### Interleaving Experiments

We hypothesize that shared layers may serve as a bottleneck to find shared features between domains, and may impact performance adversely when training domains are highly different from one another. Indeed, preliminary experiments suggest that interleaving expert layers causes large performance hits in the most distinct domains, i.e., those with lower vocabulary overlap with other domains in the corpus.

### Hyperparameter Assignments

We display hyperparameter assignments for LM pretraining in Tables 14, 15, 16, and 17. We set the total number of training steps based on this allocated runtime, set 8% of these steps to be warm-up, and use the Adam optimizer Kingma and Ba (2015) with a polynomial learning rate decay. Learning rates are tuned for each model separately over {0.0001, 0.0003, 0.0005}, taking the fastest learning rate that avoids divergence. Each worker processes two sequences of length 1,024, and gradients are accumulated over 8 updates. We clip gradients if their \(L_{2}\) norm exceeds 0.1. These settings are inspired by Lewis et al. (2021).

Figure 6: Average gating probabilities across domains (x-axis) for each expert (y-axis) in the expert layers of a GShard LM with 125M parameters per GPU. We observe high entropy of gating probabilities across experts and domains in each expert layer, with similar results in larger models.

 



### Expert Performance Ratios

We display a heatmap of expert performance ratios, using the largest DEMixLM (1.3B parameters per GPU) in Figure 7. These results suggest that experts specialize to their domain, and that leveraging the outputs of multiple experts (especially those specialized to hetereogeneous domains) at test time would lead to better language modeling performance.

### Per-Domain Results

We display the rest of the per-domain test results in the spreadsheets at the following link: [https://docs.google.com/spreadsheets/d/lymNZGSPAvhTi3JttLamiCULaoIGTJ4QGEOaj03b5kt8/edit?usp=sharing](https://docs.google.com/spreadsheets/d/lymNZGSPAvhTi3JttLamiCULaoIGTJ4QGEOaj03b5kt8/edit?usp=sharing)

#### a.10 Domain Posterior Calculations

We track calculated domain posteriors over sequences of development data in Figure 8 (training domains) and Figure 9 (novel domains). The domain posteriors are noisier for earlier sequences, stabilizing usually after around 50 sequences. For all experiments, we conservatively use 100 sequences of data to compute the domain posterior, though one may be able to accurately calcuate the domain posterior for some domains with less data.



Figure 8: Calculated domain posteriors for 8 training domains.

Figure 9: Calculated domain posteriors for 8 novel domains.

 



\begin{table}
\begin{tabular}{c c} \hline
**Computing Infrastructure** & 64 Volta 32GB GPUs \\ \hline \hline
**Hyperparameter** & **Assignment** \\ \hline architecture & GPT-3 medium \\ \hline tokens per sample & 1024 \\ \hline batch size & 2 \\ \hline number of workers & 2 \\ \hline learning rate & [5e–4, 3e–4, 1e–4] \\ \hline clip norm & 0.1 \\ \hline gradient accumulation steps & 8 \\ \hline batch size & 2 \\ \hline number of workers & 2 \\ \hline learning rate & [5e–4, 3e–4, 1e–4] \\ \hline clamp norm & 0.1 \\ \hline gradient accumulation steps & 8 \\ \hline number of steps & 300,000 \\ \hline save interval updates & 6,000 \\ \hline validation interval & 3,000 \\ \hline number of warmup steps & 24,000 \\ \hline learning rate scheduler & polynomial decay \\ \hline learning rate optimizer & Adam \\ \hline Adam beta weights & (0.9, 0.95) \\ \hline Adam epsilon & 10e-8 \\ \hline weight decay & 0.1 \\ \hline \end{tabular}
\end{table}
Table 14: Hyperparameters for pretraining the LM with 125M parameters per GPU. All hyperparameters are the same for DEMix and Dense training.

\begin{table}
\begin{tabular}{c c} \hline
**Hyperparameter** & **Assignment** \\ \hline architecture & GPT-3 small \\ \hline tokens per sample & 1024 \\ \hline batch size & 2 \\ \hline number of workers & 2 \\ \hline learning rate & [5e–4, 3e–4, 1e–4] \\ \hline clique norm & 0.1 \\ \hline gradient accumulation steps & 8 \\ \hline number of steps & 300,000 \\ \hline save interval updates & 6,000 \\ \hline validation interval & 3,000 \\ \hline number of warmup steps & 24,000 \\ \hline learning rate scheduler & polynomial decay \\ \hline learning rate optimizer & Adam \\ \hline Adam beta weights & (0.9, 0.95) \\ \hline Adam epsilon & 10e-8 \\ \hline weight decay & 0.1 \\ \hline \end{tabular}
\end{table}
Table 15: Hyperparameters for pretraining the LM with 350M parameters per GPU. All hyperparameters are the same for DEMix and Dense training.





\begin{table}
\begin{tabular}{c c} \hline
**Computing Infrastructure** & 128 Volta 32GB GPUs \\ \hline \hline
**Hyperparameter** & **Assignment** \\ \hline architecture & GPT-3 XL \\ \hline tokens per sample & 1024 \\ \hline batch size & 2 \\ \hline learning rate & [5e–4, 3e–4, 1e–4] \\ \hline clip norm & 0.1 \\ \hline gradient accumulation steps & 8 \\ \hline number of steps & 65,000 \\ \hline save interval updates & 2,000 \\ \hline validation interval & 1,000 \\ \hline number of warmup steps & 5,200 \\ \hline learning rate scheduler & polynomial decay \\ \hline learning rate optimizer & Adam \\ \hline Adam beta weights & (0.9, 0.95) \\ \hline Adam epsilon & 10e-8 \\ \hline weight decay & 0.1 \\ \hline \end{tabular}
\end{table}
Table 16: Hyperparameters for pretraining the LM with 760M parameters per GPU. All hyperparameters are the same for DEMix and Dense training.

\begin{table}
\begin{tabular}{c c} \hline
**Hyperparameter** & **Assignment** \\ \hline architecture & GPT-3 large \\ \hline tokens per sample & 1024 \\ \hline batch size & 2 \\ \hline learning rate & [5e–4, 3e–4, 1e–4] \\ \hline clip norm & 0.1 \\ \hline gradient accumulation steps & 8 \\ \hline number of steps & 50000 \\ \hline save interval updates & 2,000 \\ \hline validation interval & 500 \\ \hline number of warmup steps & 4000 \\ \hline learning rate scheduler & polynomial decay \\ \hline Adam beta weights & (0.9, 0.95) \\ \hline Adam epsilon & 10e-8 \\ \hline weight decay & 0.1 \\ \hline \end{tabular}
\end{table}
Table 17: Hyperparameters for pretraining the LM with 1.3B parameters per GPU. All hyperparameters are the same for DEMix and Dense training.