Token Dropping for Efficient BERT Pretraining (2022)

Abstract: Transformer-based models generally allocate the same amount of computation for each token in a given sequence. We develop a simple but effective “token dropping” method to accelerate the pretraining of transformer models, such as BERT, without degrading its performance on downstream tasks. In particular, we drop unimportant tokens starting from an intermediate layer in the model to make the model focus on important tokens more efficiently if with limited computational resource. The dropped tokens are later picked up by the last layer of the model so that the model still produces full-length sequences. We leverage the already built-in masked language modeling (MLM) loss to identify unimportant tokens with practically no computational overhead. In our experiments, this simple approach reduces the pretraining cost of BERT by 25% while achieving similar overall fine-tuning performance on standard downstream tasks.

## introduction nowadays, the success of neural networks in a variety of nlp tasks heavily relies on bert-type language models containing millions to billions of parameters. however, the pretraining process of these models is computationally expensive, generating significant carbon dioxide emission @xcite . in practice, there is the need to perform large-scale language model pretraining for diverse applications @xcite @xcite in different languages @xcite . in this paper, we develop a technique that significantly reduces the pretraining cost of bert models @xcite without hurting their test performance on a diverse set of fine-tuning tasks. recent efforts of efficient training involve mixedprecision training @xcite , distributed training @xcite , better modeling on rare words and phrases @xcite , designing more effective and data-efficient pretraining objectives @xcite @xcite , progressive stacking @xcite , and so on. while these approaches contribute to efficient training with reduced computational cost, most of them focus on the model architecture or the optimization process. in this paper, we focus on a simple but efficient bert-pretraining strategy that has been underexplored before, i.e., "token dropping," which removes the redundant tokens in each sequence that are less informative to training. since not all tokens contribute equally to the output or the training objective, and the computational complexity of transformer-based models grows at least linearly with respect to the sequence length, shortening the input sequences can accelerate the training effectively. among existing studies, the depth-adaptive transformer approach aims to reduce the autoregressive inference time by allocating less computation on easy-to-predict tokens @xcite . to improve the training efficiency, @xcite perform pooling on the embeddings of nearby tokens. however, directly dropping tokens during pretraining was not studied until very recently in faster depth adaptive transformer @xcite , where the important tokens are identified either through (1) mutual information-based estimation between tokens and predefined labels or through (2) a separate bert model that exhaustively computes the masked language model (mlm) loss for each token. on the contrary, we focus on accelerating the task-agnostic pretraining phase without requiring any labels or any computation by a separate language model. specifically, we identify important tokens as the ones hard to predict by the model itself through its loss during training, which is adaptive to its training process and leads to practically no computational overhead. we show examples of dropped tokens in figure 1 . recent approaches such as roberta @xcite suggest packing input sequences. in this way, there are no [pad] tokens, which makes it a non-trivial task to identify unimportant tokens. we identify unimportant tokens in each sequence with the smallest historical mlm loss (we take the running average of the mlm loss of each token throughout the pretraining process). by removing them from intermediate layers of a bert model during training, we save an enormous amount of computation and memory. we keep them in the first several layers as well as in the last layer so that they are still present in the model. therefore, the inputs and outputs of bert model are kept consistent with the conventional all-token training process. without modifying the original bert architecture or training setting, this simple tokendropping strategy trains intermediate layers mainly on a few important tokens. as demonstrated in our experiments, models pretrained in this way generalize well on diverse downstream tasks with full sequences. to summarize, our contributions are as follows. (1) we show that bert models can be pretrained with only a subset of the layers focusing on important tokens. even though the model is trained on sub-sequences of important tokens only, it generalizes well to full sequences during fine-tuning on downstream tasks. (2) we identify important tokens through the pretraining process by exploring the training dynamics, with minimal computational overhead and without modifying the model architecture. ## results ## related work several ideas are particularly relevant to tokenwise layer dropping: @xcite propose to use a fixed probability to drop an entire layer during pretraining; here, we use the more fine-grained token-wise layer dropping. the dynamic halting algorithm @xcite , motivated by the finding that transformers fail to generalize to many simple tasks, stops the processing of a token through upper layers if its representation is good enough. however, the implementation does not improve training time, as its goal is to improve performance. ## conclusion we present a simple yet effective approach to save bert pretraining time. our approach identifies unimportant tokens with practically no computational overhead and cuts unnecessary computation on these unimportant tokens for training. experiments show that bert models pretrained in this manner save 25% pretraining time, while generalizing similarly well on downstream tasks. we show that our token dropping approach performs better than average pooling along the sequence dimension. future work will involve extending token dropping to pretraining transformer models that can process a much longer context, as well as extending this algorithm to a wider range of transformer-based tasks, including translation and text generation.