Revisiting Token Dropping Strategy in Efficient BERT Pretraining

Introduction

Masked language models (MLMs), such as BERT Devlin et al. (2019) and its variants Liu et al. (2019); He et al. (2020); Zhong et al. (2023a)We refer to these models as BERT-style models., 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 Jiao et al. (2020); Hou et al. (2022). Hence, it is crucial and green to speed up the training and reduce the computational overhead for BERT-style pretraining Zhang and He (2020); Schwartz et al. (2020).

To achieve this goal, various training-efficient approaches have been developed and summarized Shoeybi et al. (2019); You et al. (2019); Zhang and He (2020); Shen et al. (2023). Among these efforts, a recently-proposed token droppingWe also refer to it as “token drop” in some cases. strategy Hou et al. (2022) 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) Yao et al. (2022); Chiang et al. (2022). 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 “semantic information is mainly encoded in the BERT’s intermediate and top layers” Jawahar et al. (2019), 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 Zhang et al. (2022); Wang et al. (2021). To verify this conjecture, we conduct a series of preliminary analyses on a representative BERT model, and find that:

❶ The training dynamics of the token dropping show a significant semantic drift.

❷ The representation of a well-trained BERT with token dropping contains less semantics.

❸ The downstream semantic-intense tasks show a clear performance degradation.

Based on these observations, we can basically conclude that (one of) the limitation of token dropping is the semantic lossAs we find that BERT models trained with token dropping are prone to losing some semantic-related polarities, e.g., less semantic knowledge in the dropped layers, we refer to this phenomenon as “semantic loss” in the paper. 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 Wang et al. (2018), SuperGLUE Wang et al. (2019) and SQuAD v1/v2 Rajpurkar et al. (2016, 2018), 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:

Our study reveals the semantic loss problem in the token dropping strategy, which limits its performance on downstream tasks, especially on semantic-intense tasks.

We propose a simple yet effective, plug-in-play approach (ScTD) to alleviate the semantic loss and further improve efficiency.

Experiments show that ScTD outperforms the vanilla token dropping with up to +1.56% average improvement and saves up to 57% of pretraining time.

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.

2 Empirical Analyses

In this part, to verify whether removing the unimportant tokens will cause the loss of semantic information and thus hinder the performance of token dropping, we conduct systematic analyses from three aspects: 1) revealing the semantic drift problem during training dynamics; 2) probing the representation of a well-trained model with token dropping; 3) evaluating the downstream performance on semantic-intense tasks. In practice, for comparison, we pre-train the representative BERT$\rm{}_{\texttt{base}}$ models with baseline training scheme and token dropping, respectively. Through the above analyses, we empirically observe that:

As suspected in §1, the corruption caused by the removal of several tokens would break the sentence structure, thus leading to semantic drift. Here, we verify this conjecture by quantitatively estimating the loss of semantic information contained in the corrupted sentence. For measuring the semantic information, we first adopt the off-the-shelf Sentence-BERT Reimers and Gurevych (2019) to capture the semantic representations. Then, suppose that the original sentence (without any corruption, such as masking or token dropping) contains full semantic information, we refer to the cosine similarity between semantic representations of the corrupted and original sentences as a metric to measure the semantic drift in the corrupted sentence.

In practice, given some sentences randomly sampled from training data, we follow the above process and measure the (average) semantic drift during the baseline/token dropping training dynamics, respectively. For reference, we also report the validation results and illustrate all results in Figure 3. It can be found that: compared to baseline training, i) sentence semantics in token dropping drifts more from the original semantics; ii) token dropping hinders the full learning of BERT, especially in the middle and later training stages (after 75K steps). To have a closer look, we show the similarity and validation gaps between both settings in the inserted figure of Figure 3. As seen, with the training going on, both gaps have a similar tendency to increaseThe curve of validation gap tends to flatten in the later training stage, as both models are going to converge., especially at the beginning of training. In general, these analyses indicate that there is a significant semantic drift during training dynamics of token dropping, which shows a correlation with the performance drop of token dropping.

❷ The representation of a well-trained BERT with token dropping contains less semantics.

In addition to the analysis during training dynamics, we then investigate the semantic properties of well-trained models. Specifically, following many prior works Conneau et al. (2018); Jawahar et al. (2019); Ding et al. (2020); Zhong et al. (2022a), we perform several semantic-aware probing tasks on the sentence representations at different layers. Taking the Tense and subject number (SubjNum) tasks as examples, we provide the comparison of semantic information between baseline and token dropping at different layers in Figure 4.

We observe that there is more semantic information in the top layers (from layer 9 to layer 12) of BERT trained with the baseline scheme, which is similar to the finding of Jawahar et al. (2019). However, when using the token dropping, the semantic information contained in BERT tends to decrease in the dropped layers (from layer 5 to layer 11). The semantic information of token dropping at 11-th layer drops dramatically, which is much lower (up to 25.2 points) than that of baseline. Moreover, due to the vulnerable and unstable training, the final representation in token dropping at the last layer is also sub-optimal. These results basically prove that the semantic drift of token dropping damages the semantic learning ability of BERT.

❸ The downstream semantic-intense tasks show a clear performance degradation.

The aforementioned analyses mainly focus on interpreting the semantic properties of models. Here, we further evaluate the downstream performance of token dropping. Specifically, several representative semantic-intenseWe chose tasks based on whether they require semantic-related information to solve. For instance, we included MRPC Dolan and Brockett (2005), a task that predicts if two sentences are semantically equivalent. tasks are used, including OntoNotes 5.0 Weischedel et al. (2013) (Onto. for short), CoNLL03 Sang and De Meulder (2003), MRPC Dolan and Brockett (2005) and SICK-Relatedness Marelli et al. (2014) (SICK-R for short). Notably, for Onto. and CoNLL03, we report the few-shot (32-shot) performance to enlarge the performance difference between different models. We measure the development performance of each task using its corresponding evaluation metrics, and report the contrastive results in Table 2.2.

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 $Fi$ be a fixed interval, ScTD first performs the vanilla token dropping training $(Fi-1)$ times and then performs once the semantic-consistent training. The overall training objective of ScTD can be formulated as:

where $t$ denotes the index of training iterations and $\lambda$ is a weight factor to balance the different objectives, which is empiricallyThe detailed analysis can be found in §4.3. set as 0.05.

To investigate the effectiveness and universality of ScTD, we follow many previous studies Zhong et al. (2022b, d) and conduct extensive experiments on various NLU tasks, covering a diversity of tasks from GLUE Wang et al. (2018), SuperGLUE Wang et al. (2019) and SQuAD benchmarks. Specifically, three semantic-intense tasks (MRPC Dolan and Brockett (2005), STS-B Cer et al. (2017) and RTE Giampiccolo et al. (2007)), five question answering tasks (BoolQ Clark et al. (2019a), COPA Roemmele et al. (2011), MultiRC Khashabi et al. (2018), SQuAD-v1 Rajpurkar et al. (2016) and -v2 Rajpurkar et al. (2018)), two natural language inference tasks (MNLI Williams et al. (2018) and CB De Marneffe et al. (2019)), and two others (CoLA Warstadt et al. (2019) and SST-2 Socher et al. (2013)) are used. For evaluation, we report the performance with Accuracy (“Acc.”) metric for most tasks, except the Pearson and Spearman correlation (“Pear./Spea.”) for STS-B, the Matthew correlation (“Mcc.”) for CoLA, the F1 score for MultiRC, and the Exact Match (“EM”) scores for SQuAD v1/v2. We report the averaged results over 5 random seeds to avoid stochasticity. The details of all tasks and datasets are shown in Appendix A.1.

Hyper-parameters

For pretraining, we train the BRET-Base and -Large models with different methodsFollowing Hou et al. (2022), we implement the token dropping and our approach under the same settings, e.g., dropping 50% of the tokens. from scratch. We basically follow the original paper Devlin et al. (2019) (e.g., the same pretraining corpus), except that we do not use the next sentence prediction (NSP) objective, as suggested in Liu et al. (2019). In practice, we train each model for 250K steps, with a batch size of 1024 and a peak learning rate of 2e-4. For fine-tuning, the learning rate is selected in {1e-5, 2e-5, 3e-5, 5e-5}, while the batch size is in {12, 16, 32} depending on tasks. The maximum length of input sentence is 384 for SQuAD v1/v2 and 256/512 for other tasks. The detailed hyper-parameters for fine-tuning are provided in Appendix A.2. We use AdamW Loshchilov and Hutter (2018) as the optimizer for both pretraining and fine-tuning processes. All experiments are conducted on NVIDIA A100 (40GB) GPUs.

2 Compared Results

Results of GLUE are shown in Table 2, while those of SuperGLUE and SQuAD are in Table 3. Based on these results, we can find that:

First, results on the semantic-intense tasks (MRPC, STS-B and RTE) show that ScTD effectively alleviates the semantic loss problem of token dropping. Specifically, for the RTE task, ScTD brings significant improvement (up to +3.4%) against the vanilla token dropping, and even outperforms the full-sequence training baseline. On the other hand, we observe that ScTD is also beneficial to the other general tasks (e.g., question answering). With the help of ScTD, token dropping strategy achieves up to +1.56% average gains among all types of tasks, proving the effectiveness and universality of ScTD.

ScTD improves performance on both model sizes.

Extensive results show that ScTD works well on both Large and Base BERT models. Specifically, compared to the vanilla token dropping, ScTD brings +1.59% and +1.37% average gains on GLUE tasks, respectively. Results on the other tasks also show a similar phenomenon. Thus, we could recommend our ScTD to speed up the training of all discriminative MLMs regardless of the regime in model capacity.

ScTD effectively improves the training efficiency.

Results in Table 2 show that, with our ScTD, BERT models can achieve comparable or even better performance with much fewer training steps, i.e., improving the training efficiencyWhile the semantic-consistent learning process in ScTD will introduce extra computation overhead, ScTD performs much better in terms of training efficiency. That is, the relatively little computation overhead is acceptable.. Specifically, compared to the full training (250K steps) BERT models, ScTD can save up to 48% pretraining time while achieving comparable performance. We attribute it to the higher data efficiency, since ScTD not only takes full advantage of the token dropping’s ability to learn important words but also alleviates the semantic loss problem in the token dropping. This can be further proved by the illustration of Figure 6, as ScTD always shows better performance against the other counterparts during the training dynamics. Furthermore, when training with the same iterations, our ScTD can even outperform the standard BERT by a clear margin. We attribute this to the regularization effect of token droppingBERT-style PLMs are often over-parameterized and prone to overfitting. Using regularization methods like token dropping and LayerDrop Fan et al. (2020) during training can improve model generalization and even boost performance..

3 Ablation Study

We evaluate the impact of each component of our ScTD, including i) semantic-consistent learning objectives, ii) coefficient $\lambda$ and iii) fixed interval $Fi$ in the hybrid training process. Notably, due to the limited computational budget, we conduct experiments on the BERT$\rm{}_{\texttt{large}}$ models trained with different methods for 5 epochs (35K steps).

As shown in §3, in addition to the original MLM objective $\mathcal{L}^{*}_{MLM}$ of token dropping, we introduce several extra training objectives ($\mathcal{L}_{MLM},\mathcal{L}_{SC_{l}},\mathcal{L}_{SC_{g}}\}$) to align the semantic information. Here, we conduct experiments to analyze the impact of different objectives and show the results in Table 4. It can be seen that all objectives are beneficial to our ScTD, where the $\mathcal{L}_{SC_{g}}$ is the most helpful. This indicates the semantic alignment in the global-level representation space is more critical. Also, we can observe that the combination of all objectives performs best, thus leaving as the default setting.

Impact of Coefficient λ𝜆\lambda.

The factor $\lambda$ in Eq. 4, which is used to balance different objectives, is an important hyper-parameters. In this study, we analyze its influence by evaluating the performance with different $\lambda$ spanning {0, 0.01, 0.05, 0.25, 0.5} on several GLUE tasks. Figure 7 illustrates the average results. Compared with the baseline, our ScTD consistently brings improvements across all ratios of $\lambda$, basically indicating that the performance of ScTD is not sensitive to $\lambda$. More specifically, the case of $\lambda=0.05$ performs best, and we thereby use this setting in our experiments.

Impact of Fixed Interval F​i𝐹𝑖Fi.

In our ScTD, we use a fixed interval $Fi$ to control the frequency for performing the semantic-align process. To verify its impact, we evaluate the performance of ScTD on different $Fi$ and show the results in Table 5. Observably, too small $Fi$ not only causes much computational overhead, but also affects the stability of hybrid training, thus leading to sub-optimal performance. On the contrary, for the larger $Fi$ (e.g., 50), it may be difficult to make full use of the semantic-consistent learning process, hindering the effect of ScTD. In the case of $Fi=10$, ScTD achieves a better trade-off between costs and performance, which we suggest as the best settingSome readers may wonder why the teacher (i.e., model with baseline training) trained with only 1/$Fi$ steps is strong enough to guide the training of student model. One possible reason for this question is that training with hard-to-learn tokens ($Fi$-1) times and training with easy-to-learn tokens once is sufficient to obtain remarkable teacher models, similar to the Lookahead Optimizer Zhang et al. (2019), which updates fast weights $k$ times before updating slow weights once..

4 Does ScTD indeed alleviate the semantic loss problem?

Here, we examine whether ScTD can alleviate the limitation of token dropping. Specifically, following the preliminary analyses in §2, we compare our ScTD with other counterparts by probing the trained BERT models (as illustrated in Figure 8) and pertinently evaluating on several semantic-intense tasks (as shown in Table 6).

It can be found that, with our ScTD, BERT learns more semantic information among most layers, especially in dropped layers. Also, ScTD brings consistent and significant performance gains on all semantic-intense tasks against the vanilla token dropping. These results can prove that ScTD is beneficial to address the semantic loss problem.

Pretraining with Transformer-based architectures like BERT Devlin et al. (2019) has achieved great success in a variety of NLP tasks Devlin et al. (2019); Liu et al. (2019); He et al. (2020); Joshi et al. (2020). Despite its success, BERT-style pretraining usually suffers from unbearable computational expenses Jiao et al. (2020); Zhang and He (2020). To this end, several training-efficient approaches are proposed to speed up the pretraining and reduce the computational overhead, such as mixed-precision training Shoeybi et al. (2019), distributed training You et al. (2019), curriculum learning Nagatsuka et al. (2021); Ding et al. (2021a) and designing efficient model architectures and optimizers Gong et al. (2019); Clark et al. (2019b); Zhang and He (2020); Zhang et al. (2023); Zhong et al. (2022c); Sun et al. (2023). These works mainly focus on efficient optimization processes or model architecture changes.

More recently, Hou et al. (2022) 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 Dehghani et al. (2018) 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 Yao et al. (2022); Chiang et al. (2022). For instance, Yao et al. (2022) 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.

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 semantic-intense scenarios. To address this limitation, we improve token dropping with a novel semantic-consistent 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 He et al. (2020); Zhong et al. (2023b). Also, it will be interesting to revisit and address the semantic loss problem in efficient training methods for generative language models (such as GPT3 Brown et al. (2020)).

Our work has several potential limitations. First, given the limited computational budget, we only validate our ScTD on the Large and Base sizes of BERT models. It will be more convincing if scaling up to the larger model size and applying ScTD to more cutting-edge model architectures. On the other hand, besides the downstream performance, we believe that there are still other properties, e.g., generalization and robustness, of MLMs that can be improved by our ScTD approach, which are not fully explored in this work.

We take ethical considerations very seriously, and strictly adhere to the ACL Ethics Policy. This paper proposes a semantic-consistent algorithm to improve the existing token dropping strategy. The proposed approach aims to speed up the pretraining of BERT-style models, instead of encouraging them to learn privacy knowledge that may cause the ethical problem. Moreover, all pretraining datasets used in this paper are publicly available and have been widely adopted by researchers. Thus, we believe that this research will not pose ethical issues.

Reproducibility

We will publicly release our code in https://github.com/WHU-ZQH/ScTD and the pretrained models in https://huggingface.co/bert-sctd-base to help reproduce the experimental results of this paper.

We are grateful to the anonymous reviewers and the area chair for their insightful comments and suggestions. This work was supported in part by the National Natural Science Foundation of China under Grants 62225113 and 62076186, and in part by the Science and Technology Major Project of Hubei Province (Next-Generation AI Technologies) under Grant 2019AEA170. Xuebo Liu was supported by Shenzhen Science and Technology Program (Grant No. RCBS20221008093121053). The numerical calculations in this paper have been done on the supercomputing system in the Supercomputing Center of Wuhan University.

Appendix A Appendix

In this work, we conduct extensive experiments on parts of tasks from GLUE and SuperGLUE. In addition, two widely-used commonsense question answering tasks are also used. Here, we introduce the descriptions of the used tasks and datasets in detail. Firstly, we present the statistics of all datasets in Table 7. Then, each task is described as:

CoLA Corpus of Linguistic Acceptability Warstadt et al. (2019) is a binary single-sentence classification task to determine whether a given sentence is linguistically “acceptable”.

MRPC Microsoft Research Paraphrase Corpus Dolan and Brockett (2005) is a task to predict whether two sentences are semantically equivalent.

STS-B Semantic Textual Similarity Cer et al. (2017) is a task to predict how similar two sentences are on a 1-5 scale in terms of semantic meaning.

RTE Recognizing Textual Entailment Giampiccolo et al. (2007), given a premise and a hypothesis, is a task to predict whether the premise entails the hypothesis.

MNLI The Multi-Genre Natural Language Inference Corpus Williams et al. (2018) is a task to predict whether the premise entails the hypothesis, contradicts the hypothesis, or neither, given a premise sentence and a hypothesis sentence.

SST-2 The Stanford Sentiment Treebank Socher et al. (2013) is a binary classification task to predict the sentiment of a given sentence.

CB CommitmentBank De Marneffe et al. (2019) is a task that can be framed as three-class textual entailment on a corpus of 1,200 naturally occurring discourses.

BoolQ Boolean Question Clark et al. (2019a) is a question answering task where each sample consists of a short passage and a yes/no question about the passage.

MultiRC Multi-Sentence Reading Comprehension Khashabi et al. (2018) is a QA task where each example consists of a context paragraph, a question about that paragraph, and a list of possible answers. The model need to predict which answers are true and which are false.

COPA Choice of Plausible AlternativesRoemmele et al. (2011) is a causal reasoning task in which a system is given a premise sentence and must determine either the cause or effect of the premise from two possible choices.

SQuAD v1 The Stanford Question Answering Dataset Rajpurkar et al. (2016) is a popular reading comprehension benchmark, where the answer to each question is a segment of text from the corresponding reading passage.

SQuAD v2 The latest version of the Stanford Question Answering Dataset Rajpurkar et al. (2018) is one of the most widely-used reading comprehension benchmarks that require the systems to acquire knowledge reasoning ability.

A.2 Hyper-parameters of Fine-tuning

For fine-tuning, we use the BERT models as the backbone PLMs and conduct experiments using the open-source toolkit fairseqhttps://github.com/facebookresearch/fairseq and transformershttps://github.com/huggingface/transformers. Notably, we apply the same hyper-parameters to all PLMs for simplicity. The training epochs/steps, batch size, and learning rate for each downstream task are listed in Table 7.