The Surprising Simplicity of the Early-Time Learning Dynamics of Neural Networks

Introduction

Modern deep learning models are enormously complex function approximators, with many state-of-the-art architectures employing millions or even billions of trainable parameters (Radford et al., 2019; Adiwardana et al., 2020). While the raw parameter count provides only a crude approximation of a model’s capacity, more sophisticated metrics such as those based on PAC-Bayes (McAllester, 1999; Dziugaite and Roy, 2017; Neyshabur et al., 2017b), VC dimension (Vapnik and Chervonenkis, 1971), and parameter norms (Bartlett et al., 2017; Neyshabur et al., 2017a) also suggest that modern architectures have very large capacity. Moreover, from the empirical perspective, practical models are flexible enough to perfectly fit the training data, even if the labels are pure noise (Zhang et al., 2017). Surprisingly, these same high-capacity models generalize well when trained on real data, even without any explicit control of capacity.

These observations are in conflict with classical generalization theory, which contends that models of intermediate complexity should generalize best, striking a balance between the bias and the variance of their predictive functions. To reconcile theory with observation, it has been suggested that deep neural networks may enjoy some form of implicit regularization induced by gradient-based training algorithms that biases the trained models towards simpler functions. However, the exact notion of simplicity and the mechanism by which it might be achieved remain poorly understood except in certain simplistic settings.

One concrete mechanism by which such induced simplicity can emerge is the hypothesis that neural networks learn simple functions early in training, and increasingly build up their complexity in later time. In particular, recent empirical work Nakkiran et al. (2019) found that, intriguingly, in some natural settings the simple function being learned in the early phase may just be a linear function of the data.

Key to our technical analysis is a bound on the spectral norm of the difference between the Neural Tangent Kernel (NTK) (Jacot et al., 2018) of the neural network at initialization and that of the linear model; indeed, a weaker result, like a bound on the Frobenius norm, would be insufficient to establish our result. Although the NTK is usually associated with the study of ultra-wide networks, our result only has a mild requirement on the width and allows the network to leave the kernel regime later in training. While our formal result focuses on two-layer fully-connected networks and data with benign concentration properties (specified in Assumption 3.1), we argue with theory and provide empirical evidence that the same linear learning phenomenon persists for more complex architectures and real-world datasets.

The early phase of neural network training has been the focus of considerable recent research. Frankle and Carbin (2019) found that sparse, trainable subnetworks – “lottery tickets" – emerge early in training. Achille et al. (2017) showed the importance of early learning from the perspective of creating strong connections that are robust to corruption. Gur-Ari et al. (2018) observed that after a short period of training, subsequent gradient updates span a low-dimensional subspace. Li et al. (2019a); Lewkowycz et al. (2020) showed that an initial large learning rate can benefit late-time generalization performance.

Implicit regularization of (stochastic) gradient descent has also been studied in various settings, suggesting a bias towards large-margin, low-norm, or low-rank solutions (Gunasekar et al., 2017, 2018; Soudry et al., 2018; Li et al., 2018; Ji and Telgarsky, 2019a, b; Arora et al., 2019a; Lyu and Li, 2019; Chizat and Bach, 2020; Razin and Cohen, 2020). These results mostly aim to characterize the final solutions at convergence, while our focus is on the early-time learning dynamics. Another line of work has identified that deep linear networks gradually increase the rank during training (Arora et al., 2019a; Saxe et al., 2014; Lampinen and Ganguli, 2018; Gidel et al., 2019).

A line of work adopted the Fourier perspective and demonstrated that low-frequency functions are often learned first (Rahaman et al., 2018; Xu, 2018; Xu et al., 2019a, b). Based on the NTK theory, Arora et al. (2019c) showed that for very wide networks, components lying in the top eigenspace of the NTK are learned faster than others. Using this principle, Su and Yang (2019); Cao et al. (2019) analyzed the spectrum of the infinite-width NTK. However, in order to obtain precise characterization of the spectrum these papers require special data distributions such as uniform distribution on the sphere.

Most relevant to our work is the finding of Nakkiran et al. (2019) that a neural network learned in the early phase of training can be almost fully explained by a linear function of the data. They supported this claim empirically by examining an information theoretic measure between the predictions of the neural network and the linear model. Our result formally proves that neural network and a corresponding linear model make similar predictions in early time, thus providing a theoretical explanation of their empirical finding.

In Section 2, we introduce notation and briefly recap the Neural Tangent Kernel. In Section 3, we present our main theoretical results on two-layer neural networks as well as empirical verification. In Section 4, we discuss extensions to more complicated architecture from both theoretical and empirical aspects. We conclude in Section 5, and defer additional experimental results and all the proofs to the appendices.

Preliminaries

We use the standard $O(\cdot)$, $\Omega(\cdot)$ and $\Theta(\cdot)$ notation to only hide universal constant factors. For $a,b\geq 0$, we also use $a\lesssim b$ or $b\gtrsim a$ to mean $a=O(b)$, and use $a\ll b$ or $b\gg a$ to mean $b\geq Ca$ for a sufficiently large universal constant $C>0$. Throughout the paper, “high probability” means a large constant probability arbitrarily close to $1$ (such as $0.99$).

Consider a single-output neural network $f({\bm{x}};{\bm{\theta}})$ where ${\bm{x}}$ is the input and ${\bm{\theta}}$ is the collection of parameters in the network. Around a reference network with parameters $\bar{{\bm{\theta}}}$, we can do a local first-order approximation:

Thus when ${\bm{\theta}}$ is close to $\bar{{\bm{\theta}}}$, for a given input ${\bm{x}}$ the network can be viewed as linear in $\nabla_{\bm{\theta}}f({\bm{x}};\bar{{\bm{\theta}}})$. This gradient feature map ${\bm{x}}\mapsto\nabla_{\bm{\theta}}f({\bm{x}};\bar{{\bm{\theta}}})$ induces a kernel $K_{\bar{{\bm{\theta}}}}({\bm{x}},{\bm{x}}^{\prime}):=\langle\nabla_{\bm{\theta}}f({\bm{x}};\bar{{\bm{\theta}}}),\nabla_{\bm{\theta}}f({\bm{x}}^{\prime};\bar{{\bm{\theta}}})\rangle$ which is called the NTK at $\bar{{\bm{\theta}}}$. Gradient descent training of the neural network can be viewed as kernel gradient descent on the function space with respect to the NTK. We use NTK matrix to refer to an $n\times n$ matrix that is the NTK evaluated on $n$ datapoints.

While in general the NTK is random at initialization and can vary significantly during training, it was shown that, for a suitable network parameterization (known as the “NTK parameterization”), when the width goes to infinity or is sufficiently large, the NTK converges to a deterministic limit at initialization and barely changes during training (Jacot et al., 2018; Lee et al., 2019; Arora et al., 2019b; Yang, 2019), so that the neural network trained by gradient descent is equivalent to a kernel method with respect to a fixed kernel. However, for networks with practical widths, the NTK does usually stray far from its initialization.

Two-Layer Neural Networks

We consider a two-layer fully-connected neural network with $m$ hidden neurons defined as:

and run vanilla gradient descent (GD) on the objective (2) starting from random initialization. Specifically, we use the following symmetric initialization for the weights $({\bm{W}},{\bm{v}})$:

Let $({\bm{W}}(0),{\bm{v}}(0))$ be a set of initial weights drawn from the symmetric initialization (3). Then the weights are updated according to GD:

where $\eta_{1}$ and $\eta_{2}$ are the learning rates. Here we allow potentially different learning rates for flexibility.

Now we state the assumption on the input distribution used in our theoretical results.

Note that a special case that satisfies Assumption 3.1 is the Gaussian distribution $\mathcal{N}({\bm{0}},{\bm{\Sigma}})$, but we allow a much larger class of distributions here. The subgaussian assumption is made due to the probabilistic tail bounds used in the analysis, and it can be replaced with a weaker bounded moment condition. The independence between $\bar{{\bm{x}}}$’s entries may also be dropped if its density is strongly log-concave. We choose to use Assumption 3.1 as the most convenient way to present our results.

We allow $\phi$ to be any of the commonly used activation functions, including ReLU, Leaky ReLU, Erf, Tanh, Sigmoid, Softplus, etc. Formally, our requirement on $\phi$ is the following:

The activation function $\phi(\cdot)$ satisfies either of the followings:

We will consider the regime where the data dimension $d$ is sufficiently large (i.e., larger than any constant) and the number of datapoints $n$ is at most some polynomial in $d$ (i.e., $n\leq d^{O(1)}$). These imply $\log n=O(\log d)<d^{c}$ for any constant $c>0$.

Under Assumption 3.1, the datapoints satisfy the following concentration properties:

Suppose $n\gg d$. Then under Assumption 3.1, with high probability we have \frac{\left\|{\bm{x}}_{i}\right\|^{2}}{d}=1\pm O\Big{(}\sqrt{\tfrac{\log n}{d}}\Big{)} ($\forall i\in[n]$), \frac{|\langle{\bm{x}}_{i},{\bm{x}}_{j}\rangle|}{d}=O\Big{(}\sqrt{\tfrac{\log n}{d}}\Big{)} ($\forall i,j\in[n],i\not=j$), and $\left\|{\bm{X}}{\bm{X}}^{\top}\right\|=\Theta(n)$.

The main result in this section is to formally prove that the neural network trained by GD is approximately a linear function in the early phase of training. As we will see, there are distinct contributions coming from the two layers. Therefore, it is helpful to divide the discussion into the cases of training the first layer only, the second layer only, and both layers together. All the omitted proofs in this section are given in Appendix D.

The width requirement in Theorem 3.2 is very mild as it only requires the width $m$ to be larger than $d^{1+\alpha}$ for some small constant $\alpha$. Note that the width is allowed to be much smaller than the number of samples $n$, which is usually the case in practice.

The agreement guaranteed in Theorem 3.2 is up to iteration $T=c\cdot\frac{d\log d}{\eta_{1}}$ (for some constant $c$). It turns out that for well-conditioned data, after $T$ iterations, a near optimal linear model will have been reached. This means that the neural network in the early phase approximates a linear model all the way until the linear model converges to the optimum. See Corollary 3.3 below.

2 Training the Second Layer

As usual, this linear model is trained with GD starting from zero:

Similar to Theorem 3.2, our main theorem for training the second layer is the following:

We remark that if the data distribution is well-conditioned, we can also have a guarantee similar to Corollary 3.3.

3 Training Both Layers

Finally we consider the case where both layers are trained, in which $\eta_{1}=\eta_{2}=\eta>0$ in (4). Since the NTK for training both layers is simply the sum of the first-layer NTK and the second-layer NTK, the corresponding linear model should have its kernel being the sum of the kernels for linear models (5) and (9), which can be derived easily:

where the constants are from (9). Note that $\langle{\bm{\psi}}({\bm{x}}),{\bm{\psi}}({\bm{x}}^{\prime})\rangle=\langle{\bm{\psi}}_{1}({\bm{x}}),{\bm{\psi}}_{1}({\bm{x}}^{\prime})\rangle+\langle{\bm{\psi}}_{2}({\bm{x}}),{\bm{\psi}}_{2}({\bm{x}}^{\prime})\rangle$.

Again, we can show that the neural network is close to the linear model (11) in early time. The guarantee is very similar to Theorems 3.2 and 3.5, so we defer the formal theorem to Appendix D; see Theorem D.1. Note that our result can be directly generalized to the case where $\eta_{1}\not=\eta_{2}$, for which we just need to redefine the linear model using a weighted combination of the kernels for (5) and (9).

4 Empirical Verification

Extensions to Multi-Layer and Convolutional Neural Networks

In this section, we provide theoretical and empirical evidence supporting that the agreement between neural networks and linear models in the early phase of training may continue to hold for more complicated network architectures and datasets than what we analyzed in Section 3.

We consider a simple 1-dimensional CNN with one convolutional layer and without pooling (generalization to the commonly used 2-dimensional CNNs is straightforward):

We have the following result concerning the NTK of this CNN:

The proof is given in Appendix E. The above result shows that the NTK of a CNN can also be close to the (scaled) data kernel, which implies the linear learning behavior in the early time of training the CNN. Our empirical results will show that this behavior can even persist to multi-layer CNNs and real data beyond our analysis.

2 Empirical Results

Conclusion

This work gave a novel theoretical result rigorously showing that gradient descent on a neural network learns a simple linear function in the early phase. While we mainly focused on two-layer fully-connected neural networks, we further provided theoretical and empirical evidence suggesting that this phenomenon continues to exist in more complicated models. Formally extending our result to those settings is a direction of future work. Another interesting direction is to study the dynamics of neural networks after the initial linear learning phase.

References

Appendices

In Appendix A, we describe additional experiment details and provide additional plots. In Appendix B, we introduce additional notation and some lemmas that will be used in the proofs. In Appendix C, we present a general result that shows how the GD trajectory of a non-linear least squares problem can be approximated by a linear one, which will be used in the proofs. Finally, in Appendices D and E we provide omitted details and proofs in Sections 3 and 4, respectively.

Appendix A Experiment Setup and Additional Plots

We provide additional plots and describe additional experiment details in this section.

In Figure 4, we repeat the same experiments in Figure 3 on the full-size ($32\times 32\times 3$) CIFAR-10 as well as MNIST datasets, using the same $4$-hidden-layer FC and CNN architectures. For both datasets we take two classes and perform binary classification. We see very good early-time agreement except for CNN on CIFAR-10, where the agreement only lasts for a shorter time.

We use the Neural Tangents Library [Novak et al., 2019] and JAX [Bradbury et al., 2018] for our experiments.

Appendix B Additional Notation and Lemmas

We introduce some additional notation and lemmas that will be used in the proofs.

For any matrix ${\bm{A}}$ and a submatrix ${\bm{A}}_{1}$ of ${\bm{A}}$, we have $\left\|{\bm{A}}_{1}\right\|\leq\left\|{\bm{A}}\right\|$.

For simplicity we assume that ${\bm{A}}_{1}$ is in the top-left corner of ${\bm{A}}$, i.e. ${\bm{A}}=\begin{bmatrix}{\bm{A}}_{1}&{\bm{A}}_{2}\\ {\bm{A}}_{3}&{\bm{A}}_{4}\end{bmatrix}$. The same proof works when ${\bm{A}}_{1}$ is any other submatrix of ${\bm{A}}$.

By the definition of spectral norm, we have

From Lemma B.1 we know that $\left|\left[{\bm{A}}\right]_{i,i}\right|\leq\left\|{\bm{A}}\right\|$ for all $i$ since $\left[{\bm{A}}\right]_{i,i}$ can be viewed as a submatrix of ${\bm{A}}$. Thus we have

For any two positive semidefinite matrices ${\bm{A}},{\bm{B}}$, we have

Appendix C General Result on the Closeness between Two Dynamics

We present a general result that shows how the GD trajectory for a non-linear least squares problem can be simulated by a linear one. Later we will specialize this result to the settings considered in the paper.

We consider an objective function of the form:

Consider another linear least squares problem:

Let ${\bm{K}}:={\bm{\Phi}}{\bm{\Phi}}^{\top}$, and let

which stand for the predictions of these two models at iteration $t$.

The linear dynamics admit a very simple analytical form, summarized below.

We make the following assumption that connects these two problems:

(boundedness of parameter movement) $\left\|{\bm{\theta}}(t)-{\bm{\theta}}(0)\right\|\leq R,\left\|{\bm{\omega}}(t)-{\bm{\omega}}(0)\right\|\leq R$.

We first prove the first two properties, and will prove the last property $\left\|{\bm{\omega}}(t)-{\bm{\omega}}(0)\right\|\leq R$ at the end.

We first prove $\left\|{\bm{\theta}}(t-1)-{\bm{\theta}}(0)\right\|\leq\frac{R}{2}$. If $t=1$, this is trivially true. Now we assume $t\geq 2$. For each $0\leq\tau<t-1$, by the fundamental theorem for line integrals we have

Let ${\bm{E}}(\tau):={\bm{J}}({\bm{\theta}}(\tau)\to{\bm{\theta}}(\tau+1)){\bm{J}}({\bm{\theta}}(\tau))^{\top}-{\bm{K}}$. Since $\left\|{\bm{\theta}}(\tau)-{\bm{\theta}}(0)\right\|\leq R$ and $\left\|{\bm{\theta}}(\tau+1)-{\bm{\theta}}(0)\right\|\leq R$, from Assumption C.1 we know that $\left\|{\bm{E}}(\tau)\right\|\leq\epsilon$. We can write

Combining the above two inequalities, we obtain

Taking sum over $\tau=0,\ldots,t-2$, we get

Then by the Cauchy-Schwartz inequality we have

Choosing $c$ sufficiently small, we can ensure $\left\|{\bm{\theta}}(t-1)-{\bm{\theta}}(0)\right\|\leq\frac{R}{2}$.

Now that we have proved $\left\|{\bm{\theta}}(t-1)-{\bm{\theta}}(0)\right\|\leq\frac{R}{2}$, to prove $\left\|{\bm{\theta}}(t)-{\bm{\theta}}(0)\right\|\leq R$ it suffices to bound the one-step deviation $\left\|{\bm{\theta}}(t)-{\bm{\theta}}(t-1)\right\|$ by $\frac{R}{2}$. Using the exact same method in (14), we have

where we have used $\eta\leq\frac{n}{\left\|{\bm{K}}\right\|}$ and $\eta\leq\eta t\leq cR^{2}$. Choosing $c$ sufficiently small, we can ensure $\left\|{\bm{\theta}}(t)-{\bm{\theta}}(t-1)\right\|\leq\frac{R}{2}$. Therefore we conclude that $\left\|{\bm{\theta}}(t)-{\bm{\theta}}(0)\right\|\leq R$.

where ${\bm{E}}(t-1)={\bm{J}}({\bm{\theta}}(t-1),{\bm{\theta}}(t)){\bm{J}}({\bm{\theta}}(t-1))^{\top}-{\bm{K}}$. Since $\left\|{\bm{\theta}}(t-1)-{\bm{\theta}}(0)\right\|\leq R$ and $\left\|{\bm{\theta}}(t)-{\bm{\theta}}(0)\right\|\leq R$, we know from Assumption C.1 that $\left\|{\bm{E}}(t-1)\right\|\leq\epsilon$. Moreover, from Claim C.1 we know

Finally, we prove the last statement in the theorem, i.e., $\left\|{\bm{\omega}}(t)-{\bm{\omega}}(0)\right\|\leq R$. In fact we have already proved this – notice that we have proved $\left\|{\bm{\theta}}(t)-{\bm{\theta}}(0)\right\|\leq R$ and that a special instance of this problem is when ${\bm{\theta}}(t)={\bm{\omega}}(t)$, i.e., the two dynamics are the same. Applying our result on that problem instance, we obtain $\left\|{\bm{\omega}}(t)-{\bm{\omega}}(0)\right\|\leq R$. ∎

Appendix D Omitted Details in Section 3

In Section D.1, we present the formal theoretical guarantee (Theorem D.1) for the case of training both layers.

In Section D.2, we calculate the formulae of various Jacobians and NTKs that will be used in the analysis.

In Section D.3, we prove Theorem 3.2 (training the first layer).

In Section D.4, we prove Corollary 3.3 (training the first layer with well-conditioned data).

In Section D.5, we prove Theorem 3.5 (training the second layer).

In Section D.6, we prove Theorem D.1 (training both layers).

In Section D.7, we prove Claim 3.1 (data concentration properties).

We remark that if the data distribution is well-conditioned, we can also have a guarantee similar to Corollary 3.3.

D.2 Formulae of Jacobians and NTKs

We first calculate the Jacobian of the network outputs at the training data ${\bm{X}}$ with respect to the weights in the network. The Jacobian for the first layer is:

After calculating the Jacobians, we can calculate the NTK matrices for the first layer, the second layer, and both layers as follows:

We also denote the expected NTK matrices at random initialization as:

These are also the NTK matrices at infinite width ($m\to\infty$).

Next, for the three linear models (5), (9) and (11) defined in Section 3, denote their feature/Jacobian matrices by:

Consequently, their corresponding kernel matrices are:

D.3 Proof of Theorem 3.2 (Training the First Layer)

For convenience we let ${\bm{v}}={\bm{v}}(0)$ which is the fixed second layer. Since we have $v_{r}\in\{\pm 1\}$ ($\forall r\in[m]$), we can write the first-layer NTK matrix as

Because it does not depend on ${\bm{v}}$, we denote ${\bm{\Theta}}_{1}({\bm{W}}):={\bm{\Theta}}_{1}({\bm{W}},{\bm{v}})$ for convenience.

Now we prove Proposition 3.4, restated below:

With high probability over the random initialization ${\bm{W}}(0)$ and the training data ${\bm{X}}$, we have

With high probability over the random initialization ${\bm{W}}(0)$ and the training data ${\bm{X}}$, we have

For convenience we denote ${\bm{W}}={\bm{W}}(0)$ and ${\bm{\Theta}}_{1}={\bm{\Theta}}_{1}({\bm{W}})={\bm{\Theta}}_{1}({\bm{W}}(0))$ in this proof.

From Claim 3.1 we know $\left\|{\bm{X}}{\bm{X}}^{\top}\right\|=O(n)$ with high probability. For the rest of the proof we will be conditioned on ${\bm{X}}$ and on Claim 3.1, and only consider the randomness in ${\bm{W}}$.

Next we will apply the matrix Bernstein inequality (Theorem 1.6.2 in Tropp ) to bound $\left\|{\bm{\Theta}}_{1}-{\bm{\Theta}}_{1}^{*}\right\|$. We will first consider the first half of independent neurons, i.e. $r\in[m/2]$. For each $r$ we have

Therefore, from the the matrix Bernstein inequality, for any $s\geq 0$ we have:

Letting $s=\frac{m}{2}\cdot\frac{n}{d^{1+\alpha}}$, we obtain

where we have used $m=\Omega(d^{1+\alpha})$ and $n=d^{O(1)}$. Therefore with high probability we have

Similarly, for the second half of the neurons we also have with high probability

Finally, by the triangle inequality we have

with high probability, completing the proof. ∎

With high probability over the training data ${\bm{X}}$, we have

We will be conditioned on the high probability events stated in Claim 3.1.

By the definition of ${\bm{\Theta}}_{1}^{*}$, we know

We consider the diagonal and off-diagonal entries of ${\bm{\Theta}}_{1}^{*}$ separately.

Now we treat the terms in (21) separately. First, we have

For the final term ${\bm{E}}$ in (21), we have

For the diagonal entries of ${\bm{\Theta}}_{1}^{*}$, we have $\left[{\bm{\Theta}}_{1}^{*}\right]_{i,i}=\frac{\left\|{\bm{x}}_{i}\right\|^{2}}{d}\cdot\Phi\left(\frac{\left\|{\bm{x}}_{i}\right\|^{2}}{d},\frac{\left\|{\bm{x}}_{i}\right\|^{2}}{d},\frac{\left\|{\bm{x}}_{i}\right\|^{2}}{d}\right)$. We denote $\bar{\Phi}(a):=\Phi(a,a,a)$ ($a\geq 0$). When $\phi$ is a smooth activation as in Assumption 3.2, we know that $\bar{\Phi}$ has bounded derivative, and thus we get

Combining the off-diagonal and diagonal approximations (23) and (25), we obtain

Finally, when $n\gtrsim d^{1+\alpha}$ ($0<\alpha<\frac{1}{4}$), we have $\left\|{\bm{I}}\right\|=1\lesssim\frac{n}{d^{1+\alpha}}$. Hence we can discard the identity component above and get

Combining Propositions D.3 and D.4 directly gives Proposition D.2.

D.3.2 Agreement on Training Data

If $\phi$ is a smooth activation as in Assumption 3.2, then with high probability over the training data ${\bm{X}}$, we have

If $\phi$ is a piece-wise linear activation as in Assumption 3.2, then with high probability over the random initialization ${\bm{W}}(0)$ and the training data ${\bm{X}}$, we have

Throughout the proof we will be conditioned on ${\bm{X}}$ and on the high-probability events in Claim 3.1.

By the definition of ${\bm{J}}_{1}({\bm{W}},{\bm{v}})$ in (15), we have

Then if $\phi$ is a smooth activation, we have with high probability,

Next we consider the case where $\phi$ is a piece-wise linear activation. From (28) and Lemma B.3 we have

Since $\phi^{\prime}$ is a step function that only depends on the sign of the input, we have

Therefore we need to bound $|M_{i}|$, i.e. how many coordinates in ${\bm{W}}{\bm{x}}_{i}$ and ${\bm{W}}(0){\bm{x}}_{i}$ differ in sign for each $i\in[n]$.

Let $\lambda>0$ be a parameter whose value will be determined later. For each $i\in[n]$, define

Taking a union bound over all $i\in[n]$, we know that with high probability,

By definition, if $r\in M_{i}$ but $r\notin N_{i}$, we must have $\left|{\bm{w}}_{r}^{\top}{\bm{x}}_{i}-{\bm{w}}_{r}(0)^{\top}{\bm{x}}_{i}\right|\geq\left|{\bm{w}}_{r}(0)^{\top}{\bm{x}}_{i}\right|>\lambda\left\|{\bm{x}}_{i}\right\|$. This leads to

Letting $\lambda=\left(\frac{\left\|{\bm{W}}-{\bm{W}}(0)\right\|^{2}}{m}\right)^{1/3}$, we get

Finally, we combine (29), (30) and (33) to obtain

The next lemma verifies Assumption C.1 for the case of training the first layer.

This proof is conditioned on all the high-probability events we have shown.

where we have used $m\gtrsim d^{1+\alpha}$. If $\phi$ is a piece-wise linear activation, from Lemma D.5 we have

Hence we always have $\left\|{\bm{J}}_{1}({\bm{W}},{\bm{v}})-{\bm{J}}_{1}({\bm{W}}(0),{\bm{v}})\right\|\leq\frac{\sqrt{n}}{d^{\frac{1}{2}+\frac{\alpha}{7}}}$. Similarly, we have $\left\|{\bm{J}}_{1}({\widetilde{\bm{W}}},{\bm{v}})-{\bm{J}}_{1}({\bm{W}}(0),{\bm{v}})\right\|\leq\frac{\sqrt{n}}{d^{\frac{1}{2}+\frac{\alpha}{7}}}$.

Note that from Proposition D.2 and Claim 3.1 we know

which implies $\left\|{\bm{J}}_{1}({\bm{W}}(0),{\bm{v}})\right\|\lesssim\sqrt{\frac{n}{d}}$. It follows that $\left\|{\bm{J}}_{1}({\bm{W}},{\bm{v}})\right\|\lesssim\sqrt{\frac{n}{d}}+\frac{\sqrt{n}}{d^{\frac{1}{2}+\frac{\alpha}{7}}}\lesssim\sqrt{\frac{n}{d}}$ and $\left\|{\bm{J}}_{1}({\widetilde{\bm{W}}},{\bm{v}})\right\|\lesssim\sqrt{\frac{n}{d}}$. Then we have

Combining the above inequality with Proposition D.2, we obtain

Finally, we can instantiate Theorem C.2 to conclude the proof of (7):

There exists a universal constant $c>0$ such that with high probability, for all $0\leq t\leq T=c\cdot\frac{d\log d}{\eta_{1}}$ simultaneously, we have:

$\left\|{\bm{W}}(t)-{\bm{W}}(0)\right\|_{F}\leq\sqrt{d\log d}$, $\left\|{\bm{\beta}}(t)\right\|\leq\sqrt{d\log d}$.

Furthermore, Theorem C.2 also tells us $\left\|{\bm{W}}(t)-{\bm{W}}(0)\right\|_{F}\leq\sqrt{d\log d}$ and $\left\|{\bm{\beta}}(t)\right\|\leq\sqrt{d\log d}$. ∎

D.3.3 Agreement on Distribution

where ${\bm{J}}_{1}({\bm{W}}(0)\to{\bm{W}}(t),{\bm{v}}):=\int_{0}^{1}{\bm{J}}_{1}({\bm{W}}(0)+x({\bm{W}}(t)-{\bm{W}}(0)),{\bm{v}})dx$. Since $\left\|{\bm{W}}(t)-{\bm{W}}(0)\right\|_{F}\leq\sqrt{d\log d}$ according to Proposition D.7, we can use Lemma D.5 in the same way as in the proof of Lemma D.6 and obtain

Since $\phi^{\prime}$ is bounded and $\frac{\left\|{\bm{x}}_{i}\right\|^{2}}{d}=O(1)$ ($\forall i\in[n]$), we can bound

Now using the standard generalization bound via Rademacher complexity (see e.g. Mohri et al. ), and noticing that the function $z\mapsto\min\{z^{2},1\}$ is $2$-Lipschitz and bounded in $$, we have with high probability, for all$t\leq T$ simultaneously,

Then letting $\delta=\frac{1}{100T}$ and taking a union bound over $t\leq T$, we obtain that with high probability, for all $t\leq T$ simultaneously,

D.4 Proof of Corollary 3.3 (Training the First Layer, Well-Conditioned Data)

According to the linear dynamics (6), we have the following relation (see Claim C.1):

Therefore we can apply the standard Rademacher complexity argument and conclude the proof of (41). ∎

D.5 Proof of Theorem 3.5 (Training the Second Layer)

Since the first layer is kept fixed in this case, we let ${\bm{W}}={\bm{W}}(0)$ for notational convenience. Similar to the proof of Theorem 3.2 in Section D.3, we still divide the proof into 3 parts: analyzing the NTK at initialization (which is also the NTK throughout training in this case), proving the agreement on training data, and proving the agreement on the distribution.

With high probability over the random initialization ${\bm{W}}$ and the training data ${\bm{X}}$, we have

With high probability over the training data ${\bm{X}}$, we have

We will be conditioned on the high probability events stated in Claim 3.1.

By the definition of ${\bm{\Theta}}_{2}^{*}$, we know

Denote $e_{i}:=\frac{\left\|{\bm{x}}_{i}\right\|}{\sqrt{d}}-1$ and $s_{i,j}:=\frac{{\bm{x}}_{i}^{\top}{\bm{x}}_{j}}{d}$. Below we consider the diagonal and off-diagonal entries of ${\bm{\Theta}}_{2}^{*}$ separately.

For $i\not=j$, we do a Taylor expansion of $\Gamma$ around $(1,1,0)$:

Here $\zeta,\vartheta_{0},\vartheta_{1},\vartheta_{2}$ are defined in (9), and $\gamma$ is the $(1,3)$-th entry in the Hessian $\nabla^{2}\Gamma(1,1,0)$ whose specific value is not important to us. Recall that $\left[{\bm{q}}\right]_{i}=\vartheta_{0}+\vartheta_{1}e_{i}+\vartheta_{2}e_{i}^{2}$.

On the other hand, by the definition (20) we have

since $n\gtrsim d^{1+\alpha}$ ($0<\alpha<\frac{1}{4}$). ∎

With high probability over the random initialization ${\bm{W}}$ and the training data ${\bm{X}}$, we have

Since $\left\|{\bm{w}}_{r}\right\|^{2}$ is a $\chi^{2}$ random variable with $d$ degrees of freedom, it has sub-exponential norm $O(d)$, which implies that the random variable $\left\|{\bm{\Theta}}_{2}^{(r)}-{\bm{\Theta}}_{2}^{*}\right\|$ has sub-exponential norm $O(n)$.

Next we bound the variance. Let $B>0$ be a threshold to be determined. We have:

Thus, letting $\frac{s^{2}}{n/d}=C\log n$ for a sufficiently large universal constant $C>0$, we know that with probability at least $1-n^{-10}$ over ${\bm{w}}\sim\mathcal{N}({\bm{0}},{\bm{I}})$,

Hence we pick the threshold $B=C^{\prime}\sqrt{n}$ which is the upper bound above, where $C^{\prime}>0$ is a universal constant.

Recall that in this case Theorem 3.5 assumes $m\gtrsim d^{2+\alpha}$.

Applying Proposition 4.1 in Klochkov and Zhivotovskiy , we know that for any $u\gg\max\{n\log m,n\sqrt{m}\}=n\sqrt{m}$,

Let $u=m\cdot\frac{n}{d^{1+\frac{\alpha}{3}}}$. We can verify $u\gg n\sqrt{m}$ since $m\gtrsim d^{2+\alpha}$. Then we have

Similarly, for the second half of the neurons we also have $\left\|\sum_{r=m/2+1}^{m}({\bm{\Theta}}_{2}^{(r)}-{\bm{\Theta}}_{2}^{*})\right\|\leq m\cdot\frac{n}{d^{1+\frac{\alpha}{3}}}$ with high probability. Therefore we have with high probability,

Recall that in this case Theorem 3.5 assumes $m\gtrsim d^{1+\alpha}$.

Applying Proposition 4.1 in Klochkov and Zhivotovskiy , we know that for any $u\gg\max\{n\log m,n\sqrt{m/d^{1-\frac{\alpha}{10}}}\}=n\sqrt{m/d^{1-\frac{\alpha}{10}}}$,

Let $u=m\cdot\frac{n}{d^{1+\frac{\alpha}{3}}}$. We can verify $u\gg n\sqrt{m/d^{1-\frac{\alpha}{10}}}$ since $m\gtrsim d^{1+\alpha}$. Then we have

Similarly, for the second half of the neurons we also have $\left\|\sum_{r=m/2+1}^{m}({\bm{\Theta}}_{2}^{(r)}-{\bm{\Theta}}_{2}^{*})\right\|\leq m\cdot\frac{n}{d^{1+\frac{\alpha}{3}}}$ with high probability. Therefore we have with high probability,

Combining Propositions D.9 and D.10 directly gives Proposition D.8.

D.5.2 Agreement on Training Data

This proves the first part in Theorem 3.5.

Note that Theorem C.2 also tells us $\left\|{\bm{v}}(t)-{\bm{v}}(0)\right\|\leq\sqrt{d\log d}$ and $\left\|{\bm{\gamma}}(t)\right\|\leq\sqrt{d\log d}$, which will be useful for proving the guarantee on the distribution $\mathcal{D}$.

D.5.3 Agreement on Distribution

Next we bound the above two traces. First, we have

with high probability for all $i\in[n]$ together. Here we have used the standard tail bound for $\chi^{2}$ random variables and a union bound over $i\in[n]$. Hence we have $\operatorname{Tr}[{\bm{\Theta}}_{2}({\bm{W}})]\lesssim n$. For the second trace, we have

with high probability. Therefore we can bound the Rademacher complexity by $\sqrt{\frac{d\log d}{n}}$. Then we can conclude the agreement guarantee on the distribution $\mathcal{D}$, i.e., for all $t\leq T$ simultaneously,

D.6 Proof of Theorem D.1 (Training Both Layers)

The proof for training both layers follows the same ideas in the proofs for training the first layer only and the second layer only. In fact, most technical components needed in the proof were already developed in the previous proofs. The only new component is a Jacobian perturbation bound for the case of training both layers, Lemma D.12 (analog of Lemma D.5 for training the first layer).

With high probability over the random initialization $({\bm{W}}(0),{\bm{v}}(0))$ and the training data ${\bm{X}}$, we have

D.6.2 Agreement on Training Data

The proof for the agreement on training data is similar to the case of training the first layer only (Section D.3.2). We will again apply Theorem C.2. For this we need a new Jacobian perturbation lemma to replace Lemma D.5, since both layers are allowed to move now.

If $\phi$ is a smooth activation as in Assumption 3.2, then with high probability over the training data ${\bm{X}}$, we have

If $\phi$ is a piece-wise linear activation as in Assumption 3.2, then with high probability over the random initialization ${\bm{W}}(0)$ and the training data ${\bm{X}}$, we have

Furthermore, with high probability over the training data ${\bm{X}}$, we have

We will be conditioned on ${\bm{X}}$ and on the high-probability events in Claim 3.1.

We first consider the first-layer Jacobian. By the definition of ${\bm{J}}_{1}({\bm{W}},{\bm{v}})$ in (15), we have

Then if $\phi$ is a smooth activation, we have with high probability,

If $\phi$ is a piece-wise linear activation, then with high probability,

For the second-layer Jacobian, we have with high probability,

Based on Lemma D.12, we can now verify Assumption C.1 for the case of training both layers:

Let $R=\sqrt{d\log d}$. With high probability over the random initialization and the training data, for all $({\bm{W}},{\bm{v}})$ and $({\widetilde{\bm{W}}},{\widetilde{{\bm{v}}}})$ such that $\left\|{\bm{W}}-{\bm{W}}(0)\right\|_{F}\leq R$, $\left\|{\widetilde{\bm{W}}}-{\bm{W}}(0)\right\|_{F}\leq R$, $\left\|{\bm{v}}-{\bm{v}}(0)\right\|\leq R$ and $\left\|{\widetilde{{\bm{v}}}}-{\bm{v}}(0)\right\|\leq R$, we have

This proof is conditioned on all the high-probability events we have shown.

Now consider $({\bm{W}},{\bm{v}})$ and $({\widetilde{\bm{W}}},{\widetilde{{\bm{v}}}})$ which satisfy the conditions stated in the lemma.

If $\phi$ is a smooth activation, from Lemma D.12 we know

where we have used $m\gtrsim d^{2+\alpha}$. If $\phi$ is a piece-wise linear activation, from Lemma D.12 we have

Hence in either case have $\left\|{\bm{J}}_{1}({\bm{W}},{\bm{v}})-{\bm{J}}_{1}({\bm{W}}(0),{\bm{v}}(0))\right\|\leq\frac{\sqrt{n}}{d^{\frac{1}{2}+\frac{\alpha}{3}}}$. Similarly, we have $\left\|{\bm{J}}_{1}({\widetilde{\bm{W}}},{\widetilde{{\bm{v}}}})-{\bm{J}}_{1}({\bm{W}}(0),{\bm{v}}(0))\right\|\leq\frac{\sqrt{n}}{d^{\frac{1}{2}+\frac{\alpha}{3}}}$.

Also, we know from Proposition D.2 that $\left\|{\bm{J}}_{1}({\bm{W}}(0),{\bm{v}}(0))\right\|\lesssim\sqrt{\frac{n}{d}}$. It follows that $\left\|{\bm{J}}_{1}({\bm{W}},{\bm{v}})\right\|\lesssim\sqrt{\frac{n}{d}}$ and $\left\|{\bm{J}}_{1}({\widetilde{\bm{W}}},{\widetilde{{\bm{v}}}})\right\|\lesssim\sqrt{\frac{n}{d}}$. Then we have

Next we look at the second-layer Jacobian. From Lemma D.12 we know $\left\|{\bm{J}}_{2}({\bm{W}})-{\bm{J}}_{2}({\bm{W}}(0))\right\|\lesssim\sqrt{\frac{n}{md}}\cdot\sqrt{d\log d}\lesssim\sqrt{\frac{n\log d}{d^{2+\alpha}}}\ll\frac{\sqrt{n}}{d^{1+\frac{\alpha}{3}}}$. Similarly we have $\left\|{\bm{J}}_{2}({\widetilde{\bm{W}}})-{\bm{J}}_{2}({\bm{W}}(0))\right\|\ll\frac{\sqrt{n}}{d^{1+\frac{\alpha}{3}}}$. Also, from Proposition D.8 we know $\left\|{\bm{J}}_{2}({\bm{W}}(0))\right\|\lesssim\sqrt{n}$, which implies $\left\|{\bm{J}}_{2}({\bm{W}})\right\|\lesssim\sqrt{n}$ and $\left\|{\bm{J}}_{2}({\widetilde{\bm{W}}})\right\|\lesssim\sqrt{n}$. It follows that

Combining the above auguments for two layers, we obtain

Combining the above inequality with Proposition D.11, the proof is finished. ∎

Finally, we can apply Theorem C.2 with $R=\sqrt{d\log d}$ and $\epsilon=O(\frac{n}{d^{1+\frac{\alpha}{3}}})$, and obtain that for all $t\leq T$:

This proves the first part in Theorem D.1.

Note that Theorem C.2 also tells us $\left\|{\bm{W}}(t)-{\bm{W}}(0)\right\|\leq\sqrt{d\log d}$, $\left\|{\bm{v}}(t)-{\bm{v}}(0)\right\|\leq\sqrt{d\log d}$ and $\left\|{\bm{\delta}}(t)\right\|\leq\sqrt{d\log d}$, which will be useful for proving the guarantee on the distribution $\mathcal{D}$.

D.6.3 Agreement on Distribution

The proof for the second part of Theorem D.1 is basically identical to the case of training the first layer (Section D.3.3), so we will only sketch the differences here to avoid repetition.

Recall that in Section D.3.3 we define an auxiliary model which is the first-order approximation of the network around initialization. Here since we are training both layers, we need to modify the definition of the auxiliary model to incorporate deviation from initialization in both layers:

There are two more minor changes to Section D.3.3:

In Sections D.3.3 and D.5.3, we have shown that the above 4 traces are all $O(n)$ with high probability. Hence we get the same Rademacher complexity bound as before.

Modulo these differences, the proof proceeds the same as Section D.3.3. Therefore we conclude the proof of Theorem D.1.

D.7 Proof of Claim 3.1

By Hanson-Wright inequality (specifically, Theorem 2.1 in Rudelson and Vershynin ), we have for any $t\geq 0$,

Let $t=C\sqrt{\log n}$ for a sufficiently large constant $C>0$. Taking a union bound over all $i\in[n]$, we obtain that with high probability, $\left\|{\bm{x}}_{i}\right\|=\sqrt{d}\pm O(\sqrt{\log n})$ for all $i\in[n]$ simultaneously. This proves the first property in Claim 3.1.

For $i\not=j$, we have $\langle{\bm{x}}_{i},{\bm{x}}_{j}\rangle=\bar{{\bm{x}}}_{i}^{\top}{\bm{\Sigma}}\bar{{\bm{x}}}_{j}$. Conditioned on $\bar{{\bm{x}}}_{j}$, we know that $\bar{{\bm{x}}}_{i}^{\top}{\bm{\Sigma}}\bar{{\bm{x}}}_{j}$ is zero-mean and $O(\left\|{\bm{\Sigma}}\bar{{\bm{x}}}_{j}\right\|^{2})$-subgaussian, which means for any $t\geq 0$,

Since we have shown that $\left\|{\bm{\Sigma}}\bar{{\bm{x}}}_{j}\right\|^{2}\lesssim\left\|{\bm{x}}_{j}\right\|^{2}\lesssim\sqrt{d}+\sqrt{\log n}\lesssim\sqrt{d}$ with probability at least $1-n^{-10}$, we have

Then we can take $t=C\sqrt{d\log n}$ and apply a union bound over $i,j$, which gives $\left|\langle{\bm{x}}_{i},{\bm{x}}_{j}\rangle\right|\lesssim\sqrt{d\log n}$ for all $i\not=j$ with high probability. This completes the proof of the second statement in Claim 3.1.

Finally, for ${\bm{X}}{\bm{X}}^{\top}$, we can use standard covariance concentration (see, e.g., Lemma A.6 in Du et al. ) to obtain $0.9{\bm{\Sigma}}\preceq\frac{1}{n}{\bm{X}}^{\top}{\bm{X}}\preceq 1.1{\bm{\Sigma}}$ with high probability. This implies $\left\|{\bm{X}}{\bm{X}}^{\top}\right\|=\left\|{\bm{X}}^{\top}{\bm{X}}\right\|=\Theta(n)$. ∎

Appendix E Omitted Details in Section 4

For two datapoints ${\bm{x}}_{i}$ and ${\bm{x}}_{j}$ ($i,j\in[n]$) and a location $k\in[d]$, we define

which is a local correlation between ${\bm{x}}_{i}$ and ${\bm{x}}_{j}$.

Now we calculate the infinite-width NTK matrix ${\bm{\Theta}}_{\mathsf{CNN}}$, which is also the expectation of a finite-width NTK matrix with respect to the randomly initialized weights $({\bm{W}},{\bm{V}})$. We divide the NTK matrix into two components corresponding to two layers: ${\bm{\Theta}}_{\mathsf{CNN}}={\bm{\Theta}}_{\mathsf{CNN}}^{(1)}+{\bm{\Theta}}_{\mathsf{CNN}}^{(2)}$, and consider the two layers separately.

Since the CNN model (12) is linear in the second layer weights, it is easy to derive the formula for the second-layer NTK:

Note that we have used the property $\left\|\left[{\bm{x}}_{j}\right]_{k:k+q}\right\|=\left\|\left[{\bm{x}}_{j}\right]_{k:k+q}\right\|=\sqrt{q}$ since the data are from the hypercube $\{\pm 1\}^{d}$.

Therefore we have shown that with high probability, for all $i\not=j$,

For the diagonal entries, we can easily see

Combining the above two equations, we obtain

We calculate the derivative of the output of the CNN with respect to the first-layer weights as:

Therefore, the entries in the first-layer NTK matrix are

Then, using the exact same analysis for the second-layer NTK, we know that with high probability,

Finally, combining the results for two layers, we conclude the proof of Proposition 4.1. ∎