Don’t Stop Pretraining: Adapt Language Models to Domains and Tasks

Language models pretrained on text from a wide variety of sources form the foundation of today's NLP. In light of the success of these broad-coverage models, we investigate whether it is still helpful to tailor a pretrained model to the domain of a target task. We present a study across four domains (biomedical and computer science publications, news, and reviews) and eight classification tasks, showing that a second phase of pretraining in-domain (domain-adaptive pretraining) leads to performance gains, under both high- and low-resource settings. Moreover, adapting to the task's unlabeled data (task-adaptive pretraining) improves performance even after domain-adaptive pretraining. Finally, we show that adapting to a task corpus augmented using simple data selection strategies is an effective alternative, especially when resources for domain-adaptive pretraining might be unavailable. Overall, we consistently find that multi-phase adaptive pretraining offers large gains in task performance.

Today's pretrained language models are trained on massive, heterogeneous corpora. For instance, RoBERTa was trained on over 160GB of uncompressed text, with sources ranging from English-language encyclopedic and news articles, to literary works and web content. Representations learned by such models achieve strong performance across many tasks with datasets of varying sizes drawn from a variety of sources. This leads us to ask whether a task's textual domain—a term typically used to denote a distribution over language characterizing a given topic or genre (such as "science" or "mystery novels")—is still relevant. Do the latest large pretrained models work universally, or is it still helpful to build separate pretrained models for specific domains?

While some studies have shown the benefit of continued pretraining on domain-specific unlabeled data, these studies only consider a single domain at a time and use a language model that is pretrained on a smaller and less diverse corpus than the most recent language models. Moreover, it is not known how the benefit of continued pretraining may vary with factors like the amount of available labeled task data, or the proximity of the target domain to the original pretraining corpus (see Figure 1). We address this question for one such high-performing model, RoBERTa. We consider four domains (biomedical and computer science publications, news, and reviews) and eight classification tasks (two in each domain). For targets that are not already in-domain for RoBERTa, our experiments show that continued pretraining on the domain (which we refer to as domain-adaptive pretraining or DAPT) consistently improves performance on tasks from the target domain, in both high- and low-resource settings.

Above, we consider domains defined around genres and forums, but it is also possible to induce a domain from a given corpus used for a task, such as the one used in supervised training of a model. This raises the question of whether pretraining on a corpus more directly tied to the task can further improve performance. We study how domain-adaptive pretraining compares to task-adaptive pretraining, or TAPT, on a smaller but directly task-relevant corpus: the unlabeled task dataset, drawn from the task distribution. Task-adaptive pretraining has been shown effective, but is not typically used with the most recent models. We find that TAPT provides a large performance boost for RoBERTa, with or without domain-adaptive pretraining.

Finally, we show that the benefits from task-adaptive pretraining increase when we have additional unlabeled data from the task distribution that has been manually curated by task designers or annotators. Inspired by this success, we propose ways to automatically select additional task-relevant unlabeled text, and show how this improves performance in certain low-resource cases. On all tasks, our results using adaptive pretraining techniques are competitive with the state of the art. In summary, our contributions include: (i) a thorough analysis of domain- and task-adaptive pretraining across four domains and eight tasks, spanning low- and high-resource settings; (ii) an investigation into the transferability of adapted LMs across domains and tasks; and (iii) a study highlighting the importance of pretraining on human-curated datasets, and a simple data selection strategy to automatically approach this performance.

§2 Background: Pretraining. Learning for most NLP research systems since 2018 consists of training in two stages. First, a neural language model (LM), often with millions of parameters, is trained on large unlabeled corpora. The word (or wordpiece) representations learned in the pretrained model are then reused in supervised training for a downstream task, with optional updates (fine-tuning) of the representations and network from the first stage. One such pretrained LM is RoBERTa, which uses the same transformer-based architecture as its predecessor BERT. It is trained with a masked language modeling objective (i.e., cross-entropy loss on predicting randomly masked tokens). The unlabeled pretraining corpus for RoBERTa contains over 160 GB of uncompressed raw text from different English-language corpora. RoBERTa attains better performance on an assortment of tasks than its predecessors, making it our baseline of choice.

Although RoBERTa's pretraining corpus is derived from multiple sources, it has not yet been established if these sources are diverse enough to generalize to most of the variation in the English language. In other words, we would like to understand what is out of RoBERTa's domain. Towards this end, we explore further adaptation by continued pretraining of this large LM into two categories of unlabeled data: (i) large corpora of domain-specific text, and (ii) available unlabeled data associated with a given task.

Our approach to domain-adaptive pretraining (DAPT) is straightforward—we continue pretraining RoBERTa on a large corpus of unlabeled domain-specific text. The four domains we focus on are biomedical (BIOMED) papers, computer science (CS) papers, news text from REALNEWS, and AMAZON reviews. We choose these domains because they have been popular in previous work, and datasets for text classification are available in each. Table 1 lists the specifics of the unlabeled datasets in all four domains (BIOMED 7.55B tokens / 47GB; CS 8.10B / 48GB; NEWS 6.66B / 39GB; REVIEWS 2.11B / 11GB), as well as RoBERTa's training corpus (~160GB).

3.1 Analyzing Domain Similarity. Before performing DAPT, we attempt to quantify the similarity of the target domain to RoBERTa's pretraining domain. We consider domain vocabularies containing the top 10K most frequent unigrams (excluding stopwords) in comparably sized random samples of held-out documents in each domain's corpus. We also sample 50K documents from sources similar to RoBERTa's pretraining corpus (i.e., BOOKCORPUS, STORIES, WIKIPEDIA, and REALNEWS) to construct the pretraining domain vocabulary, since the original pretraining corpus is not released. Figure 2 shows the vocabulary overlap across these samples. We observe that RoBERTa's pretraining domain has strong vocabulary overlap with NEWS (54%) and REVIEWS (35%), while CS (27%) and BIOMED (19%) are far more dissimilar. This simple analysis suggests the degree of benefit to be expected by adaptation of RoBERTa to different domains—the more dissimilar the domain, the higher the potential for DAPT.

3.2 Experiments. Our LM adaptation follows the settings prescribed for training RoBERTa. We train RoBERTa on each domain for 12.5K steps, which amounts to a single pass on each domain dataset, on a v3-8 TPU. This second phase of pretraining results in four domain-adapted LMs, one for each domain. We present the masked LM loss of RoBERTa on each domain before and after DAPT in Table 1: $\mathcal{L}_{\text{ROB.}}$ vs $\mathcal{L}_{\text{DAPT}}$ — BIOMED 1.32 → 0.99, CS 1.63 → 1.34, NEWS 1.08 → 1.16, REVIEWS 2.10 → 1.93. We observe that masked LM loss decreases in all domains except NEWS after DAPT, where we observe a marginal increase. Under each domain, we consider two text classification tasks. Our tasks represent both high- and low-resource ($\le 5$K labeled training examples, and no additional unlabeled data) settings. As our baseline, we use an off-the-shelf RoBERTa-base model and perform supervised fine-tuning of its parameters for each classification task; following standard practice we pass the final-layer [CLS] token representation to a task-specific feedforward layer for prediction.

Results. Test results are shown under the DAPT column of Table 3. We observe that DAPT improves over RoBERTa in all domains. For BIOMED, CS, and REVIEWS, we see consistent improvements over RoBERTa, demonstrating the benefit of DAPT when the target domain is more distant from RoBERTa's source domain. The pattern is consistent across high- and low-resource settings. Although DAPT does not increase performance on AGNEWS, the benefit we observe in HYPERPARTISAN (86.6 → 88.2) suggests that DAPT may be useful even for tasks that align more closely with RoBERTa's source domain. Notable jumps: ACL-ARC 63.0 → 75.4 (+12.4), SCIERC 77.3 → 80.8, CHEMPROT 81.9 → 84.2.

3.3 Domain Relevance for DAPT. Additionally, we compare DAPT against a setting where for each task, we adapt the LM to a domain outside the domain of interest (¬DAPT). This controls for the case in which the improvements over RoBERTa might be attributed simply to exposure to more data, regardless of the domain. For each task, DAPT significantly outperforms adapting to an irrelevant domain, suggesting the importance of pretraining on domain-relevant data. Furthermore, we generally observe that ¬DAPT results in worse performance than even RoBERTa on end-tasks. Taken together, these results indicate that in most settings, exposure to more data without considering domain relevance is detrimental to end-task performance. However, there are two tasks (SCIERC and ACL-ARC) in which ¬DAPT marginally improves performance over RoBERTa, suggesting that in some cases, continued pretraining on any additional data is useful, as noted in Baevski et al. (2019).

3.4 Domain Overlap. Our analysis of DAPT is based on prior intuitions about how task data is assigned to specific domains. For instance, to perform DAPT for HELPFULNESS, we only adapt to AMAZON reviews, but not to any REALNEWS articles. However, the gradations in Figure 2 suggest that the boundaries between domains are in some sense fuzzy; for example, 40% of unigrams are shared between REVIEWS and NEWS. As further indication of this overlap, we qualitatively identify documents that overlap cross-domain: in Table 4 we showcase reviews and REALNEWS articles that are similar to these reviews. In fact, we find that adapting RoBERTa to NEWS is not as harmful to its performance on REVIEWS tasks (DAPT on NEWS achieves 65.5 on HELPFULNESS and 95.0 on IMDB). Although this analysis is by no means comprehensive, it indicates that the factors giving rise to observable domain differences are likely not mutually exclusive. It is possible that pretraining beyond conventional domain boundaries could result in more effective DAPT; we leave this to future work.

Datasets curated to capture specific tasks of interest tend to cover only a subset of the text available within the broader domain. For example, the CHEMPROT dataset for extracting relations between chemicals and proteins focuses on abstracts of recently-published, high-impact articles from hand-selected PubMed categories. We hypothesize that in such cases where the task data is a narrowly-defined subset of the broader domain, pretraining on the task dataset itself, or data relevant to the task, may be helpful. Task-adaptive pretraining (TAPT) refers to pretraining on the unlabeled training set for a given task; prior work has shown its effectiveness. Compared to domain-adaptive pretraining (DAPT), the task-adaptive approach strikes a different trade-off: it uses a far smaller pretraining corpus, but one that is much more task-relevant (under the assumption that the training set represents aspects of the task well). This makes TAPT much less expensive to run than DAPT, and as we show in our experiments, the performance of TAPT is often competitive with that of DAPT.

4.1 Experiments. Similar to DAPT, task-adaptive pretraining consists of a second phase of pretraining RoBERTa, but only on the available task-specific training data. In contrast to DAPT, which we train for 12.5K steps, we perform TAPT for 100 epochs. We artificially augment each dataset by randomly masking different words (using the masking probability of 0.15) across epochs. As in our DAPT experiments, we pass the final-layer [CLS] token representation to a task-specific feedforward layer for classification.

Results. Our results are shown in the TAPT column of Table 5. TAPT consistently improves the RoBERTa baseline for all tasks across domains. Even on the news domain, which was part of RoBERTa's pretraining corpus, TAPT improves over RoBERTa, showcasing the advantage of task adaptation. Particularly remarkable are the relative differences between TAPT and DAPT. DAPT is more resource intensive (see Table 9), but TAPT manages to match its performance in some of the tasks, such as SCIERC. In RCT, HYPERPARTISAN, AGNEWS, HELPFULNESS, and IMDB, the results even exceed those of DAPT, highlighting the efficacy of this cheaper adaptation technique. Selected numbers (RoBERTa / TAPT): HYPERPARTISAN 86.6 → 90.4, AGNEWS 93.9 → 94.5, HELPFULNESS 65.1 → 68.5, IMDB 95.0 → 95.5.

Combined DAPT and TAPT. We investigate the effect of using both adaptation techniques together. We begin with RoBERTa and apply DAPT then TAPT under this setting. The three phases of pretraining add up to make this the most computationally expensive of all our settings. As expected, combined domain- and task-adaptive pretraining achieves the best performance on all tasks (Table 5; e.g., ACL-ARC 75.6, SCIERC 81.3, HYPERPARTISAN 90.0, AGNEWS 94.6). Overall, our results show that DAPT followed by TAPT achieves the best of both worlds of domain and task awareness, yielding the best performance. While we speculate that TAPT followed by DAPT would be susceptible to catastrophic forgetting of the task-relevant corpus, alternate methods of combining the procedures may result in better downstream performance. Future work may explore pretraining with a more sophisticated curriculum of domain and task distributions.

Cross-Task Transfer. We complete the comparison between DAPT and TAPT by exploring whether adapting to one task transfers to other tasks in the same domain. For instance, we further pretrain the LM using the RCT unlabeled data, fine-tune it with the CHEMPROT labeled data, and observe the effect. We refer to this setting as Transfer-TAPT. Our results for tasks in all four domains are shown in Table 6. We see that TAPT optimizes for single-task performance, to the detriment of cross-task transfer (e.g., CHEMPROT 82.6 → 80.4 / ↓2.2; HYPERPARTISAN 89.9 → 82.2 / ↓7.7; ACL-ARC 67.4 → 64.1 / ↓3.3; HELPFULNESS 68.5 → 65.0 / ↓3.5). These results demonstrate that data distributions of tasks within a given domain might differ. Further, this could also explain why adapting only to a broad domain is not sufficient, and why TAPT after DAPT is effective.

In §4, we continued pretraining the LM for task adaptation using only the training data for a supervised task. Inspired by the success of TAPT, we next investigate another setting where a larger pool of unlabeled data from the task distribution exists, typically curated by humans. We explore two scenarios: first, for three tasks (RCT, HYPERPARTISAN, and IMDB) we use this larger pool of unlabeled data from an available human-curated corpus (§5.1). Next, we explore retrieving related unlabeled data for TAPT, from a large unlabeled in-domain corpus, for tasks where extra human-curated data is unavailable (§5.2).

5.1 Human Curated-TAPT. Dataset creation often involves collection of a large unlabeled corpus from known sources. This corpus is then downsampled to collect annotations, based on the annotation budget. The larger unlabeled corpus is thus expected to have a similar distribution to the task's training data, and it is usually available. We simulate a low-resource setting RCT-500 by downsampling the training data of the RCT dataset to 500 examples (out of 180K available), and treat the rest of the training data as unlabeled. The HYPERPARTISAN shared task has two tracks: low- and high-resource; we use 5K documents from the high-resource setting as Curated-TAPT unlabeled data and the original low-resource training documents for task fine-tuning. For IMDB, we use the extra unlabeled data manually curated by task annotators, drawn from the same distribution as the labeled data.

We compare Curated-TAPT to TAPT and DAPT + TAPT in Table 7. Curated-TAPT further improves our prior results from §4 across all three datasets. Applying Curated-TAPT after adapting to the domain results in the largest boost in performance on all tasks; in HYPERPARTISAN, DAPT + Curated-TAPT is within standard deviation of Curated-TAPT. Moreover, Curated-TAPT achieves 95% of the performance of DAPT + TAPT with the fully labeled RCT corpus (Table 5) with only 0.3% of the labeled data. These results suggest that curating large amounts of data from the task distribution is extremely beneficial to end-task performance. We recommend that task designers release a large pool of unlabeled task data for their tasks to aid model adaptation through pretraining.

5.2 Automated Data Selection for TAPT. Consider a low-resource scenario without access to large amounts of unlabeled data to adequately benefit from TAPT, as well as absence of computational resources necessary for DAPT. We propose simple unsupervised methods to retrieve unlabeled text that aligns with the task distribution, from a large in-domain corpus. Our approach finds task-relevant data from the domain by embedding text from both the task and domain in a shared space, then selects candidates from the domain based on queries using the task data. Importantly, the embedding method must be lightweight enough to embed possibly millions of sentences in a reasonable time. Given these constraints, we employ VAMPIRE (Figure 3), a lightweight bag-of-words language model. We pretrain VAMPIRE on a large deduplicated sample of the domain (1M sentences) to obtain embeddings of the text from both the task and domain sample. We then select $k$ candidates of each task sentence from the domain sample, in embedding space, either via nearest-neighbor selection (kNN-TAPT, using a flat FAISS index with cosine similarity) or randomly (RAND-TAPT). We continue pretraining RoBERTa on this augmented corpus with both the task data (as in TAPT) and the selected candidate pool.

Results. Results in Table 8 show that kNN-TAPT outperforms TAPT for all cases. RAND-TAPT is generally worse than kNN-TAPT, but within a standard deviation on RCT and ACL-ARC. As we increase $k$, kNN-TAPT performance steadily increases and approaches that of DAPT. For ACL-ARC: TAPT 67.4 / RAND-TAPT 69.7 / 50NN 70.7 / 150NN 73.3 / 500NN 75.5, while DAPT is 75.4. Future work might consider a closer study of kNN-TAPT, more sophisticated data selection methods, and the tradeoff between the diversity and task relevance of selected examples.

The computational requirements for all our adaptation techniques on RCT-500 in the BIOMED domain are summarized in Table 9. TAPT is nearly 60 times faster to train than DAPT on a single v3-8 TPU and storage requirements for DAPT on this task are 5.8M times that of TAPT (TAPT 0.2K steps / 500 docs / 80KB; DAPT 12.5K steps / 25M docs / 47GB). Our best setting of DAPT + TAPT amounts to three phases of pretraining, and at first glance appears to be very expensive. However, once the LM has been adapted to a broad domain, it can be reused for multiple tasks within that domain, with only a single additional TAPT phase per task.

While Curated-TAPT tends to achieve the best cost-benefit ratio in this comparison, one must also take into account the cost of curating large in-domain data. Automatic methods such as kNN-TAPT are much cheaper than DAPT: e.g., 500NN-TAPT runs 9.0K steps on 185K docs / 24MB to reach F1 81.7, while DAPT runs 12.5K steps on 25M docs / 47GB to reach 82.5. The takeaway: DAPT pays off when the same domain LM will be reused across many tasks; TAPT (or kNN-TAPT / Curated-TAPT) is the right choice when budget is tight or when adaptation must be tailored per task.

Transfer learning for domain adaptation. Prior work has shown the benefit of continued pretraining in domain. We have contributed further investigation of the effects of a shift between a large, diverse pretraining corpus and target domain on task performance. Other studies have trained LMs in their domain of interest from scratch. In contrast, our work explores multiple domains, and is arguably more cost effective, since we continue pretraining an already powerful LM.

Task-adaptive pretraining. Continued pretraining of a LM on the unlabeled data of a given task (TAPT) has been shown to be beneficial for end-task performance (Howard and Ruder, 2018; Phang et al., 2018; Sun et al., 2019). In the presence of domain shift between train and test data distributions of the same task, domain-adaptive pretraining (DAPT) is sometimes used to describe what we term TAPT (Logeswaran et al., 2019; Han and Eisenstein, 2019). Related approaches include language modeling as an auxiliary objective to task classifier fine-tuning, or considering simple syntactic structure of the input while adapting to task-specific data. We compare DAPT and TAPT as well as their interplay with respect to dataset size for continued pretraining (hence, expense of more rounds of pretraining), relevance to a data sample of a given task, and transferability to other tasks and datasets.

Data selection for transfer learning. Selecting data for transfer learning has been explored in NLP. Dai et al. (2019) focus on identifying the most suitable corpus to pretrain a LM from scratch, for a single task (NER), whereas we select relevant examples for various tasks. Concurrent to our work, Aharoni and Goldberg (2020) propose data selection methods for NMT based on cosine similarity in embedding space, using DistilBERT for efficiency. In contrast, we use VAMPIRE, and focus on augmenting TAPT data for text classification tasks. Khandelwal et al. (2020) introduced kNN-LMs that allow easy domain adaptation of pretrained LMs by simply adding a datastore per domain with no further training; an alternative to integrate domain information in an LM. Our study of human-curated data is related to focused crawling for collection of suitable data, especially with LM reliance.

What is a domain? Despite the popularity of domain adaptation techniques, most research and practice seems to use an intuitive understanding of domains. A small body of work has attempted to address this question. For instance, Aharoni and Goldberg (2020) define domains by implicit clusters of sentence representations in pretrained LMs. Our results show that DAPT and TAPT complement each other, which suggests a spectrum of domains defined around tasks at various levels of granularity (e.g., Amazon reviews for a specific product, all Amazon reviews, all reviews on the web, the web).

We investigate several variations for adapting pretrained LMs to domains and tasks within those domains, summarized in Table 10. Our experiments reveal that even a model of hundreds of millions of parameters struggles to encode the complexity of a single textual domain, let alone all of language. We show that pretraining the model towards a specific task or small corpus can provide significant benefits. Our findings suggest it may be valuable to complement work on ever-larger LMs with parallel efforts to identify and use domain- and task-relevant corpora to specialize models. While our results demonstrate how these approaches can improve RoBERTa, a powerful LM, the approaches we studied are general enough to be applied to any pretrained LM. Our work points to numerous future directions, such as better data selection for TAPT, efficient adaptation of large pretrained language models to distant domains, and building reusable language models after adaptation.