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

Introduction

In order to fundamentally understand how and why deep learning works, there has been much effort to understand the learning 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 functionally complex classifiers. 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 errorTheir paper uses “Ideal World” for “Oracle World” and “Real World” for “Empirical World”. 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 relevant but mathematically tractable settings.

Thus, in the large width limit it suffices to study the dynamics on the linearization $f_{N,\rm{lin}}$. When using the squared loss, these dynamics correspond to optimizing a kernel least-squares objective with the neural tangent kernel $K_{N}({\bm{x}},{\bm{x}}^{\prime})=\langle\nabla_{\bm{\theta}}f_{N}({\bm{x}};{\bm{\theta}}^{0}),\nabla_{\bm{\theta}}f_{N}({\bm{x}}^{\prime};{\bm{\theta}}^{0})\rangle$.

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 cannot be used to explain some phenomena such as feature learning, they can still be conceptually useful for understanding other neural networks properties.

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 $f^{\rm{or}}_{t}$ and the empirical world model $\hat{f}_{t}$ are the following quantities of interest: the empirical training error $\widehat{R}_{n}(\hat{f}_{t})$, the empirical test error $R(\hat{f}_{t})$, and the oracle error $R(f^{\rm{or}}_{t})$ 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, the empirical test, and the oracle errors 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 $L^{2}$ 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 detailed mathematical characterization for two specific settings in Theorem 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.

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).

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.

where $({\bm{x}}_{\rm{new}},y_{\rm{new}})$ is i.i.d. with $({\bm{x}}_{i},y_{i})_{i\leq n}$. The test error $R(f)$ measures the fit of $f$ on the population distribution and the training error $\widehat{R}_{n}(f)$ measures the fit of $f$ to the training set.

with zero initialization $f^{\rm{or}}_{0}\equiv\hat{f}_{0}\equiv 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 $H_{d}$ similar to those in Mei et al. (2021a) that are discussed in detail in Appendix 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.

2 Dot Product Kernels

In our first example, we consider dot product kernels $H_{d}$ of the form

Then we have the following characterizations,

(Oracle World) The oracle model learns every degree component of $f_{d}$ as time progresses

(Empirical World – Train) Empirical training error follows oracle error then goes to zero

(Empirical World – Test) Empirical test error follows oracle error until the empirical model learns the degree-$\mathsf{s}$ component of $f_{d}$

The results are conceptually illustrated in an example in Fig. 3(a) which shows the stair-case phenomenon in high-dimensions and the three learning stages. We see that in both the oracle world and the empirical world, the prediction model increases in complexity over time. More precisely, the model learns the best polynomial fit to the target function (in an $L^{2}$ sense) of increasingly higher degree. In the empirical world the maximum complexity is determined by the sample size $n$, which is in contrast to the oracle world where there are effectively infinite samples.

The results imply that generally (but not always c.f. Remark 1) there will be three stages of learning. In the first stage the oracle and empirical world models are close in $L^{2}$ and fit a polynomial with degree determined by $t$. The first stage lasts from $t=0$ to $t=nd^{-\varepsilon}\ll n$ for some small $\varepsilon>0$. As $t$ approaches $n$, there is a phase transition and the empirical world training error goes to zero at $t=nd^{\varepsilon}\gg n$. From time $nd^{-\varepsilon}$ till at least $d^{\mathsf{s}+1}$ is the second stage where the empirical and oracle models remain close in $L^{2}$ but the gap between test and train error can be large. If the sample size $n$ is not large enough for $\hat{f}_{t}$ to learn the target function, then at some large enough $t$ we will enter a third stage where $f^{\rm{or}}_{t}$ improves in performance, outperforming $\hat{f}_{t}$ which remains the same. On synthetic data in finite dimensions we can see a resemblance of the staircase shape which becomes sharper with increasing $d$ (c.f. Appendix F).

3 Group Invariant Kernels

To encode invariance in our kernel we consider ${\mathcal{G}}_{d}$-invariant kernels $H_{d}$ of the form

where $\pi_{d}$ is the Haar measuare on ${\mathcal{G}}_{d}$. Such kernels satisfy the following invariance property: for all $g,g^{\prime}\in{\mathcal{G}}_{d}$ and for $H_{d}({\bm{x}}_{1},{\bm{x}}_{2})=H_{d}(g\cdot{\bm{x}}_{1},g^{\prime}\cdot{\bm{x}}_{2})$ for every ${\bm{x}}_{1},{\bm{x}}_{2}$. For the cyclic group, $\pi_{d}$ is the uniform measure. We now present our results for the group invariant setting.

Then we have the following characterizations,

(Oracle World) The oracle model learns every degree component of $f_{d}$ as time progresses

(Empirical World – Train) Empirical training error follows oracle error then goes to zero

(Empirical World – Test) Empirical test error follows oracle error until the empirical model learns the degree-$\mathsf{s}$ component of $f_{d}$

With respect to the dot product kernel setting (c.f. Theorem 1), in this setting the behavior of the oracle world is unchanged, but the empirical world behaves as if it has $d^{\alpha}$ times as many samples. This is illustrated graphically in Fig. 3(c). It can be shown that using an invariant kernel is equivalent to using a dot product kernel and augmenting the dataset to $\{(g\cdot{\bm{x}}_{i},y_{i}):g\in{\mathcal{G}}_{d},i\in[n]\}$ (c.f. Appendix E.2). Hence for the cyclic group which has size $d^{\alpha}=d$, we reach the following intriguing conclusion: if the target function is cyclically invariant, then using a dot product kernel and augmenting a training set of $n$ i.i.d. samples to $nd$ many samples is asymptotically equivalent to training with $nd$ i.i.d. samples.

Interestingly, training with an invariant kernel is slower than with a dot product kernel and takes longer to interpolate the dataset despite eventually generalizing better on invariant function estimation tasks (c.f. Fig. 3(c)). This conclusion is not an artifact of the continuous time analysis (c.f. Appendix E.3) and is observed empirically in Section 4.2 for discrete-time SGD. This example highlights the limitation of stability based analyses (e.g. Hardt et al. (2016)) which argue that faster SGD training leads to better generalization. While a faster rate leads to better generalization in the first stage when stability can control the gap between train and test error, our analysis shows that the duration length of the first stage also impacts the final generalization error.

The dynamics of high-dimensional kernel regression display behaviors that parallel some empirically observed phenomena in deep learning. For kernel regression, we have shown that the complexity of the empirical model , measured as the number of learned eigenfunctions, depends on the time optimized when $t\ll n$ and the sample size when $t\gg n$. At a high-level, we also expect a similar story for neural networks but for some other notion of complexity. It is believed that neural networks first learn simple functions and then progressively more complex ones, until the complexity saturates after interpolating at some time proportional to $n$ (Nakkiran et al., 2019b). We have also shown that in kernel regression there is a non-trivial “deep boostrap” phenomenon (Nakkiran et al., 2020) during the second learning stage: the gap between the oracle world and empirical world test errors is negligible whereas the train and test errors exhibit a substantial gap. The gradient flow results for kernel regression can also provide insight into the deep bootstrap for random feature networks trained with discrete-time SGD as these results can approximately predict their behavior (see Section 4.2).

Although we believe our results conceptually shed light on some of the interesting behaviors observed in the training dynamics of deep learning, due to our stylized setting we may not exactly see the predicted phenomena in practice. Accurately observing the three stages of kernel regression requires sufficiently high-dimensional data in order for the kernel eigenvalues to obey a staircase-like decay and for training to be sufficiently long as the time axis should be in log-scale. Our results hold for regression whereas for classification the empirical model may continue improving after classifying the train set correctly. Despite these caveats, certain conclusions can be observed in some realistic settings (c.f. Appendix E.4).

Numerical Simulations

In Section 4.1, we simulate the gradient flows of kernel least-squares with dot product kernels and cyclic kernels (Fig. 4) to reproduce the three stages as shown in Fig. 3(b). In Section 4.2 we show that SGD training of (dot product and cyclic) random-feature models (Fig. 5) exhibit similar three stages phenomena, in which the second stage behaviors are consistent with the deep bootstrap phenomena observed in deep learning experiments (Nakkiran et al., 2020). Empirical quantities are averaged over 10 trials and the shaded regions indicate one standard deviation from the mean.

for some activation function $\sigma$. We will examine a few different choices of $f_{\star}$ and kernel $H$, which are specified in the descriptions of each figure.

The oracle world error is computed analytically but the empirical world errors require sampling train and test datasets. We compute empirical world errors by averaging over $10$ trials. The results are visualized both in log-scale and linear-scale on the time axis. The log-scale plots allow for direct comparison with the cartoons in Fig. 3(b). The linear-scale plots are zoomed into the region $0<t\leq nd^{0.4}$ since: 1) plotting the full interval squeeze all curves to the left boundary which is uninformative 2) in practice one would not optimize for very long after interpolation.

In Figs. 4(a) and 4(b) we take the target function to be a polynomial of the form

where $\operatorname{He}_{i}(t)$ is the $i$th Hermite polynomial (c.f. Appendix G.4) and the approximate equality in Eq. (7) holds in high-dimensions. In panel (4(a)) we consider use the ReLU activation function $\sigma(t)=\max(t,0)$, and take $f_{\star}$ to be a quadratic polynomial with $(a_{0},a_{1},a_{2})=(1/2,1/\sqrt{2},1/\sqrt{8})$. With such a choice of parameters, we can see the three stages phenomenon. In panel (4(b)) we choose $f_{\star}$ to be a cubic polynomial with $(a_{0},a_{1},a_{2},a_{3})=(1/2,1/\sqrt{2},0,1/\sqrt{24})$ and $\sigma(t)=\max(t,0)+0.1\operatorname{He}_{3}(t)$ (we need the third Hermite coefficient of $\sigma$ to be non-zero for stage 3 to occur). This choice of coefficients for $f_{\star}$ is such that $\|\overline{\mathsf{P}}_{>1}f_{\star}\|_{L^{2}}^{2}\approx\|\overline{\mathsf{P}}_{>2}f_{\star}\|_{L^{2}}^{2}$, so that the second stage in (4(b)) is longer compared to (4(a)).

In Fig. 4(c), we take the target function to be a cubic cyclic polynomial

where the subindex addition in $x_{i+k}$ is understood to be taken modulo $d$. We compare the performance of the dot product kernel $H$ and its invariant version $H_{\text{inv}}$ (c.f. Eq. (5)) with activation function $\sigma(t)=\max(t,0)+0.1\operatorname{He}_{3}(t)$. The kernel $H_{\text{inv}}$ with $n$ samples performs equivalently to $H$ with $nd$ samples (c.f. Remark 2), but is more computationally efficient since the size of the kernel matrix is still $n\times n$. Using $H_{\text{inv}}$ elongates the first stage by a factor $d$, delaying the later stages and ensuring that the empirical world model improves longer.

Although in the simulations, the dimension $d$ is not yet high enough to see a totally sharp staircase phenomenon as in the illustrations of Fig. 3(b), even for this $d$ we are still able to clearly see the three predicted learning stages and deep bootstrap phenomenon across a range of settings. To better understand the effect of dimension we show similar plots with varying $d$ in Appendix F.

2 SGD for Two-layer Random-Feature Models

To more closely relate with deep learning practice and the deep bootstrap phenomenon (Nakkiran et al., 2020), we simulate the error curves of SGD training on random-feature (RF) models (i.e. two-layer networks with random first-layer weights and trainable second-layer weights), in the same synthetic data setup as before. In particular, we look at dot product RF models

For all following experiments we take the activation $\sigma$ to be ReLU and $N=4\times 10^{5}\approx n^{1.4}$.

For a given data distribution and RF model we train two fitted functions, one on a finite dataset (empirical world) and the other on the data distribution (oracle world). More specifically, the training of the empirical model is done using multi-pass SGD on a finite training set of size $n$ with learning rate $\eta=0.1$ and batch size $b=50$. The training of the oracle model is done using one-pass SGD with the same learning rate $\eta$ and batch size $b$, but at each iteration a fresh batch is sampled from the population distribution. Both models $\hat{f}_{t},f^{\rm{or}}_{t}$ are initialized with $a_{i}=0$ for $i\in[N]$. To speed up and stabilize optimization we use momentum $\beta=0.9$. Note that if we took $N\to\infty$, $\eta\to 0$, and $\beta=0$ we would be exactly in the dot product kernel gradient flow setting.

In Fig. 5, the data generating distributions of panels (5(a)), (5(b)), (5(c)) are respectively the same as that of panels (4(a)), (4(b)), and (4(c)) from Section 4.1. The top row of Fig. 5 shows SGD for $\{$dot product, cyclic$\}$ RF models, and the bottom row shows the corresponding gradient flow for $\{$dot product, cyclic$\}$ kernel least-squares. We see that the corresponding curves in these two rows exhibit qualitatively the same behaviors. Additionally, the results in panel (5(c)) show that as predicted, even for discrete SGD dynamics the dot product RF optimizes faster but fails to generalize, whereas the invariant RF optimizes slower but generalizes better.

Summary and Discussion

In this paper, we used precise asymptotics to study the oracle world and empirical world dynamics of gradient flow on kernel least-squares objectives for high-dimensional regression problems. Under reasonable conditions on the target function and kernel, we showed that in this setting there are three learning stages based on the behaviors of the empirical and oracle models and also connected our results to some empirical deep learning phenomena.

Although our setting already captures some interesting aspects of deep learning training dynamics, there are some limitations which would be interesting to resolve in future work. We require very high-dimensional data in order for the asymptotics to be accurate, but real data distributions have low-dimensional structure. We work in a limiting regime of neural network training where the dynamics are linear and the step-size is infinitesimal. It is an important direction to extend this analysis to the non-linear feature learning regime and to consider discrete step-size minibatch SGD, as these are considered important aspects of network training. Our results hold for the square-loss in regression problems, but many deep learning problems involve classification using cross-entropy loss. Lastly, our analysis holds specifically for gradient flow, so it would be also interesting to consider other iterative learning algorithms such as boosting.

We would like to thank Preetum Nakkiran for helpful discussions and for reviewing an early draft of the paper. This research is kindly supported in part by NSF TRIPODS Grant 1740855, DMS-1613002, 1953191, 2015341, IIS 1741340, the Center for Science of Information (CSoI), an NSF Science and Technology Center, under grant agreement CCF-0939370, NSF grant 2023505 on Collaborative Research: Foundations of Data Science Institute (FODSI), the NSF and the Simons Foundation for the Collaboration on the Theoretical Foundations of Deep Learning through awards DMS-2031883 and 814639, and a grant from the Weill Neurohub.

References

Appendix A General Setting

In this section we present our theory for training dynamics of kernel regression in an abstract setting similar to that of Mei et al. (2021a). We first introduce the setting of interest, then state the relevant assumptions, and finally we provide our theoretical results. We provide proofs of these results in Appendix B.

We are concerned with a supervised learning problem where we are given i.i.d. data $({\bm{x}}_{i},y_{i})_{i\leq n}$. The feature vectors ${\bm{x}}_{i}\sim_{iid}\nu_{d}$ are in ${\mathcal{X}}_{d}$ and the empirical-valued noisy responses $y_{i}$ are given by

for some unknown target function $f_{d}\in{\mathcal{D}}_{d}$ and $\varepsilon_{i}\sim_{iid}{\mathcal{N}}(0,\sigma_{\varepsilon}^{2})$.

where $H_{d}\in L^{2}({\mathcal{X}}_{d}\times{\mathcal{X}}_{d})$ with the property that $\int_{{\mathcal{X}}_{d}}H_{d}({\bm{x}},{\bm{x}}^{\prime})g({\bm{x}}^{\prime})\nu_{d}(\differential{{\bm{x}}^{\prime}})=0$ for $g\in{\mathcal{D}}_{d}^{\perp}$.

By the spectral theorem of compact operators, there exists an orthonormal basis $(\psi_{j})_{j\geq 1}$ such that $\operatorname{span}(\psi_{j},j\geq 1)={\mathcal{D}}_{d}\subseteq L^{2}({\mathcal{X}}_{d})$ and empirical eigenvalues $(\lambda_{d,j})_{j\geq 1}$ with nonincreasing absolute values $|\lambda_{d,1}|\geq|\lambda_{d,2}|\geq\cdots$ and $\sum_{j\geq 1}\lambda_{d,j}^{2}<\infty$ such that

where convergence holds in $L^{2}({\mathcal{X}}_{d}\times{\mathcal{X}}_{d})$.

where $({\bm{x}}_{1},y_{1}),\ldots,({\bm{x}}_{n},y_{n}),({\bm{x}}_{\rm{new}},y_{\rm{new}})$ are i.i.d. For a kernel $H_{d}\in L^{2}({\mathcal{X}}_{d}\times{\mathcal{X}}_{d})$ we will consider the oracle model $f^{\rm{or}}_{t}$ and the empirical model $\hat{f}_{t}$ which satisfy the following gradient flows

A.2 General Assumptions

We now state our assumptions on the kernel and the sequence of probability spaces $({\mathcal{X}}_{d},\nu_{d})$.

(Hypercontractivity of finite eigenspaces.) For any fixed $q\geq 1$, there exists a constant $C$ such that for any $h\in{\mathcal{D}}_{d,\leq r(d)}=\operatorname{span}(\psi_{s},1\leq s\leq r(d))$, we have

(Properly decaying eigenvalues) There exists fixed $\delta_{0}>0$, such that, for all $d$ large enough,

(Concentration of diagonal elements of kernel) For $({\bm{x}}_{i})_{i\in[n(d)]}\sim_{iid}\nu_{d}$, we have:

There exists a fixed $\delta_{0}>0$, such that

There exists a fixed $\delta_{0}>0$, such that

There exists a fixed $\delta_{0}>0$, such that

Assumptions 2(a) and 2(b) can be seen as a spectral gap assumption. This ensures a clear separation between the eigenvalues in the subspace ${\mathcal{D}}_{d,\leq\mathsf{m}(d)}$ and the subspace ${\mathcal{D}}_{d,>\mathsf{m}(d)}$. The technical requirement in Assumption 2(c) is mild.

In the asymptotic setting, we will be interested in the model learned at a time $t=t(d)$ scaling with the dimension. The following assumptions give requirements for a valid scaling.

There exists a fixed $\delta_{0}>0$, such that for $d$ large enough

If $\mathsf{u}(d)<\mathsf{m}(d)$ for infinitely many $d$, then

Assumption 3(a) is similar to the spectral gap condition Assumption 2(a), 2(b) for $(t(d),\mathsf{u}(d))_{d\geq 1}$. Assumption 3(b) relates the ordering of the indices $\mathsf{u}(d),\mathsf{m}(d)$ to the relative growth of $t(d),n(d)$. Assumption 3(c) requires that the eigenvalue tail does not decay too abruptly.

A.3 Main Results

In this section we give the main theoretical results. Recall the problem set-up and notation from Appendix A.1. We will characterize the gradient flow dynamics dynamics of the oracle model $f^{\rm{or}}_{t}$ Eq. (10) and the empirical model $\hat{f}_{t}$ Eq. (11) for a general kernel $H_{d}$.

where $\mathsf{P}_{\leq\mathsf{u}(d)}$ and $\mathsf{P}_{>\mathsf{u}(d)}$ are the projection operators onto the subspace spanned by the top $\mathsf{u}(d)$ kernel eigenfunctions and the orthogonal complement respectively, as defined in Appendix A.1.

The error of the oracle model is determined solely by optimization time $t$ through $\mathsf{u}(d)$. Due to the spectral gap assumption 3(a) learning only occurs along the top $\mathsf{u}(d)$ eigenfunctions.

The next results describe the empirical model. First we characterize the training error.

In the early-time regime $t\ll n/\kappa_{H}$ the training error may be non-zero and matches the oracle world error if also $\mathsf{u}(d)\leq\mathsf{m}(d)$. In the late-time regime $t\gg n/\kappa_{H}$ the training error is negligible and the model interpolates the training set. The quantity $\kappa_{H}$ arises since the empirical kernel matrix can be decomposed as ${\bm{H}}={\bm{H}}_{\leq\mathsf{m}}+{\bm{H}}_{>\mathsf{m}}$ and the second component is approximately a multiple of the identity: ${\bm{H}}_{>\mathsf{m}}\approx\Tr({\bm{H}}_{>\mathsf{m}})\cdot{\mathbf{I}}_{n}=\kappa_{H}\cdot{\mathbf{I}}_{n}$. This term acts as a self-induced ridge-regularizer.

Our final result characterizes the test error of the empirical model.

Appendix B Proof of General Setting

The oracle model ODE Eq. (10) with initialization $f^{\rm{or}}_{0}\equiv 0$ can be solved (c.f. Appendix E.1) to yield the solution

where $\lambda_{d,k}^{2}$ are the kernel eigenvalues and $\hat{f}_{k}:=\langle f_{d},\psi_{k}\rangle_{L^{2}}$ are the Fourier coefficients of $f_{d}$ in the kernel eigenbasis (c.f. Appendix A.1). We can control this quantity as follows,

where the last inequality follows from Assumption 3(a). This shows the first theorem statement,

where the second inequality follows from the fact that $(1-e^{-x})^{2}\leq 1-e^{-2x}$ and the final equality is from the proof of the first part of the theorem. Thus,

B.2 Empirical World – Preliminaries

We will introduce some useful notations for studying the empirical world.

B.2.2 Decompositions and Notations

In this section we recall some useful decompositions of empirical quantities from Mei et al. (2021a). As mentioned earlier the eigendecomposition of $H_{d}$ is given by

We write the orthogonal decomposition of $f_{d}$ in the basis $\{\psi_{k}\}_{k\geq 1}$ as

We have the following orthogonal basis decompositions of ${\bm{f}},{\bm{H}},{\bm{E}}$ and ${\bm{M}}$

By Lemma 6 below, under Assumptions 1 and 2(a) the matrices ${\bm{H}}$ and ${\bm{M}}$ can be written as

We will use $\alpha$ as shorthand for the scalar valued dimension dependent quantity $e^{-(t/n)\kappa_{H}}$ and take

We also introduce the shrinkage matrix defined as

B.2.3 Auxiliary Lemmas

Here we collect some lemmas which will be of use to us.

The notations in Lemmas 3-8 all follow the notations given in Appendix B.2.2.

Let Assumptions 1 and 2(a) hold. Then we can decompose the kernel matrices as follows

where the expectation is with respect to the randomness in ${\bm{\Psi}}_{S},{\bm{\Psi}}_{T}$.

Let $\iota:S\to[|S|]$ be the bijection such that ${\bm{\Psi}}_{S}=(\psi_{\iota^{-1}(k)}({\bm{x}}_{i}))_{i\in[n],k\in[|S|]}$ then we have

where the second to last inequality is by Holder’s inequality and the last inequality used the hypercontractivity assumption as in Assumption 1(a). ∎

Assume the conditions of Lemma 6 hold and assume there exists $\delta_{0}>0$, such that $(t(d),\mathsf{u}(d))$ satisfies the condition

Let $j(d)$ satisfy $\min\{\mathsf{u}(d),\mathsf{m}(d)\}\leq j(d)\leq\mathsf{m}(d)$ then

First consider the regime $t=O_{d}(n/\kappa_{H})$. Recall the decomposition,

where equality $(a)$ holds by assumption Eq. (21). Now if $t=\omega_{d}(n/\kappa_{H})$, then it is easy to see that

B.3 Empirical World - Train

From now on we focus on the case that $t=o_{d}(n/\kappa_{H})$. Let us decompose the training error as

Let us start by analysing $R_{1}$. We can write

Now we analyse the term $T_{3}$. Observe that by the inequality $1-x\leq e^{-x}\leq 1$,

Observe that by the Cauchy-Schwarz inequality,

Turning to $R_{2}$, we take the second-moment with respect to ${\bm{\varepsilon}}$

since by assumption we consider $t=o_{d}(n/\kappa_{H})$. Turning to the variance

B.4 Empirical World - Test

Recalling ${\bm{u}}(t)$ from Eq. (14), let

Recall the expansion of the test error from Eq. (15),

The proof for the test error is the most involved, but will follow a similar strategy of analysing each term in the expansion. We first begin by analysing $T_{2}$ in Appendix B.4.1, then $T_{1}$ in Appendix B.4.2, and finally terms $T_{3}$, $T_{4}$, and $T_{5}$, which are all simpler than the first two, in Appendix B.4.3. At the end of this Appendix section we present the proof of Theorem 5.

As before, we will analyze each term separately. We begin with term $T_{2}$.

where we recall the shrinkage matrix ${\bm{S}}_{\leq\mathsf{m}}$ defined in Eq. (20) and ${\bm{v}}(t)$ in Eq. (23).

To analyze $T_{2}$ we further decompose it into the following terms

Using Lemma 9 and 10 proven below, by the Cauchy-Schwarz inequality,

We will show that the dominant term is $T_{211}$ and the others are of lower order. By Lemma 3,

Using similar reasoning from Lemma 9 we can write

Let us first start with analysing $T_{231}$, defining ${\bm{B}}={\bm{\Psi}}_{\leq\mathsf{m}}^{\mathsf{T}}{\bm{\Psi}}_{\leq\mathsf{m}}/n$ we have

where the first equality is by Lemma 2 and the last line follows from Lemma 7, Markov’s inequality, and by the fact that $\mathsf{m}/n=o_{d}(1)$ by Assumption 2(c). To see that inequality $(a)$ holds, note that

Since ${\bm{S}}_{\leq\mathsf{m}}^{2}\preceq{\mathbf{I}}_{m}$,

where the last equality is by Lemma 4. Now let us turn to $T_{233}$, which we can further split into

Turning to the second term $T_{2332}$, let

By the inequality $1-e^{-x}\leq x$, in the PSD order we see that

For the second term on the right, by Assumption 3(b) and by Lemma 4,

For the first term by Lemma 7 and Assumptions 3(b), 3(c),

We break $T_{1}$ into the following terms

where $T_{23}$ is as given in Eq. (27). From the analysis of $T_{11}$ in Lemma 12 and $T_{13}$ in Lemma 13 we combine everything to get Proposition 2

Recall that ${\bm{E}}_{\leq\mathsf{m}}={\bm{\Psi}}_{\leq\mathsf{m}}{\bm{D}}_{\leq\mathsf{m}}^{2}\widehat{{\bm{f}}}_{\leq\mathsf{m}}$. By the Cauchy-Schwarz inequality

where the last equality follows from Lemma 3. ∎

where the last equality used Eq. (13) in Assumption 2(a). ∎

To analyse the terms $T_{3}$, $T_{4}$, $T_{5}$ we can adapt the corresponding steps for the proof of Theorem 4 in Mei et al. (2021a). For the following analysis we recall the definition of ${\bm{\varepsilon}}(t)$ from Eq. (24).

where $(a)$ used Lemma 3, $(b)$ used ${\bm{S}}_{\leq\mathsf{m}}\preceq{\mathbf{I}}_{m}$, and $(c)$ used Lemma 4 and Assumption 2(c). The lemma then follows from Markov’s inequality. ∎

where the last inequality follows from Lemma 3. Hence,

where $(a)$ follows from the definition of ${\bm{E}}_{\leq L}$, $(b)$ follows from the definition of operator norm, and $(c)$ follows from Eq. (35). Therefore we get

Note as in Eq. (35), that by Lemma 3 and Lemma 4,

and taking the second moment of $T_{51}$ yields

where $T_{21}$ is as given in Eq. (25). Similarly we get that

where $T_{23}$ is as given in Eq. (27). By Markov’s inequality we deduce that

Finally putting Propositions 1, 2 and Lemmas 14, 15, 16 together for terms $T_{2},T_{1},T_{3},T_{4}$ and $T_{5}$ respectively leads to the proof of Theorem 5

and as a result we obtain the first part of the theorem

Now observe that similar to Eq. (15) we have the following decomposition

where ${\bm{u}}(t)$ is given in Eq. (14). Therefore we can write

hence combining with the bound for $T_{4}$ in Lemma 15 yields

where $(a)$ used the inequality $(1-e^{-x})/x\leq 1$ and $(b)$ used Assumption 3(a). Therefore

hence by Eq. (36) and Eq. (37) we obtain the final part of the theorem

for some activation function $\sigma_{d}$. We state mild assumptions on $\sigma_{d}$ and show that under these conditions we can apply the results in Appendix A.3.

C.2 Assumptions

The assumptions given on the activations are the same as Assumption 3 of Mei et al. (2021a).

Let $\{H_{d}\}_{d\geq 1}$ be a sequence of dot product kernels with associated activation functions $\{\sigma_{d}\}_{d\geq 1}$ as in Eq. (61). We assume the following hold

We state mild assumptions on $\sigma_{d}$ and show that under these conditions we can apply the results in Appendix A.3. For additional technical background refer to Appendix G.

D.2 Assumptions

We will assume that $\sigma_{d}=\sigma$ for all $d$ and make the following assumptions on $\sigma$ which are the same as Assumption 1 in Mei et al. (2021b).

Let $\{H_{d}\}_{d\geq 1}$ be a sequence of invariant kernels with associated activation functions $\sigma_{d}=\sigma$ as in Eq. (40). We assume the following conditions hold

For ${\mathcal{G}}_{d}=\operatorname{Cyc}_{d}$, we assume $\sigma$ to be $(\mathsf{s}+1)\vee 3$ differentiable and there exists constants $c_{0}>0$ and $c_{1}<1$ such that $|\sigma^{(k)}|\leq c_{0}e^{c_{1}u^{2}/2}$ for any $2\leq k\leq(\mathsf{s}+1)\vee 3$.

For general ${\mathcal{G}}_{d}$, we assume that $\sigma$ is a (finite degree) polynomial function.

The Hermite coefficients $\mu_{k}(\sigma)$ (c.f. Appendix) verify $\mu_{k}\neq 0$ for any $0\leq k\leq\mathsf{s}$.

We assume that $\sigma$ is not a polynomial with degrees less than or equal to $\mathsf{s}$.

where $D(d,k)$ is the dimension of the subspace of invariant polynomials of degree $k$ (c.f. Appendix G.6). The verification of Assumption 1 (Kernel Concentration Property) and Assumption 2 (Eigenvalue Condition) at level $\{(n(d),\mathsf{m}(d)\}$ is exactly the same as in Theorem 1 in Mei et al. (2021b).

We must verify Assumption 3. To see part 3(a), note that $1/\lambda_{d,\mathsf{u}(d)}^{2}=\Theta_{d}(d^{\mathsf{j}})$ and $1/\lambda_{d,\mathsf{u}(d)}^{2}=\Theta_{d}(d^{\mathsf{j}+1})$. For part 3(b), the condition holds because $\mathsf{u}(d)<\mathsf{m}(d)$ for large $d$ if and only if $\mathsf{j}<\mathsf{s}$ in which case

from which it follows that for some constant $C$

Appendix E Auxiliary Results

Recall that we are interested in the following dynamics given in Eqs. (2), (3),

The gradient $\gradient V(f)\in{\mathcal{H}}_{d}$ is defined such that

exists uniquely by the Riesz representation theorem. The gradients of the risk functionals are

The oracle model ODE Eq. (41) is simply a linear differential equation which has the following solution involving the operator exponential $\exp({\bm{A}}):=\sum\limits_{k=0}^{\infty}{\bm{A}}^{k}/k!$

For the empirical model ODE Eq. (42) we first consider the system of scalar differential equations induced at the points $\{({\bm{x}}_{i},y_{i})\}_{i\in[n]}$. Letting ${\bm{u}}(t)=(\hat{f}_{t}({\bm{x}}_{1}),\ldots,\hat{f}_{t}({\bm{x}}_{n}))^{\mathsf{T}}$, ${\bm{y}}=(y_{1},\ldots,y_{n})^{\mathsf{T}}$, and ${\bm{H}}=(H_{d}({\bm{x}}_{i},{\bm{x}}_{j}))_{i,j\in[n]}$ we have

with initial condition ${\bm{u}}(0)={\bm{0}}$. As this a linear ODE, the solution is given by

which is Eq. (42) at point ${\bm{x}}$. Indeed, by the chain rule

E.2 Equivalence between Invariant Kernels and Data Augmentation

In this section we will show an equivalence between the (time rescaled) gradient flows for training invariant kernels and using an augmented dataset. Specifically consider a group ${\mathcal{G}}$ and a kernel $H$ that is ${\mathcal{G}}$-equivariant, that is

Given a ${\mathcal{G}}$-equivariant kernel $H$, we define a ${\mathcal{G}}$-invariant kernel $H_{\rm{inv}}$ as the group averaged kernel

for the Haar measure $\pi$ on ${\mathcal{G}}$ (c.f. Eq. (5)). Note that any dot product kernel is ${\mathcal{G}}$-equivariant for ${\mathcal{G}}$ a subgroup the orthogonal group e.g. the cyclic group $\operatorname{Cyc}$ (c.f. Appendix 3.3).

Given a dataset $({\bm{X}},{\bm{y}})=\{({\bm{x}}_{i},y_{i}):i\in[n]\}$ consider the augmented dataset

We consider the (rescaled c.f. Remark 5) empirical dynamics Eq. (3) of the gradient flow on $({\bm{X}},{\bm{y}})$ using $H_{\rm{inv}}$ which we denote $\hat{f}_{t,\rm{inv}}$

and the empirical dynamics of the gradient flow on $({\bm{X}}_{{\mathcal{G}}},{\bm{y}}_{{\mathcal{G}}})$ using $H$ which we denote $\hat{f}_{t,\rm{aug}}$

Let ${\mathcal{G}}$ be a finite group with $m$ elements. Given a ${\mathcal{G}}$-equivariant kernel $H$, if $\pi$ is the uniform measure on ${\mathcal{G}}$ then

Let ${\mathcal{G}}=\{g_{1},\ldots,g_{m}\}$ where $g_{1}$ is the identity. Define the output vectors

Note that by definition ${\bm{H}}_{\rm{inv}}=\frac{1}{m}\sum\limits_{j=1}^{m}{\bm{H}}_{g_{1},g_{j}}$. We will show that

From this the result follows by Theorem 4.1 in Li et al. (2019) since $\hat{f}_{t,\rm{inv}},\hat{f}_{t,\rm{aug}}$ are given by kernel regressions with targets ${\bm{u}}_{\rm{inv}}(t),{\bm{u}}_{\rm{aug}}(t)$ and kernels $H_{\rm{inv}},H$ respectively.

in order to show Eq. (45) holds. We prove the above by induction on $k$. For $k=1$, observe that

where the second equality follows from ${\mathcal{G}}$-equivariance of $H$. Assume the inductive hypothesis holds for $k$. Then

where the second equality applies the induction hypothesis and the third equality uses equivariance. Thus the inductive claim is proved and the proof is complete.

The scaling factor $m$ in the gradient flow for $\hat{f}_{t,\rm{inv}}$, leads to a natural comparison with $\hat{f}_{t,\rm{aug}}$ as elaborated in Appendix E.3. As argued in that section, in the gradient descent discretization, it is natural to take a step-size inversely proportional to the maximum kernel eigenvalue. In the case of high-dimensional invariant kernels, note that

hence the step-size for the invariant kernel flow should be $m$ times larger.

E.3 Discretizing Time

Comparing different “speeds” of optimization algorithms only makes sense for discrete-time algorithms. Consider the following gradient descent dynamics with step-size $\eta$, obtained as the discretization of the empirical gradient flow Eq. (42)

We will argue that it is natural to take $\eta\sim n/\lambda_{\max}({\bm{H}})$ where ${\bm{H}}$ is the kernel matrix.

Let $T=S^{*}S$ and define $\overline{{\bm{H}}}:=SS^{*}=\frac{1}{n}{\bm{H}}$ to be the normalized kernel matrix. Let $b:=S^{*}\overline{{\bm{H}}}^{-1}{\bm{y}}\in{\mathcal{H}}_{d}$ and note that since $Tb=S^{*}{\bm{y}}$, we can rewrite Eq. (47) as

where the second equality holds since $T$ is self-adjoint and the last equality is since $\phi_{i}$ is an eigenvector of $T$. It is easy to see that $\alpha_{i}(k)=(1-\eta\lambda_{i})^{k}\alpha_{i}(0)$. Therefore we see that gradient descent Eq. (46) is guaranteed to converge if $\eta<1/(2\lambda_{1})$ and may not otherwise. Therefore it is natural to choose the step-size $\eta$ to scale asymptotically as $\eta\sim 1/\lambda_{\max}(\overline{{\bm{H}}})$.

For a dot product kernel $H$ and its corresponding invariant kernel $H_{\rm{inv}}$ the kernel matrices have operator norms of the same order

hence no time rescaling is need to compare the corresponding optimization speeds asymptotically.

E.4 Similarities with Empirical Phenomena

In this section we elaborate upon Remark 4 and mention some connections with empirical observations in Nakkiran et al. (2020). Although the metric in our setting is the squared loss, we can still observe three stages in classification problems when measuring the soft error. In Fig. 6(a) taken from Nakkiran et al. (2020) we can observe stage 1 and stage 2. Either training has not continued long enough to observe stage 3 or $n$ is large enough so that the models have converged to the approximation error of the neural network class (c.f. Remark 1). In Fig. 6(b) taken from Nakkiran et al. (2020), although the train errors are not plotted, by extrapolating from Fig. 6(a), presumably for each $n$ stage 1 and stage 2 occur. From the dimmer curves in Fig. 6(b) we can see that for $n<50000$ stage 3 occurs as well.

In Fig. 7(a), we see a parallel between the use of cyclic versus dot product kernels and the use data augmentation versus not for a Resnet-18 trained on CIFAR-5m (note that using a cyclic kernel is equivalent to using a dot product kernel with data-augmentation c.f. Appendix E.2). In both our theoretical results and in the empirical results of Nakkiran et al. (2020) we observe that the ideal world optimization speed of augmented and non-augmented training are the same, but for augmented training the real world training speed is slowed down, eventually leading to better generalization for long enough training.

Appendix F Additional Figures

To see the effects of varying the dimension $d$ we replicate the log-scale plots of kernel gradient flow with dot product kernels in Fig. 4. We take $n=d^{1.5}$ and vary $d\in\{50,100,200,400\}$. Each plot is averaged over 10 runs with the shaded region representing one standard deviation around the mean. We can see that as $d$ increasing the standard deviation decreases and the curves approach the theoretical high-dimensional prediction.

Appendix G Technical background

G.2 Functional spaces over the sphere

The dimension of each subspace is given by

G.3 Gegenbauer polynomials

We will use the following properties of Gegenbauer polynomials

These properties imply that, up to a constant, $Q_{k}^{(d)}(\langle{\bm{x}},{\bm{y}}\rangle)$ is a representation of the projector onto the subspace of degree-$k$ spherical harmonics

then we have the following equation holds in $L^{2}([-\sqrt{d},\sqrt{d}],\tau^{1}_{d})$ sense

By rotational invariance, the space $V_{k}$ of homogeneous polynomials of degree $k$ is an eigenspace of ${\mathcal{H}}_{d}$, and we will denote the corresponding eigenvalue by $\xi_{d,k}(h_{d})$. In other words ${\mathcal{H}}_{d}f({\bm{x}})\equiv\sum_{k=0}^{\infty}\xi_{d,k}(h_{d}){\overline{\mathsf{P}}}_{k}f$. The eigenvalues can be computed via

For a dot product kernel $H_{d}({\bm{x}},{\bm{y}})=h_{d}(\langle{\bm{x}},{\bm{y}}\rangle/d)$ consider the Gegenbauer expansion of $h_{d}$ in $L^{2}([-\sqrt{d},\sqrt{d}],\tau_{d}^{1})$

Using Eq. (54) we can equivalently write the kernel as an expectation over random features for some activation $\sigma_{d}$

Note that $\sigma_{d}\in L^{2}([-\sqrt{d},\sqrt{d}],\tau_{d}^{1})$ as long as $h(1)<\infty$.

G.4 Hermite polynomials

is the standard Gaussian measure, and ${\rm He}_{k}$ has degree $k$. We will follow the classical normalization (here and below, expectation is with respect to $G\sim{\mathcal{N}}(0,1)$):

Here and below, for $P$ a polynomial, ${\rm Coeff}\{P(x)\}$ is the vector of the coefficients of $P$. As a consequence, for any fixed integer $k$, we have

where $\mu_{k}(\sigma)$ and $\xi_{d,k}(\sigma)$ are given in Eq. (64) and Eq. (57).

G.5 The invariant function class and the symmetrization operator

G.6 Orthogonal polynomials on invariant function class