Revisiting Token Dropping Strategy in Efficient BERT Pretraining (2023)

Abstract: Token dropping is a recently-proposed strategy to speed up the pretraining of masked language models, such as BERT, by skipping the computation of a subset of the input tokens at several middle layers. It can effectively reduce the training time without degrading much performance on downstream tasks. However, we empirically find that token dropping is prone to a semantic loss problem and falls short in handling semantic-intense tasks. Motivated by this, we propose a simple yet effective semantic-consistent learning method (ScTD) to improve the token dropping. ScTD aims to encourage the model to learn how to preserve the semantic information in the representation space. Extensive experiments on 12 tasks show that, with the help of our ScTD, token dropping can achieve consistent and significant performance gains across all task types and model sizes. More encouragingly, ScTD saves up to 57% of pretraining time and brings up to +1.56% average improvement over the vanilla token dropping.

## introduction masked language models (mlms), such as bert @xcite and its variants @xcite @xcite 1 , have achieved great success in a variety of natural language understanding (nlu) tasks. however, with the scaling of model size and corpus size, the pretraining of these bert-style models becomes more computationally expensive and memory intensive @xcite . hence, it is crucial and green to speed up the training and reduce the computational overhead for bert-style pretraining @xcite . figure 1 : performance of bert base on several downstream tasks. we see that: 1) despite the remarkable performance on general tasks (i.e., mnli and sst-2), token dropping leads to dramatically poor performance on the semantic-intense task (i.e., rte). 2) our sctd achieves consistent performance gains among all tasks. to achieve this goal, various training-efficient approaches have been developed and summarized @xcite @xcite . among these efforts, a recently-proposed token dropping foot_0 strategy @xcite has attracted increasing attention owing to its easy-to-implement algorithm and impressive efficiency (reducing the training cost by 25% without much average performance dropping) @xcite . different from most previous works that focus on changing model architecture or optimization process, token dropping aims to improve training efficiency by dynamically skipping the compute of the redundant (unimportant) tokens that are less informative to the current training, at some middle layers of bert during training. although achieving a remarkable speedup, the performance improvement of token dropping is usually limited and unstable, compared to the baseline training scheme. more specifically, we empirically found that token dropping falls short in handling semantic-intense tasks, as shown in figure 1 . this motivates us to explore and address the limitations of token dropping in this paper. in light of the conventional wisdom that "seman-tic information is mainly encoded in the bert's intermediate and top layers" @xcite , we suspected, apriori, that the corruption caused by the removal of unimportant tokens would break the sentence structure, and may easily lead to the semantic drift of sentence representations, as also observed in many similar scenarios @xcite . to verify this conjecture, we conduct a series of preliminary analyses on a representative bert model, and find that: based on these observations, we can basically conclude that (one of) the limitation of token dropping is the semantic loss 3 problem, which causes vulnerable and unstable training of bert models. to address this limitation, we propose a simple yet effective semantic-consistent learning method (referred to as sctd) to improve token dropping. the principle of sctd is to encourage the bert to learn how to preserve the semantic information in the representation space. specifically, sctd first introduces two semantic constraints to align the semantic information of representations between baseline-and token dropping-based models, and then adopts a novel hybrid training approach to further improve the training efficiency. we evaluate sctd on a variety of benchmarks, including glue @xcite , super-glue @xcite and squad v1/v2 @xcite , upon two typical mlms: bert-base and -large. results show that sctd can not only bring consistent and significant improvements (up to +1.56% average score among all tasks) into the token dropping strategy on both bert models, but also alleviate the semantic loss problem. moreover, compared to the standard bert models, sctd can also save up to 48% of pretraining time while achieving comparable performance, and further achieve +1.42% average gain for the same training iterations. to summarize, our contributions are as follows: ## revisiting token dropping strategy in this section, we first review the background of token dropping strategy and then present the empirical analyses of this strategy in detail. ## improving token dropping with semantic-consistent learning based on the observations in §2, we recognize that it is essential to alleviate the side effect (i.e., semantic loss problem) of token dropping. to achieve this goal, we propose a simple yet effective semantic-consistent learning (sctd) framework specifically, our sctd adopts two key techniques as follows: semantic-consistent learning. the principle of our sctd is to encourage the model to preserve the semantic information in the representation space. inspired by the success of knowledge distillation @xcite , 2020), sctd refers to the model with baseline training (containing more semantic information) as the teacher to guide the training of students (i.e., model trained with token dropping). considering that it is usually unavailable to obtain the pre-trained teacher model, we hereby recast it as a self-distillation process @xcite . given the same input x 0 , we input x 0 into the model to perform twice forwardpropagation processes, where one is for token dropping and the other is for baseline training. the outputs of baseline training (x l ) are used as the teacher distributions to teach the student (outputs xl of token dropping). as such, the student can learn how to align the semantic information with the teacher. more specifically, sctd introduces two semantic constraints in a local-to-global manner (as illustrated in figure 2 ). for the global one, we use the kl divergence to constrain the globallevel semantic distributions of baseline-and tokendropping-based models at the last l-th layer, as follows: hybrid training. since the semantic-consistent learning process requires twice forward/back-propagation, sctd would introduce much computational overhead, leading to inefficiency. to overcome this issue, sctd adopts a novel hybrid training strategy, as illustrated in figure 5 . specifically, instead of using the semantic-consistent learning method throughout the training, sctd basically follows the vanilla token dropping and adopts the semantic-consistent training intermittently. as such, sctd can combine the advantages of semantic-consistent learning (effectiveness) and token dropping (efficiency). let f i be a fixed interval, sctd first performs the vanilla token dropping training (f i -1) times and then performs once the semantic-consistent training. the overall training objective of sctd can be formulated as: ## related works pretraining with transformer-based architectures like bert @xcite has achieved great success in a variety of nlp tasks @xcite @xcite . despite its success, bert-style pretraining usually suffers from unbearable computational expenses @xcite . to this end, several training-efficient approaches are proposed to speed up the pretraining and reduce the computational overhead, such as mixed-precision training @xcite , distributed training @xcite , curriculum learning @xcite ) and designing efficient model architectures and optimizers @xcite @xcite @xcite . these works mainly focus on efficient optimization processes or model architecture changes. more recently, @xcite propose the token dropping strategy, which exposes a new mode to speed up the bert pretraining. without modifying the original bert architecture or training setting, token dropping is inspired by the dynamic halting algorithm @xcite and attempts to skip the computations on part of (unimportant) tokens in some middle bert layers during the forward-propagation process. owing to its impressive efficiency, token dropping has recently attracted increasing attention @xcite . for instance, @xcite apply the token dropping strategy to broader applications, e.g., both nlp and cv communities. along with the line of token dropping, we take a further step by exploring and addressing its limitations. to be specific, we first reveal the semantic loss problem ( §2) in the token dropping, and then propose a novel semantic-consistent learning method ( §3) to alleviate this problem and further improve performance and training efficiency. ## conclusion in this paper, we reveal and address the limitation of token dropping in accelerating language model training. based on a series of preliminary analyses, we find that removing parts of tokens would lead to a semantic loss problem, which causes vulnerable and unstable training. furthermore, experiments show such a semantic loss will hinder the performance of token dropping in most semanticintense scenarios. to address this limitation, we improve token dropping with a novel semanticconsistent learning algorithm. it designs two semantic constraints to encourage models to preserve semantic information. experiments show that our approach consistently and significantly improves downstream performance across all task types and model architectures. in-depth analyses prove that our approach indeed alleviates the problem, and further improves training efficiency. in future work, we will explore the effectiveness of our method on more advanced discriminative language models @xcite . also, it will be interesting to revisit and address the semantic loss problem in efficient training methods for generative language models (such as gpt3 @xcite ).