The Three Stages of Learning Dynamics in High-dimensional Kernel Methods

1 INTRODUCTION . In order to fundamentally understand how and why deep learning works , there has been much effort to understand the dynamics of neural networks trained by gradient descent based algorithms . This effort has led to the discovery of many intriguing empirical phenomena ( e.g . Frankle et al . ( 2020 ) ; Fort et al . ( 2020 ) ; Nakkiran et al . ( 2019a ; b ; 2020 ) ) that help shape our conceptual framework for understanding the learning process in neural networks . Nakkiran et al . ( 2019b ) provides evidence that SGD starts by first learning a linear classifier and over time learns increasingly complex functions . Nakkiran et al . ( 2020 ) introduces the “ deep bootstrap ” phenomenon : for some deep learning tasks the empirical world test error remains close to the oracle world error1 for many SGD iterations , even if the empirical training and test errors display a large gap . To better understand such phenomena , it is useful to study training dynamics in related but mathematically tractable settings . One approach for theoretical investigation is to study kernel methods , which were recently shown to have a tight connection with over-parameterized neural networks ( Jacot et al. , 2018 ; Du et al. , 2018 ) . Indeed , consider a sequence of neural networks ( fN ( x ; θ ) ) N∈N with the widths of the layers going to infinity as N → ∞ . Assuming proper parametrization and initialization , for large N the SGD dynamics on fN is known to be well approximated by the corresponding dynamics on the first-order Taylor expansion of fN around its initialization θ0 , fN , lin ( x ; θ ) = fN ( x ; θ 0 ) + 〈∇θfN ( x ; θ0 ) , θ − θ0〉 . Thus , in the large width limit it suffices to study the dynamics on the linearization fN , lin . When using the squared loss , these dynamics correspond to optimizing a kernel least-squares objective with the neural tangent kernel KN ( x , x′ ) = 〈∇θfN ( x ; θ0 ) , ∇θfN ( x′ ; θ0 ) 〉 . 1Their paper uses “ Ideal World ” for “ Oracle World ” and “ Real World ” for “ Empirical World ” . Over the past few years , researchers have used kernel machines as a tractable model to investigate many neural network phenomena including benign overfitting , i.e. , generalization despite the interpolation of noisy data ( Bartlett et al. , 2020 ; Liang & Rakhlin , 2020 ) and double-descent , i.e , risk curves that are not classically U-shaped ( Belkin et al. , 2020 ; Liu et al. , 2021 ) . Kernels have also been studied to better understand certain aspects of neural network architectures such as invariance and stability ( Bietti & Mairal , 2017 ; Mei et al. , 2021b ) . Although kernel methods can not be used to explain some phenomena such as feature learning , they can still be conceptually useful for understanding other neural networks properties . 1.1 THREE STAGES OF KERNEL DYNAMICS . Despite much classical work in the study of gradient descent training of kernel machines ( e.g . Yao et al . ( 2007 ) ; Raskutti et al . ( 2014 ) ) there has been limited work understanding the high-dimensional setting , which is the setting of interest in this paper . Although solving the linear dynamics of gradient flow is simple , the statistical analysis of the fitted model requires involved random matrix theory arguments . In our analysis we study the dynamics of the Oracle World , where training is done on the ( usually inaccessible ) population risk , and the Empirical World , where training is done on the empirical risk ( as is done in practice ) . Associated with the oracle world model fort and the empirical world model f̂t are the following quantities of interest : the empirical training error R̂n ( f̂t ) , the empirical test error R ( f̂t ) , and the oracle error R ( fort ) defined in Eqs . ( 1 ) , ( 2 ) , ( 3 ) for which we derive expressions that are accurate in high dimensions . Informally , our main results show that under reasonable conditions on the regression function and the kernel the training dynamics undergo the following three stages : • Stage one : the empirical training error , the empirical test error , and the oracle error are all close . • Stage two : the empirical training error decays to zero , but the empirical test error and the oracle error stay close and keep approximately constant . • Stage three : the empirical training error is still zero , the empirical test error stays approximately constant , but the oracle test error decays to the approximation error . We conceptually illustrate the error curves of the oracle and empirical world in Fig . 1 and provide intuition for the evolution of the learned models in Fig . 2 . The existence of the first and third stages are not unexpected : at the beginning of training the model has not fit the dataset enough to distinguish the oracle and empirical world and at the end of training an expressive enough model with infinite samples will outperform one with finitely many . The most interesting stage is the second one where the empirical model begins to “ overfit ” the training set while still remaining close to the non-interpolating oracle model in the L2 sense ( see Fig . 2 ) . In Section 2 we discuss some related work . In Section 3 we elaborate our description of the three stages and give a mathematical characterization for two particular settings in Theorems 1 and 2 . Although the three stages arise fairly generally , we remark that certain stages will vanish if the problem parameters are chosen in a special way ( c.f . Remark 1 ) . We connect our theoretical results to related empirical deep learning phenomena in Remark 3 and discuss the relation to deep learning in practice in Remark 4 . In Section 4 we provide numerical simulations to illustrate the theory more concretely and in Section 5 we end with a summary and discussion of the results . 2 RELATED LITERATURE . The generalization error of the kernel ridge regression ( KRR ) solution has been well-studied in both the fixed dimension regime ( Wainwright , 2019 , Chap . 13 ) , ( Caponnetto & De Vito , 2007 ) and the high-dimensional regime ( El Karoui , 2010 ; Liang & Rakhlin , 2020 ; Liu et al. , 2021 ; Ghorbani et al. , 2020 ; 2021 ; Mei et al. , 2021a ; b ) . Most closely related to our results is the setting of ( Ghorbani et al. , 2021 ; Mei et al. , 2021a ; b ) . Analysis of the entire KRR training trajectory has also been done ( Yao et al. , 2007 ; Raskutti et al. , 2014 ; Cao et al. , 2019 ) but only for the fixed dimensional setting . Classical non-parametric rates are often obtained by specifying a strong regularity assumption on the target function ( e.g . the source condition in Fischer & Steinwart ( 2020 ) ) , whereas in our work the assumption on the target function is mild . Another line of work directly studies the dynamics of learning in linear neural networks ( Saxe et al. , 2013 ; Li et al. , 2018 ; Arora et al. , 2019 ; Vaskevicius et al. , 2019 ) . Similar to us , these works show that some notion of complexity ( typically effective rank or sparsity ) increases in the linear network over the course of optimization . The relationship between the speed of iterative optimization and gap between population and empirical quantities has been studied before in the context of algorithmic stability ( Bousquet & Elisseeff , 2002 ; Hardt et al. , 2016 ; Chen et al. , 2018 ) . These analyses certify good empirical generalization by using stability in the first few iterations to upper bound the gap between train and test error . In contrast , our analysis directly computes the errors at an arbitrary time t ( c.f . Remark 2 ) . The relationship between oracle and empirical training dynamics has been considered before in Bottou & LeCun ( 2004 ) and Pillaud-Vivien et al . ( 2018 ) . 3 RESULTS . In this section we introduce the problem and present a specialization of our results to two concrete settings : dot product and group invariant kernels on the sphere ( Theorems 1 and 2 respectively ) . The more general version of our results is described in Appendix A.3 . 3.1 PROBLEM SETUP . We consider the supervised learning problem where we are given i.i.d . data ( xi , yi ) i≤n . The covariate vectors ( xi ) i≤n ∼iid Unif ( Sd−1 ( √ d ) ) and the real-valued noisy responses yi = fd ( xi ) + εi for some unknown target function fd ∈ L2 ( Sd−1 ( √ d ) ) and ( εi ) i≤n ∼iid N ( 0 , σ2ε ) . Given a function f ∈ L2 ( Sd−1 ( √ d ) ) , we define its test error R ( f ) and its training error R̂n ( f ) as R ( f ) ≡ E ( xnew , ynew ) { ( ynew − f ( xnew ) ) 2 } , R̂n ( f ) ≡ 1 n n∑ i=1 ( yi − f ( xi ) ) 2 ( 1 ) where ( xnew , ynew ) is i.i.d . with ( xi , yi ) i≤n . The test error R ( f ) measures the fit of f on the population distribution and the training error R̂n ( f ) measures the fit of f to the training set . For a kernel function Hd : Sd−1 ( √ d ) × Sd−1 ( √ d ) → R , we analyse the dynamics of the following two fitted models indexed by time t : the oracle model fort and the empirical model f̂t , which are given by the gradient flow on R and R̂n over the associated RKHSHd respectively d dt fort ( x ) = −∇R ( fort ( x ) ) = E [ Hd ( x , z ) ( fd ( z ) − fort ( z ) ) ] , ( 2 ) d dt f̂t ( x ) = −∇R̂n ( f̂t ( x ) ) = 1 n n∑ i=1 Hd ( x , xi ) ( yi − f̂t ( xi ) ) , ( 3 ) with zero initialization for0 ≡ f̂0 ≡ 0 . These dynamics are motivated from the neural tangent kernel perspective of over-parameterized neural networks ( Jacot et al. , 2018 ; Du et al. , 2018 ) . A precise mathematical definition and derivation of these two dynamics are provided in Appendix E.1 . For our results we make some assumptions on the spectral properties of the kernels Hd similar to those in Mei et al . ( 2021a ) that are discussed in detail in Section A.2 . At a high-level we require that the diagonal elements of the kernel concentrate , that the kernel eigenvalues obey certain spectral gap conditions , and that the top eigenfunctions obey a hyperconctractivity condition which says they are “ delocalized ” . For the specific settings of Theorems 1 and 2 we give more specific conditions on the kernels that are more easily verified and imply the required spectral properties .

Summary: This paper studies the learning dynamics of gradient flow for kernel ridge regression. The authors contrast to different setups: the "empirical world", where the model is trained on a finite data set, and the "oracle world", where the model is instead trained directly on the population loss, or in other words, by minimising the test error. Nakkiran et al. (2020) recently reported that the errors in both setups stay close to each throughout training for a range of deep neural networks. By studying the relation between empirical and oracle world in the setup of kernel ridge regression, the authors investigate this phenomenon in a setup under precise theoretical control.  The authors perform such a study by leveraging the analysis of Ghorbani et al. (2021) and Mei et al. (2021), who gave a detailed and precise analysis of the implicit bias of learning in random feature models, and in particular the learning of approximations of increasing complexity.   Contrasting the learning in the two worlds, the authors find that learning generally, but not always, proceeds in three stages. In the case of polynomial kernels, models in both the empirical and the oracle learn first the leading $\ell$-components of the target function, defined as those projections of the target function along the polynomials of small degree. After the training error of the empirical model reaches zero, a second phase ensues where the test errors of both models remain close, but the gap between training and test error in the empirical world can be large. Finally, either there is enough training data to learn the target function perfectly, or learning enters a third stage where the oracle model learns the target function perfectly, which the empirical model won't achieve.  The authors confirm this picture also for invariant kernels, say translation-invariant ones, using technical tools from Mei et al. (2021). Quantitatively, the presence of an invariance shows up in the training speed, which is slower than for the product kernel (Remark 2). Numerical experiments for kernel least-square confirm the theory, while the authors also report experiments for SGD for random feature regression.