Original: P ROGRESSIVELY S TACKING 2.0: A MULTI - STAGE LAYERWISE TRAINING METHOD FOR BERT TRAINING SPEEDUP
Revised: P ROGRESSIVELY S TACKING 2.0: A MULTI - STAGE LAYERWISE TRAINING METHOD FOR BERT TRAINING SPEEDUP

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Pre-trained language models, such as BERT, havehas achieved signiﬁcant accuracy gain in many natural language processing tasks.
Revised: Pre-trained language models, such as BERT, have achieved signiﬁcant accuracy gain in many natural language processing tasks.

Original: Despite its effectiveness, the huge number of parameters makes training a BERT model computationally very challenging.
Revised: Despite its effectiveness, the huge number of parameters makes training a BERT model computationally very challenging.

Original: In this paper, we propose an efﬁcient multi-stage layerwise training (MSLT) approach to reduce the training time of BERT.
Revised: In this paper, we propose an efﬁcient multi-stage layerwise training (MSLT) approach to reduce the training time of BERT.

Original: We decompose the whole training process into several stages.
Revised: We decompose the whole training process into several stages.

Original: 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.
Revised: 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.

Original: At each stage, we only train the top (near the output layer) few encoder layers which are newly added.
Revised: At each stage, we only train the top (near the output layer) few encoder layers which are newly added.

Original: The parameters of the other layers which have been trained in the previous stages will not be updated in the current stage.
Revised: The parameters of the other layers which have been trained in the previous stages will not be updated in the current stage.

Original: In BERT training, the backward computation is much more timeconsuming than the forward computation, especially in the distributed training setting in which the backward computation time further includes the communication time for gradient synchronization.
Revised: In BERT training, the backward computation is much more timeconsuming than the forward computation, especially in the distributed training setting in which the backward computation time further includes the communication time for gradient synchronization.

Original: In the proposed training strategy, only top few layers participate in backward computation, while most layers only participate in forward computation.
Revised: In the proposed training strategy, only top few layers participate in backward computation, while most layers only participate in forward computation.

Original: Hence both the computation and communication efﬁciencies are greatly improved.
Revised: Hence both the computation and communication efﬁciencies are greatly improved.

Original: Experimental results show that the proposed method can achieve more than 110% training speedup without signiﬁcant performance degradation.
Revised: Experimental results show that the proposed method can achieve more than 110% training speedup without signiﬁcant performance degradation

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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 ﬁeld of natural language processing (NLP).
Revised: 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 ﬁeld of natural language processing (NLP).

Original: By pre-trained on unlabeled datasets and ﬁne-tuned on small downstream labeled datasets for speciﬁc tasks, BERT achieved signiﬁcant breakthroughs in eleven NLP tasks (Devlin et al., 2018).
Revised: By pre-trained on unlabeled datasets and ﬁne-tuned on small downstream labeled datasets for speciﬁc tasks, BERT achieved signiﬁcant breakthroughs in eleven NLP tasks (Devlin et al., 2018).

Original: Due to its success, a lot of variants of
Revised: Due to its success, a lot of variants of

Original: BERT were proposed, such as RoBERTa (Liu et al., 2019b), ALBERT (Lan et al., 2019),
Revised: BERT were proposed, such as RoBERTa (Liu et al., 2019b), ALBERT (Lan et al., 2019),

Original: Structbert (Wang et al., 2019) etc., most of which yielded new state-of-the-art results.
Revised: Structbert (Wang et al., 2019) etc., most of which yielded new state-of-the-art results.

Original: 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.
Revised: 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.

Original: Hence, training these models is quite time-consuming and requires a lot of computing and storage resources.
Revised: Hence, training these models is quite time-consuming and requires a lot of computing and storage resources.

Original: Even training a BERT-Base model costs at least $7 k (Strubell et al., 2019), let alone the other larger models, such as BERT-Large.
Revised: Even training a BERT-Base model costs at least $7 k (Strubell et al., 2019), let alone the other larger models, such as BERT-Large.

Original: Such a high cost is not affordable for many researchers and institutions.
Revised: Such a high cost is not affordable for many researchers and institutions.

Original: Therefore, improving the training efﬁciency should be a critical issue to make BERT more practical.
Revised: Therefore, improving the training efﬁciency should be a critical issue to make BERT more practical.

Original: Some pioneering attempts have been made to accelerate the training of BERT.
Revised: Some pioneering attempts have been made to accelerate the training of BERT.

Original: You et al.
Revised: You et al.

Original: (2019) proposed a layerwise adaptive large batch optimization method (LAMB), which is able to train a
Revised: (2019) proposed a layerwise adaptive large batch optimization method (LAMB), which is able to train a

Original: BERT model in 76 minutes.
Revised: BERT model in 76 minutes.

Original: However, the tens of times speedup is based on the huge amount of computing and storage resources, which is unavailable for common users.
Revised: However, the tens of times speedup is based on the huge amount of computing and storage resources, which is unavailable for common users.

Original: Lan et al.
Revised: Lan et al.

Original: (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.
Revised: (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.

Original: Gong et al.
Revised: Gong et al.

Original: (2019) proposed a progressively stacking method, which trains a deep BERT network by progressively stacking from a shallow one.
Revised: (2019) proposed a progressively stacking method, which trains a deep BERT network by progressively stacking from a shallow one.

Original: Utilizing the similarity of the attention distributions across different layers, such a strategy achieves about 25% speedup without signiﬁcant performance loss.
Revised: Utilizing the similarity of the attention distributions across different layers, such a strategy achieves about 25% speedup without signiﬁcant performance loss.

Original: Progressively stacking provides a novel training strategy, namely training a BERT model from shallow to deep.
Revised: Progressively stacking provides a novel training strategy, namely training a BERT model from shallow to deep.

Original: However, progressively stacking only has a high training efﬁciency at the initial stage in which the model depth is small.
Revised: However, progressively stacking only has a high training efﬁciency at the initial stage in which the model depth is small.

Original: As the training goes on, the model depth increases and the training speed decreases.
Revised: As the training goes on, the model depth increases and the training speed decreases.

Original: The low efﬁciency of the later stages makes the overall speedup of progressively stacking limited.
Revised: The low efﬁciency of the later stages makes the overall speedup of progressively stacking limited.

Original: Note that in the progressively stacking method, the bottom layers are trained with longer time than the top layers.
Revised: Note that in the progressively stacking method, the bottom layers are trained with longer time than the top layers.

Original: However, we observe that though the bottom layers are updated all the time, they do not have signiﬁcant changes in the later stages, in terms of the attention distribution which can reﬂect the functionality of the encoder layers to some extent (Gong et al., 2019).
Revised: However, we observe that though the bottom layers are updated all the time, they do not have signiﬁcant changes in the later stages, in terms of the attention distribution which can reﬂect the functionality of the encoder layers to some extent (Gong et al., 2019).

Original: In other words, most optimization of the bottom layers has been ﬁnished in the early stage when the model is shallow.
Revised: In other words, most optimization of the bottom layers has been ﬁnished in the early stage when the model is shallow.

Original: Motivated by this observation, in this work, we propose a novel multi-stage layerwise training (MSLT) approach, which can greatly improve the training efﬁciency of BERT.
Revised: Motivated by this observation, in this work, we propose a novel multi-stage layerwise training (MSLT) approach, which can greatly improve the training efﬁciency of BERT.

Original: We decompose the training process of BERT into several stages, as shown in Fig.
Revised: We decompose the training process of BERT into several stages, as shown in Fig.

Original: We start the training from a small BERT model with only a few encoder layers and gradually add new encoder layers.
Revised: We start the training from a small BERT model with only a few encoder layers and gradually add new encoder layers.

Original: At each stage (except the ﬁrst 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 ﬁxed in the current stage.
Revised: At each stage (except the ﬁrst 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 ﬁxed in the current stage.

Original: After all the encoder layers are trained, to make the network better behaved, we further retrain the model by updating all the layers together.
Revised: After all the encoder layers are trained, to make the network better behaved, we further retrain the model by updating all the layers together.

Original: Since the whole model has already been well trained, this stage only requires a few steps (accounting for about 20% of the total steps).
Revised: Since the whole model has already been well trained, this stage only requires a few steps (accounting for about 20% of the total steps).

Original: 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-efﬁcient.
Revised: 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-efﬁcient.

Original: Experimental results demonstrate the effectiveness and efﬁciency 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.
Revised: Experimental results demonstrate the effectiveness and efﬁciency 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.

Original: According to the results, the proposed method achieves more than 110% speedup without signiﬁcant performance degradation.
Revised: According to the results, the proposed method achieves more than 110% speedup without signiﬁcant performance degradation.

Original: 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;
Revised: 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;

Original: Sanh et al., 2019) are designed for network speedup in the inference phase.
Revised: Sanh et al., 2019) are designed for network speedup in the inference phase.

Original: Namely, these methods are used after the model has been trained.
Revised: Namely, these methods are used after the model has been trained.

Original: While in this paper, we focus on the model training speedup.
Revised: While in this paper, we focus on the model training speedup.

Original: The rest of the paper is organized as follows.

Original: In Section 2, we brieﬂy introduce the related works.

Original: In Section 3, we give the description of the multi-stage layerwise training method.

Original: In Section 4, we perform experiments to demonstrate the effectiveness and efﬁciency of the proposed method.

Original: Finally, we draw the conclusion in Section 5.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: 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.
Revised: 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.

Original: Following BERT, many pre-trained models were proposed, such as RoBERTa (Liu et al., 2019b),
Revised: Following BERT, many pre-trained models were proposed, such as RoBERTa (Liu et al., 2019b),

Original: XLNet (Yang et al., 2019), KBERT (Liu et al., 2019a), StructBERT (Wang et al., 2019), and so on.
Revised: XLNet (Yang et al., 2019), KBERT (Liu et al., 2019a), StructBERT (Wang et al., 2019), and so on.

Original: Higher accuracy were achieved by these models with more training data, more training steps, or more effective loss functions.
Revised: Higher accuracy were achieved by these models with more training data, more training steps, or more effective loss functions.

Original: 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).
Revised: 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).

Original: Hence, training a BERT model is challenging in terms of both computation and storage.
Revised: Hence, training a BERT model is challenging in terms of both computation and storage.

Original: In the literature, some approaches were proposed to improve the training speed of BERT.
Revised: In the literature, some approaches were proposed to improve the training speed of BERT.

Original: 2. D ISTRIBUTED T RAINING WITH L ARGE B ATCH S IZE
Revised: 2. D ISTRIBUTED T RAINING WITH L ARGE B ATCH S IZE

Original: 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.
Revised: 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.

Original: However, traditional stochastic gradient descent (SGD) based optimization methods perform poorly in large mini-batches training.
Revised: However, traditional stochastic gradient descent (SGD) based optimization methods perform poorly in large mini-batches training.

Original: Naively increasing the batch size leads to performance degradation and computational beneﬁts reduction (You et al., 2019).
Revised: Naively increasing the batch size leads to performance degradation and computational beneﬁts reduction (You et al., 2019).

Original: An efﬁcient layerwise adaptive large batch optimization technique named LAMB was proposed in
Revised: An efﬁcient layerwise adaptive large batch optimization technique named LAMB was proposed in

Original: You et al.
Revised: You et al.

Original: (2019) to address this problem.
Revised: (2019) to address this problem.

Original: It allows the BERT model to be trained with extremely large batch size without any performance degradation.
Revised: It allows the BERT model to be trained with extremely large batch size without any performance degradation.

Original: By using 1024 TPUv3 chips, LAMB reduced the BERT training time from 3 days to 76 minutes.
Revised: By using 1024 TPUv3 chips, LAMB reduced the BERT training time from 3 days to 76 minutes.

Original: With a similar manner, by using 1472 V

Original: GPUs, NVIDIA also trained the BERT model in only 53 minutes.

Original: 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.
Revised: 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.

Original: 2. ALBERT
Revised: 2. ALBERT

Original: ALBERT (Lan et al., 2019) adopted two parameter reduction techniques, namely factorized embedding parameterization and cross-layer parameter sharing, which signiﬁcantly reduced the model size.
Revised: ALBERT (Lan et al., 2019) adopted two parameter reduction techniques, namely factorized embedding parameterization and cross-layer parameter sharing, which signiﬁcantly reduced the model size.

Original: 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.
Revised: 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.

Original: Since the communication overhead is directly proportional to the number of parameters in the model, ALBERT also improved the communication efﬁciency in distributed training setting.
Revised: Since the communication overhead is directly proportional to the number of parameters in the model, ALBERT also improved the communication efﬁciency in distributed training setting.

Original: However, since ALBERT has almost the same computational complexity as BERT, training an ALBERT model is still very time-consuming.
Revised: However, since ALBERT has almost the same computational complexity as BERT, training an ALBERT model is still very time-consuming.

Original: 2. P ROGRESSIVELY S TACKING
Revised: 2. P ROGRESSIVELY S TACKING

Original: The most related work should be progressively stacking (Gong et al., 2019). This method 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.
Revised: 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.

Original: Such a phenomenon implies that the encoder layers in the BERT model have similar functionalities.
Revised: Such a phenomenon implies that the encoder layers in the BERT model have similar functionalities.

Original: Utilizing the natural similarity characteristic, to train a N − layer
Revised: Utilizing the natural similarity characteristic, to train a N − layer

Original: BERT model, the progressively stacking method ﬁrst trains a N/ 2 − layer model, and then sticks it into N − layer by copying the parameters of the trained N/ 2 layers.
Revised: BERT model, the progressively stacking method ﬁrst trains a N/ 2 − layer model, and then sticks it into N − layer by copying the parameters of the trained N/ 2 layers.

Original: After the N − layer model is constructed, the progressively stacking method continues to train the whole model by updating all the parameters together.
Revised: After the N − layer model is constructed, the progressively stacking method continues to train the whole model by updating all the parameters together.

Original: By repeatedly using such a strategy, the deep BERT model can be trained more efﬁciently.
Revised: By repeatedly using such a strategy, the deep BERT model can be trained more efﬁciently.

Original: According to the results shown in Gong et al.
Revised: According to the results shown in Gong et al.

Original: (2019), progressively stacking can achieve the training time about 25% shorter than the original training method (Devlin et al., 2018).
Revised: (2019), progressively stacking can achieve the training time about 25% shorter than the original training method (Devlin et al., 2018).

Original: However, we can see that the speedup of progressively stacking mainly comes from the initial stage in which the model depth is small.
Revised: However, we can see that the speedup of progressively stacking mainly comes from the initial stage in which the model depth is small.

Original: As the model depth increases in the later stages, the training efﬁciency also decreases, and according to Gong et al.
Revised: As the model depth increases in the later stages, the training efﬁciency also decreases, and according to Gong et al.

Original: (2019), to guarantee the performance, more training steps should be assigned in the later stages for training the deep model.
Revised: (2019), to guarantee the performance, more training steps should be assigned in the later stages for training the deep model.

Original: Hence, the overall speedup brought by progressively stacking is limited.
Revised: Hence, the overall speedup brought by progressively stacking is limited.

Original: Such an issue is addressed in this paper.
Revised: Such an issue is addressed in this paper.

Original: In our work, we also train the BERT model from shallow to deep.
Revised: In our work, we also train the BERT model from shallow to deep.

Original: In contrast , at each stage, we only train the top few layers and we almost keep a high training efﬁciency during the whole training process.
Revised: In contrast , at each stage, we only train the top few layers and we almost keep a high training efﬁciency during the whole training process.

Original: Hence, much more signiﬁcant speedup is achieved.
Revised: Hence, much more signiﬁcant speedup is achieved.

Original:  M ETHODOLOGY
Revised:  M ETHODOLOGY

Original: In this section, we propose an efﬁcient training method to accelerate the training of BERT model.
Revised: In this section, we propose an efﬁcient training method to accelerate the training of BERT model.

Original: 3. M OTIVATION
Revised: 3. M OTIVATION

Original: The large depth should be one of the main reasons making the BERT training time-consuming.
Revised: The large depth should be one of the main reasons making the BERT training time-consuming.

Original: The original method (Devlin et al., 2018) trains all the encoder layers simultaneously.
Revised: The original method (Devlin et al., 2018) trains all the encoder layers simultaneously.

Original: 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 inefﬁcient, especially when the model is very deep.
Revised: 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 inefﬁcient, especially when the model is very deep.

Original: Inspired by progressively stacking (Gong et al., 2019), we also consider to train the BERT model from shallow to deep.
Revised: Inspired by progressively stacking (Gong et al., 2019), we also consider to train the BERT model from shallow to deep.

Original: The main problem of progressively stacking is that its training efﬁciency decreases as the training goes on.
Revised: The main problem of progressively stacking is that its training efﬁciency decreases as the training goes on.

Original: We observe that in the progressively stacking strategy, the bottom layers are trained for longer time than the top layers.
Revised: We observe that in the progressively stacking strategy, the bottom layers are trained for longer time than the top layers.

Original: For example, the ﬁrst 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.
Revised: For example, the ﬁrst 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.

Original: 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 signiﬁcant role (Khetan & Karnin, 2020).
Revised: 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 signiﬁcant role (Khetan & Karnin, 2020).

Original: In BERT model, the encoder layers are mainly used to learn the dependencies of the input elements, which can be reﬂected by the attention distribution.
Revised: In BERT model, the encoder layers are mainly used to learn the dependencies of the input elements, which can be reﬂected by the attention distribution.

Original: In Gong et al.
Revised: In Gong et al.

Original: , 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 ﬁrst CLS token.
Revised: , 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 ﬁrst CLS token.

Original: In addition, the authors also observed that the attention distribution of the top layers is very similar to that of the bottom layers.
Revised: In addition, the authors also observed that the attention distribution of the top layers is very similar to that of the bottom layers.

Original: Using a similar way, we also visualize some important attention distributions and we get some new ﬁndings when using the progressively stacking method to train a 12-layer BERT-Base model.
Revised: Using a similar way, we also visualize some important attention distributions and we get some new ﬁndings when using the progressively stacking method to train a 12-layer BERT-Base model.

Original: Speciﬁcally, we ﬁrst train a 6-layer model.
Revised: Speciﬁcally, we ﬁrst train a 6-layer model.

Original: Then we stack the trained 6-layer model into a 12-layer model and continue to train the whole model until convergence.
Revised: Then we stack the trained 6-layer model into a 12-layer model and continue to train the whole model until convergence.

Original: The attention distributions of the topencoder layers of the ﬁnal 12-layer BERT-Base model are shown in the ﬁrst row of Fig.
Revised: The attention distributions of the topencoder layers of the ﬁnal 12-layer BERT-Base model are shown in the ﬁrst row of Fig.

Original: For each layer, we randomly choose an attention head.
Revised: For each layer, we randomly choose an attention head.

Original: Then we also show the attention distributions of the corresponding heads of the bottom 6 encoder layers in the second row of Fig.
Revised: Then we also show the attention distributions of the corresponding heads of the bottom 6 encoder layers in the second row of Fig.

Original: Further, we show the attention distributions of the corresponding heads of the trained 6-layer BERT model before stacking in the third row.
Revised: Further, we show the attention distributions of the corresponding heads of the trained 6-layer BERT model before stacking in the third row.

Revised: 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.

Revised: The forth row of Fig.shows the attention distributions of the bottom 6 encoder layers of the original BERT-Base model.

Original: Combined with Fig.
Revised: Combined with Fig.

Original: 2, we ﬁnd that:1.
Revised: 2, we ﬁnd that:

Original: Except the two obvious distributions found by Gong et al.
Revised: Except the two obvious distributions found by Gong et al.

Original: (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.
Revised: (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.

Original: 2 corresponds to the position of SEP).
Revised: 2 corresponds to the position of SEP).

Original: Compared with the ﬁrstrow and second row of Fig. 3, one can see that for a trained 12layer 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.
Revised: Compared with the ﬁrst 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.

Original: In addition, there are also some bottom-top layer pairs whose attention distributions are very different.
Revised: In addition, there are also some bottom-top layer pairs whose attention distributions are very different.

Original: 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 ﬁnal 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.
Revised: 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 ﬁnal 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.

Original: The performance of the bottom encoder layers are not further improved, in terms of catching the elements’ dependencies.
Revised: The performance of the bottom encoder layers are not further improved, in terms of catching the elements’ dependencies.

Revised: 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

Original: 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.
Revised: 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.

Original: Therefore, an intuitive idea is that at each stage, let the bottom layers[Table] 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: So the time of both computation and communication can be greatly reduced.
Revised: So the time of both computation and communication can be greatly reduced.

Original: 3. M ULTI - STAGE L AYERWISE T RAINING
Revised: 3. M ULTI - STAGE L AYERWISE T RAINING

Original: Now we present our MSLT training method.
Revised: Now we present our MSLT training method.

Original: Based on the above observations, to train a deep BERT model with N encoder layers, we decompose the whole training process into k stages, as shown in
Revised: Based on the above observations, to train a deep BERT model with N encoder layers, we decompose the whole training process into k stages, as shown in

Original: Fig.
Revised: Fig.

Original: We start our training from a 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).
Revised: We start our training from a 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).

Original: At each time, we only add N/k new layers.
Revised: At each time, we only add N/k new layers.

Original: Except the ﬁrst 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.
Revised: Except the ﬁrst 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.

Original: 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.
Revised: 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.

Original: In addition, in the BERT model, the word embedding matrices in the input and output layers are shared.
Revised: In addition, in the BERT model, the word embedding matrices in the input and output layers are shared.

Original: So we only update the word embedding matrix in the ﬁrst stage, and in the later stages the word embedding matrix is ﬁxed.
Revised: So we only update the word embedding matrix in the ﬁrst stage, and in the later stages the word embedding matrix is ﬁxed.

Original: Similar to Gong et al.
Revised: Similar to Gong et al.

Original: (2019), 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.
Revised: (2019), 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.

Original: To make the ﬁnal model better behaved, after each layer is trained, we further retrain the model by updating all the parameters together for a few steps.
Revised: To make the ﬁnal model better behaved, after each layer is trained, we further retrain the model by updating all the parameters together for a few steps.

Original: In the BERT model training, backward computation generally takes much longer time than the forward computation, especially in the distributed training setting in which the backward computation also includes gradient synchronization.
Revised: In the BERT model training, backward computation generally takes much longer time than the forward computation, especially in the distributed training setting in which the backward computation also includes gradient synchronization.

Original: 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 (including communication time for gradient synchronization in the distributed setting) is almost six times as the time of forward computation.
Revised: 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 (including communication time for gradient synchronization in the distributed setting) is almost six times as the time of forward computation.

Original: The proposed MSLT method can greatly reduce the backward computation, since only top few layers participate in the backward computation in the whole training process, except the ﬁnal retraining stage.
Revised: The proposed MSLT method can greatly reduce the backward computation, since only top few layers participate in the backward computation in the whole training process, except the ﬁnal retraining stage.

Original: Hence, the total training time will be much shorter.
Revised: Hence, the total training time will be much shorter.

Original:  E XPERIMENT 4. E XPERIMENTAL S ETUP
Revised: 3. E XTENDING TO ALBERT

Revised: The proposed MSLT strategy can also be used to speed up the training of ALBERT Lan et al. 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 ﬁxed. Hence, before applying

Revised: MSLT, we make a slight modiﬁcation of ALBERT.

Revised: We decompose the whole encoder layers into k groups and we share the parameters of the encoder layers in the same group.

Revised: Then we use a similar strategy as shown in last subsection to efﬁciently train the ALBERT model.

Revised: Speciﬁcally, at each stage, we add a new group of encoder layers and only the newly added group of layers will be updated in this stage.

Revised: Compared with the original ALBERT model, the modiﬁed ALBERT model requires more storage resources since it involves more parameters.

Revised: However, the training speed is signiﬁcantly improved.

Revised: In real applications, one should decompose the encoder layers into a suitable number of groups to achieve a computation-storage efﬁciency balance, according to actual demand.

Revised:  E XPERIMENT 4. E XPERIMENTAL S ETUP

Original: In this section, we perform experiments to demonstrate the effectiveness and efﬁciency of the propose method.
Revised: In this section, we perform experiments to demonstrate the effectiveness and efﬁciency of the propose method.

Original: Following the setup in Devlin et al.
Revised: Following the setup in Devlin et al.

Original: (2018), we use English Wikipedia (2,500M words) and BookCorpus (800M words) for pre-training.
Revised: (2018), we use English Wikipedia (2,500M words) and BookCorpus (800M words) for pre-training.

Original: Though some more effective objectives were proposed in Lan et al.
Revised: Though some more effective objectives were proposed in Lan et al.

Original: ; Steinley (2006), to make the comparison as meaningful as 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).
Revised: ; Steinley (2006), to make the comparison as meaningful as 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).

Original: We set the maximum length of each input sequence to be 128.
Revised: We set the maximum length of each input sequence to be 128.

Original: 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.
Revised: 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.

Original: The models considered in this section are BERT-Base and BERT-Large, whose hyperparameters setting can be seen in Table 1.
Revised: The models considered in this section are BERT-Base and BERT-Large, whose hyperparameters setting can be seen in Devlin et al.

Original: For each model, we decompose the training process into k = 4 stages.
Revised: For each model, we decompose the training process into k = 4 stages.

Original: So at each stage, we train 3 encoder layers for BERT-Base, and 6 encoder layers for BERT-Large.
Revised: So at each stage, we train 3 encoder layers for BERT-Base, and 6 encoder layers for BERT-Large.

Original: All the experiments are performed on a distributed computing cluster consisting of32 Telsa V
Revised: All the experiments are performed on a distributed computing cluster consisting of

Original:  GPU cards with 32G memory, and the batch size is 32*32=1024.
Revised: Telsa V100 GPU cards with 32G memory, and the batch size is 32*32=1024.

Original: We use the LAMB optimizer with learning rate 0.00088 (You et al., 2019) .
Revised: We use the LAMB optimizer with learning rate 0.00088 (You et al., 2019) .

Original: All the other settings are the same as Devlin et al.
Revised: All the other settings are the same as Devlin et al.

Original: (2018), including the data pre-processing, unless otherwise speciﬁed.[Table]
Revised: (2018), including the data pre-processing, unless otherwise speciﬁed.

Original: 4. C OMPARISON FOR T RAINING WITH S AME S TEPS
Revised: 4. C OMPARISON FOR T RAINING WITH S AME S TEPS

Original: We train both BERT-Base and BERT-Large models using the MSLT method for totally 1,000, steps.
Revised: We train both BERT-Base and BERT-Large models using the MSLT method for totally 1,000, steps.

Original: Speciﬁcally, each stage uses 200,000 steps, and the ﬁnal model is retrained for the remaining 200,000 steps.
Revised: Speciﬁcally, each stage uses 200,000 steps, and the ﬁnal model is retrained for the remaining 200,000 steps.

Original: Wealso train a BERT-Base and a BERT-Large model with 1,000,000 steps from scratch using the original method (Devlin et al., 2018) as the baselines.
Revised: We train a BERT-Base and a BERT-Large model with 1,000,000 steps from scratch using the original method (Devlin et al., 2018) as the baselines.

Original: For each model, the ﬁrst 10,000 training steps are used for learning rate warmup.
Revised: For each model, the ﬁrst 10, training steps are used for learning rate warmup.

Original: We ﬁrst show the pre-training loss of all the models in Fig.
Revised: We ﬁrst show the pre-training loss of all the models in Fig.

Original: We can see that for both BERT-Base and BERT-Large, the loss of our method decreases faster than the baseline, and the ﬁnal convergence value of our method is close to that of the baseline.
Revised: We can see that for both BERT-Base and BERT-Large, the loss of our method decreases faster than the baseline, and the ﬁnal convergence value of our method is close to that of the baseline.

Original:  Finally, our method achieves more than 110% speedup (saves about 55% training time), which is a signiﬁcant improvement compared with the 25% speedup achieved by progressively stacking (Gong et al., 2019).
Revised: [Table] Finally, our method achieves more than 110% speedup (saves about 55% training time), which is a signiﬁcant improvement compared with the 25% speedup achieved by progressively stacking (Gong et al., 2019).

Original: Then we further evaluate the above models on the widely used General Language Understanding
Revised: Then we further evaluate the above models on the widely used General Language Understanding

Original: Evaluation (GLUE (Wang et al., 2018)) benchmark, which consists of nine natural language understanding tasks, namely Corpus of Linguistic Acceptability (CoLA (Warstadt et al., 2019)), Microsoft Research Paraphrase Corpus (MRPC (Dolan & Brockett, 2005)), Semantic Textual Similarity Benchmark (STS (Cer et al., 2017)), Quora Question Pairs (QQP (Iyer et al., 2017)),
Revised: Evaluation (GLUE (Wang et al., 2018)) benchmark and the Stanford Question Answering Dataset (SQuAD1.1). The sequence length of the SQuAD1.1 task is 384. So for the SQuAD1.1 task, we train the last 10% of the steps using the sequence of 512 to learn the positional embeddings.

Original: Question NLI (QNLI (Rajpurkar et al., 2016)), Multi-Genre NLI (MNLI (Williams et al., 2017)), Recognizing Textual Entailment (RTE (Bentivogli et al., 2009)), Stanford Sentiment Treebank (SST (Socher et al., 2013)), Winograd NLI (WNLI (Levesque et al., 2012)).
Revised: To better show the advantage of the proposed method, we also add the progressively stacking method

Original: Following Devlin et al.
Revised: Gong et al.

Original: (2018), the dataset WNLI is excluded in our experiment.
Revised: as the baseline.

Original: Table 2 shows the Dev set results of the GLUE tasks.
Revised: Table 1 shows the Dev set results on SQuAD1.1 and selected GLUE tasks.

Original: Similar to Liu et al.
Revised: Similar to Liu et al.

Original: (2019b), all the results are the median of ﬁve runs.
Revised: (2019b), all the results are the median of ﬁve runs.

Original: For each  task, we ﬁne-tune the model with batch size 

Revised: For each GLUE task, we ﬁne-tune the model with batch size

Original: and we perform a grid search on the learning rate set
Revised: and we perform a grid search on the learning rate set

Original: [5e-5, 4e-5, 3e-5, 2e-5], which is the same as Devlin et al.
Revised: [5e-5, 4e-5, 3e-5, 2e-5].

Original: Following Clark et al.
Revised: Following Clark et al.

Original: (2020), we ﬁne-tune the model with 10 epochs for STS and RTE, and 3 epochs for all the other tasks.
Revised: (2020), we ﬁne-tune the model with 10 epochs for STS and RTE, and 3 epochs for all the other tasks.

Revised: For SQuAD1.1 task, we ﬁne-tune the model with batch size 12, epoch 2, and learning rate 3e-5.

Original: From Table 2, we can see that for both BERT-Base and BERT-Large, the results of our method are close to those of the original BERT.
Revised: From Table 1, we can see that for both BERT-Base and BERT-Large, the results of our method are close to those of the original BERT.

Original: In addition, the results of the baselines are comparable to those shown in Devlin et al.
Revised: In addition, the results of the baselines are comparable to those shown in Devlin et al.

Original: (2018), which conﬁrms the validity of our proposed method.[Table] [Table]
Revised: (2018), which conﬁrms the validity of our proposed method.

Original:  C OMPARISON FOR T RAINING WITH S AME T IME
Revised: 4. C OMPARISON FOR T RAINING WITH S AME T IME

Original: In this section [Table] we compare the performance of all the models trained with the same time.
Revised: In this section, we compare the performance of all the models trained with the same time.

Original: Training a BERT-Base model with 1M steps using MLST method requires about 40 hours, which is similar to the time of training a BERT-Base model using the original method for 480,000 steps (abouthours).
Revised: Training a BERT-Base model with 1M steps using MLST method requires about 40 hours, which is similar to the time of training a BERT-Base model using the original method for 480,000 steps (abouthours).

Original: 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.
Revised: 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.

Original: The results  are shown in Table 3. Only 6 GLUE tasks are selected due to space limitation.
Revised: The results on selected GLUE tasks are shown in Table 2.

Original: We see that with the same training time, the models trained by our method perform much better than the baselines. 
Revised: We see that with the same training time, the models trained by our method perform much better than the baselines. [Table] [Table]

Original: 4. O NLINE T EST R ESULTS ON GLUE
Revised: A CCELERATE T RAINING OF ALBERT

Revised: As shown in Section 3. [Table] the proposed MSLT method can also be used to speed up the training of

Revised: ALBERT.

Revised: Table 3 reports the results on select GLUE tasks of the original ALBERT-Base model and the modiﬁed ALBERT-Base model trained by MSLT.

Revised: All the models are trained for 1M steps and the other settings are the same as Section 4.2.

Revised: We can see that the modiﬁed ALBERT model trained by

Revised: MSLT achieves better performance than the original BERT.

Revised: That is because the modiﬁed ALBERT model involves more parameters.

Revised: Hence, the original ALBERT has higher memory efﬁciency while the modiﬁed ALBERT trained using MSLT has higher time efﬁciency and better performance.

Revised: 4. E FFECT OF J OINTLY R ETRAINING

Revised: In the above examples, we left 200,000 steps for jointly retraining all the layers.

Revised: Here we investigate the impact of the retraining stage.

Revised: Table 4 shows the results of BERT models with/without retraining.

Revised: The results implies that the retraining stage can further improve the performance of the BERT model.

Revised: The training efﬁciency of the retraining stage is much lower than the previous stages, since in this stage all the parameters are updated.

Revised: However, since the model is already near-optimal after training by previous stage, the retraining stage only requires a few steps.

Revised: In practice, we use 10% ∼ 20% of the total steps for retraining, so the retraining stage will not make the whole training very timeconsuming.

Original: Lastly, we report the online test results on GLUE tasks (except WNLI) of all the models after ﬁnetuning, which are shown in Table 4.
Revised: Lastly, we report the online test results on GLUE tasks (except WNLI) of all the models after ﬁnetuning, which are shown in Table 5.

Original: All the submitted models are pre-trained with 1,000,000 steps.
Revised: All the submitted models are pre-trained with 1,000,000 steps.

Original: Same as Devlin et al.
Revised: Same as Devlin et al.

Original: (2018), for each task, we select the best ﬁne-tuning learning rate on the Dev set.
Revised: (2018), for each task, we select the best ﬁne-tuning learning rate on the Dev set.

Original:  C ONCLUSION
Revised:  D ISCUSSION

Revised: In last section, we empirically showed the effectiveness and efﬁciency of the proposed MSLT method.

Revised: In this section we discuss how does it work and why it can improve the convergence speed.

Revised: 5. R ELATIONSHIP B ETWEEN MSLT AND A LTERNATING O PTIMIZATION

Revised: At ﬁrst glance, MSLT is an 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.

Revised: Such a strategy is similar to alternating optimization Bezdek & Hathaway (2003), which is widely used for solving multi-variable nonconvex optimization problem, due to its simple

Revised: [Table]

Revised: [Table]

Revised: implementation, fast convergence, and superb empirical performance Li et al.

Revised: [Table]

Revised: Alternating optimization addresses the multi-variable optimization problem by updating a small set of variables and keeping the other variables ﬁxed.

Revised: 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).

Revised: According to alternating optimization, in the retraining stage, we should update the bottom layers again and keep the other layers ﬁxed.

Revised: However, when the model is stacked deep, to compute the gradient of the parameters of the bottom layers, we have to compute the gradient of the parameters of the top layers, according to the chain rule.

Revised: In this situation, updating the parameters of the bottom layers is almost as expensive as updating all the parameters.

Revised: So we update all the parameters simultaneously in the retraining stage.

Revised: Overall, in the MSLT method, we utilize the advantage of alternating optimization to quickly reach a near-optimal state.

Revised: When the model is deep enough, alternating optimization does not have efﬁciency advantage, then we use joint descent to sufﬁciently utilize the gradient information.

Revised: 5. A VOID C ONTRADICTION OF C HOICE OF L EARNING R ATE

Revised: In Section 2.3, we have shown that, though the bottom layers are updated for most steps in progressively stacking method, they do not have signiﬁcant change in the later stage.

Revised: So it is not worthwhile to spend so much the computation to updating the parameters of bottom layers.

Revised: Another issue is that for the adaptive optimizer, such as Adam or LAMB, when the variable is close to the optimal value, we should use a small learning rate, and when the variable is far from optimal, we should use a large learning rate.

Revised: In progressively stacking strategy, the bottom layers are near-optimal while the newly added top layers are far from optimal since they are not trained.

Revised: So they have completely different requirements for the learning rate.

Revised: The MSLT method address the learning rate contradiction by ﬁrst optimizing the untrained top layers and keep the trained bottom layers ﬁxed.

Revised: When all the layers are near-optimal, we further retrain all the layers with a small learning rate.

Original:  In this paper, we propose an efﬁcient multi-stage layerwise training method for accelerating the training process of BERT.
Revised: Meanwhile, such a manner saves a lot of computation.  C ONCLUSION In this paper, we propose an efﬁcient multi-stage layerwise training method for accelerating the training process of BERT.

Original: 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.
Revised: 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.

Original: We ﬁnd that the attention distributions of the bottom layers tend to be ﬁxed early, and further training in the later stages does not bring signiﬁcant changes.
Revised: We ﬁnd that the attention distributions of the bottom layers tend to be ﬁxed early, and further training in the later stages does not bring signiﬁcant changes.

Original: 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.
Revised: 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.

Original: Experimental results show that the proposed training method achieves signiﬁcant speedup without signiﬁcant performance loss, compared with the existing training method.
Revised: Experimental results show that the proposed training method achieves signiﬁcant speedup without signiﬁcant performance loss, compared with the existing training method.