DEM ix Layers: Disentangling Domains for Modular Language Modeling (2022)

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.

## 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 @xcite , assuming that models will be able to fit all domains equally well. while dense training is convenient, and densely trained lms achieve impressive results @xcite , 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 @xcite , limiting generalization to less prevalent domains. updating all parameters of the network gets substantially more expensive as model size grows @xcite , making fine-tuning or domain adaptation inference 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. 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 @xcite . 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 ; §3). this is an example of conditional computation @xcite @xcite . 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 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 demix 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 demix consistently improves model generalization, especially out-ofdomain, 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 ( §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. foot_1 ## 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 ( we divide our data into training and test domains. the training domains text from eight english corpora (top of ## in-domain performance our first set of experiments measure in-domain performance when replacing the feedforward layers in a transformer lm with demix layers. ## 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 , . . . , 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 @xcite , this approach is parameter-free and computed only at test time. ## 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; @xcite ), which involves continued dense training of the lm on the target domain. however, dapt with dense training (or dense-dapt) is expensive @xcite and may entail forgetting domains learned during earlier training phases @xcite , 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 @xcite . 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 §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). ## removing experts dense lms are also prone to unexpected behavior when deployed. for example, they may generate hatespeech @xcite , which is undesirable for user-facing tasks @xcite . 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 @xcite and away @xcite from certain behaviors, but they tend to be expensive or require retraining the model with a different objective, which 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 ( §a.8), experts assigned to unwanted domains can be disabled at test-time. foot_3 ## related work document metadata has been used to improve topic models @xcite , adapt rnnbased lms @xcite , learn document representations @xcite , and improve text generation control @xcite . other inference-time methods @xcite 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 @xcite @xcite . demix-dapt is related to model expansion techniques in reinforcement learning or vision @xcite and adapters for pretrained lms @xcite . multi-domain models have been studied in machine translation @xcite and supervised settings @xcite , and with smaller dense lms @xcite . previous studies have shown the importance of considering domains when adapting lms @xcite . our study establishes the importance of considering domains when training lms from scratch. ## 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.