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

TL;DR

This paper introduces a reinforcement learning-based method to generate adaptive text masks for better language model fine-tuning, improving performance on specific tasks like question answering and text classification.

Contribution

It presents a novel neural framework that learns to generate task- and domain-specific masks, outperforming rule-based masking strategies in language model adaptation.

Findings

Outperforms rule-based masking strategies on multiple datasets

Improves task performance of BERT and DistilBERT models

Uses a Transformer-based policy network for mask generation

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.