Speeding up Deep Learning Training by Sharing Weights and Then Unsharing
Authors: Shuo Yang
Le Hou
Xiaodan Song
qiang liu
Denny Zhou
Venue:  (2021)
Keywords: fast training
BERT
transformer
weight sharing
deep learning

Abstract: It has been widely observed that increasing deep learning model sizes often leads to significant performance improvements on a variety of natural language processing and computer vision tasks. In the meantime, however, computational costs and training time would dramatically increase when models get larger. In this paper, we propose a simple approach to speed up training for a particular kind of deep networks which contain repeated structures,  such as the transformer module. In our method, we first train such a deep network with the weights shared across all the repeated layers till some point. We then stop weight sharing and continue training until convergence. The untying point is automatically determined by monitoring gradient statistics. Our adaptive untying criterion is obtained from a theoretic analysis over deep linear networks.  Empirical results show that our method is able to reduce the training time of BERT  by 50%. 

Under review as a conference paper at ICLR 2021
SPEEDING UP DEEP LEARNING TRAINING BY SHAR- ING WEIGHTS AND THEN UNSHARING
Anonymous authors Paper under double-blind review
ABSTRACT
It has been widely observed that increasing deep learning model sizes often leads to significant performance improvements on a variety of natural language process- ing and computer vision tasks. In the meantime, however, computational costs and training time would dramatically increase when models get larger. In this paper, we propose a simple approach to speed up training for a particular kind of deep networks which contain repeated structures, such as the transformer module. In our method, we first train such a deep network with the weights shared across all the repeated layers till some point. We then stop weight sharing and con- tinue training until convergence. The untying point automatically determined by monitoring gradient statistics. Our adaptive untying criterion is obtained from a theoretic analysis over deep linear networks. Empirical results show that our method is able to reduce the training time of BERT by 50%.
INTRODUCTION
has been widely observed that increasing model size often leads to significantly better perfor- mance on various real tasks, especially natural language processing and computer vision applica- tions (Amodei et al. 2016 He et al. 60 Wu et al. 2016 Vaswani et al. 2017 Devlin et al. 2019] Brock al. 2019 Raffel et al. 2020 Brown al.
2020). However, as models getting larger, the training can become extremely resource intensive and time consuming, As consequence, there has been growing interest developing systems and algorithms for efficient distributed large-batch
training (Goyal et al. 2017 Shazeer et al. 2018 Lepikhin et al. 2020 You et al. 2020).
In this paper, we seek for speeding up deep learning training by exploiting unique network architec- tures rather than by distributed training. In particular, we are interested in speeding up the training of a special kind of deep networks which are constructed by repeatedly stacking the same layer, for example, the transformer module (Vaswani al. 2017). We propose a simple method for efficiently training such kind of networks. In our approach, we first force the weights to be shared across all the repeated layers and train the network, and then, some point, stop weight sharing and con- tinue training until convergence. The point for stopping weight sharing can be either predefined or automatically chosen by monitoring gradient statistics during training. Empirical studies show that our method can reduce the training time of BERT (Devlin et al. 2019) by 50%.
Our method is motivated by the successes of weight sharing models, in particular, ALBERT (Lan) ctt al., (020). It is variant of BERT in which the weights across all the transformer layers are shared. As long as its architecture is sufficiently large, ALBERT can be comparable with or even outper- form the original BERT on various downstream natural language processing benchmarks. However, when its architecture being the same as the original BERT, ALBERT performs significantly worse. Since the weights in the original BERT are not shared at all, it is natural to expect that ALBERT's performance will be improved if we stop its weight sharing at some point of training.
To make this idea work, however, we need to know when to untie the shared weights. A randomly chosen untying point will not work. We can see this from the two extreme cases: ALBERT which shares weights all the time, and BERT which has weight sharing at all. To find an effective solution for automatic weight untying, we turn to theoretic analysis over deep linear networks (Hardt & Ma 2017: Laurent & Brecht, 2018 Wu al. 2019). A deep linear model is constructed by stacking a series of matrix multipiccatinns. In its forward pass, a deep linear model is trivially equivalent to a single matrix. However, when being trained with backpropagation, its behavior is

---

Under review as a conference paper at ICLR 2021
analogous to the deep models with non-linearity but much easier to understand. Our theoretical analysis shows that, when learning a positive definite matrix (which admits an optimal solution with all layers having the same weights), training with weight sharing can bring significantly faster convergence. More importantly, our theoretical analysis leads to the adaptive weight untying rule that need construct our algorithm (see Algorithm 2 Empirical studies real tasks show that our adaptive untying method can be at least as effective as using the best untying point which is obtained by running multiple experiments of which each has a different point to untic weights and then choosing the best result.
The rest of this paper is organized as follows. We present our weight sharing algorithm in Section 2 It actually contains three versions, depending on how to stop weight sharing during training. In Section 3 we present our theoretical results for positive definite deep linear models. All the proofs are deferred to the Appendix. In Section 4 we discuss related work. In Section we show detailed experimental setups and results. We also provide various ablation studies on different choices in implementing our algorithm. Finally, we conclude this paper with discussions in Section 6
ALGORITHM: SHARING WEIGHTS AND THEN UNSHARING
Assume we have a deep network which obtained by repeatedly stacking the same neural module times, such the transformer module in transformer models (Vaswani et al. 2017). Denote by W1, w1,.... wn, the weights of these n layers. In our method, we first train the deep network with all the weights tied. Then, after certain number of training steps, we untie the weights and further train the network until convergence. In what follows, we first present simple version of our algorithm in which the weight untying point is predefined. Then, we move its adaptive version in which the layers are automatically gradually untied according to gradient statistics. Finally, we discuss simplified variant of this adaptive method which unties the layers all the once.
Untying weights at fixed point. This is the simplest version of our method (Algorithm I We first train the deep network with all the weights tied for a fixed number of steps, and then untie the weights and continue training until convergence.
Algorithm SHARING WEIGHTS AND THEN UNTYING AT AT A FIXED POINT
1: Input: total number of training steps T, untying point т, learning rates (a(0) 1...,, ,
2: Randomly and equally initialize weights w(") ...,,
(0)
3: for 1. 1 to T do
4: if <<< then
(t-1) (1)) x mean grad (loss, w w u)) =1, I,...,, [ i= 11,...,
5: wi = wi
6: else
(t)
(t-1) - (1)) grad [osss, w/tt 111- ; I,...,
7: W. wi
Note that, from line 1 to 5, we initialize all the weights equally, and then update them using the mean of their gradients. It easy to see that such an update equivalent to weight sharing or tying. For the sake of simplicity, in lines 5 and 7, we only show how to update the weights using the plain (stochastic) gradient descent rule. One can replace this plain update rule with any of their favorite optimization methods, for example, the Adam optimization algorithm (Kingma & Ba 2015).
While the repeated layers being the most natural units for weight sharing, that not the only choice. We may view several layers together as the weight sharing unit, and share the weights across those units. The layers within the same unit can have different weights. For example, for a 24-layer transformer model, may combine every four layers weight sharing unit. Thus, there will be six such units for weight sharing. Such flexibility of choosing weight sharing units allows for a balance between "full weight sharing" and "no weight sharing" at all.
Adaptive weight untying. The theoretical analysis in Section 3 motives us to adaptively and gradu- ally untie weights based on the gradient correlation of adjacent layers (Algorithm 2). To implement this idea, all layers are put in the same group during initialization. Then, at any time step of the train- ing, suppose we have groups G (G1, G2, ..., G For each group G,, compute the correlation
2

---

Under review as a conference paper at ICLR 2021
between the gradients of any two adjacent layers by
(0)) -99-1
c (
9 19+1 denote the gradients of two adjacent layers in Gy, (1)
(2)
Once c falls below p (0.5 as default) consecutively for certain number of times (3 as default), we split the group G2 into two subgroups. by splitting at j-th layer in Gg. The layers in the same group are updated using the gradient mean of this group as Algorithm
Algorithm 2 ADAPTIVE WEIGHT UNTYING
1: Input: total number of training steps 7", threshold learning rates {a(*) = ,.... .
2: All layers are put in the same group, and randomly initialized with the same weights
3: for = 1 to T do
4: for adjacent layers i, + in the same group do
5:
Calculate the gradient correlation Equation [ / //Motivated by Section 6: the correlation falls below μ, break the group between layer and ++1
7: Update the weights in each group using the average of gradients within each group
Simplified adaptive untying. To simplify the multi-step adaptive untying, instead of gradually splitting the groups, we may untie the weights of all layers at once. For every certain number of iteration steps. we monitor the gradient correlations as described multi-step adaptive untying. more than half of these layer correlations are less than the predefined threshold p (0.5 as default) consecutively for a certain number of times (3 as default), we stop sharing weights for all layers.
THEORETIC MOTIVATION
In this section, show how our method including the adaptive untying criterion be motivated by analyzing the dynamics of training deep linear network by gradient descent. At the first glance, single matrix. However, when being trained with backpropagation, its behavior is analogous to generic deep mod- els. Thus, analyzing deep linear models has attracted increasing interest in the theoretic research community (Hardt & Ma 2017 Laurent & Brecht 2018 Wu et al. 2019).
deep learning models may look trivial since deep linear model is just equivalent
deep linear network is series of matrix multiplication
fx; W1, Wz) W.-....W
Wi E Rdddd L I,...,, L.
The task is to train the deep linear network to learn a target matrix E IRdxd, To focus on the training dynamics, we adopt the simplified objective function
- 11 Denote V,R to be the gradient of R with respect to W.. We have
R(W,,... = FI
- W
OR
VRR ow,
WT. (WL,,WW, h)
where Wiss, = Wi, Wt - Wh+ Wi, The standard gradient update is given by
Wi(t + 1) Wi(t) R(),,, 1= L. To train with weights shared, all the layers need to have the same initialization. And the update is
Wi(t 1) wi(t)
" L V.R(),
1=
Since the initialization and updates the same for all layers, the parameters Wi W1),).... (t) equal for all t. For simplicity, we denote the weight at time t to be Wo(t). Notice that the gradients are averaged, the norm of update each layer doesn't scale with L.
Suppose that the target matrix O is positive definite matrix. It is immediate that & //n is solution to the deep linear network. Before looking into the detailed convergence analysis, worth first showing Lemma that reveals the updates in the weight sharing training.
3

---

Under review as a conference paper at ICLR 2021
Lemma 1.
With positive definite target matrix + and initializing with Wo(0) I, update the parameters according to Equation 2
we have
Wi(t + - Wi(() (W) Φ) 1= L, 42
Intuitively, the Lemma [hoows that weight sharing allows all the layers be trained "equally well", the layers that are far away from the output layer won't suffer from gradient vanishing or exploding.
In the following subsections, we first study the convergence result with continuous-time gradient descent, which demonstrates the benefit of training with weight sharing when learning positive definite matrix $ We then extend the results to the discrete-time gradient descent. We draw comparison with training with zero-asymmetric (ZAS) initialization (Wu et al. (019). To the best of our knowledge, ZAS gives the state-of-the-art convergence rate. is actually the only work showing the global convergence of deep linear network trained by gradient descent for an arbitrary target matrix.
3.1 CONTINUOUS-TIME GRADIENT DESCENT
With continuous- ciitiuuuu--tmeeeeeee gradient descent (i.e. 0), training with gradient descent and ZAS, the loss decays as R(t)= S exp( (2)) For training with weight sharing, the loss becomes R(t) s exp(-2T. min(1, Amin a (t) (()), when the target matrix + is positive definite. The extra L in the exponent demonstrates the acceleration of training with weight sharing.
With - the training dynamics of continuous time gradient descent be described as
dWi(t) wi(()) = ViR((), 1, L, di
The ZAS initializes the weights W1 = W22 = -... = W1-1 and W, = 0. It helps avoiding saddle points and has the following convergence result.
Theorem 1. |Continuous-time gradient descent without weight sharing (Wu et al.) 2019]] For the deep linear network f(x; Wi,..., WL) = W.........W the continuous time gradient descent with the zero-asymmetric initialization satisfies Rt))) S exp (-26) R(0).
Theorem shows that with the zero-asymmetric initialization, the continuous gradient descent lin- early converges the global optimal solution for general target matrix T
Next, considering the special case where the goal is to learn a positive definite matrix &. Based on Lemma 1 we have the following convergence result for training with weight sharing.
Theorem 2. /Continuous-time gradient descent with weight sharing] For the deep linear network f(x: Wi,... WL) = W, W.., initialize all W((0) with identity matrix and update ac- cording Equation With positive definite target matrix 4 the continuous-timg gradient descent satisfies R(t) I exp 22 min(1, X2 min (+))t R()..
Remark 1. The difference between convergence rates in Theorem I and Theorem 2 not an artifact of analysis. For example, when the target matrix simply
al, > can be explicitly shown that with the initialization in Theorem 1 we have R(0) (0) while training with weight sharing (Theorem 2), we have R())) (0). This implies that the convergence results in Theorem Tanoormmaaane and Theorem cannot be improved in general.
The extra L in the exponent leads to faster convergence. The key to show the acceleration is
dR(t)
di Σ tr (VTTR(()))))) < -2L22 (Wo(t)=-1)(())
where we see the I. comes from the summation. This sheds light on two important factors:
All layers need to have sufficiently large update (i.e. Wt()) is large for all 1).
2. Each layer's update needs to well correlate with its gradient (i.e. V,R(1) correlates with W())).
4

---

Under review as a conference paper at ICLR 2021
Initializing the weights be the same and using the average of gradients guarantees that all layers sufficiently trained. The high correlation of VtR(t) and Wi(t) relies on d being positive definite.
Suppose the gradients of different layers do not correlate well (e.g. tr (V.R(t) and the weights are still forced to be shared via the updates according to Equation 2
R 0, + j)
Recall that
(i(t) = + Σ V (tt,, we then have Σ1= tr V. ((t)Wa(t)) ~ - E (t)(t) [ which
loses the extra L acceleration in the convergence due to the 1/L leading factor.
When dealing with real deep learning models, there is no guarantee that all the gradients at different layers highly correlate. Thus, we may monitor gradient correlations during training: sharing weights when gradients well correlate, and break the ties when gradient correlations fall below certain threshold. This matches the adaptive untying rule we proposed in Section
3.2 DISCRETE- GRADIENT DESCENT
Here we extend the previous result to the discrete-time gradient descent with a positive constant step size 1). It can be shown that with arrmmmsymttriic initialization, training with the gradient descent will achieve R(t)= I e within O(L a Log(1/e)) steps; initializing and training with weights sharing, the deep linear network will learn positive definite matrix & to R(t)) < within OL log(//e) steps, which reduces the required iterations by a factor of T.?.
To make easy comparisons, we first repeat without proving the discrete-time gradient descent vergence result of ZAS.
Theorem 3. [Continuous-tine gradient descent without weight sharing (Wu et al. 2019] For deep linear network fx; Wi,..., Wr) W..Wr.-...W with zero-asymmetric initialization and
discrete-time gradient descent, if the learning rate satisfies 1j S min { (4 , (144L25*) }}.
where & = max [2] Lb F,, 3L
1/2 3L-1/2 .1). then we have linear convergence R(t)= S (1- R(0).
Since the learning rate is 1 = ( Theorem 3 indicates that the gradient descent can achieve (t) S c within O log steps.
In the special case of learning a positive definite matrix 4, initialize all weights W₁ to be the same and train with weights sharing, we have the following convergence result.
Theorem 4. [Discrete-time gradient descent with weight sharing] For the deep linear network
(x; WI...., WL) = WWW.W...WWxx initialize all W((0) with identity matrix I and update ac-
min(Xan (1).1)
cording to Equation With a positive definite target matrix 4 and setting S 4vdd1a max() ( (0).1)'
we have linear convergence R(t)= & exp [-(2L - 2) min (N (.). 1) my] R(0).
Take Amin (ain(()) ($) Ama (T). d as constants and focus the scaling with I,, e, we have 17 Because the extra I. in the exponent, we know that when learning a positive definite matrix t training with weight sharing can achieve R(t) S within O (L log (1/e)) steps. The dependency on L reduces from previous 33 to linear, which shows the acceleration of training by weight sharing.
OT.-2))
4 RELATED WORK
Lan et al. (2020) propose ALBERT with the weights being shared across all its transformer layers. Large ALBERT models can achieve good performance on several natural language understanding benchmarks. BBii et al. (2019)) propose trellis networks which are temporal convolution networks with shared weights and obtain good results for language modeling. This line of work is then ex- tended to deep equilibrium models (Bai et al. (019a)) which equivalent to infinite-depth weight- tied feedforward networks. Dabre Fujital (2019) show that the translation quality of a model that recurrently stacks single layer is comparable le to having the same number of separate layers. Zhang et al. 2020) also demonstrate the application of weight-sharing in neural architecture search.
Deep linear models have been widely studied for its simplicity and similarity deep learning mod- els. Baldi & Hornik (1989) show that all local minima are also global minima for two-layer linear networks. Laurent & Brecht (2018) extend the same result to deep linear networks. Hardt & Ma
5

---

Under review as a conference paper at ICLR 2021
2017) show the PL condition is satisfied within the neighbour of a global optimum. Shamir (2019)
show that, for mne-dimensional deep linear networks, with the Xavier or near-identity initialization, it requires at least exp({((L)) steps to converge. where L is the depth. [Wu et al. (2019) show that this result can be improved to ( log 1/e) with a special zero-asymmetric initialization.
5 EXPERIMENTS
In this section, we present the experimental setup and results for training the BERT Large model with the standard training procedure as in the literature well as our Sharing WEights (SWE) method. In what follows, without explicit clarification, BERT always means the BERT Large model.
5.1 EXPERIMENTAL SETUP
We use the TensorFlow official implementation of BERT tteam & contributors.. We first show experimental results with English Wikipedia and BookCorpus for pre-training in the original BERT paper (Devlin et al.] 2019). We then move to the XLNet enlarged pretraining dataset (Yang) al. 2019). We preprocess all datasets with WordPiece tokenization (Schuster & Nakajima 2012). We mask 15% tokens in each sequence. For experiments on English Wikipedia and BookCorpus, we randomly choose tokens mask. For experiments on the XLNet dataset, we do whole word masking- in case that a word is broken into multiple tokens, either all tokens are masked or not masked. For all experiments, we set both the batch size and sequence length to 512.
We use the AdamW optimizer (Loshchilov & Hutter 2019) with the weight decay rate being 0.01, 31 0.9, and 82 0.999. For English Wikipedia and BookCorpus, we use Pre-LN (He et al.) 2016b Wang et al.] (019b) instead of the original BERT's Post-LN. Note that the correct implemen- tation of Pre-LN contains a final layer-norm right before the final classification/massed language modeling layer. Unlike the claim made by Xiong et al. (2020), we notice that using Pre-LN with learning rate warmup leads to better performance. In our implementation, the learning rate starts from 0.0, linearly increases to the peak value 3 10 10-4, the learning rate used by Xiong et al. 22020) at the 10k-th iteration, and then linearly decays to 0.0. For the XLNet dataset, we apply the same Pre-LN setup except the peak learning chosen to be 2 x 10-4, The peak learning rate of 3 10-4 makes training unstable here and yields worse performance than 2 x 10-4.
After pre-training, we fine-tune the models for the Stanford Question Answering Dataset (SQuAD v1.1 and SQuAD v2.0) (Rajpurkar et al. (016) and the GLUE benchmark (Wang et al. 2019a). For all fine-tuning tasks, we follow the setting as the literature: the model fine-tuned for 3 epochs; the learning rate warms up linearly from 0.0 to peak in the first 10% of the training iterations, then linearly decay 0.0. We select the best peak learning rate based on the validation set from (1 x 10-5 1.5 x 10-5 2 x 10-5, x 10-5 4 10-5, x 10-5 7.5 x 10-5 10 x 10-5, 12 10 10-5). For the SQuAD datasets, we fine-tune each model 5 times and report the average. For the GLUE benchmark, for each training method, simply train one BERT model and submit the model's predictions over the test sets to the GLUE benchmark website to obtain test results. We observed that when using Pre-LN, the GLUE finetuning process is stable and model diverged.
Training methods. The training procedure in the TensorFlow official implementation of BERT serves as our baseline, where the baseline training takes 1 million steps (both on English Wikipedia plus BookCorpus and on the enlarged XLNet dataset). For our Sharing WEight (SWE) method, only half of the number of iterations is taken. For complete comparison, we also report the results from the baseline method with half of the number of iterations. Three versions of our method with hyperparameter settings are listed below.
SWE-F Fixed point untying. It serves as baseline of its adaptive version. We ran experiments with set of different untying points, and identified the best untying point = 50k. The full results from different values are presented in Section 5.3.
SWE-A Adaptive untying. We check the gradient correlations for every 1000 iterations. If the gradient correlation of adjacent layers below a threshold = 0.5 for three consecutive times, break the tie. The effect of different p values is studied in Section 5.3.1

---

Under review as a conference paper at ICLR 2021
SWE-S Simplified adaptive untying. We use the same setup in SWE-A, except break the tie all at once, when the majority gradient correlations are below the threshold for three consecutive times.
5.2 EXPERIMENT RESULTS
For English Wikipedia and BookCorpus, both pretraining and finetuning results of our method vs. the baseline method are shown in Table [ From the results, training iterations matches the performance of the baseline method with million training iterations, and significantly outperforms the baseline method with 500k training iterations. The results for the XLNet dataset are shown in Table baseline.
that our method with 500k
We observe similar advantages of our approach over the
Table 1: Training BERT on English Wikipedia and BookCorpus. Our method with half of million iterations matches the baseline performance with one million iteration steps, and outperforms the baseline with half of million iterations.
Baselines
Our method, 0.5m iter.
Im iter.
0.5m iter.
SWE-F
SWE-A
SWE-S
Pretrain MLM (acc. %)
74.98
73.66
73.92
74.03
74.13
SQuAD v1.1 (F-1%)
92.58
91.54
92.54
92.31
92.54
SQuAD v2.0 (F- 1%)
85.06
82.99
84.14
84.16
84.79
GLUE/AX (corr%
42.3
39.0
40.1
42.4
41.8
GLUE/MNLIInG (acc.%)
86.9
85.9
87.2
87.4
87.2
GLUE/MNLI-min (acc.
86.1
85.4
86.7
86.8
86.7
GLUE/QNLI (acc. %)
93.5
92.3
93,4
93.2
92.9
GLUE/QQP (F-1%)
72.2
72.1
71.7
71.8
71.5
GLUE/SST-2 (acc. %)
94.8
94.3
94.8
95.6
94.2
Table 2: Training BERT on the XLNer dataset. Our method with half of million iterations matches the baseline performance with one million iteration steps, and outperforms the baseline performance with half of a million iterations.
Baselines
Our method, 0.5m iter.
1m iter.
0.5m iter.
SWE-F
SWE-A
SWE-S
Pretrain MLM (ace. %)
71.75
70.06
70.18
70.92
70.60
SQuAD vI.1 (F-1%)
93.37
92.41
92.82
92.65
92.99
SQuAD v2.0 (F-1%)
86.49
85.28
85.57
85.41
86.12
GLUE/AX (corr%
44.2
43.2
43.7
44.1
43.3
GLUE/MNLI-n (acc.%)
88.6
87.3
88.4
87.7
88.2
GLUE/MNLI-min (acc. %)
87.9
87.4
87.4
87.4
87.6
GLUE/QNLI (acc.%)
92.3
91.9
92.8
92.8
93.2
GLUE/QQP (F-1%)
72.3
71.7
72.1
72.2
72.0
GLUE/SST-2 (acc.%)
95.9
95.7
95.4
95.9
96.0
5.3
ABLATION STUDIES
In this section, we study the effects of different choices in implementing method.
5.3.1 WHEN STOP WEIGHT SHARING
In this section, we study the effects of using different untying points (Algorithm I and thresholds (Algorithm()). If weights shared throughout the entire pretraining process, the final performance

---

Under review as a conference paper at ICLR 2021
will be much worse than without any form of weight sharing (Lan et al. 2020 On the other hand, without weight sharing at all yields slower convergence.
Results of using different untying point and threshold p values summarized in Table. 3 Models are trained for 500k iterations on English Wikipedia and BookCorpus. From the results, we see that for the SWE-F method, smaller value performs better than larger This means that the weight sharing stage should not be too long. We also see that the performance of the SWE-A method not sensitive the threshold p.
Table 3: Results from different untying points and thresholds p. Models trained for 500k iterations on English Wikipedia and BookCorpus.
SWE-F
SWE-A
7=25k
r=50k
T=200k
0=
/=
0=
0=
p=0,9
Pretrain MLM (acc.%) (6
71.48
73.92
72.42
73.76
74.11
74.03
74.12
73.71
SQuAD v1.1 (F-15 (-))
90.51
92.54
91.65
92.29
92.15
92.31
92.16
92.17
SQuAD v2.0 (F-1%)
81.11
84.14
83.35
84.69
84.30
84.16
83.38
84.21
GLUE/AX (corr%
37.0
40.1
42.6
40.0
41.5
42.4
40.2
40.3
GLUE/MNLI-1 (acc.%)
86.1
87.2
85.5
86.9
87.4
87.4
87.0
86.5
GLUE/MNLI-mmn (acc.%)
83.9
86.7
85.4
86.1
86.4
86.8
86.9
86.6
GLUE/QNLI (acc,% %)
92.8
93.4
91.9
93.3
93.3
93.2
93.3
92.7
GLUE/QQP (F-1%)
70.3
71.7
70.8
71.4
71.6
71.8
71.9
72.0
GLUE/SST-2 (ace. %)
93.4
94.8
93.5
93.4
94.7
95.6
95.2
94.9
5.3.2 How TO CHOOSE WEIGHT SHARING UNITS
Note that it is not necessary be restricted to share weights only across the original layers. We group several consecutive layers as a weight sharing unit. We denote x B as grouping A layers as weight sharing unit which being shared with B times. Since BERT has 24 layers, the baseline method without weight sharing can be viewed as "24x1". and our method shown in Table [ can be viewed as "1x24". We present results from more different choices of weight sharing units in Table We can see that, in order to achieve good results, the size of the chosen weight sharing unit should be larger than layers. This means that the weights of layer must be shared for least 4 times.
Table 4: We group several consecutive layers as a weight sharing unit instead sharing weights only across original layers. A B means grouping A layers unit which is being shared with B times. Models are trained for 500k iterations on English Wikipedia and BookCorpus.
Baseline
SWE-F
24x1
12x2
6x4
4x6
2x12
1x24
Pretrain MIM ())
73.66
73.82
73.99
73.90
74.16
73.92
SQuAD v1.1 (F-1%)
91.54
92.18
92.62
92.52
92.44
92.54
SQuAD v2.0 (F-1%)
82.99
83.74
84.56
85.82
85.06
84.14
GLUE/AX (corr%
39.0
40.7
40.3
40.2
43.0
40.1
GLUE/MNLIIm (ace.
85.9
86.9
87.9
87.0
87.1
87.2
GLUE/MNN.-mm (acc. %)
85.4
85.8
87.0
86.4
86.4
86.7
GLUE/QNLI (ace,90 %)
92.3
93.1
93.2
92.9
93.8
93.4
GLUE/QQP (F-1%)
72.1
72.0
72.0
71.6
71.8
71.7
GLUE/SST-2 (acc. %)
94.3
94.7
94.3
94.8
94.7
94.8
CONCLUSION
We proposed simple weight sharing method to speed up the training of deep networks with repeated layers and showed promising empirical results on BERT training. Our method motivated by the successes of weight sharing models in the literature well theoretic analysis deep linear
8

---

Under review as a conference paper at ICLR 2021
models. For future work, will extend our empirical studies to other deep learning models and tasks, and analyze under which conditions our method will be helpful.
REFERENCES
Dario Amodei, Sundaram Ananthanarayanan, Rishita Anubhai, Jingliang Bai, Eric Battenberg, Carl Case, Jared Casper, Bryan Catanzaro, Qiang Cheng, Guoliang Chen, al. Deep speech 2: End-to- end speech recognition in english and mandarin. In International conference on machine learning, pp. 173-182 2016.
Shaojie Bai, J Zico Kolter, and Vladlen Koltun. Deep equilibrium models. In Advances in Neural Information Processing Systems, pp. 690-701, 2019a.
Shaojie Bai, Zico Kolter, and Vladlen Koltun. Trellis networks for sequence modeling. In Inter- national Conference Learning Representations, 2019b.
Pierre Baldi and Kurt Horiik. Neural networks and principal component analysis: Learning from examples without local minima. Neural networks, 2(1):53-58, 1989.
Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale GAN training for high fidelity natural image synthesis. In International Conference on Learning Representations, 2019.
Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005.14166 2020.
Raj Dabre and Atsushi Fujita. Recurrent stacking of layers for compact neural machine translation models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 6292- 6299, 2019.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT:
Pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACL-HLT, pp. 4171- 4186, 2019.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: Training ImageNet in hour. arXiv preprint arXiv: 001.770067.777 2017.
Moritz Hardt and Tengyu Ma. Identity matters in deep learning. International Conference for Learning Representations, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016a.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In European conference on computer vision, pp. 630-645. Springer. 2016b.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International
Conference for Learning Representations, 2015.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Sori- cut. ALBERT: A lite BERT for self- self-supervised learning of language representations. In Interna- tional Conference on Learning Representations, 2020.
Thomas Laurent and James Brecht. Deep linear networks with arbitrary loss: All local minima global. In International conference on machine learning, pp. 2902-2907. PMLR, 2018
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu. Dehao Chen. Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. GShard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006. /6668, 2020.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. International Confer- on Learning Representations, 2019.

---

Under review as a conference paper at ICLR 2021
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1-67, 2020.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100, 000+ questions for machine comprehension of text. In EMNLP, 2016.
Mike Schuster and Kaisuke Nakajima. Japanese and Korean voice search. In IEEE International Conference on Acoustics, Speech and Signal Processing, 2012.
Ohad Shamir. Exponential convergence time of gradient descent for one-dimensional deep linear neural networks. In Conference on Learning Theory, pp. 2691-2713. PMLR, 2019.
Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, penporn Koanantakool, Peter Hawkins. HyoukJoong Lee, Mingsheng Hong. Cliff Young, et al. learning for supercomputers. In Advances Neural Information Processing Systems, pp. 10414- 10423, 2018.
Deep
TensorFlow team and contributors. TensorFlow model garden NLP. gi rithub. com/
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: multi-task benchmark and analysis platform for natural language understanding, In International Conference on Learning Representations, 2019a.
Qiang Wang, Bei Li, Tong Xiao, Jingbo Zhu, Changliang Li, Derek F Wong, and Lidia S Chao. Learning deep transformer models for machine translation. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 1810-1822, 2019b.
Lei Wu, Qingean Wang, and Chao Ma. Global convergence of gradient descent for deep linear residual networks. In Advances in Neural Information Processing Systems, pp. 13389-13398, 2019.
Yonghui Wu, Mike Schuster, Zhifeng Chen, Quoe Le, Mohammad Norouzi, Wolfgang Macherey, Maxim Krikun, Yuan Cao, Qin Gao, Klaus Macherey, et al. Google's neural machine trans- lation system: Bridging the gap between human and machine translation. arXiv preprint 1rXiv:1609.0814 2016.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie-Yar Liu. On layer normalization in the transformer architec- ture. In International Conference on Machine Learning, 2020.
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoe V Le. XLNet: Generalized autoregressive pretraining for language understanding. In Advances in neu- ral information processing systems, pp. 5753-5763, 2019.
Yang You, Jing Li, Sashank Reddi, Jonathan Hscu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large batch optimization for deep learning: Training bert in 76 minutes. In International Conference on Learning Representations, 2020.
Yuge Zhang, Zejun Lin, Junyang Jiang. Quanlu Zhang, Yujing Wang, Hui Xue. Chen Zhang, and Yaming Yang. Deeper insights into weight sharing in neural architecture search. arXiv preprint arXiv:2001.01433 2020.
10

---

Under review as a conference paper at ICLR 2021
A PROOFS OF SECTION 3
A.1 PROOF OF LEMMA
Proof. By definition, have
(tR(t) (e6())---(--0- +) (t).
To see the result Lemma] 1
it is sufficient to show that (i(t) has the same eigenvectors 4b
First, it is not hard to see that by initializing Wo = I. the requirement trivially holds for Wo()). Now, suppose Wo(to) has the same eigenvectors as o, then from the definition of update
ViR(tu) = 'to) (Wo b) Wo (ta).
We thus know that (tt( has the same cigenvectors as for all layer L. Therefore the update for Wo(fu), which is given by El= V, v,R(tu)/L L also has the same eigenvectors as 4. Inductively, it indicates that from time step ty. all subsequent weights (t) has the same eigenvectors 4.
We can thus exchange the order of matrix multiplication and for all 2 0, we have
ViR(() = Wo(1)=1 (1) +) (() - (Wo (1) +)
which directly implies that
L
Wi(t + 1) Witt) = " L
Σ R(R(t) = -nw. (t) (t) +))
A.2
PROOF OF THEOREM 2
Proof.
From Lemma
we have
Wi(t) - (w) Φ)
For the loss function R(t), we have
R(() Σ Ir (V. (VTR(I)))())) =-
-- Σ [VR(())
Σ [a((t)= (Wo(t) (Wo(t)- ②)) 11?
AL min (wo(t)=-))R(t)
By continuous gradient descent with Wo(0) = I, it is easy to see that
Amin (W--((1) 2 min (1, Amin (+)).
Therefore we have
R(t)= - min (4))R(t - R(t)= < e 2L 1.. (9))
A.3 TECHNICAL LEMMA
Lemma 2. Initializing Wo and training with weight sharing update (Equation 2}, by setting
7) < maxt X (+),,J) we huve
Amin (W)( 2 min()min ($) 11, Amaz (Wo(!) S max (Xmaaa ( 1).
11

---

Under review as a conference paper at ICLR 2021
we know that Wo(t) has the eigenvectors @. Take any eigenvector, denote (t) and X to be the corresponding eigenvalue of Wo(t) and Φ. By Lemma we have
Proof. In the proof of Lemma 1
+ 1) - n7(t)=-1 (2(t)- - x²
For À > 1, we would like to show A(t) [1, Al. V 2 0 by setting S L Since X(0) = 1, we know this claim holds trivially = 0. Then suppose we have the claim holds for t to, then
X(to + 1) = I(to) nn((to) 4-1 (2(ta) 2)
((to)'24-1)
1=
To make A(to + 1) [1, we set 17 S EX2L and have
A(to + 1) x((to) - i(to)) Σ 4-1 = ((to)'21-11
(7(to) - ->)
LADL
<S(to) - ((10)--1 Σ 4-1 7(10)*\4-1-1
(A(to) - -2) s A.
N Ei= X
And 7) 2 guarantees that A(to + 1) 2 A(to) 2 By induction, when À 1, we haveX(t) [1, 2 0.
Similarly, for À < would like to show π(t)= [A, 1] by setting n S 1. Note again the claim holds trivially when t = O. Suppose A(to)= E [A, 1], we have
(2(40) - 1) A. A(lo 1) 2 A(to)= 7(16)= L-1 \(1o)'\--11
15
And 1) 2 0 guarantees that A(to + 1) & A(to)= 2 1. By induction, when A < 1, we haveā(t) E
[X, 1]. vt 2 0.
Note that the two claims hold for all X, it then directly implies that by setting 1 S L Lmax A2 (0),1) we have
Amin (Wo(1) 2 min(Amin (2) .1), Amax (Wo(t)) S max(Amox (1) 1),
which completes the proof.
A.4 PROOF OF THEOREM
4
Proof. Denote o max(Amax (4), 1). From the proof of Lemma I
(2),
we have
R(R(() w (1) (W) - Φ),
By Lemma 2 setting 7) S L max (0).)) we have
Amin (Wo(t)) 2 min(>min( (.) 1), Amax (Wo(1)) & max(Ammx (8) 1). Denote mmx (mmux ()), 1), we immediately have R()(t) S (f.-1)// 1)/1, /2R(t). [V/R(t)] 2 min(Amin (D), () 1)/ With one step gradient update, we have
R(t 1) - R(t) - [| (Wo - (o)(()) 11 I1 113
= =(1W (o - 3 (W) 1) " 1 ())R(t) 16.l
where the @ denotes the element-wise multiplication. Let
I.
1 1
1 = n @ (Wo -4), 1
Σ
1 A o (we - 1), 13
Σ
Ax
k=2
k=1
12

---

Under review as a conference paper at ICLR 2021
where the matrix A& comes from
(Wo - (VOR(())( + +.... AL.
We have
R(t+1)-RRM $ 11 + +
Note that
|Woll2 s 001
[VoR(()] S 101-
Using the fact that for 0 S y S x/122
(or + w) S xt + 2Lx 2Lxt- y,
(x (x w) S xt + Lxt-1py y + 222
Take 77 < L207/2R(1)
have
7 Axlle S (0V/L + no 1 (2((t)
b - Φ S 2nL0² 22((1)
Σ Aulll (01/L + noo - / L - - 7 s 2
k=2
Thus we have
12 S
1/ Aull ( - + $20
I3 S 2
17 s 147
k=1
For I1, we directly have
1= -=- 3 W (W - -4) = -
Since we have [|Vo[(t]] 2 min (Amin (Ø), 1) v 2R((). Therefore by setting
1
min( (22 min ()),)
min 1)
S min
max (12 (1).1)' (2)2 2V/12
4L2
We have
1) < (2 - 2L) min (2 1)R(t).
By directly setting
min( (2
min 1 1)
A dvdL2 max (Atxx ( 1)
we can satisfy all the requirements above, which will give us
R(t)= & [1 -22 2) min (2 (2), 1) 1] R(0) z exp [-(2L - 2) min (4), 1) mt} R(t).
13


Reviews:
[{"review_id": "yU4Ykzr57s1", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "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": "##########################################################################\n\nSummary:\n\nThe paper proposed a new way for training models that stack the same basic block for multiple times -- share the weights first and then untie the weights. The author tried to provide theoretical insights on why it can be better. Also, ablation study shows that the proposed algorithm has marginal improvement over the baseline.\n\n##########################################################################\n\nReasons for score: \n\n\nThe improvement over the baseline method, which just trains BERT for more steps, is quite marginal. In addition, after the weights are untied, the training process becomes exactly the same as that in BERT. It's highly likely that both methods can have similar performance after being trained for a large number of steps. Also, I do not think the theoretical results on a deep linear network can help explain the  phenomena we see in BERT training. BERT uses a much more complicated building block that involves layer normalization, attention and non-linearity. Due to these reasons, I vote for rejection.\n\n##########################################################################\n\nPros: \n \n\n1. The idea of first train the model with shared weights and then untie the weights is interesting.\n \n\n##########################################################################\n\nCons: \n \n\n1. The improvement over the baseline method is not very substantial. Basically, it is possible that both methods perform similarly after being trained for a longer number of steps (e.g., 1M, 1.5M).\n\n2. I'm not convinced by the theoretical analysis. Building block in BERT is much more complicated that the building block discussed in Section 3.\n\n3. The criteria of changing from sharing to unsharing is largely heuristic.\n\n\n##########################################################################\n\nQuestions during rebuttal period: \n \n\nPlease address and clarify the cons above \n \n\n#########################################################################\n\nTypos: \n\n(1) Page 3, \"To be best of\" should be \"To the best of\"\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "ZBu0w1eIxJO", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "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 proposes a simple yet effective method to train transformer-like networks for NLP tasks. Weight sharing idea is not new, it has been investigated in ALBERT as mentioned in the paper and CV areas (see refs [1],[2]).  Though novelty of this paper is not strong enough, their experimental results are adequate to show the effectiveness of proposed method.  The writing and description are quite clear, but you need to clarify the definition of 'iters' you used in  SWE exps, does it mean iters_of_weight_share + iters_of_wight_unshare?\n\nTo show the advantage of weight sharing, the authors take deep linear model as an example to show that with weight sharing, layers far away from output will not suffer from gradient vanishing or explosion, which is true for deep LINEAR models. But for deep learning based models, non-linear operations is required, which still leads to gradient vanishing or explosion even with weight sharing. For example, recurrent neural network is a typical kind of weight-sharing network and its weight is shared through time. Even with single hidden layer, RNN still suffers from gradient vanishing or explosion, and that is why LSTM is proposed in 1990s. So in my opinion, you'd better take a TRUE deep learning model as an example to analyze the effectiveness of weight sharing during training. For example, you can try to analyze deep feed forward neural networks with ReLU activation, which can be treated as a deep piece-wise linear model.\n\n[1] J. Wang and X. Hu, “Gated recurrent convolution neural network for ocr,” in NIPS, 2017.\n[2] M. Liang and X. Hu, “Recurrent convolutional neural network for object recognition,” in CVPR, 2015, pp. 3367–3375", "extra_scores": {}, "extra_text": {}}, {"review_id": "Nz_cfFfxJym", "reviewer": "AnonReviewer4", "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": "This paper provides a simple way to speed up training deep networks which contain repeated structures, such as the transformer module, by sharing the weights first and then unsharing. Three methods are shown to stop weight sharing: 1) stop weight sharing at a predefined point; 2) stop weight sharing at a point using an adaptive way, in which unsharing will be triggered when the majority of the gradient correlations between layers is below a threshold; 3) gradually untie weights between layers by splitting groups based on the gradient correlations. The authors did experiments on BERT large for SQuAD and the GLUE benchmark, and showed that their approach outperform the baseline at the same iteration of 500k, and is on par w/ the baseline training at 1M iterations.\n\nBasically I think this approach could be a practical way to train deep networks w/ repeated layers. It's like between BERT large and ALBERT, the results are not that surprising. The main contribution of this work is that it provides the empirical results of this training method on several tasks.\n\nMy questions regarding to this paper are:\n\n1. In Section 3, the authors provide analysis of weighting sharing can achieve better convergence rate and for good weight sharing each layer's update needs to well correlate with its gradients for deep linear network. How about any theoretical insights for networks w/ activation functions? And How could different activation functions matter (ReLU, Swish et al)? And will these activation functions affect the parameters of when to stop weight sharing?\n\n2. For untie weights between adjacent layers, SWE-A and SWE-M, what's the motivation of setting a unified threshold for all layers? Basically I'm wondering what's the gradient correlation between adjacent layers of top layers and bottom layers, are the correlations on the same level? Will we get better accuracy if we do not limit to a unified threshold? Say, for the top 12 layers, we set the threshold to X and for the bottom 12 layers we set the threshold to Y. Need further information on this.\n\n3. In Table 2, why is SWE-A missing?\n\n4. For Table 3, I think the authors need to add the experiments when \\tau is smaller than 50k. Basically the experiments will show that \\tau should be carefully designed, an arbitrary number may not work, we still need to train some iterations w/ weight sharing for the good model performance.\n\n5. This paper may need to cite some other papers discussing some insights of weight sharing, such as  DEEPER INSIGHTS INTO WEIGHT SHARING IN NEURAL ARCHITECTURE SEARCH (https://arxiv.org/pdf/2001.01431.pdf).", "extra_scores": {}, "extra_text": {}}, {"review_id": "G5_5a3tA3Ax", "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": "This work introduces a method to accelerate the training speed of deep learning models. The key idea is to start training with shared weights and unroll the shared weights in the middle of the training. The authors report that this strategy accelerates the convergence speed. The paper introduces heuristics on when to stop weight sharing and how many layers to share weights. It further provides an analysis via the view of deep linear models on why weight sharing helps improve the convergence speed. In the evaluation, the paper evaluates their approach against the training of BERT, and shows that their method can obtain comparable and sometimes even better accuracy on downstream tasks while with 50% faster training speed.\n \nStrengths:\n+ The paper aims to address an important problem in large model training: slow training speed.\n+ The paper proposes an easy-to-implement approach to accelerate the convergence speed of BERT-like models. \n \nWeakness:\n- The technical contribution seems rather incremental. The main difference between this work and the prior work [1], which also train Transformers via shared weights,  seems to be switching from sharing weights to unsharring weights in the middle of the training.\n- Important references are missing, making it not clear the advantage of this work as compared with existing approaches that accelerate the training of Transformer networks. \nThe comparison with existing work is inadequate. Important ablation studies are needed. \n \nComments:\n \nPrior work [1] uses weight sharing to train a smaller Transformer model to obtain similar accuracy.  However, weight sharing does not improve training speed per batch, because the training still needs to perform the same amount of forward and backward computation for each batch. Training may actually be slowed down, since the model may need to train with more batches to reach the same accuracy. This paper aims to speed up the training process by switching from shared to unshared weights in the middle of the training, and it observes faster convergence -- achieving similar downstream accuracy with less number of pretraining samples. This is an interesting empirical observation and can potentially become useful in practice. However, there are some major concerns about the paper.\n \nIt is still unclear whether this faster convergence comes from switching from sharing to unsharring weights or is an effect of the model or hyperparameter changes. First, the stop condition (e.g., the switching threshold) cannot be known as a prior. Therefore, it is controlled with an additional hyperparameter. From the text, it is unclear how this hyperparameter has been chosen or will be chosen when training a new model. It would be better to test the sensitivity of the hyperparameter on another model such as GPT-2 to verify the effectiveness of the proposed method.\n \nSecond, the paper adopts Pre-LN in its evaluation (as briefly mentioned in Section 5.1). However, from the text description, it seems it employs the original BERT as the baseline (\"We first show experiment results English Wikipedia and BookCorpus for pretraining as in the original BERT paper\"). As Pre-LN has been studied in several prior work [2,3] and has been demonstrated to also have the effect of accelerating the convergence speed, it is unclear whether the observed speedup in this paper is an effect of PreLN or weight sharing/unsharing. An ablation study with BERT-PreLN is needed if not already included.\n \nThird, the analysis on deep linear models appears to be over-simplified, where important characteristics of the DNNs such as non-linear activations and residual branches are not represented, making it difficult to connect it with the actual observations in practice.\n \nFinally, importance reference [4] is missing. [4] starts with a shallow BERT and progressively stacks Transformer blocks to accelerate training speed, which is in some sense similar to the proposed technique, which starts with shallow BERT with shared weights and switches to full-length BERT in the middle of the training. The paper might need to highlight the difference between this work and [4]. \n \nSeveral places in the paper are vague or inconsistent: \n\n1. The paper claims that it uses the same BERT implementation and training procedure as the one used by Devlin et. al.. However, the accuracy reported in Table 2 seems to be consistently lower than what was reported in the original BERT paper. For example, QQP in the original paper reaches 72.1, whereas this paper reports 71.4, QNLI was 92.7 in the BERT paper and 91.7 in this paper. If we use the original BERT reported results, the proposed technique seems to incur low accuracy on most tested datasets. Some clarification on the accuracy results is needed. \n\n2. The paper claims that \"it sounds natural for us to expect that ALBERT's performance will be improved if we stop its weight sharing at some point of training. The optimal models are supposed to not be far from weight sharing.\" However, this explanation actually creates some confusion. First, the paper does not provide an analysis of how the weight distribution between the weights trained with and without weight sharing, so it is unclear what \"not be far\" means. Second, what does it mean by \"optimal model\"?  Does it refer to models trained not through weight sharing? If so, prior work [5] identified that model weights and gradients at different layers can exhibit very different characteristics, which seems to contradict the argument that \"The optimal models are supposed to not be far from weight sharing\". \n\n3. I find it challenging to claim the proposed technique generic for models with repeatedly layers, whereas only BERT is evaluated in the experiments.\n\n4. The paper says \"This means that the weight sharing stage should not be too long\", but it is unclear how long is considered as not too long. \n \n[1] Lan et. al. \"ALBERT: A Lite BERT for Self-supervised Learning of Language Representations\", https://arxiv.org/abs/1909.11942\n\n[2] Xiong et. al. \"On Layer Normalization in the Transformer Architecture\", https://arxiv.org/abs/2002.04745\n\n[3] Shoeybi et. al. \"Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism\", https://arxiv.org/abs/1909.08053\n\n[4] Gong et. al. \"Efficient Training of BERT by Progressively Stacking\", http://proceedings.mlr.press/v97/gong19a/gong19a.pdf\n\n[5] You et. al. \"Large Batch Optimization for Deep Learning: Training BERT in 76 minutes\", https://arxiv.org/abs/1904.00962", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: 
The paper proposed a new way for training models that stack the same basic block for multiple times -- share the weights first and then untie the weights. Ablation study shows that the proposed algorithm has marginal improvement over the baseline. The authors also provide some theoretical justifications to how the proposed idea works.
The proposed idea is straightforward and intuitive. Weight sharing has been used in previous works, and what’s new in this paper is to unshare the weights in the middle (with a heuristic rule). The hope is that by doing so, one can achieve a better tradeoff between speedup and accuracy. However the experimental supports are somehow weak and incomplete. For example, in order to show the real speedup, one should provide the full training curve (until convergence) under different settings, instead of just showing one data point (at 500K). It is very common that one can get some speedup at 500K, but the speedup totally disappears after another 500K steps. 
Furthermore, the theoretical analysis is conducted in a simplified setting, and it is not very clear whether it can be used to explain what really happened during BERT training.
The reviewers conducted some lengthy discussions after the author rebuttal was available. As a final consensus, we think that there are still concerns on the paper, which makes us hesitate to give an ACCEPT recommendation.

Decision: Reject