E arly BERT : Efficient BERT Training via Early-bird Lottery Tickets (2021)

Abstract: Heavily overparameterized language models such as BERT, XLNet and T5 have achieved impressive success in many NLP tasks. However, their high model complexity 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 models, but only focusing on reducing inference time while still requiring an expensive training 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 computer vision tasks, we propose EarlyBERT, a general computationally-efficient training algorithm 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 tickets in the early stage of BERT training. We apply those tickets towards efficient BERT training, and conduct comprehensive pre-training and fine-tuning experiments on GLUE and SQuAD downstream tasks. Our results show that EarlyBERT achieves comparable performance to standard BERT, with 35 45% less training time. Code is available at https://github.com/VITA-Group/EarlyBERT .

## introduction large-scale pre-trained language models (e.g., bert @xcite , xlnet @xcite , t5 @xcite ) have significantly advanced the state of the art in the nlp field. despite impressive empirical success, their computational 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 resources and extremely long training time. model compression is one approach to alleviating this issue. recently, many methods have been proposed to encode large nlp models compactly @xcite @xcite . however, the focus is solely on reducing inference time or resource costs, leaving the process of searching for the right compact model ever more costly. furthermore, most model compression methods start with a large pretrained model, which may not be available in practice. recent work @xcite 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 resourceefficient training algorithm for large nlp models, which can be applied to both pre-training and finetuning stages. our goal is to trim down the training time and avoid more costs of the total training resources (e.g., taking large-batch or distributed training). to meet this challenge demand, we draw inspirations from recent work @xcite that explores the use of lottery ticket hypothesis (lth) for efficient training of computer vision models. lth was first proposed in @xcite as an exploration to understand the training process of deep networks. the original lth substantiates a trainable sparse sub-network at initialization, but it cannot be directly utilized for efficient training, since the subnetwork itself has to be searched through a tedious iterative process. in addition, most lth works discussed only unstructured sparsity. the study of @xcite presents discoveries that structured lottery tickets can emerge in early stage of training (i.e., early-bird ticket), and therefore a structurally sparse subnetwork can be identified with much lower costs, leading to practical efficient training algorithms. inspired by the success of lth and early-bird ticket, we propose earlybert, a general efficient training algorithm based on structured early-bird tickets. due to the vast differences between the architectures and building blocks of computer vision models and bert, directly extending the method of @xcite does not apply to our work. by instead using network slimming @xcite on the self-attention and fully-connected sub-layers inside a transformer, we are the first to introduce an effective approach that can identify structured winning tickets in the early stage of bert training, that are successfully applied for efficient language modeling pre-training and fine-tuning. extensive experiments on bert demonstrate that earlybert can save 35∼45% training time with minimal performance degradation, when evaluated on glue and squad benchmarks. ## related work efficient nlp models it is well believed that bert and other large nlp models are considerably overparameterized @xcite . this explains the emergence of many model compression works, which can be roughly categorized into quantization @xcite , knowledge distillation @xcite @xcite , dynamic routing @xcite , and pruning @xcite @xcite . almost all model compression methods focus on reducing inference time, while their common drawback is the reliance on fully-trained and heavily-engineered dense models, before proceeding to their compact, sparse versions -which essentially transplants the resource burden from the inference to the training stage. pruning is the mainstream approach for compressing bert so far @xcite . @xcite proposed to greedily and iteratively prune away attention heads contributing less to the model. @xcite proposed to structurally prune bert models using low-rank factorization and augmented lagrangian ℓ 0 norm regularization. @xcite pruned less important self-attention heads and slices of mlp layers by applying ℓ 0 regularization to the coefficient corresponding to each head/mlp layer. others aim to reduce the training time of transformer-based models via large-batch training and gpu model parallelism @xcite . our work is orthogonal to these works, and can be readily combined for further efficiency boost. ## the earlybert framework in this section, we first revisit the original lottery ticket hypothesis (lth) @xcite and its variant early-bird ticket @xcite , then describe our proposed earlybert. ## conclusion in this paper, we present earlybert, an efficient framework for large-scale language model pretraining and fine-tuning. based on lottery ticket hypothesis, earlybert identifies structured winning tickets in an early stage, then uses the pruned network for efficient training. experimental results demonstrate that the proposed method is able to achieve comparable performance to standard bert with much less training time. future work includes exploring more data-efficient strategies to enhance the current training pipeline.