On the Explicit Role of Initialization on the Convergence and Generalization Properties of Overparametrized Linear Networks

1 INTRODUCTION . Neural networks have shown excellent empirical performance in many application domains such as vision ( Krizhevsky et al. , 2012 ; Rawat & Wang , 2017 ) , speech ( Hinton et al. , 2012 ; Graves et al. , 2013 ) and video games ( Silver et al. , 2016 ; Vinyals et al. , 2017 ) . Among the many unexplained mysteries behind this success is the fact that gradient descent with random initialization and without explicit regularization enjoys good generalization performance despite being highly overparametrized . A promising attempt to explain this phenomena is the Neural Tangent Kernel ( NTK ) ( Jacot et al. , 2018 ) , which characterizes the implicit regularization effect of gradient flow/descent on infinitely wide neural networks with random initialization . Precisely , under this infinite width assumption , a proper initialization , together with gradient flow training , can be understood as a kernel gradient flow ( NTK flow ) of a functional that is constrained on a manifold that guarantees good generalization performance . The analysis further admits extensions to the finite width ( Arora et al. , 2019b ; Buchanan et al. , 2020 ) . The core of the argument , illustrated in Figure 1 , amounts to showing that ( properly scaled ) random initialization of networks with sufficiently large width , leads to trajectories that are , in some sense initialized close to the aforementioned manifold . Thus , approximately good initialization , together with acceleration , ensures that the training stays close to the NTK Flow . Such analysis , however , leads to bounds on the network width that significantly exceed practical values ( Buchanan et al. , 2020 ) , and seems to suggest the need of acceleration to achieve generalization . This motivates the following questions : • Is the kernel regime , which requires impractical bounds on the network width , necessary to achieve good generalization ? • Does generalization depends explicitly on acceleration ? Or is acceleration required only due to the choosing an initialization outside the good generalization manifold ? For the simplified , yet certainly non-trivial , single-hidden layer linear network setting , this paper finds an answer to these questions . Contributions . We present a novel analysis of the gradient flow dynamics of overparametrized single-hidden layer linear networks , which provides disentangled conditions on initialization that lead to acceleration and generalization . Specifically , we show that imbalanced initialization ensures acceleration , while orthogonal initialization ensures that trajectories remain close to the generalization manifold . Interestingly , properly scaled random initialization of moderately wide networks is sufficient to ensure that initialization is approximately imbalanced and orthogonal , yet it is not necessary for either . More specifically , as illustrated in Figure 1 , this paper makes the following contributions : 1 . We show first that gradient flow on the squared-l2 loss preserves a certain matrix-valued quantity , akin to constants of motion in mechanics or conservation laws of physics , that measures the imbalance of the network weights . Notably , some level of imbalance , measured by certain singular value of the imbalance matrix and defined at initialization , is sufficient to guarantee the exponential rate of convergence of the loss . Our analysis is non-probabilistic and valid under very mild assumptions , satisfied by moderately wide single-hidden-layer linear networks . 2 . We characterize the existence of a low-dimensional manifold defined by a specific orthogonality condition on the parameters , which is invariant under the gradient flow . All trajectories within this manifold lead to a unique ( w.r.t the end-to-end function ) minimizer with good generalization performance , which corresponds to the min-norm solution . As a result , initializing the network within this manifold guarantees good generalization . 3 . We further show that by randomly initializing the network weights using N ( 0 , 1/h ) ( where h is the hidden layer width ) , an initialization setting related to the kernel regime ( see Appendix E ) , one can approximately satisfy both our sufficient imbalance and orthogonality conditions with high probability . Notably , the inaccurate initialization relative to the good generalization manifold requires acceleration to control the generalization error . In the context of NTK , our result further provide for linear networks a novelO ( h−1/2 ) upper-bound on the operator norm distance between the trained network and the min-norm solution . To the best of our knowledge , this is the first non-asymptotic bound regarding the generalization property of wide linear networks under random initialization in the global sense . Notation . For a matrix A , we let AT denote its transpose , tr ( A ) denote its trace , λi ( A ) and σi ( A ) denote its i-th eigenvalue and i-th singular value , respectively , in decreasing order ( when adequate ) . We let [ A ] ij , [ A ] i , : , and [ A ] : ,j denote the ( i , j ) -th element , the i-th row and the j-th column of A , respectively . We also let ‖A‖2 and ‖A‖F denote the spectral norm and the Frobenius norm of A , respectively . For a scalar-valued or matrix-valued function of time , F ( t ) , we let Ḟ = Ḟ ( t ) = ddtF ( t ) denote its time derivative . Additionally , we let In denote the identity matrix of order n andN ( µ , σ2 ) denote the normal distribution with mean µ and variance σ2 . 2 RELATED WORK . Wide neural networks . There is a rich line of research that studies the convergence ( Du et al. , 2019b ; a ; Du & Hu , 2019 ; Allen-Zhu et al. , 2019b ) and generalization ( Allen-Zhu et al. , 2019a ; Arora et al. , 2019a ; b ; Li & Liang , 2018 ; Cao & Gu , 2019 ; Buchanan et al. , 2020 ) of wide neural networks with random initialization . The behavior of such networks in their infinite width limit can be characterized by the Neural Tangent Kernel ( NTK ) ( Jacot et al. , 2018 ) . With the concept of NTK , heuristically , training wide neural networks can be approximately viewed as kernel regression under gradient flow/descent ( Arora et al. , 2019b ) , hence the convergence and generalization can be understood by studying the non-asymptotic results regarding the equivalence of finite width networks to their infinite limit ( Du et al. , 2019b ; a ; Du & Hu , 2019 ; Allen-Zhu et al. , 2019b ; Arora et al. , 2019a ; b ; Buchanan et al. , 2020 ) . More generally , such non-asymptotic results are related to the “ lazy training ” ( Chizat et al. , 2019 ; Du et al. , 2019a ; Allen-Zhu et al. , 2019b ) , where the network weights do not deviate too much from its initialization during training . Our results for wide linear networks presented in Section 4.2 do not follow the NTK analysis , but provide an alternative , presumably more general view on the effect of random initialization when the hidden layer is sufficiently wide . Convergence of linear networks . Convergence in overparametrized linear networks has been studied for both gradient flow ( Saxe et al. , 2014 ) and gradient descent ( Bartlett et al. , 2018 ; Arora et al. , 2018a ; b ) . In the kernel regime , Du & Hu ( 2019 ) applied the analysis of convergence of wide neural networks ( Du et al. , 2019b ) to deep linear networks . Aside from the large hidden layer width and random initialization assumptions , Saxe et al . ( 2014 ) analyzed the trajectory of network parameters under spectral initialization , while Bartlett et al . ( 2018 ) studied the case of identity initialization . Although the fact that the imbalance is conservative under gradient flow has been exploited in Arora et al . ( 2018b ; a ) , they consider the case of balanced initialization only to simplify the learning dynamics , and additional conditions are required for convergence . Most works mentioned above considered specific , often data-dependent , types of initialization that makes the learning dynamics tractable . Our result , based on an imbalance measure , is data-agnostic and satisfied under a wide range of random initialization schemes , see lemmas 1 and F.2 . Min-norm solution in high-dimension linear regression . For high-dimensional under-determined linear regression , the asymptotic generalization error of min-norm solution has been studied in Hastie et al . ( 2019 ) . In Bartlett et al . ( 2020 ) ; Mei & Montanari ( 2019 ) , the min-norm solution was proved to have near-optimal generalization performance under mild assumptions on the data model . For our purpose , we study the generalization of a trained linear network as its distance to the min-norm solution . In this way , we refer to solutions with good generalization performance , as those with small distance to the min-norm solution . 3 CONVERGENCE RATE OF GRADIENT FLOW FOR SINGLE-HIDDEN-LAYER LINEAR NETWORKS . In this section we study the convergence of gradient flow on squared l2-loss for single-hidden-layer linear networks . Given training data of n samples { x ( i ) , y ( i ) } ni=1 with x ( i ) ∈ RD , y ( i ) ∈ Rm , we aim to solve the linear regression problem min Θ∈RD×m L = 1 2 n∑ i=1 ( y ( i ) −ΘTx ( i ) ) 2 , ( 1 ) by training a single-hidden-layer linear network y = f ( x ; V , U ) = V UTx , V = Rm×h , U ∈ RD×h with gradient flow , or equivalently , gradient descent with “ infinitesimal step size ” , here h is the hidden layer width . We consider an overparametrized model such that h ≥ min { m , D } . We rewrite the loss function with respect to our parameters V , U as L ( V , U ) = 1 2 n∑ i=1 ( y ( i ) − V UTx ( i ) ) 2 = 1 2 ‖Y −XUV T ‖2F , ( 2 ) where Y = [ y ( 1 ) , · · · , y ( n ) ] T , X = [ x ( 1 ) , · · · , x ( n ) ] T are concatenations of the training data in rows . Assuming that the input data X has full rank , we consider the under-determined case D > n for our regression problem , i.e. , there are infinitely many solutions Θ∗ to achieve zero loss of ( 1 ) . With minor a reformulation , our convergence result works for the case where the input data X is rank deficient . We refer the reader to Appendix B . We will show that under certain conditions , the trajectory of the loss function L ( t ) = L ( V ( t ) , U ( t ) ) under gradient flow of ( 2 ) , i.e. , V̇ ( t ) = − ∂ ∂V L ( V ( t ) , U ( t ) ) , U̇ ( t ) = − ∂ ∂U L ( V ( t ) , U ( t ) ) , ( 3 ) converges to 0 exponentially , and that proper initialization of U ( 0 ) , V ( 0 ) controls the convergence rate via a time-invariant matrix-valued term , the imbalance of the network .

Summary: This paper studies the optimization and generalization properties of a two-layer linear network. The considered setting is over-parameterized linear regression where the input dimension is D, number of samples is n<D, and the target dimension is m. The hidden width is h. The paper has two main results. The first result is exponential convergence of gradient flow to global minimum, where the convergence rate depends on the (m+n-1)-th singular value of an "imbalance" matrix. The second result shows that the solution found is close to the minimum L2 norm solution if certain orthogonality assumption is approximately satisfied at initially; then it was shown that if the width h is sufficiently large, then under a random initialization scheme, the solution found is close to the minimum L2 norm solution with a distance $1/\sqrt{h}$.