DEM IX Layers: Disentangling Domains for Modular Language Modeling
Abstract
We introduce a new domain expert mixture (DEM IX ) layer that enables conditioning a language model (LM) on the domain of the input text.
A DEM IX 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 DEM IX layers reduce testtime 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.
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 [Figure] Layer [Figure]
Opinions
Training
[Figure]
Inference 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, 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 (DEM IX ) 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 DEM IX 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; §3).
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 DEM IX layer routing mechanism follows from a natural, observable segmentation of the data.
We identify domains using coarse provenance categories (e.g., whether a document is a medical research paper or a Reddit post; §2).
Training on data from eight different domains, we find that replacing every feedforward layer in the transformer with a DEM IX layer consistently improves in-domain performance (§4).
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 (§5).
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 DEM IX consistently improves model generalization, especially out-ofdomain, while enabling many new modular capabilities.
Because DEM IX forces experts to specialize to domains, the overall model can be (partially) disentangled after training.
Beyond mixing, we can add (§6) or remove (§7) domain experts, predictably changing model behavior at inference time.
Adding experts allows for model adaptation without updating all parameters (hence avoiding forget- ting), 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.
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 multidomain 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 whitespaceseparated 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).
DEM IX Layer 3. 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:
[Equation]
[Table]
where h t,ℓ is the vector for the t th token produced by layer ℓ .
Shazeer et al.
(2017) propose to replace dense feedforward layers with an ensemble of n experts
FFN 1 , . . . , FFN n , assigned weights respectively by functions g 1 , . .
. , g n :
[Equation]
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.
3. DEM IX 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 D denote the set of domain labels (i.e., the eight labels in Table 1).
If we index the experts by D and d ∈ D is the domain label for the current training instance, then
[Equation]
We assume that each training document is associated with a single domain label.
However, we relax this requirement at inference time (§5), to model unseen or heterogeneous domains.
We perform a detailed comparison of DEM IX routing with GS HARD (Lepikhin et al., 2020), a mixture-of-experts transformer LM with learned token-level routing, in §A.5.
Our results suggest that learned token-level routing does not enable modularity, underperforms DEM IX at similar computational budgets (especially on novel domains), and is much less efficient to train and evaluate.
3. DEM IX Architecture
Our design results in one expert in a DEM IX 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 DEM IX layer, in contrast to previous work (Fedus et al., 2021;
Lepikhin et al., 2020) that interleaves shared and expert layers. [Table]
Preliminary experiments showed that interleaving led to worse in-domain performance (see §A.6 for more details).
Future work may comprehensively compare different architectural choices.
Each expert FFN j is a two-layer MLP with the same dimensions as the original FFN layer of the transformer.
This means that the effective number of parameters in the overall DEM IX LM increases (Table 2), without increasing inference runtime.
3. DEM IX Training
To train an LM with DEM IX 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, DEM IX 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).
DEM IX 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.
3. Comparison with Adapters
DEM IX 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.
In-Domain Performance
Our first set of experiments measure in-domain performance when replacing the feedforward layers in a transformer LM with DEM IX layers.
4. Experimental Setup
Architecture, Input and Hyperparameters
We 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 §A.7 for training hyperparameters.
Evaluation
We 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).
4. Models D ENSE
The first baseline is a dense LM, implemented with distributed data parallel (Li, 2021).
There is no explicit conditioning on domain.
D ENSE (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.
[Equation]
This model is trained identically to D ENSE (balanced), but we prepend a token to every sequence indicating its domain (during training and test time).
We ignore the domain token when computing perplexity during evaluation.
Parameters per GPU 125M 350M 760M 1.3B
D ENSE 20. 13.8 D ENSE (balanced) 19. 13.6 +D OMAIN -T OKEN 19. 13.4 DEM IX (naive) 18.
13.8 DEM IX (cached; §5.4) 17. DEM IX (naive) We replace every feedforward layer in the transformer with a DEM IX layer, and use DEM IX training (§3).
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.
4. 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 DEM IX 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 DEM IX (naive) model, at least when considering the average perplexity across domains.
4. 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 DEM IX layers in a naive setting (where the domain label is revealed at test time), but some do not; W EB T EXT , R EAL N EWS , and R EDDIT 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 (§A.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.
[Table]
Mixing Experts at Inference Time
The previous section established that incorporating
DEM IX 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 , . . .
, g n as mixture coefficients, transforming the domain membership of an input text into a matter of probabilistic belief.
Unlike previously proposed mixtureof-experts formulations (Shazeer et al., 2017; Lepikhin et al., 2020), this approach is parameter-free and computed only at test time.
5. Dynamic Domain Mixtures
Consider the probabilistic view of language modeling, where we estimate p ( x t | x <t ) .
We introduce a domain variable, D t , alongside each word.
We assume that this hidden variable depends on the history, x <t , so that:
[Equation]
This model is reminiscent of class-based n -gram
LMs (Brown et al., 1992; Saul and Pereira, 1997).
We have already designed the DEM IX LM to condition on a domain label, giving a form for p ( X t | x <t , D t = j ) .
We now further treat g 1 , . .
g n as a posterior probability over domains, calculated either globally or at each timestep. [Figure] x t
P ( D t | x < t ) D t [Figure] For a
, we estimate a posterior domain probabilities p ( D t | 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 x t . To do this, we apply Bayes’ rule:
[Equation]
The conditional probabilities of word sequences given a domain label, as noted above, are already defined by the DEM IX LM.
For the prior over domain labels, we consider three alternatives:
Uniform
Set a uniform prior across domains.
Updating
Set the prior at timestep t to be an exponentially weighted moving average of the posteriors from previous timesteps:
[Equation]
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 (§A.10).
We performed a small grid search to set the value λ , and found that λ
= 0.3 worked well.
1 B C S Lega l M ed
W eb t e x t
R ea l ne w s
R edd i t
R e v i e w s x <102, 1BCS
LegalMed WebtextRealnewsRedditReviews D 102 , Training Domains C O RDG i t hub
G u t enbe r g
B r e a k
i ng N e w s
C on t r a c t s A C L
T w ee t s Y e l p x <102, 1BCS
LegalMed
WebtextRealnewsRedditReviews D 102 ,
Novel Domains 0.
P ( D t | x < t ) 0.
P ( D t | x < t )
Cached
We 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 §A.10 for more details, and Figure 2 for an illustration of expert mixing.
5. Visualizing Domain Membership
In Figure 3, we plot domain posteriors calculated using the largest DEM IX LM from §4 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.
5. Experimental Setup
Here, we experiment with the corpus of novel domains (Table 1).
We evaluate the three mixture treatments of DEM IX layers (§5.1) against five Parameters [Table] baselines.
No new models are trained for this experiment beyond those used in §4.
D ENSE and D ENSE (balanced)
These are the basic baselines from §4.
+D OMAIN -T OKEN
Here test data is evaluated using each domain label token, and we choose the lowest among these perplexity values per test set.
DEM IX (naive)
Similar to +D OMAIN -T
OKEN , we evaluate the data separately with each of the eight experts, and report the lowest among these perplexity values per test set.
DEM IX (average)
At every timestep, we take a simple average of the eight experts’ predictions.
5. Results Novel Domain Performance Ensembling
DEM IX 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 DEM IX 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.
In-Domain Performance
We 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 §A.9; across all model sizes, R EDDIT improves on average 10.7%, W EB T EXT 2.4%,
R EAL N EWS 1.9%), confirming that domain labels may not align with the most effective boundaries.
5. Summary
As opposed to other token-level routing mechanisms (e.g., Lepikhin et al. 2020), expert mixing in
DEM IX is introduced at test-time and is parameterfree; 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 (§A.5).
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
D ENSE -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).
DEM IX layers allow for cheap adaptation without forgetting through a technique we call DEM IXDAPT (Figure 4).
To adapt to a new domain, we initialize a new expert in each DEM IX layer using the parameters of the nearest pretrained expert, which we identify using domain posteriors from §5.
We then train the added expert on target data, updating only the new expert parameters .
For inference, we mix experts with a cached prior (§5).
6. Experimental Setup
We compare DEM IX -DAPT to D ENSE -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 [Figure] <
[Figure] Posteriors [Figure] ( D t | x < t )
D t
[Figure]
each novel domain with the same hyperparameters as §4, except with a 10x smaller learning rate.
DEM IX -DAPT updates about 10% of the total parameters in the DEM IX LM, while D ENSE -DAPT updates all parameters of the dense LM [Figure]
6. Results Adding One Expert
We display examples of
DEM IX -DAPT and D ENSE -DAPT on a single domain in Figure 5.
As D ENSE -DAPT proceeds, its performance on the training domains progressively worsens (see §A.12 for results with larger LMs).
In contrast, DEM IX -DAPT reduces perplexity on the novel domain without forgetting.
Adding Eight Experts
We find that adding all eight experts adapted to novel domains to the
DEM IX model from §4 significantly reduces perplexity on novel and previously seen domains (Table 6) while also helping in-domain for smaller
[Figure]
[Table]
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 M ED , 1.8% reduction on R EAL N EWS , and 2% reduction on R EDDIT (see §A.9 for details).
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 [Table] 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
[Table]
becomes less feasible as the LM grows in size.
DEM IX layers offer a simple mechanism for cheap, lightweight control of large LMs: since domain experts specialize (§A.8), experts assigned to unwanted domains can be disabled at test-time.
7. 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: +E XPERT , a DEM IX LM with all experts active, –E XPERT , a DEM IX LM with a domain expert deactivated, and –D OMAIN , a DEM IX LM trained from scratch without a particular domain.
For all settings, we use a DEM IX LM (125M parameters per GPU) from §4 and expert mixing with a cached prior (§5) for inference.
7. 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., W EB T EXT and R EAL N EWS ), –E XPERT even underperforms –D OMAIN .
This leads us to conjecture that most domain-specific learning happens within the DEM IX layer.
Our preliminary analysis here suggests that
DEM IX 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.
Related Work
Document metadata has been used to improve topic models (Mimno and McCallum, 2012), adapt RNNbased 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 DEM IX experts.
Future work may explore applying DEM IX to multilingual settings, as multilingual models benefit from language-specific parameters (Fan et al., 2020;
Pfeiffer et al., 2020; Chau et al., 2020).
DEM IX -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 Schütze, 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.
Conclusion
We introduce DEM IX layers, which provide modularity to an LM at inference time, addressing limitations of dense training by providing a rapidly adaptable system.
DEM IX 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 Bhosale, 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. Ethical Considerations While DEM IX 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, DEM IX 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 A Appendix A. Image Attribution
Images retrieved from emojipedia.org or istockphoto.com .
A. 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.
R EDDIT 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 G UTENBERG , we use the scraping tool provided in https:// github.com/aparrish/gutenberg-dammit .
For
B REAKING N EWS , we identify a list of factually reliable English news sources, using the list curated by Baly et al.
Specifically, we filter on "high" factuality in the data provided in this repository: https://github. com/ramybaly/News-Media-Reliability .
We then use Newspaper3K ( 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 .
A. 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.
A. 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: // /spreadsheets/
d/1NO-Lz_
VqZGF2fpJTFxtXyjhmaoYi6qnz50Xr8W8hgGw/ edit?usp=sharing .
A. Gshard Comparison
Here we describe empirical comparisons between
DEM IX and GS HARD , the token-level mixture of experts transformer proposed by Lepikhin et al.
As opposed to DEM IX , which uses domain labels to route data to experts, GS HARD 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 GS HARD is emblematic of an learned routing procedure, we are generally interested if GS HARD naturally learns to specialize experts to domains, whether its experts are modular, and how GS HARD
LM generally performs compared to DEM IX and
D ENSE models on our multi-domain corpus.
Experimental Setup
We aim to make minimal changes to the overall architecture of the model, to focus on the differences afforded by token-level routing (vs. DEM IX routing).
As such, we keep all architecture and computational budgets the same as our DEM IX and D ENSE baselines (we generally display results for the 125M, 350M, and 760M parameter LMs).
We only add the GS HARD 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 GS HARD experts are of the same size as our DEM IX 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 GS HARD in §A.7.
Model Scale
In DEM IX , 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 GS HARD implementation, on the other hand, allocates one expert per GPU.
This means that GS HARD adds many more experts to the system, which results in a substantially larger increase in model size (Table 9).
Unlike DEM IX , GS HARD results in an increase in FLOP count relative to the D ENSE model, due to a variety of additional computation during training, like load balancing and routing to two experts for every token, which
DEM IX does not need.
Training efficiency
However, unlike DEM IX , which increases model size while maintaining or improving GPU throughput, GS HARD in fact reduces GPU throughput during training (Table 9).
This is due to the necessity of expensive all-to-all operations in GS HARD which mediate communication between experts on different GPUs that are ac- [Table] Link [Table]
Dummy [Table] Email https://regex101.com/r/ZqsF9x/ <EMAIL> [Table]
https://regex101.com/r/0tQ6EN/ <DART>
[Table]
https://regex101.com/r/GZl5EZ/ <FB_USERID>
[Table]
https://regex101.com/r/YrDpPD/ <PHONE_NUMBER>
[Table]
https://
[Table]
com/r/
[Table]
<CREDIT
[Table]
NUMBER
[Table]
https
[Table]
//regex
[Table]
com/r/V5GPNL
[Table]
<
SSN>
[Table]
https://regex101.com/r/vpey04/ <USER>
[Table]
tivated for different tokens of the same document.
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 GS HARD ’s load balancing, since some experts are not used at test time.
DEM IX 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.
Evaluation efficiency
Another benefit to
DEM IX is that its experts specialize to their domain, and only a sparse subset of them are [Table] [Table] activated at test time.
Does token-level routing via GS HARD also result in a modular model?
We explore this question by computing the average gating probabilities in the GS HARD router across all experts for all test data in each domain.
We generally find that gating probabilities in GS HARD 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 DEM IX , we only need
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.
Layer Layer
Layer1 B C S Lega l M ed
W eb t e x t
R ea l ne w s
R edd i t
R e v i e w s
Layer1 B C S Lega l M ed
W eb t e x t
R ea l ne w s
R edd i t
R e v i e w s
Layer1 B C S Lega l M ed
W eb t e x t
R ea l ne w s
R edd i t
R e v i e w s
Layer0.
Model performance
As noted earlier, our GShard implementation substantially increases the effective parameter count of the model relative to DEM IX (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,
GS HARD underperforms DEM IX 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.
Summary
Our results suggest that while
GS HARD 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 GS HARD suffers from similar downstream issues as D ENSE models, e.g., forgetting after adaptation and lack of lightweight controllability, though we leave a close exploration of those phenomena to future work.
Overall, DEM IX LMs are substantially simpler and more efficient for training and evaluation, and even outperform GS HARD (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.
A. 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.
A. 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 warmup, 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.
1 B Lega l C S M ed
W eb t e x t
R ea l ne w s
R edd i t
R e v i e w s
Domain 1B Legal CS Med
Webtext
Realnews
Reddit
Reviews E x pe r t 1x 2x 4x 6x 8x pe r p l e x i t y i n c r ea s e [Figure]
A. Expert Performance Ratios
We display a heatmap of expert performance ratios, using the largest DEM IX LM (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. A. 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/ 1yNMZGSPAvhTi3JttLamiCULaOIGTJ4QGEOajO3b5kt8/ edit?usp=sharing
A. 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, [Table] 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.
A. GPT-3 Da-Vinci Comparison
We conduct an experiment comparing our largest
DEM IX LM with GPT-3 Da-Vinci from Brown et al.
(2020), using the zero-shot language modeling evaluation they report: Penn TreeBank (Marcus et al. 1993; Table 12).
We observe that the largest
DEM IX LM achieves competitive results with the
GPT-3 Da-Vinci result with a fraction of the computation, and gives large performance boosts on this benchmark over our other D ENSE baselines.
These results further suggest the importance of domain modularity as a mechanism to improve generalization performance, in addition to model scaling.
A. Perplexity Changes after D ENSE -DAPT
In Table 13, we display the average perplexity change after performing D ENSE -DAPT on a new domain.
We observe that across all model sizes,
D ENSE -DAPT improves performance in the novel domain, at the cost of a large performance hit in the training domains.
[Figure]
1B 1b
[Figure]
reviews cs legal med reddit
[Figure]
Gutenberg
Breaking News
[Figure]
ACL Papers
[Figure]
Tweets
[Figure]
Yelp Reviews 1b openwebtext realnews reviews cs legal med reddit
Computing Infrastructure
[Table] Volta [Table]
[Table]
[Table]
[Table]