S PEEDING UP D EEP L EARNING T RAINING BY S HARING W EIGHTS AND T HEN U NSHARING
Anonymous authors
Paper under double-blind review
A BSTRACT
It has been widely observed that increasing deep learning model sizes often leads to signiﬁcant 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 ﬁrst 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%.
I NTRODUCTION
It has been widely observed that increasing model size often leads to signiﬁcantly better performance on various real tasks, especially natural language processing and computer vision applications (Amodei et al., 2016;
He et al., 2016a; Wu et al., 2016; Vaswani et al., 2017; Devlin et al., 2019; Brock et al., 2019; Raffel et al., 2020; Brown et al., 2020).
However, as models getting larger, the training can become extremely resource intensive and time consuming.
As a consequence, there has been a growing interest in developing systems and algorithms for efﬁcient 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 architectures 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 et al., 2017).
We propose a simple method for efﬁciently training such kind of networks.
In our approach, we ﬁrst force the weights to be shared across all the repeated layers and train the network, and then, at some point, we stop weight sharing and continue training until convergence.
The point for stopping weight sharing can be either predeﬁned 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 et al., 2020).
It is a variant of BERT in which the weights across all the transformer layers are shared.
As long as its architecture is sufﬁciently large, ALBERT can be comparable with or even outperform the original BERT on various downstream natural language processing benchmarks.
However, when its architecture being the same as the the original BERT, ALBERT performs signiﬁcantly 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.
The other motivation for our method comes from our theoretical analysis on deep linear models.
A deep linear model is constructed by a series of matrix multiplication (Hardt & Ma, 2017; Laurent & Brecht, 2018; Wu et al., 2019).
In its forward pass, a deep linear model is just equivalent to a single matrix.
However, when being trained with backpropagation, its behavior is analogous to the deep models with non-linearity but much easier to understand.
Our theoretical analysis shows that, when learning a positive deﬁnite matrix (which admits an optimal solution with all layers having the same weights), training with weight sharing can bring signiﬁcantly faster convergence.
Moreover, our theoretical analysis on deep linear models also provides insights in implementing the adaptive weight untying rule of our algorithm.
The rest of this paper is organized as follows.
We present our training 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 deﬁnite deep linear models.
All the proofs are deferred to the Appendix.
In Section 4, we discuss related work.
In Section 5, we show detailed experimental setup and results.
We also provide various ablation studies on different choices in implementing our algorithm.
Finally, we conclude this paper with some discussions in Section 6.
A LGORITHM : S HARING W EIGHTS AND T HEN U NSHARING
Assume we have a deep network which is obtained by repeatedly stacking the same neural module n times, such as the transformer module in transformer models (Vaswani et al., 2017).
Denote by w 1 , . .
w n the weights of these n layers.
In our method, we ﬁrst train the deep network with all the weights tied.
Then, after a certain number of training steps, we untie the weights and further train the network until convergence.
In what follows, we ﬁrst present a simple version of our algorithm in which the weight untying point is predeﬁned. Then, we move to its adaptive version in which the layers are automatically gradually untied according to gradient statistics. Finally, we discuss a simpliﬁed variant of this adaptive method which unties the layers all the once.
Untying weights at a ﬁxed point .
This is the simplest version of our method (Algorithm 1).
We ﬁrst train the deep network with all the weights tied for a ﬁxed number of steps, and then untie the weights and continue training until convergence.
if t < τ then 5: [Equation]
6: else 7:
[Equation]
Note that, from line 1 to 5, we initialize all the weights equally, and then update them using the mean of their gradients.
It is easy to see that such an update is equivalent to weight sharing or tying.
For the sake of simplicity, in line 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 is 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, we may combine every four layers as a weight sharing unit.
Thus, there will be six such units for weight sharing.
Such a ﬂexibility 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 gradually 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 training, suppose we have groups G = { G 1 , G 2 , ..., G k } .
For each group G i , we compute the correlation between the gradients of any two adjacent layers by correlation ( i )
j = (cid:68) g
( i
[Equation]
g
[Equation]
1 denote the gradients of two adjacent layers in G i . Once correlation ( i ) j falls below ρ (0.5 as default) consecutively for a certain number of times (3 as default), we split the group G i into two subgroups, by splitting at j -th layer in G i . The layers in the same group are updated using the gradient mean of this group as in Algorithm 1.
Algorithm Calculate the correlation as //Motivated by Section6: If the correlation falls below ρ , break the group between layer
i and i +7: Update the weights in each group using the average of gradients within each group Simpliﬁed adaptive untying .
To simplify the multi-step adaptive untying, instead of gradually splitting the groups, we may untie the weights [-shared (i.e. existing W ∗ 
-]of all layers at once. For every certain number of iteration steps, we monitor the gradient correlations as described in multi-step adaptive untying.
If more than half of these layer correlations are less than the predeﬁned threshold ρ (0.5 as default) consecutively for a certain number of times (3 as default), we stop sharing weights for all layers.  T HEORETIC M OTIVATION
In this section, we show how our method including the adaptive untying criterion can be motivated from analyzing the dynamics of training a deep linear network by gradient descent. At the ﬁrst glance, deep learning models may look trivial since a deep linear model is just equivalent to a single matrix. However, when being trained with backpropagation, its behavior is analogous to generic deep models. Thus, analyzing deep linear models has attracted increasing interest in the theoretic research community (Hardt & Ma, 2017; Laurent & Brecht, 2018; Wu et al., 2019).
A deep linear network is a series of matrix multiplication
[Equation]
The task is to train the deep linear network to learn a target matrix Φ ∈ R d × d .
To focus on the training dynamics, we adopt the simpliﬁed objective function
[Equation]
Denote ∇ l R to be the gradient of R with respect to W l .
We have
[Equation]
where W l 2 : l 1 = W l 2 W l 2 − 1 ...
W l 1 +1 W l 1 .
The standard gradient update is given by
[Equation]
To train with weights shared, all the layers need to have the same initialization.
And the update is
[Equation]
Since the initialization and updates are the same for all layers, the parameters W 1 ( t ) , ..., W L ( t ) are equal for all t .
For simplicity, we denote the weight at time t to be W 0 ( t ) .
Notice that the gradients are averaged, the norm of update to each layer doesn’t scale with L .
Suppose that the target matrix Φ is a positive deﬁnite matrix.
It is immediate that Φ 1 /n is a solution to the deep linear network.
Before looking into the detailed convergence analysis, it worth ﬁrst showing a Lemma that reveals the updates in the weight sharing training.
Lemma 1.
With a positive deﬁnite target matrix Φ and initializing with W 0 (0)
=
I , update the parameters according to Equation 2, we have
[Equation]
Intuitively, the Lemma 1 shows that weight sharing allows all the layers to 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 ﬁrst study the convergence result with continuous-time gradient descent, which demonstrates the beneﬁt of training with weight sharing when learning a positive deﬁnite 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., 2019).
To be best of our knowledge, ZAS gives the state-of-the-art convergence rate.
It is actually the only work showing the global convergence of deep linear network trained by gradient descent for an arbitrary target matrix.
3. C ONTINUOUS - TIME GRADIENT DESCENT
With continuous-time gradient descent (i.e. η → 0 ), training with gradient descent and ZAS, the loss decays as R ( t ) ≤ exp( − 2 t ) R (0) .
For training with weight sharing, the loss becomes R ( t ) ≤ exp( − 2 L min(1 , λ min (Φ)) t ) R
(0) , when the target matrix Φ is positive deﬁnite.
The extra L in the exponent demonstrates the acceleration of training with weight sharing.
With η → 0 , the training dynamics of continuous-time gradient descent can be described as
[Equation]
The ZAS initializes the weights W
= W 2 = · · · = W L − 1 =
I and W L = 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 ; W 1 , ..., W L ) =
W L W L − 1 ...
W 1 x , the continuous time gradient descent with the zero-asymmetric initialization satisﬁes R ( t ) ≤ exp ( − 2 t ) R (0) .
Theorem 1 shows that with the zero-asymmetric initialization, the continuous gradient descent linearly converges to the global optimal solution for general target matrix Φ .
Next, considering the special case where the goal is to learn a positive deﬁnite 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 ; W 1 , ..., W L ) =
W L W L − 1 ...
W 1 x , initialize all W l (0) with identity matrix I and update according to Equation 2.
With a positive deﬁnite target matrix Φ , the continuous-time gradient descent satisﬁes R (
[Equation]
Remark 1.
The difference between convergence rates in Theorem 1 and Theorem 2 is not an artifact of analysis.
For example, when the target matrix is simply Φ = αI, α > 1 .
It can be explicitly shown that with the initialization in Theorem 1, we have ˙ R (0) =
− 2 R
(0) while training with weight sharing (Theorem 2), we have ˙ R (0)
=
− 2 L R (0) .
This implies that the convergence results in Theorem 1 and Theorem 2 cannot be improved in general.
The extra L in the exponent leads to faster convergence.
The key to show the acceleration is
[Equation]
where we see the L comes from the summation.
This sheds light on two important factors:
All layers need to have sufﬁciently large update (i.e. ˙ W l ( t ) is large for all l ).
Each layer’s update needs to well correlate with its gradient (i
e. ∇ l R ( t ) correlates with ˙ W l ( t ) ).
Initializing the weights to be the same and using the average of gradients guarantees that all layers are sufﬁciently trained.
The high correlation of ∇ l R ( t ) and ˙ W l ( t ) relies on Φ being positive deﬁnite.
Suppose the gradients of different layers do not correlate well (e.g. tr ( ∇ i R ( t ) ∇ j R ( t ))
≈ 0 , i (cid:54) =
j ) and the weights are still forced to be shared via the updates according to Equation 2.
Recall that ˙ W l ( t ) =
1 L (cid:80) Li =1 ∇ i R ( t ) , we then have (cid:80) Ll
=1 tr (cid:16) ∇ (cid:62) l R ( t ) ˙
W l ( t ) (cid:17)
≈ − 1 L (cid:80) Ll =1 (cid:107)∇R ( t ) (cid:107) 2 F , 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 a certain threshold.
This matches the adaptive untying rule we proposed in Section 2.
3. D ISCRETE - TIME GRADIENT DESCENT
Here we extend the previous result to the discrete-time gradient descent with a positive constant step size η .
It can be shown that with zero-asymmetric initialization, training with the gradient descent will achieve R ( t ) ≤ (cid:15) within O ( L 3 log(1 /(cid:15) ))
steps; initializing and training with weights sharing, the deep linear network will learn a positive deﬁnite matrix Φ to R ( t ) ≤ (cid:15) within O ( L log(1 /(cid:15) )) steps, which reduces the required iterations by a factor of L 2 .
To make easy comparisons, we ﬁrst repeat without proving the discrete-time gradient descent convergence result of ZAS.
Theorem 3.
[Continuous-time gradient descent without weight sharing (Wu et al., 2019)]
For deep linear network f ( x ; W 1 , ..., W L ) =
W L W L − 1 ...
W 1 x with zero-asymmetric initialization and discrete-time gradient descent, if the learning rate satisﬁes η ≤ min (cid:110)(cid:0) 4 L 3 ξ 6 (cid:1) − 1 , (cid:0) 144 L 2 ξ 4 (cid:1) − 1 (cid:111) , where ξ = max (cid:8) 2 (cid:107) Φ (cid:107) F , 3 L − 1 / 2 , 1 (cid:9) , then we have linear convergence R ( t ) ≤ (cid:0) 1 − η 2 (cid:1) t R (0) .
Since the learning rate is η = O
( L − 3 ) , Theorem 3 indicates that the gradient descent can achieve
R ( t ) ≤ (cid:15) within O ( L 3 log(1 /(cid:15) )) steps.
In the special case of learning a positive deﬁnite matrix Φ , initialize all weights W l 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 f ( x ; W 1 , ..., W L ) =
W L W L − 1 ...
W 1 x , initialize all W l (0) with identity matrix I and update according to Equation 2.
With a positive deﬁnite target matrix Φ , and setting η ≤ min( λ 2 min (Φ) , 1) 4 √ dL 2 max( λ 4 max (Φ) , 1) , we have linear convergence R ( t ) ≤ exp (cid:2) − (2 L − 2) min (cid:0) λ 2 min (Φ) , 1 (cid:1) ηt (cid:3) R (0) .
Take λ min (Φ) /λ max (Φ) , d as constants and focus on the scaling with L, (cid:15) , we have η = O
( L − 2 ) .
Because of the extra L in the exponent, we know that when learning a positive deﬁnite matrix Φ , training with weight sharing can achieve R ( t ) ≤ (cid:15) within O
( L log (1 /(cid:15) )) steps.
The dependency on
L reduces from previous L 3 to linear, which shows the acceleration of training by weight sharing.
R ELATED W ORK
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.
Bai et al.
(2019b) propose trellis networks which are temporal convolution networks with shared weights and obtain good results for language modeling.
This line of work is then extended to deep equilibrium models (Bai et al., 2019a) which are equivalent to inﬁnite-depth weighttied feedforward networks.
Dabre & Fujita (2019) show that the translation quality of a model that recurrently stacks a single layer is comparable to having the same number of separate layers.
Deep linear models have been widely studied for its simplicity and similarity to deep learning models.
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 (2017) show the PL condition is satisﬁed within the neighbour of a global optimum.
Shamir (2019) show that, for one-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 O ( L 3 log 1 /(cid:15) ) with a special zero-asymmetric initialization.
E XPERIMENTS
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 as well as our Sharing WEights (SWE) method.
In what follows, without explicit clariﬁcation, BERT always means the BERT Large model.
5. E XPERIMENTAL S ETUP
We use the TensorFlow ofﬁcial implementation of BERT (team & contributors).
We ﬁrst show experimental results with English Wikipedia and BookCorpus for pre-training as in the original BERT paper (Devlin et al., 2019).
We then move to the XLNet enlarged pretraining dataset (Yang et 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 to 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 , β 1 = 0 .
9 , and β 2 = 0 .
For English Wikipedia and BookCorpus, we use Pre-LN
(He et al., 2016b; Wang et al., 2019b) instead of the original BERT’s Post-LN.
Note that we the correct implementation of Pre-LN contains a ﬁnal layer-norm right before the ﬁnal classiﬁcation/masked language modeling layer.
Unlike the claim made by Xiong et al.
, 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 − 4 at the 10k-th iteration, and then linearly decays to 0 .
For the XLNet dataset, we apply the same Pre-LN setup except the peak learning chosen to be 2 × 10 − 4 .
After pre-training, we ﬁne-tune the models for the Stanford Question Answering Dataset (SQuAD v1.1 and SQuAD v2.0) (Rajpurkar et al., 2016) and the GLUE benchmark (Wang et al., 2019a).
For all ﬁne-tuning tasks, we follow the setting as in the literature: the model is ﬁne-tuned for 3 epochs; the learning rate warms up linearly from 0.0 to peak in the ﬁrst 10% of the training iterations, then linearly decay to 0 .
We select the best peak learning rate based on the validation set from { 1 × 10 − 5 , 1 . 5 × 10 − 5 , 2 × 10 − 5 , 3 × 10 − 5 , 4 × 10 − 5 , 5 × 10 − 5 , 7 . 5 × 10 − 5 , 10 × 10 − 5 , 12 × 10 − 5 } .
For the SQuAD datasets, we ﬁne-tune each model 5 times and report the average.
For the GLUE benchmark, for each training method, we 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 ﬁnetuning process is stable and no model diverged.
Training methods.
The training procedure in the TensorFlow ofﬁcial implementation of BERT serves as our baseline, where training on English Wikipedia plus BookCorpus takes 1 million steps.
Training on the enlarged XLNet dataset takes 2 million steps.
For our Sharing WEight (SWE) method, only half of the numbers of iterations are taken.
For a complete comparison, we also report the results from the baseline method with half of the numbers of iterations.
Three versions of our method with hyperparamter settings are listed below.
SWE-F F ixed point untying.
By default, we set the untying point τ
=
50 k. We study the effect of different τ values in Section 5.3.1.
SWE-A A daptive untying.
We check the gradient correlations for every 1k iterations.
If the gradient correlation of adjacent layers is below a threshold ρ = 0 .
5 for three consecutive times, we break the tie.
The effect of different ρ values is studied in Section 5.3.1.
SWE-S S impliﬁed adaptive untying.
We use the same setup as in SWE-A, except we break the tie all at once, when the majority of gradient correlations are below ρ for 3 consecutive times.
5. E XPERIMENT R ESULTS
For English Wikipedia and BookCorpus, both pretraining and ﬁnetuning results of our method vs. the baseline method are shown in Table 1.
From the results, we see that our method with 500k training iterations matches the performance of the baseline method with 1 million training iterations, and signiﬁcantly outperforms the baseline method with 500k training iterations.
The results for the
XLNet dataset are shown in Table 2.
We observe similar advantages of our approach over the baseline.
[Table] AX
[Table] MNLI-m
[Table]
acc.
[Table]
/MNLI-mm (acc
[Table]
QQP
[Table]
SST
[Table]
[Table]
SQuAD v1.1 (F-1%) 93.
SQuAD v2.0 (F-1%) 86.
GLUE/AX (corr%) 44.
GLUE/MNLI-m (acc.%) 88.
GLUE/MNLI-mm (acc.%) 87.
GLUE/QNLI (acc.%) 92.
GLUE/QQP (F-1%) 72.
GLUE/SST-2 (acc.%) 95.
A BLATION STUDIES
In this section, we study the effects of different choices in implementing our method.
5.3. W HEN TO S TOP W EIGHT S HARING
In this section, we study the effect of using different untying points and thresholds.
If weights are shared throughout the entire pretraining process, the ﬁnal performance 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 ρ values are summarized in Table.
Models are trained for 500k iterations on English Wikipedia and BookCorpus.
From the results, we see that for the SWE-F method, a smaller τ value performs better than a larger τ .
This means that the weight sharing stage should not be too long.
We also see that the performance of the SWE-A method is not sensitive to the threshold ρ . [Table] Pretrain MLM ( [Table]
[Table] [Table] [Table] [Table]
[Table]
[Table]
[Table]
[Table]
SQuAD v1.1 (F-1%) 90.
[Table]
SQuAD v2.0 (F-1%) 81.
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
GLUE/AX (corr%) 37.
[Table]
[Table]
[Table]
[Table]
GLUE/MNLI-m (acc.%) 86.
GLUE/MNLI-mm (acc.%) 83.
[Table]
[Table]
[Table]
[Table]
[Table]
GLUE/QNLI (acc.%) 92.
[Table]
GLUE/QQP (F-1%) 70.
GLUE/SST-2 (acc.%) 93.
H OW TO CHOOSE WEIGHT SHARING UNITS
Note that it is not necessary to be restricted to share weights only across the original layers.
We can group several consecutive layers as a weight sharing unit.
We denote A × B as grouping A layers as a weight sharing unit which is 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 1 can be viewed as “1x24”.
We present results from more different choices of weight sharing units in Table 4.
We can see that, in order to achieve good results, the size of the chosen weight sharing unit should not be lager than 6 layers.
This means that the weights of a layer must be shared for at least 4 times.
[Table]
SQuAD v1.1 (F-1%)
[Table]
SQuAD v2.0 (F-
[Table]
[Table] GLUE/AX (corr%) 39. GLUE/MNLI-m (acc.%) 85.
[Table]
GLUE/MNLI-mm (acc.%) 85. [Table] [Table]
[Table]
GLUE/QNLI (acc.%) 92.
GLUE/QQP (F-1%) 72.
[Table]
C ONCLUSION AND D ISCUSSION
We proposed a simple method to speed up the training of deep networks with repeated layers and showed promising empirical results.
Our method is motivated by the successes of weight sharing models in the literature as well as our theoretic analysis on deep linear models.
For the future work, we will extend our empirical studies to more deep learning models and tasks, and analyze under which conditions our method will be efﬁcient.
In addition, the presented theoretic analysis is fairly preliminary.
It even cannot be directly applied to illustrate a slightly more realistic case where the truth may somewhat deviate from weight sharing, not to mention dealing with the deep models with non-linearity.
We believe a deep theoretic understanding will help further improve this approach.
A P ROOFS OF S ECTIONA. P ROOF OF L EMMA
Proof.
By deﬁnition, we have
[Equation]
To see the result in Lemma 1, it is sufﬁcient to show that W 0 ( t ) has the same eigenvectors as Φ .
First, it is not hard to see that by initializing W
=
I , the requirement trivially holds for W 0 (0) .
Now, suppose W 0 ( t 0 ) has the same eigenvectors as Φ , then from the deﬁnition of update
[Equation]
We thus know that ∇ l R ( t 0 ) has the same eigenvectors as Φ for all layer l .
Therefore the update for
W 0 ( t 0 ) , which is given by (cid:80) Li =1 ∇ i R ( t 0 ) /L also has the same eigenvectors as Φ .
Inductively, it indicates that from time step t 0 , all subsequent weights W 0 ( t ) has the same eigenvectors as Φ .
We can thus exchange the order of matrix multiplication and for all t ≥ 0 , we have
[Equation]
which directly implies that
[Equation]
A. P ROOF OF T HEOREM
Proof.
From Lemma 1, we have
[Equation]
For the loss function R ( t ) , we have
[Equation]
By continuous gradient descent with W 0 (0)
=
I , it is easy to see that
[Equation]
Therefore we have
[Equation]
A. T ECHNICAL L EMMA
Lemma 2.
Initializing W
=
I and training with weight sharing update (Equation 2), by setting η ≤ 1 L max( λ 2 max (Φ) , 1) , we have λ min ( W 0 ( t )) ≥ min( λ min (Φ) 1 /L , 1) , λ max ( W 0 ( t )) ≤ max( λ max (Φ) 1 /L , 1)
[Equation]
Proof.
In the proof of Lemma 1, we know that W 0 ( t ) has the same eigenvectors as Φ .
Take any eigenvector, denote (cid:98) λ ( t ) and λ L to be the corresponding eigenvalue of W 0 ( t ) and Φ .
By Lemma 1, we have
[Equation]
For λ > 1 , we would like to show (cid:98) λ ( t ) ∈ [1 , λ ] , ∀ t ≥ 0 by setting η ≤ 1 Lλ 2 L .
Since (cid:98) λ (0) = 1 , we know this claim holds trivially at t = 0 .
Then suppose we have the claim holds for t = t 0 , then
[Equation]
L −
[Equation]
i = (cid:98) λ
[Equation]
t 0 )
i λ L − 1 − i .
To make (cid:98) λ ( t 0 + 1) ∈
[1 , λ ] , we set η ≤ 1 Lλ 2 L and have (cid:98)
[Equation] [Equation]
λ
[Equation] [Equation]
[Equation]
λ
[Equation]
− 1 − i
Lλ 2 L (cid:16)(cid:98)
[Equation] [Equation]
≤ (cid:98)
[Equation]
( t 0 ) − (cid:98) λ ( t 0 )
L
[Equation]
1 (cid:80) L
[Equation] [Equation]
0 (cid:98) λ ( t 0 )
[Equation]
λ L − 1 − i λ L − 1 (cid:80)
[Equation]
And η ≥ 0 guarantees that (cid:98) λ
( t 0 + 1) ≥ (cid:98) λ
( t 0 ) ≥ 1 .
By induction, when λ > 1 , we have (cid:98) λ ( t ) ∈ [1 , λ ] , ∀ t ≥ 0 .
Similarly, for λ < 1 , we would like to show (cid:98) λ ( t ) ∈ [ λ, 1] by setting η ≤ 1 L .
Note again the claim holds trivially when t = 0 .
Suppose (cid:98) λ ( t 0 ) ∈
[ λ, 1] , we have
[Equation]
And η ≥ 0 guarantees that (cid:98) λ
( t 0 + 1) ≤ (cid:98) λ ( t 0 ) ≥ 1 .
By induction, when λ < 1 , we have (cid:98) λ ( t ) ∈
[ λ, 1] , ∀ t ≥ 0 .
Note that the two claims hold for all λ , it then directly implies that by setting η ≤
[Equation]
max
[Equation]
Φ) , 1) , we have
[Equation]
which completes the proof.
[Equation]
A. P ROOF OF T HEOREM
Proof.
Denote φ = max( λ max (Φ) , 1) .
From the proof of Lemma 1, we have
[Equation]
Φ
[Equation]
By Lemma 2, setting η ≤
[Equation] [Equation]
max (Φ) , 1) , we have
[Equation]
Denote φ = max( λ max (Φ) , 1) , we immediately have
[Equation]
With one step of gradient update
[Equation]
we have
[Equation]
where the (cid:12) denotes the element-wise multiplication.
Let
[Equation]
where the matrix A k comes from
[Equation]
We have
[Equation]
Note that
[Equation]
Using the fact that for 0 ≤ y ≤ x/L 2 ,
[Equation]
Take η ≤
[Equation]
, we have
[Equation]
Thus we have
[Equation]
For I 1 , we directly have
[Equation]
Since we have (cid:107)∇ 0 R ( t ) (cid:107) F ≥ min ( λ min (Φ) , 1) (cid:112) 2 R ( t ) .
Therefore by setting η ≤ min
[Equation]
We have
[Equation]
By directly setting
[Equation]
we can satisfy all the requirements above, which will give us
[Equation]