SlowGuess/NAACL_2022_OCR huggingface dataset:
MinhQuang Pham, Josep Maria Crego, and François Yvon. 2021. Revisiting multi-domain machine translation. Transactions of the Association for Computational Linguistics, 9:17-35.
Project Gutenberg. Project Gutenberg.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. 2019. Language models are unsupervised multitask learners.
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.
Stephen Roller, Sainbayar Sukhbaatar, Arthur Szlam, and Jason Weston. 2021. Hash layers for large sparse models.
Andrei A. Rusu, Neil C. Rabinowitz, Guillaume Desjardins, Hubert Soyer, James Kirkpatrick, Koray Kavukcuoglu, Razvan Pascanu, and Raia Hadsell. 2016. Progressive neural networks.
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.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, 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, Toulouse, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net.
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.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc.
Lucy Lu Wang, Kyle Lo, Yoganand Chandrasekhar, Russell Reas, Jiangjiang Yang, Doug Burdick, Darrin Eide, Kathryn Funk, Yannis Katsis, Rodney Kinney, Yunyao Li, Ziyang Liu, William Merrill, Paul Mooney, Dewey Murdick, Devvret Rishi, Jerry Sheehan, 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.
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.
Jing Xu, Da Ju, Margaret Li, Y-Lan Boureau, Jason Weston, and Emily Dinan. 2020. Recipes for safety in open-domain chatbots.
Yelp Reviews. Yelp reviews.
Rowan Zellers, Ari Holtzman, Hannah Rashkin, Yonatan Bisk, Ali Farhadi, Franziska Roesner, and Yejin Choi. 2019. Defending against neural fake news. In NeurIPS.
A.1 Image Attribution
Images retrieved from emojipedia.org or istockphoto.com.
A.2 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. 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. 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.3 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.4 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_VqZGF2fpJTFxtXyjhmaoYi6qnz50Xr8W8hgGw/edit?usp=sharing.
A.5 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 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. 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 §A.7.
Model Scale In 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 efficiency However, 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
Category	Link to Regex	Dummy Token
Email	https://regex101.com/r/Zqsf9x/1	<EMAIL>
DART	https://regex101.com/r/0tQ6EN/1	<DART>
FB User ID	https://regex101.com/r/GZ15EZ/1	<FB_USERID>
Phone Number	https://regex101.com/r/YrDpPD/1	<PHONE_NUMBER>
Credit Card Number	https://regex101.com/r/9NTO6W/1	<CREDIT_CARD_number>
Social Security Number	https://regex101.com/r/V5GPNL/1	<SSN>
User handles	https://regex101.com/r/vpey04/1	<USER>
Table 8: Anonymization schema. We anonymize text using the regexes provided in the above links for the categories listed.
Parameters per GPU
125M	350M	760M	1.3B
DENSE	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
DEMIX	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
TFLOPs/update	556	3279	13,637	23,250
TFLOPs/GPU	31	37	48	55
GSHARD	GPUs	32	64	128	128
Total Experts	32	64	128	128
GPUs/expert	1	1	1	1
Total params	1B	6.7B	29.5B	52.5B
TFLOPs/update	675	4120	17,400	30,000
TFLOPs/GPU	15	16	19	13
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 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.
Evaluation efficiency Another benefit to DEMIX is that its experts specialize to their domain, and only a sparse subset of them are
Table 9: Our specifications for training DENSE, DEMIX, and GSHARD LMs. All models are trained for about 48 hours on V100 GPUs. DEMIX 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.
Parameters per GPU
125M	350M	760M	1.3B
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
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 §A.9 for per-domain results.
Parameters per GPU
125M	350M	760M	1.3B
DENSE (balanced)	25.9	21.4	18.4	17.8
DEMIX	21.4	18.3	17.4	17.0
GSHARD	24.0	19.5	18.9	17.2
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 §A.9 for per-domain results.
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.
$^{7}$ https://images.nvidia.com/events/sc15/pdfs/NCCL-Woolley.pdf
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.
Model performance As 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  $4\mathrm{x}$  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 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.
A.6 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.7 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 7: Heatmap of expert performance ratios, using the largest DEMIX LM (1.3B parameters per GPU). The diagonal indicates that expert specialization to their own domain. While some experts (e.g., 1B, MED) do not transfer well to most domains in the training corpus, WEBTEXT and REALNEWS experts transfer much better, confirming the heterogeneity of those domains.
DENSE (1.3B params per GPU)	29.4
DEMIX (cached; 1.3B params per GPU)	21.8
GPT-3 Da-Vinci	20.5
A.8 Expert Performance Ratios
We display a heatmap of expert performance ratios, using the largest DEMIX 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 heterogeneous domains) at test time would lead to better language modeling performance.
A.9 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/
1yNMZGSPAvhTi3JttLamiCULaOIGTJ4QGEOaj03b5kt8/
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,
Table 12: Zero-shot perplexity on the Penn TreeBank Corpus (Marcus et al., 1993), comparing our largest DENSE and DEMIX baselines with GPT-3 Da-Vinci, the largest Brown et al. (2020). Our largest DEMIX LM gains a large boost in performance over DENSE baseline, approaching the performance of GPT-3 Da-Vinci with a fraction of the compute budget.
Parameters
125M	350M	760M	1.3B
DENSE-	T	+70.1%	+21.4%	+16.7%	+20.6%
DAPT	N	-55.1%	-46.6%	-38.3%	-44.4%
Table 13: Average change in perplexity in training (T) and novel (N) domains after DENSE-DAPT. Negative values indicate better performance relative to the original DENSE LM. While average perplexity in the novel domains decreases more for DENSE-DAPT, this comes at the cost of a significant deterioration in performance in training domains.
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 calculate the domain posterior for some domains with less data.
A.11 GPT-3 Da-Vinci Comparison
We conduct an experiment comparing our largest DEMIX 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 DEMIX 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 DENSE baselines. These results further suggest the importance of domain modularity as a mechanism to improve generalization performance, in addition to model scaling.
A.12 Perplexity Changes after DENSE-DAPT
In Table 13, we display the average perplexity change after performing DENSE-DAPT on a new domain. We observe that across all model sizes, DENSE-DAPT improves performance in the novel domain, at the cost of a large performance hit in the training domains.
1b openwebtext realnews reviews legal med redis
Figure 8: Calculated domain posteriors for 8 training domains.
Figure 9: Calculated domain posteriors for 8 novel domains.
Computing Infrastructure	32 Volta 32GB GPUs
Hyperparameter	Assignment
architecture	GPT-3 small
tokens per sample	1024
batch size	2
number of workers	2
learning rate	[5e-4, 3e-4, 1e-4]
clip norm	0.1
gradient acculumption steps	8
number of steps	300,000
save interval updates	6,000
validation interval	3,000
number of warmup steps	24,000
learning rate scheduler	polynomial decay
learning rate optimizer	Adam
Adam beta weights	(0.9, 0.95)
Adam epsilon	10e-8
weight decay	0.1
Table 14: Hyperparameters for pretraining the LM with 125M parameters per GPU. All hyperparameters are the same for DEMIX and DENSE training.
Computing Infrastructure	64 Volta 32GB GPUs
Hyperparameter	Assignment
architecture	GPT-3 medium
tokens per sample	1024
batch size	2
number of workers	2
learning rate	[5e-4, 3e-4, 1e-4]
clip norm	0.1
gradient acculumption steps	8
number of steps	120,000
save interval updates	3,000
validation interval	2,000
number of warmup steps	9,600
learning rate scheduler	polynomial decay
learning rate optimizer	Adam
Adam beta weights	(0.9, 0.95)
Adam epsilon	10e-8
weight decay	0.1
Table 15: Hyperparameters for pretraining the LM with 350M parameters per GPU. All hyperparameters are the same for DEMIX and DENSE training.
Computing Infrastructure	128 Volta 32GB GPUs
Hyperparameter	Assignment
architecture	GPT-3 large
tokens per sample	1024
batch size	2
number of workers	2
learning rate	[5e-4, 3e-4, 1e-4]
clip norm	0.1
gradient accumulation steps	8
number of steps	65,000
save interval updates	2,000
validation interval	1,000
number of warmup steps	5,200
learning rate scheduler	polynomial decay
learning rate optimizer	Adam
Adam beta weights	(0.9, 0.95)
Adam epsilon	10e-8
weight decay	0.1
Table 16: Hyperparameters for pretraining the LM with 760M parameters per GPU. All hyperparameters are the same for DEMIX and DENSE training.
Computing Infrastructure	128 Volta 32GB GPUs
Hyperparameter	Assignment
architecture	GPT-3 XL
tokens per sample	1024
batch size	2
number of workers	2
learning rate	[5e-4, 3e-4, 1e-4]
clip norm	0.1
gradient accumulation steps	8
number of steps	50000
save interval updates	2,000
validation interval	500
number of warmup steps	4000
learning rate scheduler	polynomial decay
learning rate optimizer	Adam
Adam beta weights	(0.9, 0.95)
Adam epsilon	10e-8
weight decay	0.1
Table 17: Hyperparameters for pretraining the LM with 1.3B parameters per GPU. All hyperparameters are the same for DEMIX and DENSE training.