Effects of Depth, Width, and Initialization: A Convergence Analysis of Layer-wise Training for Deep Linear Neural Networks
Yeonjong Shin
Introduction
Deep learning has drawn a lot of attention from both academia and industry due to its tremendous empirical success in various applications . One of the key components in the success of deep learning is the intriguing ability of gradient-based optimization methods. Despite of the non-convex and non-smooth nature of the loss function, it somehow finds a local (or global) minimum, which performs well in practice. Mathematical analysis of this phenomenon has been undertaken. There are several theoretical works, which show that under the assumption of over-parameterization, more precisely, very wide networks, the (stochastic) gradient descent algorithm finds a global minimum . These theoretical progresses have its own importance, however, it does not directly help practitioners to have better training results. This is mainly because there are still many parameters to be determined a priori; learning rate, the depth of network, the width of intermediate layers, optimization algorithms with its own internal parameters, to name just a few. The learning rates from existing theoretical works are not applicable in practice. For example, when a fully-connected ReLU network of depth 10 is trained over 1,000 training data, theoretically guaranteed learning rate is either or . Thus, practitioners typically choose these aforementioned parameters by either a grid search or trial and error.
Despite its expressive power, training deep neural networks is not an easy task. It has been widely known that the deeper the network is, the harder it is to be trained . Empirical success of deep learning heavily relies on numerous engineering tricks used in the training process. These includes but not limited to dropout , dropconnect , batch-normalization , weight-normalization , pre-training , and data augmentation . Although these techniques are shown to be effective in many machine learning applications, it lacks rigorous justifications and hinders a thorough mathematical understanding of the training process of deep learning. The layer-wise training is an alternative to the standard end-to-end back-propagation, especially for training deep neural networks. The underlying principle is to train only a few layers (or a single layer) at a time, rather than train the whole layers simultaneously. This approach is not new and has been proposed in several different contexts. One stream of layer-wise training is adaptive training. At each stage, only a few layers (or a single) are trained. Once training is done, new layers are added. By fixing all the previously trained layers for the rest of the training, only newly added layers are trained. This procedure is repeated. The works of this direction include . Another stream of layer-wise training is the block coordinate descent (BCD) method . The BCD is a Gauss-Seidel type of gradient-free methods, which trains each layer at a time by freezing all other layers, in a sequential order. Thus, all layers are updated once in every sweep of training. This paper concerns with the layer-wise training in this line of approach. In , layer-wise training is employed as a pretraining strategy.
Deep linear network (DLN) is a neural network that uses linear activation functions. Although DLN is not a popular choice in practice, it is an active research subject as it is a class of decent simplified models for understanding the deep neural network with non-linear activation functions . DLN has a trivial representation power (product of weight matrices), however, its training process is not trivial at all. It has been studied the loss surface of DLNs and it is shown that although the loss surface is not convex, there are no spurious local minima. The works of studied a convergence analysis of gradient descent for DLNs, under various settings. showed that under some assumptions, the gradient descent finds a global optimum. The learning rate from the analysis, however, is not applicable in practice as it requires prior knowledge of the global minimizer. The theoretically guaranteed learning rate of should meet , where is the global minimizer, is a constant related to the initial error, and is the depth. showed that under the assumptions of Gaussian random initialization, and severely wide networks, the gradient descent finds a global optimum. The learning rate from the analysis of does not require any prior knowledge of and can be applied in practice. However, we found that it leads divergence of GD in all of our tests. The theoretically guaranteed width of is too large to be used. For example, if the condition number of the input data matrix (full rank) is , the width should be at least .
In this paper, we study a layer-wise training for DLNs using a block coordinate gradient descent (BCGD) . Similar to BCD, the BCGD trains each layer at a time in a sequential order by freezing all other layers at their last updated values. However, a key difference is the use of gradient descent in every update. We aim to identify the effects of depth, width, and initialization in the training process through the lens of DLNs. We first establish a general convergence analysis and found the optimal learning rate, which leads to the fastest decrease in the loss for the next iterate. More importantly, the optimal learning rate can directly be applied in practice. Neither trial and error nor a grid search for tuning parameters are required. To illustrate the performance of BCGD with the optimal learning rate, we consider a learning task of fitting 600 data (see Section 4.1 for details) and plot the training loss trajectories by BCGD and GD in Figure 1. Despite the fact that BCGD updates only a single matrix per iteration, while GD updates all the weight matrices per iteration, we clearly see that BCGD converges drastically faster than GD. The learning rates for GD are found by trial-and-error.
Next, we show that when the orthogonal-like initialization is employed, as long as the width of intermediate layers is greater than or equal to both the input and output dimensions, the width plays no role in any gradient-based training. Also, we rigorously prove that when (i) the orthogonal-like initialization is used, (ii) the initial loss is sufficiently small, whenever the depth is sufficiently large, the convergence to the global optimum (within machine accuracy) is guaranteed by updating each weight matrix only once. Furthermore, we found that a well-chosen depth could result in a significant acceleration in convergence when it is compared to those of a single layer, even when the computational cost is considered. This clearly demonstrates the benefit of using deep networks (over-parameterization via depth). Similar behavior was empirically reported in as implicit acceleration.
Lastly, we establish a convergence analysis of the block coordinate stochastic gradient descent (BCSGD). Our analysis indicates that the BCSGD cannot reach the global optimum, however, the converged loss will be staying close to the global optimum. This can be understood as an implicit regularization, which avoids over-fitting.
The rest of paper is organized as follows. In Section 2, we present the mathematical setup and introduce the block coordinate (stochastic) gradient descent. We then present a general convergence analysis and the optimal learning rate in Section 3. In Section 4, several numerical examples using both synthetic and real data sets are presented to demonstrate the effectiveness of the layer-wise training by BCGD and justify our theoretical findings.
Setup and Preliminary
Given a set of training data , the goal is to learn the parameters which minimize the loss function defined by
respectively. Here is the Euclidean norm, is the Frobenius norm, is the largest singular value, and is the -th largest singular value. Also, we denote the -th largest singular value by . When , we simply write the condition number as .
2 Block Coordinate Gradient Descent
If for all , we write for . For notational completeness, we set whenever . Also, we simply write as .
3 Initialization
Let be a matrix of size and be of size where . We say is equivalent to upto zero-valued padding if
Here is the identity matrix of size . We consider the following weight initialization schemes.
Orthogonal Initialization : for all , where is an orthogonal matrix of size .
Orth-Identity Initialization: for . This is a special case of orthogonal initialization that is proposed in the present work.
Identity Initialization : for .
Random Initialization: for all . Often is chosen to so that the expected value of the square norm of each row is 1.
The orth-indentity initialization can be viewed as a hybrid initialization between the orthogonal and the identity initialization schemes. This paper primarily concerns with the orth-indentity initialization.
Convergence Analysis
In this section, we present a convergence analysis of BCGD and establish the optimal learning rate. The optimality is defined to be the learning rate which results in the fastest decrease in the loss at the current parameters. The standard -loss will be mainly discussed. However, we also present a convergence result for general differentiable convex loss functions whose gradient are Lipshitz continuous in a bounded domain, such as -loss where is even. We measure the approximation error in terms of the distance to the global optimum. For example, when the -loss is employed, the error is .
We first identify the effects of width in DLNs in gradient-based training under either the orth-identity or the balanced initialization (Section 2.3).
Theorem 3.1 implies that the width does not play any role in gradient-based training if the condition of (5) is met and the weight matrices are initialized in a certain manner.
However, the same conclusion does not follow if the random initialization is employed. This indicates that the role of width highly depends on how the weight matrices are initialized. With a proper initialization, over-parameterization by the width can be avoided.
We first focus on the standard loss function and present a general convergence analysis of BCGD. We do not make any assumptions other than . We follow the convention of .
where , , , and
Furthermore, the optimal learning rate is
and with the optimal learning rate of (8), we obtain
The optimality of (8) should be understood in the sense that it gives the smallest loss for the next iterate.
In order for an iteration of BCGD to strictly decrease the approximation error, it is important to guarantee the condition of
In what follows, we show that if the initial approximation error is sufficiently close to the global optimum under the orth-identity initialization (Section 2.3), the convergence to the global optimum is guaranteed at a linear rate by the layer-wise training (BCGD).
and . Then, with the learning rates of (6), the -th sweep of BCGD satisfies
where and .
By Lemma 3.7, the proof readily follows from Theorem 3.3.
Under the same conditions of Theorem 3.5, we have
We remark that the rate of convergence for a single sweep is . When the speed of convergence is measured against the number of sweeps, this implies that the deeper the network is, the faster convergence is obtained. Thus, if the depth of a linear network is sufficiently large, the global optimum can be reached (within machine accuracy) by the layer-wise training (BCGD) after updating each weight matrix only once. Also, we note that the work of also has a similar initialization condition.
Theorem 3.5 relies on the assumption that the initial approximation is sufficiently close to the global optimum in terms of , and the depth . As a special case of , a similar result can be obtained without this restriction.
where is defined in (10) and . Then, the -th sweep of descending BCGD with the learning rate of (6) satisfies
where .
2 Convergence of BCGD for general convex loss function
We present a general convergence analysis of the layer-wise training (BCGD) for convex differentiable loss functions. For general loss functions, let be the solution to . For a matrix , the matrix norm is defined by
and the max norm is .
Note that when , the above is identical to the optimal learning rate of (8).
3 Convergence of BCSGD
In this subsection, a convergence analysis of BCSGD (16) is presented with the standard -loss.
Given a discrete random variable on , we denote the expectation with respect to conditioned on all other previous random variables by .
Then, the approximation by BCSGD (16) with the learning rates of
This indicates that unlike the BCGD, if a randomly chosen datum is used to update a weight matrix, an extra term, which is proportional to , is introduced in both upper and lower bounds of the expected error. Therefore, the BCSGD would not achieve the global optimum, unless . However, the expected loss by BCSGD will be within the distance proportional to from . In practice, will almost never be zero. This indicates that the stochasticity introduced by the random selection of mini-batch (of size 1) results in an implicit regularization effect, which avoids over-fitting. We defer further characterization of BCSGD to future work.
Remark: The proposed stochastic gradient-descent in Theorem 3.13 can be viewed as a generalized version of the sampling used in .
Numerical Examples
In what follows, we employ the layer-wise training by BCGD for deep linear neural networks. The learning rate is chosen to be (near) optimal according to (14). We emphasize that the (near) optimal learning rate of (14) does not require any prior knowledge, and can completely be determined by the loss function, the current weight matrices and the input data matrix. This allows us to avoid a cumbersome grid-search over learning rate. When the -loss is employed, the optimal learning rate of (8) is identical to the one of (14).
1.2 Big Condition number
We now consider the input data matrix whose condition number is rather big. To do this, we first generate as in the above and conduct the singular value decomposition. We then assign randomly generated numbers from to the singular values. In our experiment, the condition number of was 236. The output data matrix is generated in the same way as before. In Figure 3, the approximation errors are plotted with respect to the number of (left) sweeps and (right) iterations of the descending BCGD at different depths . When the speed of convergence is measured against the number of sweeps, we see that the deeper the network is, the faster the convergence is obtained. When the amount of computation is considered, unlike the case where has a good condition number, we now see that the errors by deep linear networks decay drastically faster than those by a shallow network of depth 1. This demonstrates that over-parameterization by the depth can indeed accelerate convergence, even when the computational cost is considered. We note that from Theorem 3.1, the width plays no role in gradient-based training, as the width of intermediate layers is . Furthermore, the optimal learning rate is employed and adding more layers does not increase any representational power. Therefore, this acceleration is solely contributed by the depth and this clearly demonstrates the benefit of using deep networks. We also observe that the error decrease per iteration does not grow proportionally to the depth. In this case, either depth 5 or 7 performs the best among others.
1.3 Comparison with GD
Next, we compare the performance between BCGD and the standard gradient descent (GD) on the two same tasks of Section 4.1.1 and 4.1.2. By trial-and-error, we choose constant learning rates for GD that leads the fastest convergence. We tried the learning rate of from , however, we observed that it makes GD diverge within few iterations. Despite the fact that GD updates all the weight matrices in a single iteration, while BCGD updates only a singe matrix, we compare the performance with respect to the number of iterations to emphasize the performance of BCGD. Figure 4 shows the approximation errors by GD and BCGD. The results for the task with a small (big) condition number are presented on the left (right). In both cases, we observe that GD converges linearly when learning rate is chosen properly. However, choosing an appropriate learning rate requires a time consuming fine tuning. It is also clear that GD is highly sensitive with respect to learning rate. For example, on the left of Figure 4, we see that GD with the learning rate of produces a linear convergence, however, GD with the learning rate of leads a highly oscillatory behavior in the error. When it is compared to the results by BCGD and also considering the fact that BCGD updates only a single matrix per iteration, it is clear that BCGD converges significantly faster than GD, especially when the model matrix has a big condition number. Furthermore, BCGD does not require one to put any efforts on finding a proper learning rate. This clearly demonstrate superior performances of BCGD over GD in these cases.
1.4 Effect of Width
1.5 Ascending versus Descending
2 Real Data Experiments
We employ the dataset from UCI Machine Learning Repository’s “Gas Sensor Array Drift at Different Concentrations” . Specifically, we used the dataset’s “Ethanol” problem — a scalar regression task with 2565 examples, each comprising 128 features (one of the largest numeric regression tasks in the repository). The input and output data sets are normalized to have zero mean and unit variance. After the normalization, the condition number of the input data matrix is 70,980. We note that this is the same data set used in . The width of intermediate layers is set to and the identity initialization (Section 2.3) is employed. On the left of Figure 7, we show the errors by the descending BCGD with respect to the number of iterations at five different depths . We use the optimal learning rate (8), which does not require any prior knowledge. We clearly see that the over-parameterization by depth significantly accelerates convergence. We remark that in the work of , although a different optimization method is used, the same problem is considered and the learning rate is chosen by a grid search. Similar implicit acceleration was demonstrated only for -loss, not -loss. In our experiment, by exploiting the layer-wise training and the optimal learning rate, we demonstrate implicit acceleration for -loss. On the right of Figure 7, we show the results by -loss, i.e,
The near optimal learning rate of (15) is employed. We observe that updating a single layer multiple-times results in the fastest error convergence than updating multiple layers once. In this case, there is no advantages of using deep networks. For reference, we also plot the best error shown at after 1,000,000 iterations as the dashed line. Unlike the conclusion of , we found that the depth leads to acceleration for the -loss, but not for the -loss.
We now train DLNs on the MNIST handwritten digit classification dataset. For an input image, its corresponding output vector contains a 1 in the index for the correct class and zeros elsewhere. The input and output dimensions are and , respectively. In order to strictly compare the effect of depth, we employ the identity initialization to completely remove the randomness from the initialization. Also, we set the width to according to Theorem 3.1. The networks are trained over the entire MNIST training dataset of 60,000 samples. The input data matrix is not full rank. Figure 8 shows the distances to the global optimum by -loss with respect to the number of iterations of the descending BCGD at ten different depths . Thus, the speed of convergence is measured against the amount of computation. We observe the accelerated convergence by the network whose depth is even but not odd. We also see that the results by DLNs of odd-depth are very similar so that the lines are overlapping each other. In this case, the depth 2 network performs the best among others. We suspect that there is a connection between the parity of depth and the acceleration in convergence. We defer such further investigation to future work.
Lastly, we compare the performance of BCGD to GD on the same real data sets. Again, the learning rates for GD are chosen by trial-and-error. Figure 9 shows the error trajectories for the same learning tasks. On the left and right, the results for the UCI and the MNIST datasets are presented, respectively. In all cases, we see that BCGD converges faster than GD while GD with a well-chosen learning rate converges linearly. We emphasize that GD updates all the weights matrices per iteration, while BCGD updates only a single matrix per iteration.
Conclusion
In this paper, we studied a layer-wise training for deep linear networks using the block coordinate gradient descent (BCGD). We established a convergence analysis and found the optimal learning rate which results in the fastest decrease in the loss for the next iterate. More importantly, the optimal learning rate can directly be applied in practice as no prior knowledge is required. Also, we identified the effects of depth, width, and initialization in the training process. Firstly, we showed that when the orthogonal-like initialization is employed and the width of the intermediate layers is great than or equal to both the input and output dimensions, the width plays no roles in gradient-based training. Secondly, under some assumptions, we proved that the deeper the network is, the faster the convergence is guaranteed (when the speed is measured against the number of sweeps). In an extreme case, the global optimum (within machine accuracy) is achieved after updating each weight matrix only once. Thirdly, we empirically demonstrated that adding more layers could drastically accelerate convergence, when it is compared to those of a single layer, even when the computational cost is considered. Lastly, we establish a convergence analysis of the block coordinate stochastic gradient descent (BCSGD). Our analysis indicates that the BCSGD cannot reach the global optimum, however, the converged loss will be staying close to the global optimum. This can be understood as an implicit regularization, which avoids over-fitting. Numerical examples were provided to justify our theoretical findings and demonstrate the performance of the layer-wise training by BCGD.
Acknowledgments
The author would like to thank Dr. Pual Dupuis for his helpful discussion in the early stages of this work, Dr. Mark Ainsworth for his helpful comments and suggestions on both analysis and examples, and Dr. Nadav Cohen for sharing code for numerical experiments.
Appendix A Least square solution
Without the rank constraint, the solution of (3) is
where is the identity matrix of size and is the Moore-Pensore pseudo-inverse of . Assuming is a full row rank matrix, we have , which allows an explicit formula . If is not a full row rank matrix, (3) allows infinitely many solutions. In this case, the least norm solution is often sought and it is . Also, for any , the following holds:
Thus, the minimizing -loss is equivalent to minimizing . Furthermore, for whitened data, the least norm solution is simply .
With the rank constraint, we consider two cases. If , the rank constrain plays no role in the minimization. Thus, the global minimizer is (19). Let us consider the case of . Let , and be a compact singular value decomposition (SVD) of where only left-singular vectors and right-singular vectors corresponding to the non-zero singular values are calculated. Then, . and it can be checked that . Let be a compact SVD of . It then can be shown that the problem (3) is equivalent to
To be more precise, if is a solution (the best -rank approximation to ) to the above, is a solution of (3), which can be explicitly written as
where and is the principal submatrix consisting of the first rows and columns of . We remark that in general, (20) and the best -rank approximation to are not the same.
Appendix B Gradient of the loss
Appendix C Proof of Theorem 3.1
For a matrix of size and a matrix of size where , we say is equivalent to upto zero-valued padding if
It then follows from the gradient descent update
which completes the proof for the balanced initialization.
Let be a matrix of size and for all . Suppose
Appendix D Proof of Theorem 3.3
For notational convenience, for , let
By definition, it follows from the update rule that
where is used in the 4th equality.
Thus, with the optimal learning rate of (23), we obtain
Note that from the relation of , we have
By recursively applying the above, we obtain
Appendix E Proof of Lemma 3.7
Then for any satisfying , we have
From Theorem 3.3, since for any , we obtain .
that . Similarly, .
By applying the induction on the number of iterations of the BCGD, we claim that there exists such that
The recursive relation with respect to gives
where , we have
Therefore, under the above assumption on , we have
that and . Also, since and are decreasing functions of , we have
Appendix F Proof of Theorem 3.9
where is defined in (10), it follows from Lemma 3.7 and Theorem 3.5 that for all and , and
Note that is from the fact that for all . Hence, we have
Appendix G Proof of Theorem 3.11
For notational convenience, for , let
It follows from the BCGD update rule that
It then can be checked that the learning rate which minimizes the above upper bound is
Appendix H Proof of Theorem 3.13
For notational convenience, for , let
By definition, it follows from the update rule that
Let and . Then
By summing up with respect to , we have