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 proves the convergence rate of gradient flow for training two-layer linear networks. In particular, this paper discusses the connection between initialization, optimization, generalization, and overparameterization. The results show that gradient flow can converge to the global minimum at a rate depending on the level of imbalance of the initialization. Moreover, the authors show that random initialization and overparameterization can implicitly constrain the gradient flow trajectory to converge to a point lying in a low-dimensional manifold, thus guarantees good generalization ability.