Neural Mask Generator: Learning to Generate Adaptive Word Maskings for Language Model Adaptation (2020)

Abstract: We propose a method to automatically generate a domain- and task-adaptive maskings of the given text for self-supervised pre-training, such that we can effectively adapt the language model to a particular target task (e.g. question answering). Specifically, we present a novel reinforcement learning-based framework which learns the masking policy, such that using the generated masks for further pre-training of the target language model helps improve task performance on unseen texts. We use off-policy actor-critic with entropy regularization and experience replay for reinforcement learning, and propose a Transformer-based policy network that can consider the relative importance of words in a given text. We validate our Neural Mask Generator (NMG) on several question answering and text classification datasets using BERT and DistilBERT as the language models, on which it outperforms rule-based masking strategies, by automatically learning optimal adaptive maskings.

## introduction the recent success of the language model pretraining approaches @xcite @xcite @xcite , which train language models on diverse text corpora with self-supervised or multitask learning, have brought up huge performance improvements on several natural language understanding (nlu) tasks @xcite . the key to this success is their ability to learn generalizable text embeddings that achieve near optimal performance on diverse tasks with only a few additional steps of fine-tuning on each downstream task. most of the existing works on language model aim to obtain a universal language model that can address nearly the entire set of available natural language tasks on heterogeneous domains. although this train-once and use-anywhere approach has been shown to be helpful for various natural language tasks @xcite @xcite , there have been considerable needs on adapting the learned language models to domain-specific corpora (e.g. healthcare or legal). such domains may contain new entities that are not included in the common text corpora, and may contain only a small amount of labeled data as obtaining annotation on them may require expert knowledge. some recent works @xcite @xcite suggest to further pre-train the language model with selfsupervised tasks on the domain-specific text corpus for adaptation, and show that it yields improved performance on tasks from the target domain. masked language models (mlms) objective in bert @xcite has shown to be effective for the language model to learn the knowledge of the language in a bi-directional manner @xcite . in general, masks in mlms are sampled at random @xcite , which seems reasonable for learning a generic language model pre-trained from scratch, since it needs to learn about as many words in the vocabulary as possible in diverse contexts. however, in the case of further pre-training of the already pre-trained language model, such a conventional selection method may lead a domain adaptation in an inefficient way, since not all words will be equally important for the target task. repeatedly learning for uninformative instances thus will be wasteful. instead, as done with instance selection @xcite @xcite , it will be more effective if the masks focus on the most important words for the target domain, and for the specific nlu task at hands. how can we then obtain such a masking strategy to train the mlms? several works @xcite @xcite propose rule-based masking strategies which work better than random masking @xcite when applied to language model pre-training from scratch. based on those works, we assume that adaptation of the pre-trained language model can be improved via a learned masking policy which selects the words to mask. yet, existing models are inevitably suboptimal since they do not consider the target domain and the task. to overcome this limitation, in this work, we propose to adaptively generate mask by learning the optimal masking policy for the given task, for the task-adaptive pre-training @xcite of the language model. as described in figure 1 , we want to further pretrain the language model on a specific task with a task-dependent masking policy, such that it directs the solution to the set of parameters that can better adapt to the target domain, while task-agnostic random policy leads the model to an arbitrary solution. to tackle this problem, we pose the given learning problem as a meta-learning problem where we learn the task-adaptive mask-generating policy, such that the model learned with the masking strategy obtains high accuracy on the target task. we refer to this meta-learner as the neural mask generator (nmg). specifically, we formulate mask learning as a bi-level problem where we pre-train and fine-tune a target language model in the inner loop, and learn the nmg at the outer loop, and solve it using renforcement learning. we validate our method on diverse nlu tasks, including question answering and text classification. the results show that the models trained using our nmg outperforms the models pre-trained using rule-based masking strategies, as well as finds a proper adaptive masking strategy for each domain and task. our contribution is threefold: target domain original lm domain * * * universally pre-trained language model * context a myocardial infraction, also known as a heart attack, occurs when blood flow.. masked context … domain-adaptive language model pre-trained model 𝑷𝑷𝑷𝑷𝑷𝑷𝑷𝑷𝑷𝑷𝑷𝑷 𝑽𝑽𝑽𝑽𝑷𝑷𝑽𝑽𝑽𝑽 neural mask generator a [mask] [mask], also known as a [mask] attack, occurs when [mask] flow… learning we assume that adjusting the masking policy of the mlm objective affects the training trajectory of the language model, such that it moves towards a better solution space for the target domain. this illustration of the solution spaces for the two domains is motivated by @xcite . a popular approach is to use self-supervised pre-training tasks for learning the contextualized embedding from large unannotated text corpora using auto-regressive @xcite or auto-encoding @xcite language modeling. following the success of the masked language model (mlm) from @xcite , several works have proposed different model architecture @xcite @xcite and pre-training objectives @xcite @xcite @xcite , to improve upon its performance. some works have also proposed alternative masking policies for the mlm pre-training over random sampling, such as spanbert @xcite and ernie @xcite ). yet, none of the existing approaches have tried to learn the task-adaptive mask in a context-dependent manner which is the problem we target in this work. language model adaptation pre-training the language model on the target domain, then finetuning on downstream tasks, is the most simple yet successful approach for adapting the language model to a specific task. some studies @xcite @xcite have shown the advantage of further pre-training the language model on a large unlabeled text corpus collected from a specific domain. moreover, @xcite and @xcite investigate the effectiveness of further pre-training of the language model on small domain-related text corpora. recently, @xcite integrates prior works and defines domain-adaptive pre-training and task-adaptive pre-training, showing that domain adaptation of the language model can be done with additional pre-training with the mlm objective on a domain-related text corpus, as well as a smaller but directly task-relevant text corpus. meta-learning meta-learning @xcite aims to train the model to generalize over a distribution of tasks, such that it can generalize to an unseen task. there exist large number of different approaches to train the metalearner @xcite @xcite @xcite . however, existing meta-learning approaches do not scale well to the training of large models such as masked language models. thus, instead of the existing meta-learning method such a gradient based approach, we formulate the problem as a bi-level problem @xcite of learning the language model in the inner loop and the mask at the outer loop, and solve it using reinforcement learning. such optimization of the outer objective using rl is similar to the formulation used in previous works on neural architecture search @xcite . ## problem statement we now describe how to formulate the problem of learning to generate masks for the masked language model (mlm) as a bi-level optimization problem. then, we describe how we can reformulate it as a reinforcement learning problem. ## neural mask generator in this section, we describe our model, neural mask generator (nmg), which learns the masking policy to generate an optimal mask on an unseen context using deep rl with the detailed descriptions of the framework setup. the overview of the meta mask generator framework is shown in figure 2. for detailed descriptions of the approaches, the procedures of both training and test phase, and algorithm, please see appendix a. ## experiment we now experimentally validate our neural mask generator (nmg) model on multiple nlu tasks, including question answering and text classification tasks using two different language models, and analyze its behaviors. in section 5.1, we evaluate the nmg with several baselines. then, we evaluate the effect of the specific design choices made for our model through ablation studies in section 5.2. finally, we analyze how the policy learned by the nmg model works in section 5.3. tasks and datasets for question answering, we use three datasets, namely squad v1.1 @xcite , newsqa @xcite , and emrqa @xcite to validate our model. we use the mrqa foot_0 version for both squad and newsqa to sustain a coherency. we also preprocess emrqa to fit the format of other datasets. we use a standard evaluation metric named exact-matchs (em) and f1 score for question answering task. for text classification, we use imdb @xcite and chemprot @xcite , following the experimental settings of @xcite . baselines according to @xcite and @xcite , training the language models with difficult objectives is much more beneficial when pre-training from scratch. to test whether it is also the case for task-adaptive pre-training, we experiment with two heuristic masking strategies, which we refer to as whole-random and spanrandom. in addition, we also tested the named entity masking proposed by @xcite . below is the complete list of the heuristic baselines we compare against in the experiments. ## conclusion we proposed a novel framework which automatically generates an adaptive masking for masked language models based on the given context, for language model adaptation to low-resource domains. to this end, we proposed the neural mask generator (nmg), which is trained with reinforcement learning to mask out words that are helpful for domain adaptation. we performed an empirical study of various rule-based masking strategies on multiple datasets for question answering and text classification tasks, which shows that the optimal masking strategy depends on both the language model and the domain. we then validated nmg against rule-based masking strategies, and the results show that it either outperforms, or obtains comparable performance to the best heuristic. further qualitative analysis suggests that such good performance comes from its ability to adaptively mask meaningful words for the given task.