Mask More and Mask Later: Efficient Pre-training of Masked Language Models by Disentangling the [ MASK ] Token (2022)

Abstract: The pre-training of masked language models (MLMs) consumes massive computation to achieve good results on downstream NLP tasks, resulting in a large carbon footprint. In the vanilla MLM, the virtual tokens, [MASK]s, act as placeholders and gather the contextualized information from unmasked tokens to restore the corrupted information. It raises the question of whether we can append [MASK]s at a later layer, to reduce the sequence length for earlier layers and make the pre-training more efficient. We show: (1) [MASK]s can indeed be appended at a later layer, being disentangled from the word embedding; (2) The gathering of contextualized information from unmasked tokens can be conducted with a few layers. By further increasing the masking rate from 15% to 50%, we can pre-train RoBERTa-base and RoBERTa-large from scratch with only 78% and 68% of the original computational budget without any degradation on the GLUE benchmark. When pre-training with the original budget, our method outperforms RoBERTa for 6 out of 8 GLUE tasks, on average by 0.4%.

## introduction large-scale pre-trained mlms, like bert @xcite and its variants @xcite @xcite @xcite , have achieved great success in various nlp tasks, such as machine translation @xcite , general language understanding @xcite , question answering @xcite , summarization @xcite and claim verification @xcite . to make the pre-trained model generalize well to a wide range of tasks, mlms tend to have a large number of parameters, even in the billion scale @xcite , and are trained with plenty of data. this is prohibitively expensive and generates significant amounts of co 2 emissions 1 code at https://github.com/baohaoliao/3ml figure 1 : information flows of vanilla mlm. a sentence, {x 1 , x 2 , x 3 , x 4 , x 5 }, is corrupted by replacing x 2 and x 3 with a virtual token m indexed with the corresponding positions. the information transferred among tokens can be split into: position information and nonpositional information (termed as token information in this paper). as shown in figure 1 , the information flows for a corrupted sentence during training consist of: (1) position and token information among unmasked tokens; (2) position and token information from unmasked tokens to @xcite s; intuitively, the amount of information transferred in each flow is not at the same level. the 3 the information flow assumption. position and token information among unmasked tokens (1 st flow) are more difficult to learn than the transfer of this knowledge to [mask]s (2 nd flow). this assumption is empirically proven later ( §4.2). since the difficulty of information transfer varies among different flows, it makes sense to divide the flows into multiple stages, forcing the model to spend more computation on the more complex one. we propose a two-stage learning method. for the early layers of an mlm, we detach [mask]s and only input the embedding of unmasked tokens. so the model firstly focuses on the most complex ( 1 st ) information flow. at an intermediate layer, we append the embedding of [mask]s with their corresponding position information back to the sequence. then the remaining layers of the mlm fuse all information. in this way, the sequence length for the early layers becomes shorter due to excluding [mask]s. we further reduce the sequence length by increasing the masking rate for higher efficiency. we call our method mask more and mask later (3ml), since [mask]s are appended later and we have a larger masking rate. in this work, we introduce two models designed for 3ml ( §2), empirically show two prerequisites of efficiency for 3ml hold ( §4), conduct extensive experiments to select an optimal setting for high performance and efficiency ( §5), and finally compare 3ml's results to strong baselines ( §6). our main contributions are summarized as: (1) we introduce a simple, intuitive but effective method for the efficient pre-training of mlm; ## model in this section, we first discuss the information flows in a vanilla mlm, i.e. bert, then introduce two architectures designed for 3ml (figure 2 ). ## two prerequisites for efficiency in this section, we show that the two prerequisites, masking more and masking later, for efficiency hold for 3ml, which is also an empirical test of our information flow assumption. ## efficient setting section 4 shows that two prerequisites for higher efficiency hold for 3mls. in this section, we further explore different choices of 3ml architecture and the masking rate, trying to select an optimal setting with the trade-off between performance and efficiency. ## comparison with previous work in this section, we compare 3ml with strong baselines on the development set of glue. we list the results of both efficient and longer pre-training recipes in ## related work the most related work to this paper is mae @xcite from computer vision. 3ml self shares almost the same architecture as mae, but with the additional prediction of unchanged tokens and randomly replaced tokens, while mae only reconstructs the masked patches. we also have different findings from mae: a small decoder works better for 3ml self . while mae applied a deeper decoder for better performance. hou et al. ( 2022 ) do a concurrent work, dropping the representations of unimportant tokens for some intermediate layers to reduce the sequence length for efficiency. we don't include any prior information and only drop the [mask] token at the very beginning. in addition, 3ml achieves better results. @xcite and our work have the same finding, better performance for a middle-level masking rate. although albert @xcite and electra @xcite ) make the pretraining of mlm more efficient, their studies are orthogonal to ours. ## conclusion we propose a two-stage learning method for efficient masked language modeling and design two models, 3ml self and 3ml cross , for our method. two prerequisites for our efficient method are: we can have a higher masking rate and append [mask]s at a later layer. our experiments show that both, masking more and masking later, are possible and favorable. this allows us to reduce the sequence length of the encoder during pre-training by a factor depending on the masking rate. by conducting extensive experiments, we observe that 3ml self performs better on downstream tasks than 3ml cross . it can speed up the pre-training by a factor of 1.5x for our efficient pre-training recipe without any performance degradation. using roughly the same computation budget, 3ml self outperforms all of our strong baselines like electra and roberta.