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

Abstract: We propose a novel task-agnostic in-domain pre-training method that sits between generic pre-training and fine-tuning. Our approach selectively masks in-domain keywords, i.e., words that provide a compact representation of the target domain. We identify such keywords using KeyBERT (Grootendorst, 2020). We evaluate our approach using six different settings: three datasets combined with two distinct pre-trained language models (PLMs). Our results reveal that the fine-tuned PLMs adapted using our in-domain pre-training strategy outperform PLMs that used in-domain pre-training with random masking as well as those that followed the common pre-train-then-fine-tune paradigm. Further, the overhead of identifying in-domain keywords is reasonable, e.g., 7-15% of the pre-training time (for two epochs) for BERT Large (Devlin et al., 2019).

## introduction employing large pre-trained language models (plms) is currently a common practice for most natural language processing (nlp) tasks @xcite . a two-stage pre-train-then-fine-tune framework is usually used to adapt/fine-tune plms to downstream tasks @xcite . however, motivated by ulmfit @xcite and elmo @xcite , @xcite showed that incorporating in-domain pre-training (also known as domain-adaptive pretraining) between generic pre-training and finetuning 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 @xcite . here, we argue that this intermediate pretraining 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" @xcite . by applying token masking only to indomain 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 @xcite , which may overlook domain-meaningful information, or the in-domain pre-training methods that selectively mask tokens deemed important given the downstream task @xcite , 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: ## related work bidirectional encoder representations from transformers (bert) brought pre-training to transformer networks @xcite through masked language modeling (mlm) @xcite . they showed that a simple two-step paradigm of generic pre-training followed by finetuning 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 @xcite , @xcite found that the same advantage applies to plms as well. since then, several efforts proposed different domainadaptive 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 @xcite @xcite @xcite . following this, @xcite introduced the fully-explored mlm in which random masking is applied to specific nonoverlapping 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: @xcite let an additional neural model learns important tokens given the task at hand; @xcite 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, @xcite designed a task-dependent objective for dialogue adaptation, and @xcite 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 @xcite @xcite . note that other mlm-based pre-training strategies focused on training a language model from scratch @xcite @xcite . 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." ## experimental setup 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. 4 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 @xcite , and amazon pet product reviews dataset (from a kaggle competition). 5 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 @xcite to extract up to 10 unigram keywords per input document utilizing contextualized word embeddings of bert base @xcite , stratified by the maximal marginal relevance (mmr) @xcite 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 @xcite . 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 @xcite . 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 @xcite , 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 @xcite . the models that obtained the highest f1 score in the validation partition are then chosen and evaluated on the test split of the datasets. ## results and discussion 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 pretraining on target domain unlabeled data. although the benefits of continual pre-training vary depending on the domain and the task at hand @xcite , our adaptation strategy always has a greater impact on plms in capturing domainspecific 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" @xcite (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). ## conclusion 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.