Token Dropping for Efficient BERT Pretraining

Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 3774 - 3784 May 22-27, 2022 c⃝2022 Association for Computational Linguistics Token Dropping for Efficient BERT Pretraining Le Hou1∗ Richard Yuanzhe Pang2§∗ Tianyi Zhou13§ Yuexin Wu1 Xinying Song1 Xiaodan Song1 Denny Zhou1 1 Google 2 New York University 3 University of Maryland, College Park lehou@google.com, yzpang@nyu.edu Abstract Transformer-based models generally allocate the same amount of computation for each to- ken in a given sequence. We develop a sim- ple but effective “token dropping” method to accelerate the pretraining of transformer mod- els, such as BERT, without degrading its per- formance 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 effi- ciently 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 model- ing (MLM) loss to identify unimportant to- kens with practically no computational over- head. In our experiments, this simple approach reduces the pretraining cost of BERT by 25% while achieving similar overall fine-tuning per- formance on standard downstream tasks. 1 Introduction Nowadays, the success of neural networks in a va- riety of NLP tasks heavily relies on BERT-type lan- guage models containing millions to billions of pa- rameters. However, the pretraining process of these models is computationally expensive, generating significant carbon dioxide emission (Strubell et al., 2019; Patterson et al., 2021). In practice, there is the need to perform large-scale language model pre- training for diverse applications (Lee et al., 2020; Chalkidis et al., 2020; Zou et al., 2021; Rogers et al., 2020) in different languages (Antoun et al., 2020; Sun et al., 2021). In this paper, we develop a technique that significantly reduces the pretraining cost of BERT models (Devlin et al., 2019) without hurting their test performance on a diverse set of fine-tuning tasks. ∗ Equal contribution. § Work done at Google Brain. Recent efforts of efficient training involve mixed- precision training (Shoeybi et al., 2019), distributed training (You et al., 2020), better modeling on rare words and phrases (Wu et al., 2021), designing more effective and data-efficient pretraining ob- jectives (Lan et al., 2020; Clark et al., 2020; Raf- fel et al., 2020), progressive stacking (Gong et al., 2019), and so on. While these approaches con- tribute to efficient training with reduced compu- tational 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 under- explored before, i.e., “token dropping,” which re- moves the redundant tokens in each sequence that are less informative to training. Since not all to- kens contribute equally to the output or the train- ing objective, and the computational complexity of transformer-based models grows at least linearly with respect to the sequence length, shortening the …

Content truncated to 500 words