Dynamic Masking Rate Schedules for MLM Pretraining (2024)

Abstract: Most works on transformers trained with the Masked Language Modeling (MLM) objective use the original BERT model’s fixed masking rate of 15%. We propose to instead dynamically schedule the masking rate throughout training. We find that linearly decreasing the masking rate over the course of pretraining improves average GLUE accuracy by up to 0.46% and 0.25% in BERT-base and BERT-large, respectively, compared to fixed rate baselines. These gains come from exposure to both high and low masking rate regimes, providing benefits from both settings. Our results demonstrate that masking rate scheduling is a simple way to improve the quality of masked language models, achieving up to a 1.89x speedup in pretraining for BERT-base as well as a Pareto improvement for BERT-large.

## introduction bert @xcite is a popular encoderonly transformer @xcite architecture that is pretrained using a cloze-inspired @xcite ) masked language modeling (mlm) objective. during mlm training, we mask out a subset of the input tokens and train the model to reconstruct the missing tokens. the proportion of tokens to be masked out is determined by the masking rate hyperparameter. most practitioners use a fixed masking rate of 0.15 @xcite , but @xcite found that the standard 15% masking rate is suboptimal for a variety of model settings and recommended a higher rate. we build on their work by studying the impact of dynamically scheduled masking rates. correspondence to ankner@mit.edu. hyperparameter scheduling-i.e., changing the learning rate, dropout rate, batch size, sequence length, etc., during training-is a common practice in deep learning @xcite @xcite @xcite . masking rate is a good candidate for hyperparameter scheduling for a number of reasons. first, a high masking rate, like a high dropout rate, directly reduces the amount of feature information available during a training step. this information removal may smooth the loss landscape, which permits simulated annealing if performed earlier in training. furthermore, a higher masking rate adds training signal, as loss is computed for a larger portion of tokens, similar to a larger sequence length or batch size. we therefore study whether scheduling the masking rate during training could lead to model quality improvements, as scheduling these other hyperparameters does. we present a series of experiments to assess the effects of masking rate scheduling on the quality of bert-base @xcite . we evaluate our masking rate scheduled models on mlm loss and downstream tasks. our contributions are: ## methods we perform typical mlm pretraining, with the key difference that a scheduler sets the masking rate dynamically. ## experiments and results in this section, we evaluate the performance of masking rate scheduling on a collection of downstream tasks and determine why our schedule is successful. we pretrain all models on the colossal cleaned common crawl (c4) dataset @xcite , and then fine-tune and evaluate on the glue benchmark @xcite . we use bert-base and bert-large models as implemented in hug-gingface @xcite , and train models with the composer library @xcite . we list further details of our experimental setup in appendix a. ## related work masked language modeling since elmo (peters et al., 2018), self-supervised pretraining has become the dominant paradigm for many nlp tasks, and bert has been established as a basic standard for transfer learning. many works have changed the bert model architecture while retaining the original mlm objective, including the 15% constant masking rate @xcite @xcite . other encoder-only models have modified the mlm objective itself to mask out spans of tokens instead of individual tokens @xcite @xcite . we note that both architectural changes and span masking are compatible with our masking rate scheduling. electra @xcite proposes an alternate denoising objective to masking; using a separate "generator" encoder language model, they replace a subset of tokens in the input sequence. while the gradual improvement of the generator may implicitly parallel a masking rate schedule, explicit scheduling may still be beneficial since accuracy can be sensitive to masking rate (appendix g). additionally, the generator is trained using an mlm objective, and as such could benefit from masking rate scheduling. there has also been previous work exploring whether the standard 15% masking rate is optimal. @xcite empirically investigate the optimal fixed masking rate and demonstrate that for larger bert models higher masking rates are more performant. most closely related to our method is @xcite , which also examines dynamic masking rates for mlm pretraining. although there is significant overlap in the proposed methodologies, their work sets the final masking rate to be close to 0%, while we found that maintaining a higher final masking rate of 15% was necessary for performance improvements. additionally, our analysis differs significantly from theirs. while both their work and ours evaluate downstream performance improvements, @xcite also investigates how dynamic masking rates affect performance when the training duration is extended and study nonrandom token masks. our analysis, by contrast, focuses on why masking rate scheduling improves performance. to this end, we investigating whether dynamic masking rates must follow a decaying scheduling (section 3.3), whether the observed gains are due to the additional training signal or the added noise (section 3.4), the impact of differing masking rate schedules on grammatical capabilities (section 3.5), and the impact of dynamic masking rates on the pre-training objective itself (section 3.6). hyperparameter scheduling although learning rate is the most commonly-scheduled hyperparameter @xcite @xcite , other hyperparameter schedules are common. our approach is also not the first to schedule a hyperparameter that removes information content from the model; prior work has suggested scheduling dropout @xcite and input resolution @xcite . scheduling has also been applied to hyperparameters that control the training signal to the model such as batch size @xcite and sequence length @xcite . masking rate combines both of these properties, making it a particularly good candidate for scheduling. ## discussion and conclusions in addition to our method's improvement on the average final downstream performance, we find that scheduling is a pareto improvement for all examined pretraining durations over the typical constant masking rate baselines on glue. our analysis suggests that this benefit comes from the combined advantages of higher and lower masking rates. we also demonstrate that our approach generalizes to other pretraining objectives (appendix h). our method of beginning with a larger masking ratio and decaying, which we found necessary (section 3.3), parallels the motivation behind simulated annealing @xcite . simulated annealing is a general method for avoiding local minima by smoothing the loss surface early in training through the addition of noise early in training. however, we found that the increasing noise early in training is not the only source of advantage. we also benefit from increasing the signal by predicting more masked tokens (section 3.4). overall, our work demonstrates that masking rate scheduling is a simple and reliable way to improve the quality and efficiency of mlm pretraining.