Training Multi-Layer Over-Parametrized Neural Network in Subquadratic Time

1 INTRODUCTION . Deep neural networks have been playing a central role in both practical ( such as computer vision ( LeCun et al. , 1998 ; Krizhevsky et al. , 2012 ; Szegedy et al. , 2015 ; He et al. , 2016 ) , natural language processing ( Collobert et al. , 2011 ; Devlin et al. , 2018 ) , automatic driving system , game playing ( Silver et al. , 2016 ; 2017 ) ) and theoretical machine learning community ( Li & Liang ( 2018 ) ; Jacot et al . ( 2018 ) ; Du et al . ( 2019b ) ; Allen-Zhu et al . ( 2019a ; b ) ; Du et al . ( 2019a ) ; Song & Yang ( 2019 ) ; Brand et al . ( 2021 ) ; Zou et al . ( 2018 ) ; Cao & Gu ( 2019 ) ; Lee et al . ( 2019a ) ; Liu et al . ( 2020 ; 2021 ) ; Chen et al . ( 2021 ) ) . In order to analyze the dynamics of neural networks and obtain provable guarantees , using over-parametrization has been a growing trend . In terms of understanding the convergence behavior of over-parametrized networks , most of the attentions have been directed to the study of first-order method such as gradient descent or stochastic gradient descent . The widespread use of first-order method is explained , to a large degree , by its computational efficiency , since computing the gradient of the loss function at each iteration is usually cheap and simple , let alone with its compatibility with random sampling-based method such as minibatch . One of the major drawbacks of first-order methods is their convergence rate is typically slow in many non-convex settings ( poly ( n , L , log ( 1/ ) ) iterations , where n is the number of training samples , L is the number of layers and is the precision of training ) , e.g. , deep neural network with ReLU activation , as shown in Allen-Zhu et al . ( 2019a ) ) , which is often the case of a deep over-parameterized neural network . Second-order method ( which employs the information of the Hessian matrix ) on the other hand enjoys a much faster convergence rate ( only log ( 1/ ) iterations ( Zhang et al. , 2019 ) , but not poly ( n ) log ( 1/ ) iterations ) and exploits the local geometry of the loss function to overcome the pathological curvature issues that are critical in first-order method . Another clear advantage of second-order method over first-order method is it does not require the tuning of learning rate . The expense of using second-order method is its prohibitive cost per iteration , as it is imperative to invert a dynamically-changing Hessian matrix or equivalently , solving a regression task involving Hessian . Given any weight matrix of size m × m ( m is the width of network ) , its Hessian matrix has size m2 ×m2 , which makes any naive implementation of second-order algorithm takes at least O ( m4 ) time since one needs to write down the Hessian . This explains the scarcity of deploying large-scale second-order method in non-convex setting , such as training deep neural networks , in contrast to their popular presence in convex optimization setting ( Vaidya ( 1989 ) ; Daitch & Spielman ( 2008 ) ; Lee et al . ( 2015 ) ; Cohen et al . ( 2019 ) ; Lee et al . ( 2019b ) ; Jiang et al . ( 2020b ; a ) ; Song & Yu ( 2021 ) ) . Recent works ( Cai et al . ( 2019 ) ; Zhang et al . ( 2019 ) ) improved the practicality of second-order method on training deep networks and presented algorithms to train one-hidden layer overparametrized networks with smooth activation functions . Specifically , they achieve a running time ofO ( mn2 ) per iteration . Their methods are essentially variants of Gauss-Newton method combined with using neural tangent kernels ( Jacot et al . ( 2018 ) ) to prove the convergence . By using the idea of randomized linear algebra ( Clarkson & Woodruff ( 2013 ) ; Woodruff ( 2014 ) ) and a clever sketching matrix as a preconditioner , Brand et al . ( 2021 ) further improves the running time to Õ ( mn ) per iteration . However , all of these algorithms are for training a shallow network with one-hidden layer and fall short on deep networks — First , it is not clear that their algorithms can be generalized to multi-layer setting , due to the presence of gradient vanishing or exploding . In the seminal work of Allen-Zhu et al . ( 2019a ) , they showed that as long as the networks are over-parametrized , first-order methods such as gradient descent and stochastic gradient descent won ’ t encounter such problem . But does this still hold for second-order method ? Can we provably show that second-order method has a good performance when training deep over-parametrized networks ? Second , even the fastest of them ( Brand et al . ( 2021 ) ) would incur a running time of Õ ( m2nL ) per iteration , which seems unavoidable due to the size of intermediate weight matrices is m×m . In this work , we take the first step to tame the beast — We propose a second-order method that achieves subquadratic cost per iteration with respect to m , and show that it has linear convergence rate in training deep over-parametrized neural networks . We emphasize the importance of obtaining subquadratic algorithm , since in multi-layer settings , the network width is typically much larger than in one-hidden layer setting ( m ≥ n8L12 , ( Zou & Gu , 2019 ) ) . Our work can be decomposed into two parts : algorithmically and analytically . From an algorithmic perspective , our method builds upon a variant of Gauss-Newton method ( Björck ( 1996 ) ) called Gram-Gauss-Newton method ( Cai et al . ( 2019 ) ; Brand et al . ( 2021 ) ) . In order to achieve a feasible running time , we exploit two features of the gradient , which is the key ingredient to form the Jacobian matrix : 1 ) . The gradient is low rank ( rank n ) . 2 ) . The gradient can be formulated as the outer product of two vectors . From an analytical perspective , our work is inspired by Allen-Zhu et al . ( 2019a ) . In contrast to their proof which is a straightforward analysis of the gradient , we make use of the multi-layer neural tangent kernels ( Du et al . ( 2019a ) ) and establish a connection between a Gram matrix we compute at each iteration and the NTK matrix . Our Contributions . We summarize our technical contributions below . • We develop an analytical framework for the convergence behavior of second-order method on training multi-layer over-parametrized neural network . To facilitate the analysis , we exploit the equivalence between neural tangent kernels and our over-parametrized network . • We design a second-order algorithm to train such networks , and achieve a cost per iteration of õ ( m2 ) . Our algorithm makes use of Gram-Gauss-Newton method , tensor-based sketching techniques , and data structures that maintains a low rank representation efficiently . • By combining fast tensor algebra techniques and sketching-based preconditioning , we devise an algorithm to efficiently solve a regression problem where the involved matrix has its rows being tensor product of vectors . 1.1 OUR RESULT . Our main result can be summarized in the following three theorems , with one analyzing the convergence behavior of a general Gram-based optimization framework , one designing an efficient algorithm to realize this second-order optimization scheme , and the other is a novel algorithm to solve tensor-based regression in high precision and fast , which is a key step in our second-order method . Throughout this paper , we will use n to denote the number of training data points , d to denote the dimension of input data points , m to denote the width of the network and L to denote the number of layers of the network . We use ft ∈ Rn to denote the prediction of neural network at time t. Our first theorem demonstrates the fast convergence rate of our algorithm . Theorem 1.1 ( Convergence , informal version of Theorem F.19 ) . Suppose the width of the neural network satisfies m ≥ poly ( n , L ) , then there exists an algorithm ( Algorithm 1 ) such that , over the randomness of initialization of the network and the algorithm , with probability at least 1 − e−Ω ( log 2m ) , we have ‖ft+1 − y‖2 ≤ 1 2 ‖ft − y‖2 , where ft ∈ Rn is the the prediction produced by neural network at time t. The above theorem establishes the linear convergence behavior of our second-order method , which is a standard convergence result for second-order method , as well as the same behavior as in onehidden layer over-parametrized networks ( Brand et al . ( 2021 ) ) . However , compared to one-hidden layer case , our analysis is much more sophisticated since we have to carefully control the probability so that it does not blow up exponentially with respect to the number of layers . The next theorem concerns the cost per iteration of our second-order algorithm . Theorem 1.2 ( Runtime , informal version of Theorem B.1 ) . There exists a randomized algorithm ( Algorithm 1 ) that trains a multi-layer neural network of widthm with the cost per training iteration being O ( m2−Ω ( 1 ) ) . We improve the overall training time of multi-layer over-parametrized networks with second-order method from Tinit +T ·O ( m2 ) to Tinit +T · o ( m2 ) , where Tinit is the initialization time of training , typically takes O ( m2 ) . As we have argued before , multi-layer over-parametrized networks require m to be in the order of n8 , hence improving the cost per iteration from quadratic to subquadratic is an important gain in speeding up training . Its advantage is even more evident when one seeks a high precision solution , and hence the number of iterations T is large . We highlight that it is non-trivial to obtain a subquadratic running time per iteration : If not handled properly , computing the matrix-vector product with weight matrices will take O ( m2 ) time ! This means that even for first-order methods such as gradient descent , it is not clear how to achieve a subquadratic running time , since one has to multiply the weight matrix with a vector in both forward evaluation and backpropagation . In our case , we have also a Jacobian matrix of size n × m2 , so forming it naively will cost O ( nm2 ) time , which is prohibitively large . Finally , note that the update matrix is also an m×m matrix . In order to circumvent these problems , we exploit the fact that the gradient is of low rank ( rank n ) , hence one can compute a rank-n factorization and use it to support fast matrix-vector product . We also observe that each row of the Jacobian matrix can be formulated as a tensor product of two vectors , therefore we can make use of fast randomized linear algebra to approximate the tensor product efficiently . As a byproduct , we have the following technical theorem : Theorem 1.3 ( Fast Tensor Regression , informal version of Theorem D.14 ) . Given two n×m matrices U and V with m n and a target vector c ∈ Rn . Let J = [ vec ( u1v > 1 ) > , . . . , vec ( unv > n ) > ] ∈ Rn×m2 where ui is the i-th row of matrix U ∀i ∈ [ n ] . There is a randomized algorithm that takes Õ ( nm+ n2 ( log ( κ/ ) + log ( m/δ ) ) + nω ) time and outputs a vector x̂ ∈ Rn such that ‖JJ > x̂− c‖2 ≤ ‖c‖2 holds with probability at least 1− δ , and κ is the condition number of J . From a high level , the algorithm proceeds as follows : given matrices U and V , it forms an approximation J̃ ∈ Rn×n log ( m/δ ) , where each row is generated by applying fast tensor sketching technique to ui and vi ( Ahle et al . ( 2020 ) ) . Then , it uses another sketching matrix for J̃ to obtain a good preconditioner R for J̃ . Subsequently , it runs a gradient descent to solve the regression . To understand this runtime better , we note that nm term is the size of matrices U and V , hence reading the entries from these matrices will take at least O ( nm ) time . The algorithm then uses tensor-based sketching techniques ( Ahle et al . ( 2020 ) ) to squash length m2 tensors to length O ( n log ( m/ δ ) ) . All subsequent operations are performed on these much smaller vectors . Finally , computing the preconditioner takes Õ ( nω ) time , and running the gradient descent takes Õ ( n2 log ( κ/ ) ) time .

Summary: 1. This paper proved that the second-order method can minimize the training loss in linea rate on multi-layer over-parameterized neural networks. This analysis relies on the connection between neural tangent kernel and over-parameterized neural networks. 2. This paper also reduced the per-iteration cost of second-order methods to $\widetilde{o}(m^2)$ where $m$ is the hidden layer width, by combining Gram-Gauss-Newton method, tensor-based sketching, and efficient data structures that maintain low-rank representations.  3. This paper also designed an algorithm to efficiently solve a regression problem in which the matrix has its rows being tensor products of two vectors. This algorithm may be of general interest.