SlowGuess/NAACL_2022_OCR huggingface dataset:
+DOMAIN-TOKEN Here 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.
Novel Domain Performance Ensembling 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
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 heterogeneous domains (Table 4 and §A.9; across all model sizes, REDIT improves on average  $10.7\%$ , WEBTEXT  $2.4\%$ , REALNEWS  $1.9\%$ ), confirming that domain labels may not align with the most effective boundaries.
5.5 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 (§A.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 §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).
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
4We 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.
3. Adapt new expert, freezing all other parameters
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.
each novel domain with the same hyperparameters as §4, 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.
Adding One Expert We 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 §A.12 for results with larger LMs). In contrast, DEMIX-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 DEMIX model from §4 significantly reduces perplexity on novel and previously seen domains (Table 6) while also helping in-domain for smaller
Figure 5: Adapting an LM (125M parameters per GPU) to CORD-19 or GUTENBERG. Top row: with DENSE-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.
Domains	# Experts	Parameters per GPU
125M	350M	760M	1.3B
TRAINING	8	17.8	14.7	13.9	13.4
16	17.7	14.6	13.7	13.4
NOVEL	8	21.4	18.3	17.4	17.0
16	16.0	14.0	13.5	12.5
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.
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  $\S A.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
Domain	125M Parameters per GPU
+EXPERT	-EXPERT	-DOMAIN
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
REIEWS	20.5	39.0	43.0
Average	18.2	28.5	30.6
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.
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.<sup>5</sup>
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
For all settings, we use a DEMIX LM (125M parameters per GPU) from §4 and expert mixing with a cached prior (§5) for inference.
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.
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 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.
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.
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.
$^{5}$ Removing an expert offers no guarantee of having fully forgotten content from the removed domain, since there are shared parameters in the model.  
 $^{6}$ We replace the removed domain with GUTENBERG, since our cluster allocates training jobs via 8-GPU nodes.
While DEMIX 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.
Armen Aghajanyan, Akshit 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.
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.
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.
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.
Jason Baumgartner, Savvas Zannettou, Brian Keegan, Megan Squire, and Jeremy Blackburn. 2020. The pushshift reddit dataset.
Emily M. Bender, Timnit Gebru, Angelina McMillan-Major, 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.
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.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. 2020. Language models are few-shot learners.
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).
Caselaw Access Project. Caselaw access project.
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.
Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillip Koehn, and Tony Robinson. 2014. One billion word benchmark for measuring progress in statistical language modeling.
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.
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.
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. 2017 International Joint Conference on Neural Networks (IJCNN), pages 526-533.
Angela Fan, Shruti Bhosale, Holger Schwenk, Zhiyi Ma, Ahmed El-Kishky, Siddharth Goyal, Mandep Baines, Onur Celebi, Guillaume Wenzek, Vishrav Chaudhary, Naman Goyal, Tom Birch, Vitaliy Liptchinsky, Sergey Edunov, Edouard Grave, Michael Auli, and Armand Joulin. 2020. Beyond english-centric multilingual machine translation.
William Fedus, Barret Zoph, and Noam Shazeer. 2021. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity.
Leo Gao, Stella Biderman, Sid Black, Laurence 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.
Samuel Gehman, Suchin Gururangan, Maarten Sap, Yejin Choi, and Noah A. Smith. 2020. RealToxicityPrompts: Evaluating neural toxic degeneration in language models. In Findings of the Association for Computational Linguistics: EMNLP 2020, pages 3356-3369, Online. Association for Computational Linguistics.
GibHub Archive Project. Github archive project.
Aaron Gokaslan and Vanya Cohen. 2019. Openwebtext corpus.
Suchin Gururangan, Ana Marasovic, Swabha Swayamdipta, Kyle Lo, Iz Beltagy, Doug Downey, and Noah 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, pages 8342-8360, Online. Association for Computational Linguistics.
Dan Hendrycks, Collin Burns, Anya Chen, and Spencer Ball. 2021. Cuad: An expert-annotated nlp dataset for legal contract review.
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin de Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. 2019. Parameter-efficient transfer learning for nlp.
Aaron Jaech and Mari Ostendorf. 2018. Low-rank rnn adaptation for context-aware language modeling.
Nitish Shirish Keskar, Bryan McCann, Lav R. Varshney, Caiming Xiong, and Richard Socher. 2019. Ctrl: A conditional transformer language model for controllable generation.
Diederik P. Kingma and Jimmy 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.
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, Tony Lee, Etienne David, Ian Stavness, Wei Guo, Berton Earnshaw, Imran Haque, Sara M Beery, Jure Leskovec, Anshul Kundaje, Emma Pierson, Sergey Levine, Chelsea Finn, and Percy Liang. 2021. WILDS: A benchmark of in-the-wild distribution shifts. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 5637-5664. PMLR.
Angeliki Lazaridou, Adhiguna Kuncoro, Elena Gribovskaya, Devang Agrawal, Adam Liska, Tayfun Terzi, Mai Gimenez, Cyprien de Masson d'Autume, Sebastian Ruder, Dani Yogatama, Kris Cao, Tomas Kocisky, Susannah Young, and Phil Blunsom. 2021. Pitfalls of static language modelling.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan First, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. 2020. Gshard: Scaling giant models with conditional computation and automatic sharding.
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. 2021. Base layers: Simplifying training of large, sparse models. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 6265-6274. PMLR.
Shen Li. 2021. Getting started with distributed data parallel.
Alisa Liu, Maarten Sap, Ximing Lu, Swabha Swayamdipta, Chandra Bhagavatula, Noah A. Smith, and Yejin Choi. 2021. Dexperts: Decoding-time controlled text generation with experts and anti-experts.
Kyle Lo, Lucy Lu Wang, Mark Neumann, Rodney Kinney, and Daniel Weld. 2020. S2ORC: The semantic scholar open research corpus. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 4969-4983, Online. Association for Computational Linguistics.
Li Lucy and David Bamman. 2021. Characterizing English Variation across Social Media Communities with BERT. Transactions of the Association for Computational Linguistics, 9:538-556.
Kelvin Luu, Daniel Khashabi, Suchin Gururangan, Karishma Mandyam, and Noah 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.
Antonios Maronikolakis and Hinrich Schütze. 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.
M. McCloskey and N. Cohen. 1989. Catastrophic interference in connectionist networks: The sequential learning problem. Psychology of Learning and Motivation, 24:109-165.
David Mimno and Andrew McCallum. 2012. Topic models conditioned on arbitrary features with dirichlet-multinomial regression.
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 ongpu clusters using megatron-lm.
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.
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.
Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. 2019. *fairoseq: 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.
Jonas Pfeiffer, Ivan Vulic, Iryna Gurevych, and Sebastian Ruder. 2020. MAD-X: An Adapter-Based Framework for Multi-Task Cross-Linguual Transfer. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 7654–7673, Online. Association for Computational Linguistics.