Do not Mask Randomly:
Effective Domain-adaptive Pre-training by Masking In-domain Keywords

Employing large pre-trained language models (PLMs) is currently a common practice for most natural language processing (NLP) tasks (Tunstall et al., 2022). A two-stage pre-train-then-fine-tune framework is usually used to adapt/fine-tune PLMs to downstream tasks (Devlin et al., 2019). However, motivated by ULMFiT (Howard and Ruder, 2018) and ELMo (Peters et al., 2018), Gururangan et al. (2020) showed that incorporating in-domain pre-training (also known as domain-adaptive pre-training) between generic pre-training and fine-tuning stages can lead to further performance improvements in downstream tasks because it “pulls” the PLM towards the target domain. At this intermediate stage, the domain adaptation for PLMs is typically handled by continuing pre-training in the same way, i.e., using randomly-masked tokens on unstructured in-domain data Devlin et al. (2019). Here, we argue that this intermediate pre-training should be performed differently, i.e., masking should focus on words that are representative of target domain to streamline the adaptation process.

We propose a novel task-independent in-domain pre-training approach for adapting PLMs that increases domain fit by focusing on keywords in the target domain, where keywords are defined as “a sequence of one or more words that offers a compact representation of a document’s content” (Rose et al., 2010). By applying token masking only to in-domain keywords, the meaningful information in the target domain is more directly captured by the PLM. This is in contrast to the classic pre-training strategy that randomly masks tokens (Devlin et al., 2019), which may overlook domain-meaningful information, or the in-domain pre-training methods that selectively mask tokens deemed important given the downstream task (Gu et al., 2020, inter alia), which require incorporating information from the downstream task into the pre-training stage. We empirically show that our method offers a better transmission of high-quality information from the target domain into PLMs, yielding better generalizability for the downstream tasks.

The key contributions of this paper are:

(1)

We propose the first task-agnostic selective masking technique for domain adaptation of PLMs that relies solely on in-domain keywords. In particular, we first extract contextually-relevant keywords from each available document in the target domain using KeyBERT Grootendorst (2020) and keep the most frequently occurring keywords to be masked during the adaptation phase.

(2)

We evaluate our proposed strategy by measuring the performance of fine-tuned PLMs in six different settings. We leverage three different datasets for text classification from multiple domains: IMDB movie reviews Maas et al. (2011), Amazon pet product reviews from Kaggle,²²2https://www.kaggle.com/datasets/kashnitsky/exploring-transfer-learning-for-nlp and PUBHEALTH (Kotonya and Toni, 2020). Our experiments show that the classifiers trained on top of two PLMs—in our case, Bidirectional Encoder Representations from Transformers (BERT) Base and Large (Vaswani et al., 2017; Devlin et al., 2019)—that are adapted based on our suggested approach outperform all baselines, including the fine-tuned BERT with no in-domain adaptation, and fine-tuned BERT adapted by random masking. Further, the overhead of identifying in-domain keywords is reasonable, e.g., 7–15% of the pre-training time (for two epochs of data) for BERT Large.

Bidirectional Encoder Representations from Transformers (BERT) brought pre-training to transformer networks Vaswani et al. (2017) through masked language modeling (MLM) Devlin et al. (2019). They showed that a simple two-step paradigm of generic pre-training followed by fine-tuning to the target domain can significantly improve performance on a variety of tasks.

However, after showing that infusing an intermediate pre-training stage (commonly known as in-domain pre-training) can help pre-trained Long Short-Term Memory models learn domain-specific patterns better (Howard and Ruder, 2018; Peters et al., 2018), Gururangan et al. (2020) found that the same advantage applies to PLMs as well. Since then, several efforts proposed different domain-adaptive pre-training strategies.

Unsurprisingly, one of the most extensively utilized in-domain pre-training methodologies has been to employ classic random masking to adapt PLMs into several domains (Lee et al., 2020; Beltagy et al., 2019; Alsentzer et al., 2019; Tavabi et al., 2022b, a; Araci, 2019). Following this, Zheng et al. (2020) introduced the fully-explored MLM in which random masking is applied to specific non-overlapping segments of the input sequence. The limitation of random masking that we aim to address is that it may put unnecessary focus on tokens that are not representative of the target domain.

In contrast, task-specific selective masking methods mask tokens that are important to the downstream task. For each task, “importance“ is defined differently: Gu et al. (2020) let an additional neural model learns important tokens given the task at hand; Ziyadi et al. (2020) defined importance by masking entities for the named entity recognition task, and Feng et al. (2018) found important tokens by input reduction—maintaining model’s confidence in the original prediction by reducing input—and they were left with a few (potentially nonsensical) tokens that were treated as important to model. Similarly, Li et al. (2020) designed a task-dependent objective for dialogue adaptation, and Ke et al. (2019) proposed label-aware MLM for a sentiment analysis task. In the same vein, token selection in certain domains, e.g., biomedical and clinical domains, was performed based on the entities relevant to the domain Lin et al. (2021); Zhang et al. (2020b); Pergola et al. (2021).

Note that other MLM-based pre-training strategies focused on training a language model from scratch (Zhang et al., 2020a; Joshi et al., 2020; Sun et al., 2019, inter alia). However, since our work focuses on in-domain pre-training, we skip this part for brevity.

In this study, we propose an information-based domain-adaptive pre-training that, without being aware of the downstream task, selectively masks words that are information-dense with respect to the target domain. As a result, PLMs adapted using our mechanism outperform baselines adapted with random masking or fine-tuned directly. In the following sections, we refer to our approach as “keyword masking pre-training.“

Data: In our experiments, we chose tasks and datasets with sufficient amounts of unlabeled data for the domain adaptation stage in order to observe the effects of keyword selection.⁴⁴4For example, we did not use the GLUE dataset Wang et al. (2018) because the included texts are short. In particular, we evaluate our method on three text classification datasets: PUBHEALTH Kotonya and Toni (2020), which contains public health claims associated with veracity labels, IMDB movie reviews dataset Maas et al. (2011), and Amazon pet product reviews dataset (from a Kaggle competition).⁵⁵5https://www.kaggle.com/datasets/kashnitsky/exploring-transfer-learning-for-nlp

Based on the thresholds we studied for filtering out the noisy keywords (see Section 3.2), we gathered 2,116, 7,274, and 6,881 domain-specific keywords from the PUBHEALTH dataset, IMDB dataset, and Amazon dataset, respectively.

Settings:

We use KeyBERT Grootendorst (2020) to extract up to 10 unigram keywords per input document utilizing contextualized word embeddings of BERT Base Devlin et al. (2019), stratified by the Maximal Marginal Relevance (MMR) Carbonell and Goldstein (1998) with a threshold of 0.8.

To perform keyword masking pre-training, we set the masking probability to 0.75 with a constant learning scheduler. The other hyperparameters are left at their default values from the Hugging Face data collator for whole word masking Wolf et al. (2020). For random masking pre-training, we set the masking probability to 0.15, which is a standard value for continual MLM pre-training, and left the remaining hyperparameters at the values provided by the Hugging Face data collator for language modeling Wolf et al. (2020). Note that the default learning rate scheduler is linear. Further, in all settings, pre-training is limited to two epochs, and the batch size of 16 is adopted during both the adaptation and fine-tuning stages.

With the learning rate set to 2e-5 and the weight decay set to 0.01 Devlin et al. (2019), we fine-tune the whole network for all of our adapted models and baselines for up to four epochs in all datasets, while keeping the other hyperparameters at the default value of Hugging Face Wolf et al. (2020). The models that obtained the highest F1 score in the validation partition are then chosen and evaluated on the test split of the datasets.

Table 1 and 2 report the performance of fine-tuned models that used multiple domain-adaptive pre-training methods for each of our settings: three different datasets and two distinct PLMs. Table 1 contains the results for BERT Base as underlying PLM; Table 2 uses BERT Large.

In particular, each table contrasts the performance of two fine-tuned baselines—one without adaptation/in-domain pre-training and one with random masking in-domain pre-training—to a fine-tuned model adapted using our keyword masking.

Both tables show that our approach outperforms all other baselines in all six settings. The improvements are statistically significant in four out of six settings (Appendix E). This highlights the importance of selecting information-carrying keywords for masking during the in-domain pre-training.

The results reveal that our suggested in-domain pre-training technique outperforms alternative settings with or without standard in-domain pre-training on target domain unlabeled data. Although the benefits of continual pre-training vary depending on the domain and the task at hand (Gururangan et al., 2020), our adaptation strategy always has a greater impact on PLMs in capturing domain-specific patterns compared to typical random masking when in-domain adaptation has a positive impact on downstream tasks. This indicates that our pre-training method indeed exposes the PLMs to relevant in-domain representations.

Given the superior outcomes seen in our six different experiments, we can argue that our selective masking strategy, which is task-agnostic as random masking yet more effective, could potentially widely replace random masking in the intermediate pre-training stage for a variety of NLP tasks. Other than performance, our method is simple and has no “pathological behavior“ (Feng et al., 2018) (see Appendix C). Additionally, our method takes 2 to 10 minutes of computational overhead to extract keywords. This accounts for 7% to 39% of pre-training time of only two epochs (Appendix A).

We proposed the first task-agnostic selective masking pre-training approach, dubbed “keyword masking,“ to adapt PLMs to the target domains. For keyword masking, we first extract in-domain keywords from the target domain using KeyBERT (Grootendorst, 2020), and after excluding the noisy ones, we only mask the selected keywords during adaptation.

We evaluated our methodology using six different settings. The results revealed that when in-domain pre-training is conducted using our approach, all fine-tuned PLMs outperform those with no adaptation or adapted using random masking. Further, we observed that our pre-training approach was superior for difficult tasks, i.e., datasets with many labels and more complexity. Lastly, keyword masking pre-training can be widely substituted with random masking during shift domain in NLP tasks since it is task-independent, as simple to use as random masking, and more effective.

Although all pre-training approaches require a sufficient amount of data, given how we defined keywords, longer sequences suit our approach better than short ones for studying the effects of keyword selection. Further, as shown in this study, our findings strongly imply that the strategy we suggested for adapting PLMs can effectively enhance their performance on text classification as the downstream task. To determine whether these findings can translate to other NLP applications, however, further experiments are required.

Although keyword extraction may amplify bias depending on the input documents and the way it extracts keywords, KeyBERT (Grootendorst, 2020) has not been reported to exhibit this behavior. Further work may be necessary to thoroughly explore the potential of introducing undesired bias.