Optimal Rates for Averaged Stochastic Gradient Descent under Neural Tangent Kernel Regime

1 INTRODUCTION . Recent studies have revealed why a stochastic gradient descent for neural networks converges to a global minimum and why it generalizes well under the overparameterized setting in which the number of parameters is larger than the number of given training examples . One prominent approach is to map the learning dynamics for neural networks into function spaces and exploit the convexity of the loss functions with respect to the function . The neural tangent kernel ( NTK ) ( Jacot et al. , 2018 ) has provided such a connection between the learning process of a neural network and a kernel method in a reproducing kernel Hilbert space ( RKHS ) associated with an NTK . The global convergence of the gradient descent was demonstrated in Du et al . ( 2019b ) ; Allen-Zhu et al . ( 2019a ) ; Du et al . ( 2019a ) ; Allen-Zhu et al . ( 2019b ) through the development of a theory of NTK with the overparameterization . In these theories , the positivity of the NTK on the given training examples plays a crucial role in exploiting the property of the NTK . Specifically , the positivity of the Gram-matrix of the NTK leads to a rapid decay of the training loss , and thus the learning dynamics can be localized around the initial point of a neural network with the overparameterization , resulting in the equivalence between two learning dynamics for neural networks and kernel methods with the NTK through a linear approximation of neural networks . Moreover , Arora et al . ( 2019a ) provided a generalization bound of O ( T−1/2 ) , where T is the number of training examples , on a gradient descent under the positivity assumption of the NTK . These studies provided the first steps in understanding the role of the NTK . However , the eigenvalues of the NTK converge to zero as the number of examples increases , as shown in Su & Yang ( 2019 ) ( also see Figure 1 ) , resulting in the degeneration of the NTK . This phenomenon indicates that the convergence rates in previous studies in terms of generalization are generally slower than O ( T−1/2 ) owing to the dependence on the minimum eigenvalue . Moreover , Bietti & Mairal ( 2019 ) ; Ronen et al . ( 2019 ) ; Cao et al . ( 2019 ) also supported this observation by providing a precise estimation of the decay of the eigenvalues , and Ronen et al . ( 2019 ) ; Cao et al . ( 2019 ) proved the spectral bias ( Rahaman et al. , 2019 ) for a neural network , where lower frequencies are learned first using a gradient descent . By contrast , several studies showed faster convergence rates of the ( averaged ) stochastic gradient descent in the RKHS in terms of the generalization ( Cesa-Bianchi et al. , 2004 ; Smale & Yao , 2006 ; Ying & Zhou , 2006 ; Neu & Rosasco , 2018 ; Lin et al. , 2020 ) . In particular , by extending the results in a finite-dimensional case ( Bach & Moulines , 2013 ) , Dieuleveut & Bach ( 2016 ) ; Dieuleveut et al . ( 2017 ) showed convergence rates of O ( T −2rβ 2rβ+1 ) depending on the complexity r ∈ [ 1/2 , 1 ] of the target functions and the decay rate β > 1 of the eigenvalues of the kernel ( a.k.a . the complexity of the hypothesis space ) . In addition , extensions to the random feature settings ( Rahimi & Recht , 2007 ; Rudi & Rosasco , 2017 ; Carratino et al. , 2018 ) , to the multi-pass variant ( Pillaud-Vivien et al. , 2018b ) , and to the tail-averaging and mini-batching variant ( Mücke et al. , 2019 ) have been developed . Motivation . The convergence rate of O ( T −2rβ 2rβ+1 ) is always faster than O ( T−1/2 ) and is known as the minimax optimal rate ( Caponnetto & De Vito , 2007 ; Blanchard & Mücke , 2018 ) . Hence , a gap exists between the theories regarding NTK and kernel methods . In other words , there is still room for an investigation into a stochastic gradient descent due to a lack of specification of the complexities of the target function and the hypothesis space . That is , to obtain faster convergence rates , we should specify the eigenspaces of the NTK that mainly contain the target function ( i.e. , the complexity of the target function ) , and specify the decay rates of the eigenvalues of the NTK ( i.e. , the complexity of the hypothesis space ) , as studied in kernel methods ( Caponnetto & De Vito , 2007 ; Steinwart et al. , 2009 ; Dieuleveut & Bach , 2016 ) . In summary , the fundamental question in this study is Can stochastic gradient descent for overparameterized neural networks achieve the optimal rate in terms of the generalization by exploiting the complexities of the target function and hypothesis space ? In this study , we answer this question in the affirmative , thereby bridging the gap between the theories of overparameterized neural networks and kernel methods . 1.1 CONTRIBUTIONS . The connection between neural networks and kernel methods is being understood via the NTK , but it is still unknown whether the optimal convergence rate faster than O ( T−1/2 ) is achievable by a certain algorithm for neural networks . This is the first paper to overcome technical challenges of achieving the optimal convergence rate under the NTK regime . We obtain the minimax optimal convergence rates ( Corollary 1 ) , inherited from the learning dynamics in an RKHS , for an averaged stochastic gradient descent for neural networks . That is , we show that smooth target functions efficiently specified by the NTK are learned rapidly at faster convergence rates than O ( 1/ √ T ) . Moreover , we obtain an explicit optimal convergence rate of O ( T −2rd 2rd+d−1 ) for a smooth approximation of the ReLU network ( Corollary 2 ) , where d is the dimensionality of the data space and r is the complexity of the target function specified by the NTK of the ReLU network . 1.2 TECHNICAL CHALLENGE . The key to showing a global convergence ( Theorem 1 ) is making the connection between kernel methods and neural networks in some sense . Although this sort of analysis has been developed in several studies ( Du et al. , 2019b ; Arora et al. , 2019a ; Weinan et al. , 2019 ; Arora et al. , 2019b ; Lee et al. , 2019 ; 2020 ) , we would like to emphasize that our results can not be obtained by direct application of their results . A naive idea is to simply combine their results with the convergence analysis of the stochastic gradient descent for kernel methods , but it does not work . The main reason is that we need the L2-bound weighted by a true data distribution on the gap between dynamics of stochastic gradient descent for neural networks and kernel methods if we try to derive a convergence rate of population risks for neural networks from that for kernel methods . However , such a bound is not provided in related studies . Indeed , to the best of our knowledge , all related studies make this kind of connection regarding the gap on training dataset or sample-wise high probability bound ( Lee et al. , 2019 ; Arora et al. , 2019b ) . That is , a statement “ for every input data x with high probability |g ( t ) nn ( x ) − g ( t ) ntk ( x ) | < ” can not yield a desired statement “ ‖g ( t ) nn − g ( t ) ntk‖L2 ( ρX ) < ” where g ( t ) nn and g ( t ) ntk are t-th iterate of gradient descent for a neural network and corresponding iterate described by NTK , and ‖·‖L2 ( ρX ) is the L2-norm weighted by a marginal data distribution ρX over the input space . Moreover , we note that we can not utilize the positivity of the Gram-matrix of NTK which plays a crucial role in related studies because we consider the population risk with respect to ‖ · ‖L2 ( ρX ) rather than the empirical risk . To overcome these difficulties we develop a different strategy of the proof . First , we make a bound on the gap between two dynamics of the averaged stochastic gradient descent for a two-layer neural network and its NTK with width M ( Proposition A ) , and obtain a generalization bound for this intermediate NTK ( Theorem A in Appendix ) . Second , we remove the dependence on the width of M from the intermediate bound . These steps are not obvious because we need a detailed investigation to handle the misspecification of the target function by an intermediate NTK . Based on detailed analyses , we obtain a faster and precise bound than those in previous results ( Arora et al. , 2019a ) . The following is an informal version of Proposition A providing a new connection between a two-layer neural networks and corresponding NTK with width M . Proposition 1 ( Informal ) . Under appropriate conditions we simultaneously run averaged stochastic gradient descent for a neural network with width of M and for its NTK . Assume they share the same hyper-parameters and examples to compute stochastic gradients . Then , for arbitrary number of iterations T ∈ Z+ and > 0 , there exists M ∈ Z+ depending only on T and such that ∀t ≤ T , ‖g ( t ) nn − g ( t ) ntk‖L∞ ( ρX ) ≤ , where g ( t ) nn and g ( t ) ntk are iterates obtained by averaged stochastic gradient descent . This proposition is the key because it connects two learning dynamics for a neural network and its NTK through overparameterization without the positivity of the NTK . Instead of the positivity , this proposition says that overparameterization increases the time stayed in the NTK regime where the learning dynamics for neural networks can be characterized by the NTK . As a result , the averaged stochastic gradient descent for the overparameterized two-layer neural networks can fully inherit preferable properties from learning dynamics in the NTK as long as the network width is sufficiently large . See Appendix A for detail . 1.3 ADDITIONAL RELATED WORK . Besides the abovementioned studies , there are several works ( Chizat & Bach , 2018b ; Wu et al. , 2019 ; Zou & Gu , 2019 ) that have shown the global convergence of ( stochastic ) gradient descent for overparameterized neural networks essentially relying on the positivity condition of NTK . Moreover , faster convergence rates of the second-order methods such as the natural gradient descent and GaussNewton method have been demonstrated ( Zhang et al. , 2019 ; Cai et al. , 2019 ) in the similar setting , and the further improvement of Gauss-Newton method with respect to the cost per iteration has been conducted in Brand et al . ( 2020 ) . There have been several attempts to improve the overparameterization size in the NTK theory . For the regression problem , Song & Yang ( 2019 ) has succeeded in reducing the network width required in Du et al . ( 2019b ) by utilizing matrix Chernoff bound . For the classification problem , the positivity condition can be relaxed to a separability condition using another reference model ( Cao & Gu , 2019a ; b ; Nitanda et al. , 2019 ; Ji & Telgarsky , 2019 ) , resulting in mild overparameterization and generalization bounds of O ( T−1/2 ) or O ( T−1/4 ) on classification errors . For an averaged stochastic gradient descent on classification problems in RKHSs , linear convergence rates of the expected classification errors have been demonstrated in Pillaud-Vivien et al . ( 2018a ) ; Nitanda & Suzuki ( 2019 ) . Although our study focuses on regression problems , we describe how to combine their results with our theory in the Appendix . The mean field regime ( Nitanda & Suzuki , 2017 ; Mei et al. , 2018 ; Chizat & Bach , 2018a ) that is a different limit of neural networks from the NTK is also important for the global convergence analysis of the gradient descent . In the mean field regime , the learning dynamics follows the Wasserstein gradient flow which enables us to establish convergence analysis in the probability space . Moreover , several studies ( Allen-Zhu & Li , 2019 ; Bai & Lee , 2019 ; Ghorbani et al. , 2019 ; Allen-Zhu & Li , 2020 ; Li et al. , 2020 ; Suzuki , 2020 ) attempt to show the superiority of neural networks over kernel methods including the NTK . Although it is also very important to study the conditions beyond the NTK regime , they do not affect our contribution and vice versa . Indeed , which method is better depends on the assumption on the target function and data distribution , so it is important to investigate the optimal convergence rate and optimal method in each regime . As shown in our study , the averaged stochastic gradient descent for learning neural network achieves the optimal convergence rate if the target function is included in RKHS associated with the NTK with the small norm . It means there are no methods that outperform the averaged stochastic gradient descent under this setting .

Summary: This paper considers the optimization of a wide two layers neural network (for a regression task) using averaged SGD. The authors consider the Neural Tangent Kernel (NTK) regime. The NTK is a kernel defined using the activation function and the initial distribution of the parameters of the input layer. The RKHS H associated to this NTK is assumed to contain the Bayes predictor. Based on this, the authors derive a convergence rate for the predictor constructed from the T-th iterate of averaged SGD and the Bayes predictor in terms of the L2 distance wrt the distribution of the features. By specifying this bound in terms of the decay of the eigenvalues of the integral operator in H, they obtain an explicit generalization error bound, which is optimal for the class of problems considered in the paper.