SlowGuess/NAACL_2022_OCR huggingface dataset:
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
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.
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 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.
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,
Human Language Technologies, pages 5557 - 5576
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 (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 (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 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-of-domain, 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 ( $\S 6$ ) or remove ( $\S 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. $^2$
2 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 whitespace-separated 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.
$\S 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).
3 DEMIX Layer
3.1 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:
$$
\mathbf {h} _ {t, \ell} = \operatorname {F F N} \left(\mathbf {h} _ {t, \ell - 1}\right), \tag {1}
$$
<sup>1</sup>We perform a detailed comparison of learned and DEMIX routing in §A.5.
2github.com/kernelmachine/demix
$^{3}$ github.com/kernelmachine/demix-data
Domain	Corpus	# Train (Eval.) Tokens
TRAINING	1B	30M NewsWire sentences (Chelba et al., 2014)	700M (10M)
CS	1.89M full-text CS papers from S2ORC (Lo et al., 2020)	4.5B (10M)
LEGAL	2.22M U.S. court opinions, 1658 to 2018 (Caselaw Access Project)	10.5B (10M)
MED	3.2M full-text medical papers from S2ORC (Lo et al., 2020)	9.5B (10M)
WEBTEXT†	8M Web documents (Gokaslan and Cohen, 2019)	6.5B (10M)
REALNEWS†	35M articles from REALNEWS (Zellers et al., 2019)	15B (10M)
REDDIT	Reddit comments from pushshift.io (Baumgartner et al., 2020)	25B (10M)
REVIEWS†	30M Amazon product reviews (Ni et al., 2019)	2.1B (10M)
Total 73.8B (80M)
Domain	Corpus	# Train (Eval.) Tokens
NOVEL	ACL PAPERS	1.5K NLP papers from ACL (Dasigi et al., 2021)	1M (1M)
BREAKING NEWS†	20K latest articles from 400 English news sites (Baly et al., 2018)	11M (1M)
CONTRACTS†	500 commercial legal contracts (Hendrycks et al., 2021)	1.5M (1M)
CORD-19	400K excerpts from COVID-19 research papers (Wang et al., 2020)	60M (10M)
GITHUB	230K public Github repository contents (Github Archive Project)	200M (10M)
GUTENBERG	3.2M copyright-expired books (Project Gutenberg)	3B (10M)
TWEETS†	1M English tweets from 2013-2018	8M (1M)
YELP REVIEWS†	6M Yelp restaurant reviews (Yelp Reviews)	600M (10M)
Table 1: Domains that make up our multi-domain training corpus, including the size of our training and evaluation (i.e. validation and test) data, in whitespace-separated tokens.  $\dagger$  indicates datasets that we (partially) anonymize ( $\S 2$ ). See Appendix  $\S A.2$  for more details on how these data were collected.
where  $\mathbf{h}_{t,\ell}$  is the vector for the  $t$ th token produced by layer  $\ell$ .
Shazeer et al. (2017) propose to replace dense feedforward layers with an ensemble of  $n$  experts  $\mathrm{FFN}_1,\ldots ,\mathrm{FFN}_n$  assigned weights respectively by functions  $g_{1},\dots,g_{n}$ :
$$
\operatorname {F F N} \left(\mathbf {h} _ {t, \ell - 1}\right) = \sum_ {j = 1} ^ {n} g _ {j} \left(\mathbf {h} _ {t, \ell - 1}\right) \cdot \operatorname {F F N} _ {j} \left(\mathbf {h} _ {t, \ell - 1}\right) \tag {2}
$$
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.2 DEMIX 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  $\mathcal{D}$  denote the set of domain labels (i.e., the eight labels in Table 1). If we index the experts by  $\mathcal{D}$  and  $d\in \mathcal{D}$  is the domain label for the current training instance, then
$$
g _ {j} \left(\mathbf {h} _ {t, \ell}\right) = \left\{ \begin{array}{l l} 1 & \text {i f} j = d \\ 0 & \text {o t h e r w i s e} \end{array} \right. \tag {3}
$$
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 DEMIX routing with GSHARD (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 DEMIX at similar computational budgets (especially on novel domains), and is much less efficient to train and evaluate.
3.3 DEMix Architecture
Our design results in one expert in a DEMIX 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 DEMIX layer, in contrast to previous work (Fedus et al., 2021; Lepikhin et al., 2020) that interleaves shared and expert layers.
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
Table 2: Our specifications for training DENSE and DEMIX 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.
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  $\mathrm{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 DEMix LM increases (Table 2), without increasing inference runtime.
3.4 DEMIX Training
To train an LM with DEMIX 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, DEMix 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). DEMix 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.5 Comparison with Adapters
DEMIX experts are related to adapters (Bapna and First, 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.
4 In-Domain Performance
Our first set of experiments measure in-domain performance when replacing the feedforward layers in a transformer LM with DEMIX layers.
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).
DENSE The first baseline is a dense LM, implemented with distributed data parallel (Li, 2021). There is no explicit conditioning on domain.
DENSE (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.
+DOMAIN-TOKEN This model is trained identically to DENSE (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
DENSE	20.6	16.5	14.5	13.8
DENSE (balanced)	19.9	15.8	14.3	13.6
+DOMAIN-TOKEN	19.2	15.9	14.3	13.4
DEMIX (naive)	18.4	15.5	14.2	13.8
DEMIX (cached; §5.4)	17.8	14.7	13.9	13.4
DEMIX (naive) We replace every feedforward layer in the transformer with a DEMIX layer, and use DEMIX 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.
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 DEMIX 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 DEMIX (naive) model, at least when considering the average perplexity across domains.
4.4 Domain Heterogeneity
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 DEMIX layers in a naive setting (where the domain label is revealed at test time), but some do not; WEBTEXT, REALNEWS, and REDDIT fare worse than the dense baseline. We hypothesize that dense training is advantageous for heterogenous 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 3: 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.
Domain	1.3B parameters per GPU
DENSE	DEMIX (naive)	DEMIX (cached prior; §5.4)
1B	11.8	11.5	11.3
CS	13.5	12.2	12.1
LEGAL	6.8	6.7	6.7
MED	9.5	9.2	9.1
WEBTEXT	13.8	14.6	14.3
REALNEWS	12.5	13.3	13.1
REDDIT	28.4	30.6	28.1
REIEWS	14.0	12.6	12.5
Average	13.8	13.8	13.4
Table 4: Test perplexity by domain for the largest models. We discuss the last column in §5.4.
5 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}, \ldots, g_{n}$  as mixture coefficients, transforming the domain membership of an input text into a matter of probabilistic belief. Unlike previously proposed mixture-of-experts formulations (Shazeer et al., 2017; Lepikhin et al., 2020), this approach is parameter-free and computed only at test time.
5.1 Dynamic Domain Mixtures
Consider the probabilistic view of language modeling, where we estimate  $p(\pmb{x}_t \mid \pmb{x}_{<t})$ . We introduce a domain variable,  $D_t$ , alongside each word. We assume that this hidden variable depends on the history,  $\pmb{x}_{<t}$ , so that:
$$
p \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {<   t}\right) = \sum_ {j = 1} ^ {n} p \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {<   t}, D _ {t} = j\right) \cdot \underbrace {p \left(D _ {t} = j \mid \boldsymbol {x} _ {<   t}\right)} _ {g _ {j}} \tag {4}
$$
This model is reminiscent of class-based  $n$ -gram LMs (Brown et al., 1992; Saul and Pereira, 1997).
We have already designed the DEMIX LM to condition on a domain label, giving a form for  $p(X_{t} \mid x_{<t}, D_{t} = j)$ . We now further treat  $g_{1}, \ldots, g_{n}$  as a posterior probability over domains, calculated either globally or at each timestep.
Figure 2: Illustration of inference with domain expert mixing. For a given input text  $\boldsymbol{x}_{<t}$  from CORD-19, we estimate a posterior domain probabilities  $p(D_t \mid \boldsymbol{x}_{<t})$ , informed by a prior that is either iteratively updated during inference, or is precomputed and cached on held-out 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  $\boldsymbol{x}_t$ .
To do this, we apply Bayes' rule:
$$
\begin{array}{l} p \left(D _ {t} = j \mid \boldsymbol {x} _ {<   t}\right) = \frac {p \left(\boldsymbol {x} _ {<   t} \mid D _ {t} = j\right) \cdot p \left(D _ {t} = j\right)}{p \left(\boldsymbol {x} _ {<   t}\right)} (5) \\ = \frac {p \left(\boldsymbol {x} _ {<   t} \mid D _ {t} = j\right) \cdot p \left(D _ {t} = j\right)}{\sum_ {j ^ {\prime} = 1} ^ {n} p \left(\boldsymbol {x} _ {<   t} \mid D _ {t} = j ^ {\prime}\right) \cdot p \left(D _ {t} = j ^ {\prime}\right)} (6) \\ \end{array}
$$
The conditional probabilities of word sequences given a domain label, as noted above, are already defined by the DEMIX 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:
$$
p (D _ {t} = j) \propto \sum_ {t ^ {\prime} = 1} ^ {t - 1} \lambda^ {t - t ^ {\prime}} \cdot p (D _ {t ^ {\prime}} = j \mid \boldsymbol {x} _ {<   t ^ {\prime}}) \quad (7)
$$
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  $\lambda$ , and found that  $\lambda = 0.3$  worked well.
Figure 3: Estimates of posteriors  $p(D_t \mid x_{<t})$  with a DEMIX LM (1.3B parameters per GPU), after 100 sequences (i.e., 102,400 tokens) of data in training (top heatmap) and novel domains (bottom heatmap).
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.2 Visualizing Domain Membership
In Figure 3, we plot domain posteriors calculated using the largest DEMIX 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 heterogeneous 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.
Here, we experiment with the corpus of novel domains (Table 1). We evaluate the three mixture treatments of DEMix layers (§5.1) against five
Parameters per GPU
125M	350M	760M	1.3B
DENSE	25.9	21.4	18.4	17.8
DENSE (balanced)	25.3	19.6	18.3	17.1
+DOMAIN-TOKEN	24.8	20.4	18.4	18.0
DEMUX (naive)	28.8	23.8	21.8	21.1
DEMUX (average)	27.2	22.4	21.5	20.1
DEMUX (uniform)	24.5	20.5	19.6	18.7
DEMUX (updating)	21.9	18.7	17.6	17.1
DEMUX (cached)	21.4	18.3	17.4	17.0
Table 5: 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.
basielines. No new models are trained for this experiment beyond those used in  $\S 4$
DENSE and DENSE (balanced) These are the basic baselines from  $\S 4$