DEMix Layers: Disentangling Domains for Modular Language Modeling

Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 5557 - 5576 July 10-15, 2022 ©2022 Association for Computational Linguistics DEMIX Layers: Disentangling Domains for Modular Language Modeling Suchin Gururangan†♢ Mike Lewis♢ Ari Holtzman† Noah A. Smith†♠ Luke Zettlemoyer†♢ †Paul G. Allen School of Computer Science & Engineering, University of Washington ♠Allen Institute for AI ♢Meta AI Seattle, WA, USA sg01@cs.washington.edu Abstract We introduce a new domain expert mixture (DEMIX) layer that enables conditioning a lan- guage 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: ex- perts can be mixed, added, or removed after initial training. Extensive experiments with autoregressive transformer LMs (up to 1.3B pa- rameters) show that DEMIX layers reduce test- time perplexity (especially for out-of-domain data), increase training efficiency, and enable rapid adaptation. Mixing experts during infer- ence, using a parameter-free weighted ensem- ble, enables better generalization to heteroge- neous or unseen domains. We also show it is possible to add experts to adapt to new domains without forgetting older ones, and remove ex- perts to restrict access to unwanted domains. Overall, these results demonstrate benefits of domain modularity in language models. 1 Introduction Most language models (LMs) are trained with data homogeneity: all parameters are updated to min- imize 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 discov- ered (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 re- spect to generalization, efficiency, and flexibility. Even if training data is sourced from many do- mains, 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. Updat- ing all parameters of the network gets substantially more expensive as model size grows (Strubell et al., 2019), making fine-tuning or domain adaptation x0 h0 DEmix Layer Self Attn Github Code FFN 1 FFN 2 FFN 3 FFN 4 FFN 5 Medical Papers U.S. Court Opinions Training COVID-19 Papers FFN 1 FFN 2 FFN 3 FFN 4 Inference Figure 1: Illustration of a DEMIX layer in a single transformer block. During training, expert feedforward networks are conditionally activated based on the do- main (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 re- duce the influence of unwanted domains (e.g., social media). Image attribution in §A.1. hard to perform with smaller computational bud- …

Content truncated to 500 words