Progressively Stacking 2.0: A Multi-stage Layerwise Training Method for BERT Training Speedup
Authors: Cheng Yang
Shengnan Wang
Chao Yang
Yuechuan Li
Ru He
Jingqiao Zhang
Venue:  (2021)
Keywords: BERT
Training speedup
Multi-stage training
Natural language processing

Abstract: Pre-trained language models, such as BERT, have achieved significant accuracy gain in many natural language processing tasks. Despite its effectiveness, the huge number of parameters makes training a BERT model  computationally very challenging.  In this paper, we propose an efficient multi-stage layerwise training (MSLT) approach to reduce the training time of BERT.  We decompose the whole training process into several stages. The training is started from a small model with only a few encoder layers and we gradually increase the depth of the model by adding new encoder layers. At each stage, we only train the top  (near the output layer)  few encoder layers which are newly added. The parameters of the other layers which have been trained in the previous stages will not be updated in the current stage. In  BERT training, the backward calculation is much more time-consuming than the forward calculation, especially in the distributed training setting in which the backward calculation time further includes the communication time for gradient synchronization. In the proposed training strategy,  only  top few layers  participate  backward calculation, while most  layers  only participate forward calculation. Hence both  the computation  and communication efficiencies are greatly improved.  Experimental results show that the proposed method can greatly reduce the training time without significant performance degradation.

Under review as a conference paper at ICLR 2021
PROGRESSIVELY STACKING 2.0: A MULTI-STAGE LAYERWISE TRAINING METHOD FOR BERT TRAIN- ING SPEEDUP
Anonymous authors Paper under double-bling blind review
ABSTRACT
Pre-trained language models, such as BERT, have achieved significant accuracy gain in many natural language processing tasks. Despite its effectiveness, the huge number of parameters makes training BERT model computationally very challenging. In this paper, we propose efficient multi-stage layerwise training (MSLT) approach to reduce the training time of BERT. We decompose the whole training process into several stages. The training is started from small model with only a few encoder layers and we gradually increase the depth of the model by adding new encoder layers. At each stage, we only train the top (near the out- put layer) few encoder layers which are newly added. The parameters of the other layers which have been trained in the previous stages will not be updated in the current stage. In BERT training, the backward computation is much more time- consuming than the forward computation, especially in the distributed training setting in which the backward computation time further includes the communi- cation time for gradient synchronization. In the proposed training strategy, only top few layers participate in backward computation, while most layers only par- ticipate forward computation. Hence both the computation and communication efficiencies greatly improved. Experimental results show that the proposed method can achieve more than 110% training speedup without significant perfor- mance degradation.
INTRODUCTION
In recent years, the pre-trained language models, such as BERT (Devlin et al. 2018), XLNet (Yang ct al. 2019), GPT (Radford al. 2018), have shown their powerful pertormance various areas, especially in the field of natural language processing (NLP). By pre-trained on unlabeled datasets and fine-tuned 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 proposed, such RoBERTa (Liu al. 2019b). ALBERT (Lan et al. 2019). Structbert Wang et al. 2019) etc., most of which yielded new state-of-the thr results.
Despite the accuracy gains, these models usually involve a large number of parameters (e.g. BERT- Base has more than 110M parameters and BERT-Large has more than 340M parameters), and they 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 high cost is not affordable for many researchers and institutions. Therefore, improving the training efficiency should be a critical issue 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
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 ALBERT model, which shares parameters across all the hidden layers, 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 deep BERT

---

Under review as a conference paper at ICLR 2021
cage
stage
Slipe
Clesslar
Grea
MA Dreode Lapez
+>
codd Joar
Grosdr Layer
Erowiing
Paco Toper
Eanoday Laper
Gitadey Lapero
Mik
Eheode Love!
Kils Eacoale Laper
Ipp
Irpa doa
The framework of MSLT method. The green blocks mean that these layers only partici- pate in forward computation and are not updated. The red blocks mean that these layers are trained in this stage and they participate in both forward and backward computations.
Figure I:
network by progressively stacking from shallow one. Utilizing the similarity of the attention distributions across different layers, such strategy achieves about 25% speedup without significant performance loss.
Progressively stacking provides a novel training strategy, namely training a BERT model from shal- low to deep. However, progressively stacking only has high training efficiency 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 pro- gressively 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 functionalityy 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 novel multi-stage layerwise training (MSLT) approach, which can greatly improve the training efficiency of BERT. We decompose the training process BERT into several stages, as shown in Fig. We start the training from 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 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, further retrain the model by updating all the layers together. Since the whole model has already been well trained, this stage only requires few steps (accounting for about 20% of the total steps). Compared with the progressively stacking method, which requires 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 -eeeeefficeen..
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. should be mentioned that some widely-known methods such 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

---

Under review as a conference paper at ICLR 2021
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 al., 2019b), XLNet (Yang 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 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 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 em- bedding 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 down- stream performance. Since the communication overhead is directly proportional to the number of parameters the model, ALBERT also improved the communication efficiency in distributed train- ing setting. However, since ALBERT has almost the same computational complexity 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 bot- tom layers, shown in Fig. 2 Such phenomenon implies that the encoder layers in the BERT model have similar functionalities. Utilizing the natural similarity characteristic, to train a N N-layer BERT model, the progressively stacking method first trains a layer model, and then sticks into N -llayer by copying the parameters of the trained N/2 layers. After the N 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 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, also train the BERT model from shallow to deep. In contrast, each stage, 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

---

Under review as a conference paper at ICLR 2021
L7
LB
L10
L11 L12 H6
L1 H10
15
13
14 L5 L6
The attention distribution trained BERT-Base model given a sample sentence. We compare the attention distribution of the same heads from different layers. For example, "L5 H4" denotes the attention distribution from the forth heads of the fifth layer.
Figure 2:
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 very inefficient, especially when the model 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 progressively stacking is that its training efficiency decreases as the training goes 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 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 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 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 6-layer model. Then we stack the trained 6-layer model into 12-layer model and continue to train the whole model until convergence. The attention distributions of the top encoder layers of the final 12-layer BERT-Base model are shown in the first row of Fig. 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 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. shows the attention distributions of the bottom 6 encoder layers of the original BERT-Bas Base model.
Combined with Fig. we find that:
1. Except the two obvious distributions found by Gong et al. (201)), namely the distribution fo- cusing local positions and the distribution focusing 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 the position of SEP).
4

---

Under review as a conference paper at ICLR 2021
From before
The first row is the attention distributions for randomly chosen sample sentence on random 6 heads of the top 6 layers from the final trained 12-layer BERT-Base model. The second row the attention distributions of the bottom 6 layers from the final trained 12-layer BERT-Base model. The third row is the attention distributions of the encoder layers from the trained 6-layer model before stacking. The fourth row is the attention distributions the bottom 6 layers from the BERT-Bas Base model trained from scratch using the original method.
Figure 3:
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 lay- ers, 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
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 lay- ers 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 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 the forward computation, and only the newly added top layers 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 distributed training setting. So the time of both computation and communication can be greatly reduced.
3.2 MULTI-STAGE LAYERWISE TRAINING
Now we present our MSLT training method. Based on the above observations, train a deep BERT model with N encoder layers. we decompose the whole training process into & stages, as shown in Fig. We start our training from shallow model with only N/k encoder layers and we gradually add new encoder layers on the top of the model (below the output layer). At each time, we only add N/k new layers. Except the first stage in which we update all the parameters, in the other stages, only the parameters of the newly added top layers and the output layer are updated. Namely, the bottom layers which have been trained in the previous stages only participate in the forward computation to provide the training input for the top layers. In addition, in the BERT model, the word embedding matrices in the input and output layers are shared. So we only update the word embedding matrix in the first stage, and in the later stages the word embedding matrix is fixed.
5

---

Under review as a conference paper at ICLR 2021
Similar to Gong et al. 2 at each stage, we initialize the parameters of the newly added N/k: layers by copying the parameters of the top N/k layers of the model trained in the previous stage. To make the final model better behaved, after each layer is trained, we further retrain the model by updating all the parameters together for a few steps.
In the BERT model training, backward computation generally takes much longer time than the for- ward computation, especially in the distributed training setting in which the backward computation also includes gradient synchronization. For example, in our experiment, when using the original method (Devlin et al., 2018) to train a BERT-Base model, the time of backward computation (in- cluding communication time for gradient synchronization in the distributed setting) is almost six times as the time of forward computation. The proposed MSLT method can greatly reduce the back- ward computation, since only top few layers participate in the backward computation the whole training process, except the final retraining stage. Hence, the total training time will be much shorter.
3.3 EXTENDING TO ALBERT
The proposed MSLT strategy can also be used to speed up the training of ALBERT Lan et al. (2019). The only problem is that ALBERT shares the parameters across all the encoder layers, so we are not able to only update the top layers while keeping the other layers fixed. Hence, before applying MSLT, we make a slight modification of ALBERT. We decompose the whole encoder layers into k groups and we share the parameters of the encoder layers in the same group, Then we use similar strategy as shown in last subsection to efficiently train the ALBERT model. Specifically, at each stage, add new group of encoder layers and only the newly added group of layers will be updated in this stage.
Compared with the original ALBERT model, the modified ALBERT model requires more storage resources since it involves more parameters. However, the training speed is significantly improved. In real applications, should decompose the encoder layers into suitable number of groups to achieve a computation-storag efficiency balance, according to actual demand.
4
EXPERIMENT
4.1 EXPERIMENTAL SETUP
In this section, we perform experiments to demonstrate the effectiveness and efficiency of the pro- pose method. Following the setup in Devlin et (218)). use English Wikipedia (2,500M words) and BookCorpus (800M words) for pre-training. Though some more effective objectives were pro- posed in Lan al. 2019 Steinley 2006), to make the comparison as meaningful possible, in this experiment we still adopt the same Masked Language Model (MLM) and Next Sentence Prediction (NSP) objectives used by the original BERT (Devlin et al. 2018). We set the maximum length of each input sequence to be 128. We focus on the training time speed up of the proposed method and whether similar performance is achieved compared with the original BERT with the same setting. The models considered in this section are BERT-Base and BERT-Large. whose hyperparameters setting can be seen in Devlin al. (2018). For each model, we decompose the training process into k = 4 stages. So at each stage, we train 3 encoder layers for BERT-Base, and 6 encoder layers for BERT-Large, All the experiments are performed on a distributed computing cluster consisting of 32 Telsa V100 GPU cards with 32G memory, and the batch size is 32#32=1024. We use the LAMB optimizer with learning rate 0.00088 (You et al. 2019). All the other settings are the same as Devlin ct al. (2018), including the data pre-processing, unless otherwise specified.
4.2 COMPARISON FOR TRAINING WITH SAME STEPS
We train both BERT-Base and BERT-Large models using the MSLT method for totally 1,000,000 steps. Specifically, each stage uses 200,000 steps, and the final model is retrained for the remaining 200,000 steps. We train a BERT-Base and BERT-Large model with 1,000,000 steps from scratch using the original method DDellin et al. 2018) as the baselines. For each model, the first 10,000 training steps used for learning rate warmup. We first show the pre-training loss of all the models in Fig. 4 We can see that for both BERT-Bare and LERTeLarge, the loss of our method decreases faster than the baseline, and the final convergence value of our method is close to that of the baseline.

---

Under review as a conference paper at ICLR 2021
BERT-Base
BERT-Large
baseline MSLT
baseline
MSLT
loss
4-
0
20
40
60
80
0
50
100
150
200
traing time (hours)
training time (hours)
Figure 4: Pre-training loss curves of the BERT models.
Dev set results of SQuAD and selected GLUE tasks for the pre-trained models. FI/EM scores are reported for SQuAD tasks. F1 scores are reported for QQP, MRPC, and Spearman corre- lations are reported for STS-B. The accuracy scores are reported for the other tasks. The result of MNLI the average the scores of MNLI-m and MNLI-mm tasks. The Avg result is the average of the scores of the downstream tasks its left, where the Fl and EM scores of SQuAD tasks are
Table I:
firstly averaged.
Model
time
MNLI
QNLI
SST-2
MRPC
SQuAD1.1
SQuAD2.0
Avg
BERT-Bas Base (ours)
40h
84,3
91.2
92.2
91.2
89.2/82.3
77.8/74,1
86.8
BERT-Base (original)
85h
84.4
91.5
92.7
89.7
89.5/82.7
77.3/73.9
86.7
BERT-Base (stacking)
63h
84.6
91.6
92.1
91.2
89.4/82.5
77.3/74.2
86.9
BERT- BERT-Large (ours)
84h
86.2
92.6
93.2
92.2
91.6/84.8
81.6/79.2
88.8
BERT-Larg (original)
188h
86.5
93.0
93.6
91.7
91.9/85.2
81.5/79.1
88.9
BERT-Large (stacking)
136h
86.7
93.1
93.8
92.4
92.1/85.2
81.7/79.2
89.2
Finally, our method achieves more than 110% speedup (saves about 55% training time). which is significant improvement compared with the 25% speedup achieved by progressively stacking (Gong jet al. 2019).
Then we further evaluate the above models on the widely used General Language Understanding Evaluation (GLUE (Wang et al., 2018)) benchmark and the Stanford Question Answering Dataset (SQuADI.1 and SQuAD2.0). The sequence length of the SQuAD tasks is 384. So for the SQuAD tasks, we train the last 10% of the steps using the sequence 512 to learn the positional embeddings. To better show the advantage of the proposed method, also add the progressively stacking method Gong et al. (2019) as the baseline. Table shows the Dev set results on SQuAD and selected GLUE tasks. Similar to Liu et al. (2019b). all the results
the median of five runs. For each GLUE task, we fine-tune the model with batch size 32 and we perform a grid search on the learning rate set [5e-5, 4e-5, 3e-5, 2e-5]. Following Clark al. (2020), fine-tune the model with 10 epochs for STS and RTE, and 3 epochs for all the other tasks. For SQuAD1.1 task, we fine-tune the model with batch size 12, epoch and learning rate 3e-5. For SQuAD2.0 task, we fine-tune the model with batch size 48, epoch 2, and learning rate 5e-5.
and BERT-Large, the results of our method are close to those of the original BERT. In addition, the results of the baselines are comparable to those shown Devlin al. 2018). which confirms the validity of our proposed method.
From Table I we can see that for both
4.3 COMPARISON FOR TRAINING WITH SAME TIME
In this section, we compare the performance of all the models trained with the same time. Training a BERT-Base model with IM steps using MLST method requires about 40 hours, which similar to the time of training a BERT-Base model using the original method for 480,000 steps (about 41 hours). In addition, we further train two BERT-Large models using MSLT with 500,000 steps (about 42 hours) and the original method with 230,000 steps (about 42 hours), respectively. The results selected GLUE tasks are shown in Table 2 We see that with the same training time, the models trained by method perform much better than the baselines.

---

Under review as a conference paper at ICLR 2021
Table 2: Dev results of selected GLUE tasks for the pre-trained models by controlling training time.
Model
steps
time
MNLI
QNLI
SST-2
STS-B
MRPC
RTE
Avg
BERT- Base (ours)
1M
40h
84.3
91.2
92.2
89.3
91.2
69.3
86.3
BERT-Base (original)
480k
41h
83.3
90.6
91.2
89.2
89.2
68.9
85.4
BERT- Large (ours)
500k
42h
84.7
91.5
92.1
89.9
91.6
70.7
86.8
BERT-Large (original)
230k
42h
84.1
90.8
91.1
89.2
88.9
68.5
85.4
Table 3: Dev set results of selected GLUE tasks for ALBERT-Base models.
Model
time
parameters
MNLI
QNLI
SST-2
STS-B
MRPC
Avg
ALBERT-Bass (ours)
32h
32M
82.4
89.9
90.9
88.8
90.7
88.5
ALBERT-Bas (original)
62h
12M
81.8
89.2
89.4
88.1
87.4
87.2
4.4 ACCELERATE TRAINING OF ALBERT
As shown in Section the proposed MSLT method can also be used speed up the training of ALBERT. Table reports the results on select GLUE tasks of the original ALBERT-Bas model and the modified ALBERT Base model trained by MSLT. All the models are trained for 1M steps and the other settings are the same as Section 4.2 We can see that the modified ALBERT model trained by MSLT achieves better performance than the original BERT. That is because the modified ALBERT model involves more parameters. Hence, the original ALBERT has higher memory efficiency while the modified ALBERT trained using MSLT has higher time efficiency and better performance.
4.5 EFFECT OF JOINTLY RETRAINING
In the above examples, we left 200,000 steps for jointly retraining all the layers. Here we investigate the impact of the retraining stage. Table shows the results of BERT models with/without retraining. The results implies that the retraining stage can further improve the performance of the BERT model. The training efficiency of the retraining stage much lower than the previous stages, since this stage all the parameters are updated. However, since the model is already near-optimal after training by previous stage, the retraining stage only requires few steps. In practice, we use 10% 20% of the total steps for retraining, so the retraining stage will not make the whole training very time- consuming.
4.6 ONLINE TEST RESULTS ON GLUE
Lastly, we report the online test results on GLUE tasks (except WNLI) all the models after fine- tuning, which are shown in Table All the submitted models are pre pre-trained with 1,000,000 steps. Same as Devlin et al. (2018), for each task, we select the best fine-tuning learning rate on the Dev
DISCUSSION
In last section, we empirically showed the effectiveness and efficiency of the proposed MSLT method. this section discuss how does it work and why can improve the convergence speed.
5.1 RELATIONSIIIP BETWEEN MSLT AND ALTERNATING OPTIMIZATION
At first glance, MSLT is extension of the progressively stacking method, and the main difference between MSLT and progressively stacking is that at each time progressively stacking updates all the parameters, while MSLT freezes most parameters and only updates the parameters from the
few top layers. Such strategy is similar to alternating optimization Bezdek & Hathaway (2003), which is widely used for solving multi-variable nonconvex optimization problem, due to its simple
8

---

Under review as a conference paper at ICLR 2021
Table 4: Results selected GLUE tasks of BERT models with/without retraining.
Model
steps
time
MNLI
QNLI
SST-2
STS-B
MRPC
Avg
BERT-Base without retraining
800K
26h
82.1
88.7
90.2
87.9
89.5
87.7
BERT-Base with retraining
IM
40h
84.4
91.5
92.7
89.7
89.7
89.6
BERT-L without retraining
SOOK
54h
83.3
89.7
91.2
88.2
89.9
88.5
BERT-Large with retraining
1M
84h
86.2
92.6
93.2
90.2
92.2
90.9
Table 5: Online test results of GLUE tasks for the pre-trained models.
Model
MNLI
QQP
QNLI
SST-2
CoLA
STS-B
MRPC
RTE
Avg
BERT-Bas (ours)
83.9
71.0
90.4
93.0
56.6
84.1
89.2
68.0
79.5
BERT-Basc Bass (baseline)
84.2
71.4
90.7
93.5
57.2
84.4
88.8
68.1
79.8
BERT-Large (ours)
85.9
71.9
92.4
94.1
60.5
85,8
89.9
69.8
81.3
BERT-Largeebaaeelne)
86.4
72.1
92.8
94.7
60.6
86.1
89.1
70.2
81.5
implementation, fast convergence, and superb empirical performance Li et al. (2019). Alternating optimization addresses the multi-variable optimization problem by updating a small set of variables and keeping the other variables fixed. The alternating optimization method is quite suitable for dealing with the problems in which updating a part of variables is much easier than updating all variables Bezdek & Hathaway (2002).
According to alternating optimization, in the retraining stage, we should update the bottom layers again and keep the other layers fixed. However, when the model is stacked deep, to compute the gradient of the parameters the bottom layers, we have to compute the gradient of the parameters of the top layers, according to the chain rule. In this situation, updating the parameters of the bottom layers is almost as expensive as updating all the parameters. So we update all the parameters simultaneously in the retraining stage.
Overall, in the MSLT method, we utilize the advantage of alternating optimization to quickly reach near-optimal state. When the model is deep enough, alternating optimization does have effi- ciency advantage, then we use joint descent to sufficiently utilize the gradient information.
5.2 AVOID CONTRADICTION OF CHOICE OF LEARNING RATE
In Section 2.3, we have shown that, though the bottom layers are updated most steps progres- sively stacking method, they do not have significant change in the later stage. So it is not worthwhile to spend so much the computation to updating the parameters of bottom layers. Another issue is that for the adaptive optimizer, such as Adam LAMB, when the variable is close to the optimal value, we should use a small learning rate, and when the variable far from optimal, we should use large learning rate. In progressively stacking strategy, the bottom layers near-optimal while the newly added top layers are far from optimal since they are not trained. So they have completely different requirements for the learning rate. The MSLT method address the learning rate contradiction by first optimizing the untrained top layers and keep the trained bottom layers fixed. When all the layers are oetrao, we further retrain all the layers with small learning rate. Meanwhile, such manner saves lot of computation.
6 CONCLUSION
In this paper, we propose an efficient multi-stage layerwise training method for accelerating the training process of BERT. We decompose the whole training process into several stages and we adopt the progressively stacking strategy that trains a BERT model from shallow to deep by gradually adding new encoder layers. We find that the attention distributions of the bottom layers tend to be fixed early, and further training in the later stages does not bring significant changes. So in our method, at each stage, we only train the top few layers which are newly added, while the bottom layers only participate in the forward computation. Experimental results show that the proposed training method achieves significant speedup without significant performance loss, compared with the existing training method.

---

Under review as a conference paper at ICLR 2021
REFERENCES
James C Bezdek and Richard J Hathaway. Some notes on alternating optimization. In AFSS Inter- national Conference on Fuzzy Systems, pp. 288-300. Springer, 2002.
James C Bezdek and Richard J Hathaway, Convergence alternating optimization. Neural, Parallel & Scientific Computations, 11(4):351-368 2003.
Kevin Clark, Minh-Thang Luong, Quoe V Le, and Christopher D Manning. Electra: Pre-training text encoders as discriminators rather than generators. arXiv preprint arXiv:2003. 10555, 2020.
Jacob Devlin, Ming- Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:/810.0450 2018.
Linyuan Gong, Di He, Zhuohan Li, Tao Qin, Liwei Wang, and Tieyan Liu. Efficient training of bert by progressively stacking. In International Conference on Machine Learning, pp. 2337-2346, 2019.
Song Han, Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. arXiv preprint arXiv:1510.0014 2015a.
Song Han, Jeff Pool, John Tran, and William Dally. Learning both weights and connections for efficient neural network. In Advances in neural information processing systems, pp. 1135-1143, 2015b.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv preprint arXiv:1503.0255 2015.
Ashish Khetan and Zohar Karnin. schubert: Optimizing elements of bert. arXiv preprint 0rXiv:2005.06628, 2020.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Sori- cut. Albert: lite bert for self-supervised learning of language representations. arXiv preprint arXiv: 999.0 1942, 2019.
Qiuwei Li, Zhihui Zhu, and Gongguo Tang. Alternating minimizations converge to second-order optimal solutions. In International Conference on Machine Learning, pp. 3935-3943, 2019.
Weijie Liu, Peng Zhou, Zhe Zhao, Zhiruo Wang, Qi Ju, Haotang Deng, and Ping Wang. K-bert: En- abling language representation with knowledge graph. arXiv preprint arXiv: 1909007606, 2019a.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1907. 1692, 2019b.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language derstanding by generative pre-training. URI. https://s3-us-wed amazonaws. com/openai- understanding paper. pdf. 2018.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint
2019.
Douglas Steinley. K-means clustering: half-century synthesis. British Journal of Mathematical and Statistical Psychology, 59(1):1-34 2006.
Emma Strubell, Ananya Ganesh, and Andrew McCallum. Energy and policy considerations for deep learning in nlp. arXiv preprint 1rXiv:1906.02243 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint / 007411 2018.
10

---

Under review as a conference paper at ICLR 2021
Wei Wang, Bin Bi, Ming Yan, Chen Wu, Zuyi Bao, Liwei Peng, and Luo Si. Structbert: Incor- porating language structures into pre- training for deep language understanding. arXiv preprint arXiv:1908.04577, 2019.
Zhilin Yang, Zihang Dai. Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. In Advances in neural information processing systems, pp. 5754-57656 5764, 2019.
Junho Yim, Donggyu Joo, Jihoon Bae, and Junmo Kim. A gift from knowledge distillation: Fast optimization, network minimization and transfer learning. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 4133-4141, 2017.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint (00dv::904400000 1(5), 2019.


Reviews:
[{"review_id": "uoyYDZhX3Zo", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary\nThis paper proposes a simple method, ie multi-stage layerwise training (MSLT), to speedup BERT training. Specifically, the authors progressively stack layers. The bottom layers are fixed and only the new added top layers are trained. The proposed method can achieve more than 110% training speedup and achieve comparable (slightly worse) performance. Although compared with other speedup training methods like ELECTRA [1], the idea of this paper is not novel. However, the proposed method is simple and effective to some extend. Thus, I give a marginal score to this paper.\n\n### Strengths\n* The proposed training method is simple and easy to implement. \n* It can achieve 110% training speedup without significant performance degradation.\n\n### Weaknesses and Questions\n* The most related method to this paper is [2]. Thus, it is better to give more comparisons and discussions in the Experiment, so that the author can know the advantages of MSLT compared with [2].\n* From Figure 4, it seems that BERT-base (baseline) needs to train around 80 hours. If training with MSLT using the same time as baselines (like train 2M steps), how about the performance compared with baseline? [3] shows that training with more steps can help improve performance. Does this phenomenon still remain in your method?\n\n\n\n\n[1] Clark, Kevin, et al. \"Electra: Pre-training text encoders as discriminators rather than generators.\" ICLR. 2020.\n\n[2] Gong, Linyuan, et al. \"Efficient training of bert by progressively stacking.\" ICML. 2019.\n\n[3] Liu, Yinhan, et al. \"Roberta: A robustly optimized bert pretraining approach.\" arXiv preprint arXiv:1907.11692 (2019).", "extra_scores": {}, "extra_text": {}}, {"review_id": "tMQs-arME6V", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose a multi-stage layerwise training (MSLT) approach to reduce the training time of BERT.   Experimental results show that the proposed method can achieve  110%+ training speedup without significant performance degradation.\n\nOverall the idea of multi-stage layerwise training is reasonable and the results look promising. \n\nMy major concern about the work is the empirical comparisons. ALBERT, a closely related work, is not compared in this work.  This paper claims that \"ALBERT has almost the same computational complexity as BERT, training an ALBERT model is still very time-consuming.\" \nI checked the ALBERT paper, which claims 1.7x training speedup. \n\"An ALBERT configuration similar to BERT-large has 18x fewer parameters and can be trained about 1.7x faster. The parameter reduction techniques also act as a form of regularization that stabilizes the training and helps with generalization.\"\nSeems the proposed method in this work (with 1.1x speedup) is not as good as ALBERT. I suggest to take ALBERT as backbone and check whether the MSLT can speed up the training of ALBERT.\n\nBesides, it is better to also test on more complex downstream tasks, such as SQuAD1.1/2.0 and RACE.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "SzJ5b--Jbyy", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presents a training strategy to progressively adding top transformer layers, which results in training time speedup. Usually when training transformers, all layers are updated simultaneously. The author found out doing a multi-stage layer-wise schedule helps convergence speed without significant performance degradation.\n\nThe paper is written very clearly. The idea is conveyed with sufficient background and related work. \nMy biggest concern is from the originality and significance. It looks to me the approach that the paper proposed is a very straightforward extension from the work of (Gong et al., 2019).  Instead of updating all parameters in all transformer layers, this paper freezes bottom layers and only update top transformer layer (newly added). With experiments performed on base and large BERT models on GLUE dataset, this training strategy is approved to have slightly drop of quality but faster convergence speed. \n\nI would argue that this is a great investigation and experimentation but it does not meet the criteria of acceptance. In section 3.1, the author has used the attention distribution (before and after) to motivate the work in this paper. However, if the author could discuss/explains why this change affects the convergence speed in analytical/mathematical solution, that would definitely gives more credit to originality and significance.  I would be more than willing to re-evaluate my ratings if that happens. \n\nMinor comments/typos:\n\n* abstract: \"have has achieved\" -> \"have achieved\".", "extra_scores": {}, "extra_text": {}}, {"review_id": "QaRiexqwvbT", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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. \n\nOne problem with the current paper is the empirical evaluation is only conducted on the GLUE benchmark, which is sequence-level and relatively simple. I think an experience on a slightly more difficult task such as SQuAD, which also requires token-level prediction, would be necessary to test the capacity boundary of the proposed approach.\n\nAnother question is what would happen or what the performance would be if the entire model is not jointly trained for the last 20% steps. This information will help to better understand this method.\n\nIn addition, the original motivation of the work comes from the fact that the attention patterns in the bottom layers do not change much after jointly trained with more higher layers. However, this does not mean the lower-layer attention patterns don't change much if the entire network is jointly trained from scratch. To truly establish the validity of motivation, it would be good to monitor and evaluate how much the lower-layer attention patterns change when jointly trained.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposed a variant of the existing training method "progressive stacking", and showed good empirical results comparing with the normal training procedure for BERT models. It contains some interesting points on the technical side, e.g, freezing the bottom layers when training newly added layers, but there are several concerns:  (1) This proposed method is called progressive stacking 2.0 but there is no comparison against the original progressive stacking in experiments. We had to check the empirical results in the original paper of progressive stacking, and did not notice any performance improvement of this new method over the original one; (2) The proposed method even introduced one more hyperparameter to tune: the number of top layers to copy.  This hyperparameter seems hard to choose. Different choices of this hyperparameter may dramatically impact the performance of this new method.  An ablation study on this should be conducted, e.g.,  what will the results look like if we only copy the last layer?  (3) The original progressive stacking method does not provide any practical guidance on how to determine the time to stack when the training goes. This severely limits the practical value of progressive stacking. If one stacks layers too early or too late, the stacking method may have no advantage at all. Unfortunately, this method called 2.0 still leaves this most critical issue away. 
Decision: Reject