Progressively Stacking 2.0: A Multi-stage Layerwise Training Method for BERT Training Speedup

1 INTRODUCTION . In recent years , the pre-trained language models , such as BERT ( Devlin et al. , 2018 ) , XLNet ( Yang et al. , 2019 ) , GPT ( Radford et al. , 2018 ) , have shown their powerful performance in various areas , especially in the field of natural language processing ( NLP ) . By pre-trained on unlabeled datasets and fine-tuned on small downstream labeled datasets for specific tasks , BERT achieved significant breakthroughs in eleven NLP tasks ( Devlin et al. , 2018 ) . Due to its success , a lot of variants of BERT were proposed , such as RoBERTa ( Liu et al. , 2019b ) , ALBERT ( Lan et al. , 2019 ) , Structbert ( Wang et al. , 2019 ) etc. , most of which yielded new state-of-the-art results . Despite the accuracy gains , these models usually involve a large number of parameters ( e.g . BERTBase has more than 110M parameters and BERT-Large has more than 340M parameters ) , and they are generally trained on large-scale datasets . Hence , training these models is quite time-consuming and requires a lot of computing and storage resources . Even training a BERT-Base model costs at least $ 7k ( Strubell et al. , 2019 ) , let alone the other larger models , such as BERT-Large . Such a high cost is not affordable for many researchers and institutions . Therefore , improving the training efficiency should be a critical issue to make BERT more practical . Some pioneering attempts have been made to accelerate the training of BERT . You et al . ( 2019 ) proposed a layerwise adaptive large batch optimization method ( LAMB ) , which is able to train a BERT model in 76 minutes . However , the tens of times speedup is based on the huge amount of computing and storage resources , which is unavailable for common users . Lan et al . ( 2019 ) proposed an ALBERT model , which shares parameters across all the hidden layers , so the memory consumption is greatly reduced and training speed is also improved due to less communication overhead . Gong et al . ( 2019 ) proposed a progressively stacking method , which trains a deep BERT network by progressively stacking from a shallow one . Utilizing the similarity of the attention distributions across different layers , such a strategy achieves about 25 % speedup without significant performance loss . Progressively stacking provides a novel training strategy , namely training a BERT model from shallow to deep . However , progressively stacking only has a high training efficiency at the initial stage in which the model depth is small . As the training goes on , the model depth increases and the training speed decreases . The low efficiency of the later stages makes the overall speedup of progressively stacking limited . Note that in the progressively stacking method , the bottom layers are trained with longer time than the top layers . However , we observe that though the bottom layers are updated all the time , they do not have significant changes in the later stages , in terms of the attention distribution which can reflect the functionality of the encoder layers to some extent ( Gong et al. , 2019 ) . In other words , most optimization of the bottom layers has been finished in the early stage when the model is shallow . Motivated by this observation , in this work , we propose a novel multi-stage layerwise training ( MSLT ) approach , which can greatly improve the training efficiency of BERT . We decompose the training process of BERT into several stages , as shown in Fig . 1 . We start the training from a small BERT model with only a few encoder layers and gradually add new encoder layers . At each stage ( except the first stage ) , only the output layer and the newly added top encoder layers are updated , while the other layers which have been trained in the previous stages will be fixed in the current stage . After all the encoder layers are trained , to make the network better behaved , we further retrain the model by updating all the layers together . Since the whole model has already been well trained , this stage only requires a few steps ( accounting for about 20 % of the total steps ) . Compared with the progressively stacking method , which requires a lot of steps ( accounting for about 70 % of the total steps ( Gong et al. , 2019 ) ) to train the whole model , our method is much more time-efficient . Experimental results demonstrate the effectiveness and efficiency of the proposed method in two aspects : 1 ) with the same data throughput ( same training steps ) , our method can achieve comparable performance , compared with the original training method , but consumes much less training time ; 2 ) with the same training time , our method can achieve better performance than the original method . According to the results , the proposed method achieves more than 110 % speedup without significant performance degradation . To avoid misunderstanding , it should be mentioned that some widely-known methods such as model compression ( Han et al. , 2015a ; b ) and knowledge distillation ( Yim et al. , 2017 ; Hinton et al. , 2015 ; Sanh et al. , 2019 ) are designed for network speedup in the inference phase . Namely , these methods are used after the model has been trained . While in this paper , we focus on the model training speedup . 2 RELATED WORK . Based on the bidirectional Transformer ( Vaswani et al. , 2017 ) encoder , BERT has shown its great representational ability and it achieved state-of-the-art results in eleven NLP tasks . Following BERT , many pre-trained models were proposed , such as RoBERTa ( Liu et al. , 2019b ) , XLNet ( Yang et al. , 2019 ) , KBERT ( Liu et al. , 2019a ) , StructBERT ( Wang et al. , 2019 ) , and so on . Higher accuracy were achieved by these models with more training data , more training steps , or more effective loss functions . However , the BERT models are generally large-scale and they need to be trained on massive datasets ( e.g . BERT-base is trained on BooksCorpus and Wikipedia with totally 3.3 billion word corpus ) . Hence , training a BERT model is challenging in terms of both computation and storage . In the literature , some approaches were proposed to improve the training speed of BERT . 2.1 DISTRIBUTED TRAINING WITH LARGE BATCH SIZE . A direct way to reduce the training time is to increase the training batch size by using more machines and train the model in a distributed manner . However , traditional stochastic gradient descent ( SGD ) based optimization methods perform poorly in large mini-batches training . Naively increasing the batch size leads to performance degradation and computational benefits reduction ( You et al. , 2019 ) . An efficient layerwise adaptive large batch optimization technique named LAMB was proposed in You et al . ( 2019 ) to address this problem . It allows the BERT model to be trained with extremely large batch size without any performance degradation . By using 1024 TPUv3 chips , LAMB reduced the BERT training time from 3 days to 76 minutes . Though tens of times speedup is achieved , these methods require a huge amount of computing and storage resources , which are far beyond the reach of common users . 2.2 ALBERT . ALBERT ( Lan et al. , 2019 ) adopted two parameter reduction techniques , namely factorized embedding parameterization and cross-layer parameter sharing , which significantly reduced the model size . In addition , ALBERT adopted the sentence-order prediction loss instead of the next-sentence prediction loss during pre-training , which is demonstrated to be more effective in terms of downstream performance . Since the communication overhead is directly proportional to the number of parameters in the model , ALBERT also improved the communication efficiency in distributed training setting . However , since ALBERT has almost the same computational complexity as BERT , training an ALBERT model is still very time-consuming . 2.3 PROGRESSIVELY STACKING . The most related work should be progressively stacking ( Gong et al. , 2019 ) , which is mainly based on the observation that in a trained BERT model , the attention distributions of many heads from top layers are quite similar to the attention distributions of the corresponding heads from the bottom layers , as shown in Fig . 2 . Such a phenomenon implies that the encoder layers in the BERT model have similar functionalities . Utilizing the natural similarity characteristic , to train a N−layer BERT model , the progressively stacking method first trains a N/2−layer model , and then sticks it into N−layer by copying the parameters of the trained N/2 layers . After the N−layer model is constructed , the progressively stacking method continues to train the whole model by updating all the parameters together . By repeatedly using such a strategy , the deep BERT model can be trained more efficiently . According to the results shown in Gong et al . ( 2019 ) , progressively stacking can achieve the training time about 25 % shorter than the original training method ( Devlin et al. , 2018 ) . However , we can see that the speedup of progressively stacking mainly comes from the initial stage in which the model depth is small . As the model depth increases in the later stages , the training efficiency also decreases , and according to Gong et al . ( 2019 ) , to guarantee the performance , more training steps should be assigned in the later stages for training the deep model . Hence , the overall speedup brought by progressively stacking is limited . Such an issue is addressed in this paper . In our work , we also train the BERT model from shallow to deep . In contrast , at each stage , we only train the top few layers and we almost keep a high training efficiency during the whole training process . Hence , much more significant speedup is achieved . 3 METHODOLOGY . In this section , we propose an efficient training method to accelerate the training of BERT model . 3.1 MOTIVATION . The large depth should be one of the main reasons making the BERT training time-consuming . The original method ( Devlin et al. , 2018 ) trains all the encoder layers simultaneously . At each training step , the parameters need to wait for the cost function gradients to propagate backwards across all the layers before update , which is very inefficient , especially when the model is very deep . Inspired by progressively stacking ( Gong et al. , 2019 ) , we also consider to train the BERT model from shallow to deep . The main problem of progressively stacking is that its training efficiency decreases as the training goes on . We observe that in the progressively stacking strategy , the bottom layers are trained for longer time than the top layers . For example , the first encoder layer ( near the input layer ) is updated from beginning to end , while the last encoder layer ( near the output layer ) is only trained at the last stage . We doubt whether it is necessary to spend much more time training the bottom layers , since some research implies that the top encoder layers play a much more significant role ( Khetan & Karnin , 2020 ) . In BERT model , the encoder layers are mainly used to learn the dependencies of the input elements , which can be reflected by the attention distribution . In Gong et al . ( 2019 ) , the authors showed that the distributions of most attention heads are mixtures of two distributions : one distribution focuses on local positions , and another focuses on the first CLS token . In addition , the authors also observed that the attention distribution of the top layers is very similar to that of the bottom layers . Using a similar way , we also visualize some important attention distributions and we get some new findings when using the progressively stacking method to train a 12-layer BERT-Base model . Specifically , we first train a 6-layer model . Then we stack the trained 6-layer model into a 12-layer model and continue to train the whole model until convergence . The attention distributions of the top 6 encoder layers of the final 12-layer BERT-Base model are shown in the first row of Fig . 3 . For each layer , we randomly choose an attention head . Then we also show the attention distributions of the corresponding heads of the bottom 6 encoder layers in the second row of Fig . 3 . Further , we show the attention distributions of the corresponding heads of the trained 6-layer BERT model before stacking in the third row . As a comparison , we also train a 12-layer BERT-Base model from scratch using the original method , where the parameters of the bottom 6 encoder layers use the same initialization as the above BERT model trained by progressively stacking . The forth row of Fig . 3 shows the attention distributions of the bottom 6 encoder layers of the original BERT-Base model . Combined with Fig . 2 , we find that : 1 . Except the two obvious distributions found by Gong et al . ( 2019 ) , namely the distribution focusing on local positions and the distribution focusing on the CLS token , the attention distributions of many heads also focus on the SEP token ( for example , the dark vertical line in “ L11 H4 ” in Fig . 2 corresponds to the position of SEP ) . 2 . Compared with the first and second rows of Fig . 3 , one can see that for a trained 12-layer BERT model , some bottom layers have similar attention distributions to the corresponding top layers , which is in line with the observation in Gong et al . ( 2019 ) . In addition , there are also some bottom-top layer pairs whose attention distributions are very different . On the other hand , compared with the second and third rows of Fig . 3 , we can see that the attention distributions of the bottom layers from the final 12-layer model are almost the same as those of the corresponding layers from the trained 6-layer model before stacking , which implies that the further training of the bottom layers after stacking does not bring substantial optimization . The performance of the bottom encoder layers are not further improved , in terms of catching the elements ’ dependencies . Compared with the third and forth rows of Fig . 3 , we see that the attention distributions of the trained 6-layer model are also very similar to those of the bottom layers of the BERT-Base model with all the layers jointly trained from scratch . Therefore , it is not worth spending too much time training the bottom layers , especially updating the bottom layers is generally much more expensive than updating the top layers , since backward computation is top-down . An intuitive idea is that at each stage , let the bottom layers having been trained in the previous stages only participate in the forward computation , and only the newly added top layers as well as the output layer participate in the backward computation . Then the gradient information of the parameters from the bottom layers will not be computed and also will not be communicated in distributed training setting . So the time of both computation and communication can be greatly reduced .

Summary: The work proposes a simple enough idea to speed up the training of BERT by progressively stacking new layers while fixing older layers. Empirically, with the same number of training steps (and less time), the proposed method can achieve a comparable performance to the original BERT. When the same amount of running time (more steps) is used, the proposed strategy can further improve the performance.