EarlyBERT: Efficient BERT Training via Early-bird Lottery Tickets

Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing, pages 2195–2207 August 1–6, 2021. ©2021 Association for Computational Linguistics 2195 EarlyBERT: Efficient BERT Training via Early-bird Lottery Tickets Xiaohan Chen1∗ Yu Cheng2 Shuohang Wang2 Zhe Gan2 Zhangyang Wang1 Jingjing Liu2 1University of Texas at Austin, 2Microsoft Corporation {xiaohan.chen, atlaswang}@utexas.edu {yu.cheng, shuowa, zhe.gan, jingjl}@microsoft.com Abstract Heavily overparameterized language mod- els such as BERT, XLNet and T5 have achieved impressive success in many NLP tasks. However, their high model complex- ity requires enormous computation resources and extremely long training time for both pre- training and fine-tuning. Many works have studied model compression on large NLP mod- els, but only focusing on reducing inference time while still requiring an expensive train- ing process. Other works use extremely large batch sizes to shorten the pre-training time, at the expense of higher computational resource demands. In this paper, inspired by the Early- Bird Lottery Tickets recently studied for com- puter vision tasks, we propose EarlyBERT, a general computationally-efficient training al- gorithm applicable to both pre-training and fine-tuning of large-scale language models. By slimming the self-attention and fully- connected sub-layers inside a transformer, we are the first to identify structured winning tick- ets in the early stage of BERT training. We ap- ply those tickets towards efficient BERT train- ing, and conduct comprehensive pre-training and fine-tuning experiments on GLUE and SQuAD downstream tasks. Our results show that EarlyBERT achieves comparable perfor- mance to standard BERT, with 35∼45% less training time. Code is available at https: //github.com/VITA-Group/EarlyBERT. 1 Introduction Large-scale pre-trained language models (e.g., BERT (Devlin et al., 2018), XLNet (Yang et al., 2019), T5 (Raffel et al., 2019)) have significantly advanced the state of the art in the NLP field. De- spite impressive empirical success, their computa- tional inefficiency has become an acute drawback in practice. As more transformer layers are stacked ∗Work was done when the author interned at Microsoft. with larger self-attention blocks, model complexity increases rapidly. For example, compared to BERT- Large model with 340 million parameters, T5 has more than 10 billion to learn. Such high model complexity calls for expensive computational re- sources and extremely long training time. Model compression is one approach to allevi- ating this issue. Recently, many methods have been proposed to encode large NLP models com- pactly (Sanh et al., 2019; Jiao et al., 2019; Sun et al., 2019, 2020b). However, the focus is solely on reducing inference time or resource costs, leav- ing the process of searching for the right com- pact model ever more costly. Furthermore, most model compression methods start with a large pre- trained model, which may not be available in prac- tice. Recent work (You et al., 2020b) proposes to use large training batches, which significantly shortens pre-training time of BERT-Large model but demands daunting computing resources (1,024 TPUv3 chips). In contrast, our quest is to find a general resource- efficient training algorithm for large NLP models, which can …

Content truncated to 500 words