Source-linked AI summary

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

Suchin Gururangan, Ana Marasović, Swabha Swayamdipta, Kyle Lo, Iz Beltagy, Doug Downey, Noah A. Smith

arXiv:2004.10964v3cs.CLcs.LG

TL;DR

Broad pretrained language models may still benefit from specialization to a task’s domain. This paper studies domain- and task-adaptive pretraining across four domains and eight classification tasks, finding consistent performance gains, including in low-resource settings.

  • Problem

    The paper asks whether a task’s textual domain remains relevant when using broad-coverage pretrained language models.

  • Method

    The study continues pretraining RoBERTa on unlabeled domain-specific text and task datasets across four domains and eight classification tasks.

  • Results

    Domain-adaptive pretraining consistently improves performance across high- and low-resource settings, while task-adaptive pretraining adds gains with or without domain adaptation.

  • Takeaways & Limitations

    Multi-phase adaptive pretraining offers substantial benefits, and task- or domain-relevant corpora can help specialize pretrained language models.

  • Takeaways & Limitations

    Task-adaptive pretraining involves three pretraining phases and can initially appear computationally expensive.

Abstract

from arXiv · show

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.

1 Introduction

The paper evaluates continued pretraining for domain and task adaptation across four domains and eight classification tasks. Domain-adaptive and task-adaptive pretraining consistently improve performance, while curated or automatically selected task-relevant data can further help, especially in low-resource settings.

  • Scope: The study evaluates domain- and task-adaptive pretraining across four domains and eight classification tasks in low- and high-resource settings.The domains are biomedical publications, computer science publications, news, and reviews, with two tasks in each domain.
  • Domain adaptation: Domain-adaptive pretraining consistently improves performance on target-domain tasks in both high- and low-resource settings.The experiments focus on targets that are not already in-domain for ROBERTA.
  • Task adaptation: Task-adaptive pretraining uses the unlabeled task dataset to adapt models to a corpus more directly tied to the supervised task.The paper compares task-adaptive pretraining with domain-adaptive pretraining on a smaller, directly task-relevant corpus.
  • Data selection: Benefits from task-adaptive pretraining increase with additional manually curated unlabeled data from the task distribution.The study also examines automatically selecting additional task-relevant unlabeled text.
  • Data selection: Automatically selected task-relevant unlabeled text improves performance in certain low-resource cases, and adaptive-pretraining results are competitive with state of the art on all tasks.The proposed selection strategies are motivated by the gains from human-curated datasets.

2 Background: Pretraining

Modern NLP systems typically pretrain a language model on large unlabeled corpora, then reuse and optionally fine-tune its representations for downstream supervised tasks. This paper examines continued pretraining of RoBERTa on domain-specific and task-associated unlabeled data because its broad corpus may not cover the variation of English sufficiently.

  • Pretraining paradigm: NLP systems commonly use two stages: language-model pretraining on large unlabeled corpora followed by supervised downstream training with reused representations.The pretrained representations and network may be optionally updated during fine-tuning.
  • RoBERTa: RoBERTa uses a Transformer architecture and masked language modeling, trained on over 160 GB of raw English-language text from multiple corpora.Its objective is cross-entropy loss for predicting randomly masked tokens.
  • Motivation for adaptation: The study asks whether RoBERTa’s diverse pretraining sources sufficiently generalize across English-language variation and investigates continued pretraining on additional unlabeled data.The additional data include large domain-specific corpora and unlabeled data associated with the task.

3 Domain-Adaptive Pretraining

Domain-adaptive pretraining continues ROBERTA’s pretraining on unlabeled, domain-specific text across biomedical, computer science, news, and review corpora. It improves classification performance across domains and resource settings, while comparisons and overlap analyses show that domain relevance matters and boundaries can be fuzzy.

  • Method: DAPT continues pretraining ROBERTA on unlabeled text from biomedical papers, computer science papers, news, and reviews.The resulting domain-adapted language models are trained for 12.5K steps on each domain.
  • Domain similarity: Vocabulary overlap suggests that NEWS and REVIEWS are closest to ROBERTA’s pretraining domain, whereas CS and BIOMED are more dissimilar.The analysis proposes that greater domain dissimilarity may correspond to greater expected adaptation benefit.
  • Results: DAPT improves over ROBERTA in all domains, with consistent gains for BIOMED, CS, and REVIEWS across high- and low-resource settings.AGNEWS does not improve, but HYPERPARTISAN results suggest DAPT can still help tasks aligned with ROBERTA’s source domain.
  • Domain relevance: DAPT significantly outperforms adaptation to an irrelevant domain for every task, indicating that domain-relevant pretraining matters beyond exposure to additional data.Irrelevant-domain adaptation generally performs worse than even ROBERTA on end tasks.
  • Domain overlap: Cross-domain vocabulary and document overlap indicate fuzzy boundaries, with NEWS adaptation transferring positively to REVIEWS tasks.The paper notes that pretraining beyond conventional domain boundaries may yield more effective DAPT, but leaves this for future work.

4 Task-Adaptive Pretraining

Task-adaptive pretraining (TAPT) further pretrains RoBERTa on a task’s unlabeled training data, using a smaller but more task-relevant and less expensive corpus than domain-adaptive pretraining (DAPT). TAPT consistently improves the RoBERTa baseline, can match DAPT on some tasks, and achieves the best performance when applied after DAPT.

  • Task-Adaptive Pretraining: TAPT pretrains RoBERTa on only the task-specific unlabeled training data, using a smaller but more task-relevant corpus than DAPT.TAPT uses 100 epochs with randomly masked words at masking probability 0.15, whereas DAPT is trained for 12.5K steps.
  • Task-Adaptive Pretraining: TAPT consistently improves the RoBERTa baseline across all tasks and domains, including news despite news being part of RoBERTa’s pretraining corpus.This demonstrates an advantage from adapting to task-specific data rather than only to the broader domain.
  • Task-Adaptive Pretraining: TAPT is often competitive with DAPT and matches its performance on some tasks, such as SCIERC, despite requiring fewer resources.The paper identifies DAPT as more resource intensive while describing TAPT as much less expensive to run.
  • Combined DAPT and TAPT: Combining DAPT followed by TAPT achieves the best performance on all tasks, although it is the most computationally expensive setting.The authors characterize this sequence as providing both domain and task awareness.
  • Cross-Task Transfer: Adapting to one task with TAPT optimizes single-task performance but harms cross-task transfer within the same domain.The authors evaluate Transfer-TAPT by pretraining on one task’s unlabeled data and fine-tuning on another task’s labeled data.

5 Augmenting Training Data for Task-Adaptive Pretraining

The section shows that task-adaptive pretraining benefits from larger pools of unlabeled task-distribution data, whether human-curated or automatically retrieved from an in-domain corpus. Curated-TAPT improves performance across tasks, while kNN-TAPT offers a cheaper alternative when curated data or DAPT resources are unavailable.

  • 5 Augmenting Training Data for Task-Adaptive Pretraining: The study evaluates Curated-TAPT with larger human-curated unlabeled pools and automatically retrieved task-relevant data from large in-domain corpora.The scenarios cover RCT, HYPERPARTISAN, and IMDB with available curated data, plus tasks lacking extra human-curated data.
  • 5.1 Curated-TAPT: Curated-TAPT further improves prior TAPT results across RCT, HYPERPARTISAN, and IMDB, with DAPT + Curated-TAPT producing the largest boost on every task.For HYPERPARTISAN, DAPT + Curated-TAPT is within standard deviation of Curated-TAPT.
  • 5.1 Curated-TAPT: 95% of DAPT + TAPT’s performance is achieved on RCT using only 0.3% of the labeled data.The result supports releasing large pools of unlabeled task data to aid adaptation through pretraining.
  • 5.2 Automated Data Selection: kNN-TAPT selects task-relevant domain sentences by embedding task and domain text with lightweight VAMPIRE and retrieving nearest neighbors.The method uses a deduplicated 1M-sentence domain sample and selects k candidates per task sentence; RAND-TAPT instead samples randomly.
  • 5.2 Automated Data Selection: kNN-TAPT outperforms TAPT in all evaluated cases, improves steadily as k increases, and approaches DAPT, while automatic selection is much cheaper than DAPT.TAPT is nearly 60 times faster to train than DAPT, and DAPT storage requirements are 5.8M times those of TAPT on RCT-500.

6 Related Work

Related work establishes continued pretraining for domain adaptation and task-adaptive pretraining as beneficial, while this paper compares their interaction across dataset size, task relevance, and transferability. It also situates the work within data-selection research and views domains as task-centered spectra at varying granularity.

  • Transfer learning for domain adaptation: Prior studies report benefits from continued pretraining in a target domain, whereas this work investigates shifts from diverse pretraining corpora across multiple domains.Other work has trained language models in-domain from scratch; this paper instead examines multiple domains and the effects of pretraining-corpus shifts on task performance.
  • Task-adaptive pretraining: Task-adaptive pretraining continues language-model pretraining on a task’s unlabeled data and has been shown to improve end-task performance.The paper compares DAPT and TAPT, including their interplay, with respect to continued-pretraining dataset size, task-data relevance, and transferability.
  • Data selection for transfer learning: Prior data-selection research includes choosing a corpus for language-model pretraining from scratch, while this work selects relevant examples for various tasks.The paper distinguishes its task-oriented selection from Dai et al.’s focus on a single NER task.
  • What is a domain?: The complementary effects of DAPT and TAPT suggest a spectrum of domains centered on tasks at different levels of granularity.Examples range from Amazon reviews for a specific product to all Amazon reviews, all web reviews, and the web.

7 Conclusion … A.2 State of the Art

The paper finds that adapting pretrained language models to domain- and task-relevant corpora provides significant benefits, while the supplementary material documents implementation details, related work, corpus composition, domain overlap, and state-of-the-art comparisons. The authors argue that these general adaptation approaches complement scaling efforts and suggest several directions for future work.

  • 7 Conclusion: Domain- and task-specific continued pretraining can provide significant benefits, even though hundreds-of-millions-parameter models struggle to encode a single textual domain.The authors frame specialization as a complement to efforts focused on ever-larger language models.
  • 7 Conclusion: The adaptation approaches improve ROBERTA and are general enough to apply to any pretrained language model.The conclusion presents domain- and task-relevant corpus use as broadly applicable beyond the model studied.
  • 7 Conclusion: Future directions include better data selection for TAPT, efficient adaptation to distant domains, and reusable language models after adaptation.These directions are explicitly identified as extensions of the study.
  • Appendix Overview: The supplementary material provides additional information for producing the paper’s results and results omitted from the main body.It organizes this material across appendices covering methods, related work, corpus details, domain overlap, and state-of-the-art references.
  • A Related Work: The related-work overview compares continued-pretraining strategies and their interplay using a language model trained on a much more heterogeneous corpus.Table 11 identifies which strategies had already been explored in prior work.
  • A.1 ROBERTA’s Pretraining Corpus: ROBERTA’s pretraining data comprise BOOKCORPUS, WIKIPEDIA, CCNEWS, OPENWEBTEXT, and a CommonCrawl subset resembling the story-like style of STORIES.These sources are listed in the description of ROBERTA’s pretraining corpus.
  • A.2 State of the Art: The state-of-the-art reference models vary by task, including SCIBERT for ACL-ARC, S2ORC-BERT for CHEMPROT and SCIERC, XLNet-large for AGNEWS and IMDB, Cohan et al. (2019) for RCT, and LONGFORMER for HYPERPARTISAN.The appendix directs readers to Table 5 for the reported performance of these models.

B Experimental Setup · C Development Set Results

The experiments adapt RoBERTa through domain- and task-adaptive pretraining using controlled preprocessing, training schedules, and optimization settings. Development-set results are reported in Tables 15, 17, and 18 following replication standards.

  • B Experimental Setup: Biomedical and computer science corpora used sentence-split abstracts and body paragraphs, while news used summaries and full text, and reviews used entire bodies.scispaCy handled biomedical and computer science articles; spaCy handled news and reviews.
  • B Experimental Setup: RoBERTa underwent DAPT for 12.5K steps per domain, with matched dataset exposure and an effective batch size of 2048.Amazon reviews contained more documents, but each document was shorter.
  • B Experimental Setup: TAPT reused DAPT hyperparameters, randomly masked tokens with probability 0.15 across 100 epochs, and used batch size 256 for tasks below 5K examples.Different tokens were masked across epochs through artificial dataset augmentation.
  • B Experimental Setup: Optimization used Adam, a linear learning-rate scheduler with 6% warm-up, and maximum learning rates of 0.0005 or 0.0001 depending on batch size.The 0.0001 maximum applied when batch size was 256.
  • B Experimental Setup: Prior continued-pretraining studies used narrower corpora, whereas RoBERTa was pretrained on a substantially more diverse corpus.The prior corpora included English Wikipedia, English tweets, 1BWORDBENCHMARK, and BOOKCORPUS.
  • B Experimental Setup: The implementation used HuggingFace Transformers and PyTorch XLA, ran adaptive pretraining on a single Google Cloud v3-8 TPU, and used AllenNLP for classification.Classification passed the final-layer [CLS] representation to a task-specific feedforward prediction layer.
  • C Development Set Results: Development-set results were reported in Tables 15, 17, and 18 in accordance with replication standards suggested by Dodge et al. (2019).The paper separately reports development-set results for the evaluated experiments.

D Analysis of Domain Overlap

Table 20 provides additional examples illustrating overlap between IMDB reviews and REALNEWS articles, supporting the analysis in §3.1.

  • D Analysis of Domain Overlap: Table 20 displays additional examples highlighting overlap between IMDB reviews and REALNEWS articles.These examples are relevant to the analysis in §3.1.

E Analysis of Cross-Domain Masked LM Loss

The analysis evaluates ROBERTA’s masked language-model loss across domains before and after domain-adaptive pretraining. Cross-domain losses largely reflect corpus similarities, while NEWS adaptation transfers better to REVIEWS than the reverse.

  • Cross-Domain Masked LM Loss: Table 12 evaluates cross-domain masked-LM loss on samples from other domains after DAPT.The analysis compares ROBERTA’s loss before and after domain-adaptive pretraining.
  • Cross-Domain Masked LM Loss: Cross-domain masked-LM loss mostly follows expected domain similarities: ROBERTA’s pretraining corpus is closer to NEWS, and BIOMED is closer to CS.These relationships hold relative to the other evaluated domains.
  • Cross-Domain Masked LM Loss: NEWS-adapted ROBERTA has lower loss on REVIEWS samples than REVIEWS-adapted ROBERTA has on NEWS samples.The result supports similarities between REVIEWS and NEWS while indicating asymmetric transfer between them.

F k-Nearest Neighbors Data Selection

The section presents k-nearest-neighbor selection as an approach for task-adaptive pretraining, selecting domain documents near the task data. For RCT, it identifies nearest-neighbor documents from the BIOMED domain and compares kNN-TAPT with random candidate selection.

  • F k-Nearest Neighbors Data Selection: kNN-TAPT selects neighbors of the task data from the domain data.The comparison uses RAND-TAPT with 50 candidates as the alternative selection strategy.
  • F k-Nearest Neighbors Data Selection: The evaluation compares RAND-TAPT and kNN-TAPT using development-set macro-F1 for HYPERPARTISAN and IMDB and micro-F1 for RCT.Results are averaged across five random seeds with standard deviations as subscripts.
  • F k-Nearest Neighbors Data Selection: For RCT, the method identifies five nearest-neighbor sentences from the BIOMED domain.Table 21 lists the RCT sentences as sources and the selected BIOMED documents as neighbors 0–4.
Loading 2004.10964v3…