# A Theory of Neural Tangent Kernel Alignment and Its Influence on Training

Haozhe Shan  Blake Bordelon

###### Abstract

The training dynamics and generalization properties of neural networks (NN)
can be precisely characterized in function space via the neural tangent kernel
(NTK). Structural changes to the NTK during training reflect feature learning
and underlie the superior performance of networks outside of the static kernel
regime. In this work, we seek to theoretically understand kernel alignment, a
prominent and ubiquitous structural change that aligns the NTK with the target
function. We first study a toy model of kernel evolution in which the NTK
evolves to accelerate training and show that alignment naturally emerges from
this demand. We then study alignment mechanism in deep linear networks and two
layer ReLU networks. These theories provide good qualitative descriptions of
kernel alignment and specialization in practical networks and identify factors
in network architecture and data structure that drive kernel alignment. In
nonlinear networks with multiple outputs, we identify the phenomenon of kernel
specialization, where the kernel function for each output head preferentially
aligns to its own target function. Together, our results provide a mechanistic
explanation of how kernel alignment emerges during NN training and a normative
explanation of how it benefits training.

##  1 Introduction

Deep learning provides a flexible framework to solve difficult statistical
inference problems across a variety of application areas (LeCun et al., 2015).
During optimization of the statistical objective, useful features are often
extracted by the neural network (NN) as the weights in the network evolve.
Though feature learning appears to be crucial to the success of neural
networks on large-scale problems (Geiger et al., 2020; Ghorbani et al., 2021;
Yang & Hu, 2021) (as well as enabling transfer learning on new, related
problems), the precise way that neural network features evolve to benefit
learning is not well understood theoretically.

One powerful framework for studying feature learning in NNs is through the
neural tangent kernel (NTK) (Jacot et al., 2020a; Lee et al., 2020b). This
framework grew out of the observation that in the limit of large widths and
small learning rates, NNs with certain parameterization behave like linear
models in their parameters. In this case, NN training is equivalent to kernel
gradient descent (KGD) with a static neural tangent kernel. While this limit
allows a precise characterization of training and generalization dynamics to
be obtained (Bietti & Mairal, 2019; Yang & Salman, 2020; Bordelon et al.,
2020; Bahri et al., 2021), the time stationarity of the kernel indicates that
feature learning does not occur. In practical NNs, however, widths are finite
and the NTK evolves during training (Dyer & Gur-Ari, 2020; Geiger et al.,
2020). In this work, we consider feature learning in NNs as the evolution of
the NTK during training.

Recent empirical works have identified that a ubiquitous feature of the NTK
evolution (Baratin et al., 2021; Fort et al., 2020; Geiger et al., 2020;
Atanasov et al., 2021; Paccolat et al., 2021) in practical settings is that
the kernel aligns with the target function over time, a phenomenon we
hereafter refer to as “kernel alignment”. The prominence and ubiquity of
kernel alignment suggest that it may play an important role in NN feature
learning. In particular, it has been speculated that this underlies the
superior generalization performance of practical NNs, when compared to their
infinite-width counterparts (Baratin et al., 2021; Fort et al., 2020; Geiger
et al., 2020; Paccolat et al., 2021). In this work, we aim to provide a
theoretical understanding of the dynamics of kernel alignment during NN
training and how it affects learning. Specifically, our main contributions are

  * •

We demonstrate that kernel alignment accelerates training by showing that a
kernel aligning over time accelerates convergence of the training loss. In
particular, we study a toy model of kernel evolution where the NTK features
explicitly evolve towards the direction that accelerates training and find
that kernel alignment occurs as a result. The strength of the acceleration is
controlled by a single parameter, the feature learning rate ($\gamma$), which
determines not only the acceleration in training but also the final alignment
of the NTK.

  * •

We provide an analytical theory of how kernel alignment emerges during
gradient descent learning in deep linear networks(Sec.5) and an approximate
treatment of two layer ReLU networks (Sec.6). The theory captures key
qualitative features of kernel alignment in real-world networks and makes the
novel prediction that kernel alignment is stronger in deeper networks, which
we validate numerically in ReLU networks.

  * •

We report the novel empirical finding of “kernel specialization” (Sec.4.3),
which occurs in NNs with multiple output heads (e.g., those for multiclass
classification). For these NNs, the NTK has different components for each pair
of output heads. We found that the diagonal components corresponding to each
head becomes aligned with its specific target function. Our theory shows how
this emerges from the interaction between network architecture and data
structure.

##  2 Related Work

Characterizing and understanding the time evolution of the NTK for during NN
training has been the subject of considerable interest in the deep learning
theory community. This is motivated by the empirical finding that the NTK
evolution underlies a significant performance gap between practical NNs and
infinite-width NNs (Lee et al., 2020a; Geiger et al., 2020). General
expressions for the leading order corrections to NN dynamics for finite width
networks with NTK parameterization can be characterized within the framework
of perturbation theory through truncation of an infinite set of ODEs known as
the neural tangent hierarchy (Dyer & Gur-Ari, 2020; Huang & Yau, 2020; Aitken
& Gur-Ari, 2020; Roberts et al., 2021). This leading order truncation contains
corrections which scale as $1/\text{width}$ and depends on data through the
targets and a third and fourth order tensor of the inputs. The alternative
mean-field infinite-width scaling reveal finite evolution of hidden features
which depend on the supervised training signal which persist even at infinite
width (Mei et al., 2019; Sirignano & Spiliopoulos, 2020; Yang & Hu, 2021;
Nguyen, 2019; Nguyen & Pham, 2020). These results, however, generally involve
nonlinear partial differential equations for the parameter or feature
distributions.

Instead of characterizing evolution of the entire kernel, recent empirical and
theoretical works have focused on describing and analyzing specific features
of the evolution. In particular, kernel alignment has emerged as a prominent
and ubiquitous phenomenon in this type of analysis. (Fort et al., 2020) find
an initial transient period of kernel evolution followed by a period where the
kernel can be approximated as static. (Baratin et al., 2021) showed that the
early dynamics align NTK eigenvectors with task relevant directions. (Atanasov
et al., 2021) provide conditions for a “silent alignment” effect in networks
with small initialization in which the NTK aligns to the task relevant
subspace before scale growth, giving network output which is a kernel
regression solution with the final, rather than initial, NTK.

Figure 1: Feature evolution alters the structure of the NTK and accelerates
learning.
A The training loss for a two-layer $N=500$ MLP when trained on a subset of
MNIST (NN) is compared to kernel gradient descent with the initial kernel
(KGD) and the initial kernel using a rescaled learning rate to account for the
difference in norm of the NN’s NTK and the initial NTK (aKGD). We see that
even the optimistically rescaling of the learning rate by the final NN’s NTK
norm does not account for the gap in the loss. B The norm of the kernel
increases non-monotonically throughout training. C The alignment between the
NTK and the task kernel throughout training increases to an asymptote for the
neural network but remains constant for the static kernel dynamics (KGD,
aKGD). Average and standard deviation over five different initializtions are
plotted.

##  3 Preliminaries: NTK and Kernel Alignment

In this section, we briefly review the NTK definition (Jacot et al., 2020b)
and give a precise definition of kernel alignment during NN training. For
simplicity, we will first consider NNs with scalar output functions and will
extend our discussion to multiple class outputs in a later section. Let
$f(\bm{x},\bm{\theta})$ represent the output of a neural network with
parameters $\bm{\theta}$ and input vector $\bm{x}$. We optimize the parameters
$\bm{\theta}$ with gradient flow on a loss function
$\mathcal{L}=\sum_{\mu}\ell(f(\bm{x}^{\mu},\bm{\theta}),y^{\mu})$ with $P$
examples $\mathcal{D}=\\{(\bm{x}^{\mu},y^{\mu})\\}_{\mu=1}^{P}$ in the
training set. Throughout this work, we assume batch gradient descent dynamics,
which give

$\frac{d\bm{\theta}}{dt}=-\eta\sum_{\mu=1}^{P}\frac{\partial
f(\bm{x}^{\mu},\bm{\theta})}{\partial\bm{\theta}}\frac{\partial\ell(f(\bm{x}^{\mu},\bm{\theta}),y^{\mu})}{\partial
f(\bm{x}^{\mu},\bm{\theta})}$ (1)

where $\eta$ is the learning rate. Rather than studying the dynamics of the
parameters $\bm{\theta}$, the NTK formulation focuses on dynamics of the
network output,

$\frac{df(\bm{x}^{\mu},\bm{\theta})}{dt}=-\eta\sum_{\nu}K(\bm{x}^{\mu},\bm{x}^{\nu};\bm{\theta})\frac{\partial\ell(f(\bm{x}^{\mu},\bm{\theta}),y^{\mu})}{\partial
f(\bm{x}^{\nu},\bm{\theta})}$ (2)

where the $K$ represents the NTK
$K(\bm{x},\bm{x}^{\prime};\bm{\theta})=\frac{\partial
f(\bm{x},\bm{\theta})}{\partial\bm{\theta}}\cdot\frac{\partial
f(\bm{x}^{\prime},\bm{\theta})}{\partial\bm{\theta}}$. On the training set, it
is completely described by the gram matrix
$\bm{K}(\bm{\theta})\in\mathbb{R}^{P\times P}$.

On the training set, the NTK evolution, including kernel alignment, is
completely reflected by the dynamics of this matrix over time. Importantly,
this evolution entirely characterizes the difference from the infinite-width
limit, where the NTK is static, and the regime of practical NNs. As a simple
demonstration of this phenomenon, we trained MLPs with two hidden layers with
500 hidden ReLU units per layer to do odd-even binary classification of MNIST
dataset. We compared the dynamics of its training loss (“NN”) against the
expected loss dynamics if the NTK were frozen at its initial value (“KGD”) and
found that the NN training is significantly faster(Fig.1A). Importantly, while
the NN $\bm{K}(\bm{\theta})$ increases in its Frobenius norm(Fig.1B), as
reported previously (Baratin et al., 2021), this alone does not account for
the acceleration of learning. To see this, we simulated the training loss
dynamics of KGD using the initial NTK but scaled up to match the amplitude of
the final NN NTK and found that it still could not match NN performance
(Fig.1A). Therefore, understanding changes to the structure of $\bm{K}$,
rather than simply its amplitude, is essential.

To quantify kernel alignment, it is useful to introduce the kernel alignment
metric (Cortes et al., 2012).

$A(t)=\frac{\left<\bm{y}\bm{y}^{\top},\bm{K}(\bm{\theta})\right>_{F}}{||\bm{K}(\bm{\theta})||_{F}||\bm{\bm{y}}\bm{y}^{\top}||_{F}}=\frac{\bm{y}^{\top}\bm{K}(\bm{\theta})\bm{y}}{||\bm{K}(\bm{\theta})||_{F}||\bm{\bm{y}}||^{2}}.$
(3)

Indeed, $A(t)$ increases substantially in our experiment before
stabilizing(Fig.1C).

##  4 How Kernel Alignment Influences Learning

###  4.1 Kernel Alignment and Optimal Feature Evolution

We first investigate how kernel alignment influences learning. In particular,
we would like to understand if and how kernel alignment is responsible for the
accelerated decrease of training loss seen in Fig.1. To do so, we first step
away from NNs for a moment and consider the general case of KGD with an
evolving kernel $\bm{K}(t)=\bm{\Psi}(t)^{\top}\bm{\Psi}(t)$. The very
important issue of how kernel alignment emerges during NN training is
addressed in Sec.5 and Sec.6.

The matrix $\bm{\Psi}(t)\in\mathbb{R}^{P\times Q}$, where $Q$ is the number of
parameters in the model, defines the feature map used by the kernel. Assuming
a mean squared error loss, the loss dynamics of KGD is given by

$\mathcal{L}_{t+1}(\bm{\Psi})=||(\bm{I}-\eta\bm{\Psi}_{t}^{\top}\bm{\Psi}_{t})\bm{\Delta}_{t}||^{2},$
(4)

where $\bm{\Delta}_{t}=\bm{f}_{t}-\bm{y}\in\mathbb{R}^{P}$.

In order to formulate a toy model of “optimal” feature evolution for
accelerating learning, we consider the case where the feature map is
explicitly updated via gradient descent on $\mathcal{L}_{t+1}$ to accelerate
learning, i.e.

$\bm{\Psi}_{t+1}=\bm{\Psi}_{t}-\gamma\frac{\partial
L_{t+1}}{\partial\bm{\Psi}}|_{\bm{\Psi}_{t}}.$ (5)

Here, $\gamma$ is a scalar which we hereafter refer to as the “feature
learning rate”. The fact that it is finite reflects the constraint that the
feature map cannot evolve infinitely fast during KGD. The limit
$\gamma\rightarrow 0$ recovers the static kernel limit. In order to analyze
what this optimal feature evolution (OFE) would look like, we go to the
continuous time limit (Appendix B) where

$\displaystyle\dot{\bm{\Delta}}(t)$
$\displaystyle=-\eta\bm{\Psi}(t)^{\top}\bm{\Psi}(t)\bm{\Delta}(t)$
$\displaystyle\dot{\bm{\Psi}}(t)$
$\displaystyle=\gamma\eta\bm{\Psi}(t)\bm{\Delta}(t)\bm{\Delta}(t)^{\top}.$ (6)

In this limit, one finds that the matrix
$\bm{C}=\gamma\bm{\Delta}(t)\bm{\Delta}(t)^{\top}+\bm{\Psi}(t)^{\top}\bm{\Psi}(t)$
is static in time (Appendix B). Exploiting this fact and noting that the
training loss $\lVert\bm{\Delta}(t)\rVert^{2}$ must eventually reach zero, one
can identify the final kernel

$\bm{K}_{\infty}=\gamma\bm{y}\bm{y}^{T}+\bm{K}_{0}.$ (7)

Further, we can verify that increasing $\gamma$ is indeed beneficial to
training dynamics by noting that
$\frac{1}{2}\frac{d}{dt}{|\bm{\Delta}|^{2}}=-\eta\bm{\Delta}^{\top}\left[\bm{K}_{0}+\gamma(\bm{y}\bm{y}^{\top}-\bm{\Delta}\bm{\Delta}^{\top})\right]\bm{\Delta}\leq-\eta\bm{\Delta}^{\top}\bm{K}_{0}\bm{\Delta}$,
which shows that positive $\gamma$ accelerates the convergence of the loss.
The analysis of “optimal feature evolution” here suggests that if one is to
simultaneously optimize the features $\bm{\Psi}$ and the training errors
$\bm{\Delta}$, the kernel will become aligned with the task functions, which
is exactly kernel alignment. This expression bears remarkable resemblance the
final NTK in the deep linear network case as we discuss in Sec.5(Eq.12).

While the OFE model does not necessarily capture the trajectory of the NTK, it
suggests the heuristic that network depth $L$ acts like $\gamma$ to control
how quickly the NTK evolves. We test this heuristic in the next section where
we compare the training and alignment dynamics of real NNs with different
feature learning rates.

###  4.2 Enhancing Feature Evolution Through Rescaling

As was explored in the prior works of (Chizat et al., 2019; Geiger et al.,
2020), the rate at which the NTK evolves can be altered through rescaling of
the trainable network function: $g(\bm{x})=\frac{1}{\gamma}f(\bm{x})$ and
learning rate $\eta=\eta_{0}\gamma^{2}$. Letting,
$L=|\bm{\Delta}(t)|^{2}=|\bm{y}-\bm{g}(t)|^{2}$, gradient flow
$\frac{d\bm{\theta}}{dt}=-\nabla_{\bm{\theta}}L$ gives
$\frac{dL}{dt}=O_{\gamma}\left(1\right)$ and $\frac{d}{dt}\frac{\partial
f}{\partial\bm{\theta}}=O_{\gamma}(\gamma)$ (see Appendix C). Recognizing that
for NN dynamics, the relevant features are the parameter gradients
$\bm{\psi}_{\mu}=\frac{\partial f_{\mu}}{\partial\bm{\theta}}$, we see that
increasing $\gamma$ increases the relative rate at which the network gradients
evolve compared to the loss at initialization, leading to faster kernel
evolution. We illustrate that such rescaling indeed alters the kernel and
training dynamics in a Wide ResNet (Zagoruyko & Komodakis, 2017) with fixed
widening factor of $k=3$ and block size $b=2$ on a subset of $100$ CIFAR-10
images in Figure 2 (details in AppendixE.1). For $\gamma=1$, the network is
sufficiently wide that no significant feature evolution occurs, but at large
$\gamma$, the NTK reaches a high final alignment value. Consistent with our
OFE toy model, the NNs trained with highest $\gamma$ not only align more
quickly but also reach a higher final alignment value. Further, increasing
$\gamma$ accelerates training convergence due to superior alignment, again
consistent with OFE. Estimating the feature learning rate for the OFE from the
initial and final alignment values, we obtain qualitatively similar
loss(Fig.2Avs.C) and alignment curves(Fig.2Bvs.D).

Figure 2: Kernel and loss evolution in Wide Res-Net on CIFAR-10 subsample. A
Though the initial loss dynamics are similar for different $\gamma$, the drops
more quickly after the NTK starts to aligning to the targets. B The alignment
increases more quickly and reaches a higher final value for large $\gamma$,
consistent with the toy OFE model. C-D The OFE model with $\bm{K}_{0}$
extracted as the initial NTK, and the feature learning rate $\hat{\gamma}_{3}$
estimated from the initial and final alignment of the green curve in panel B.
Both loss and alignment dynamics are qualitatively similar to the real NN
experiments.

###  4.3 Kernel Specialization in Multiclass NNs

So far, we have focused the analysis on NNs with a scalar output. We now
consider NNs with multiclass outputs, which are common in practical
applications. For networks with $C$ output nodes (e.g. $C$-class
classification networks), the NTK $\bm{K}$ is a 4D tensor of dimensions
$P\times P\times C\times C$. We define a matrix-valued ”subkernel” as

$\bm{K}^{c,c^{\prime}}\in\mathbb{R}^{P\times
P}:(\bm{K}^{c,c^{\prime}})_{\mu\mu^{\prime}}=\nabla_{\bm{\theta}}f_{c}^{\mu}\cdot\nabla_{\bm{\theta}}f_{c^{\prime}}^{\nu}$
(8)

We now consider a typical setting where the scalar activation of the $c$th
output node, $f_{c}(\bm{x})$, is fitted to a separate target function (e.g.,
an indicator function of whether the input belongs to a certain class),
$\bm{y}_{c}\in\mathbb{R}^{P}$. In addition, the loss function is decomposed
into terms depending on individual nodes, e.g.,
$\mathcal{L}=C^{-1}P^{-1}\sum_{\mu=1}^{P}\sum_{c=1}^{C}(f_{c}(\bm{x}_{\mu})-y_{c}^{\mu})^{2}$.
In this case, the dynamics of the loss for the $c$th output head only depend
on $\bm{K^{c,c}}$.

At infinite network width, $\bm{K}^{c,c^{\prime}}=\bm{0}$ if $c\neq
c^{\prime}$ and all “diagonal” subkernels, $\bm{K}^{c,c}$, are identical. This
suggests that each output node evolves towards its own target function under
KGD dynamnics governed by the same subkernel. Our analysis of the OFE suggests
that to accelerate learning, it is best for each output node to evolve using a
separate subkernel. Its own subkernel should learn features specific to its
own target function, such that $\bm{K}^{c,c}$ becomes aligned with
$\bm{y}_{c}$ but not $\bm{y}_{d\neq c}$. We call this phenomenon kernel
specialization.

Does kernel specialization occur in neural network learning? Previous
empirical studies of multi-output networks studies the “traced kernel”,
$\bm{K}_{\text{tr}}=C^{-1}\sum_{c=1}^{C}\bm{K}^{c,c}$, and found that it
becomes aligned with all $C$ target functions (Baratin et al., 2021) during
training. While this observation is consistent with kernel specialization, it
can also arise from indiscriminate alignment where $\bm{K}^{c,c}$ becomes
aligned with all $\\{\bm{y}_{d}\\}_{d=1,...,C}$. To see whether kernel
specialization occurs, we first define a kernel specialization matrix (KSM),
defined as

$\text{KSM}(c,d)\equiv\frac{A(\bm{K}^{c,c},\bm{y}_{d}{\bm{y}_{d}}^{T})}{C^{-1}\sum_{d^{\prime}=1}^{C}A(\bm{K}^{d^{\prime},d^{\prime}},\bm{y}^{d}{\bm{y}^{d}}^{T})}.$
(9)

If kernel specialization occurs, KSM should be higher when $c=d$ (i.e.
diagonal elements of the KSM). We computed the KSM for a two-layer $N=500$ MLP
trained on 10-class classification of MNIST digits and found that diagonal
elements are indeed higher than off-diagonal ones, demonstrating kernel
specialization (Fig.3 A, details in AppendixE.2). The same qualitative result
was replicated in a convolutional network trained on classification of
CIFAR-10 images (Fig.3 B, details in AppendixE.3; only two classes were used).

Figure 3: NTK becomes specialized during training in NNs with multiclass
outputs and nonlinear activation. A. KSM (defined in Eq.9) of a two-layer
$N=500$ ReLU MLP trained on 10-class classification of MNIST digits. B. KSM of
a convolutional neural network trained on CIFAR-10.

##  5 Mechanisms of Alignment in Linear NNs

Having shown that kernel alignment can indeed accelerate training, we now turn
to the important question of how it mechanistically arises from training NNs
with gradient descent. In particular, we would like to identify important
features of network architectures and data structures that give rise to kernel
alignment. We begin by analyzing kernel alignment in deep linear networks in
Sec.5.1 before moving on to two-layer ReLU networks in Sec.6.

###  5.1 Kernel Alignment in Deep Linear Networks

Deep linear networks have supplied rich theoretical insights on NN training
dynamics which generalize well to nonlinear counterparts (Advani et al., 2020;
Arora et al., 2018; Du et al., 2018). They are therefore a natural starting
point for developing a theory of kernel alignment in NNs. First, we examined
numerically whether kernel alignment occurs at all in linear networks. We
trained a deep linear network with two hidden layers and a scalar output to
learn a linear mapping $\mathbb{R}^{N}\rightarrow\mathbb{R}$,
$y^{\mu}=\bm{\beta}\bm{x^{\mu}}$, where $\\{\bm{x^{\mu}}\\}$ are $P$ i.i.d.
Gaussian vectors as the training set and $\bm{\beta}$ are the weights of a
linear teacher. We found that the NTK becomes aligned with the target function
in a similar fashion to the kernel alignment observed in our MLP experiment
(Fig.1) and previous empirical work (Fort et al., 2020; Geiger et al., 2020;
Baratin et al., 2021; Atanasov et al., 2021). This suggests that kernel
alignment can indeed occur in deep linear networks.

Figure 4: Kernel alignment in linear networks during gradient descent.
A Alignment between network weights and teacher weights during learning, as
measured by cosine similarity. B NTK becomes more aligned with the task kernel
$\bm{yy}^{T}$. C The top eigenfunction of the NTK becomes more aligned with
the target function, as measured by cosine similarity. D The kernel
specialization matrix of the linear network shows no specialization.

We consider deep linear networks of arbitrary depth with $L$ hidden layers and
a scalar output, given by
$f(\bm{x})=\bm{w}^{L+1\top}\bm{W}^{L}...\bm{W}^{1}\bm{x}$. We assume a small
initialization and gradient flow training dynamics. Our results rest on the
conservation law identified in prior works in the literature on linear neural
networks(Advani et al., 2020; Arora et al., 2018; Du et al., 2018)

$\displaystyle\frac{d}{dt}\left[\bm{W}^{\ell}\bm{W}^{\ell\top}-\bm{W}^{\ell+1\top}\bm{W}^{\ell+1}\right]=0$
(10)

Under the assumption that the weights are all initialized with small variance
this conservation law implies
$\bm{W}^{\ell}\bm{W}^{\ell\top}\approx\bm{W}^{\ell+1\top}\bm{W}^{\ell+1}$.
Starting from the last layer, we infer that
$\bm{W}^{L}\bm{W}^{L\top}\approx\bm{w}^{L+1}\bm{w}^{L+1\top}$ so that
$\bm{W}^{L}$ is approximately rank-one
$\bm{W}^{\ell}=u(t)\hat{\bm{w}}^{L+1}(t)\bm{r}_{L}(t)$ where $\bm{r}_{L}(t)$
and $\hat{\bm{w}}^{L+1}(t)$ are unit vectors. Repeating this argument
inductively generates the conclusion that each layer’s weight matrix is rank-
one $\bm{W}^{\ell}=u(t)\bm{r}_{\ell+1}(t)\bm{r}_{\ell}(t)^{\top}$ and thus the
kernel has the form

$\displaystyle
K(\bm{x},\bm{x}^{\prime})=u(t)^{2L-2}\bm{x}^{\top}\left[L\bm{r}_{1}(t)\bm{r}_{1}(t)^{\top}+\bm{I}\right]\bm{x}^{\prime}$
(11)

We provide a derivation in Appendix A. The fixed point for $\bm{r}_{1}$ is the
direction of the linear teacher $\bm{\beta}$ , since the network must
interpolate the data. Thus if one were to evaluate the kernel on the training
data, they would obtain

$\bm{K}_{\infty}\propto L\bm{y}\bm{y}^{\top}+\bm{K}_{0}.$ (12)

This is the central result of our theory for linear networks and provides
several insights. First, nonlinear activation functions are not necessary for
kernel alignment. Second, Even at infinite time, the kernel is not fully
aligned with the target function, $\bm{y}\bm{y}^{T}$, as observed in empirical
studies(Bahri et al., 2021). Finally, since the first term is linear in
network depth, this expression predicts that kernel alignment is more
prominent in deeper networks. To test whether this prediction generalizes to
nonlinear networks, we trained ReLU two-layer MLPs of different depths on the
same MNIST task and found that deeper networks indeed have stronger kernel
alignment(Fig.5; details of the experiments are in Appendix.E.2).

Figure 5: Kernel alignment in ReLU MLPs of different depths. A While all
networks start with comparable levels of initial alignment, deeper networks
reach higher alignment values at the end of training. B Estimated feature
learning rates ($\gamma$) from networks of different depths.

###  5.2 Linear NNs Cannot Develop Specialized Kernels

We next analyze whether kernel specialization may occur in deep linear
networks. We assume a typical architecture used for multiclass classification,
where the network has $L$ shared hidden layers and $C$ linear readouts from
the last hidden layer, trained under the setting described in Sec.4.3. In this
case and under a weak assumption about symmetry of target functions with
respect to $c$, it can be shown that linear networks of arbitrary depth cannot
develop specialized kernels(see C.1). To test this prediction, we trained a
four-layer linear MLP with $10$ output nodes on 10 linear target functions
(input vectors were randomly drawn from the unit sphere and the teacher
weights for each target function were i.i.d. Gaussian) and computed its KSM.
As predicted, the KSM does not show specialization (Fig.4D).

##  6 Dynamics of Kernel Alignment in Two-Layer ReLU Networks

We next studied how kernel alignment emerges in nonlinear NNs by considering
the case of two-layer ReLU networks, a common toy model for studying NN
training dynamics in nonlinear networks (Ergen & Pilanci, 2020; Safran &
Shamir, 2018; Luo et al., 2021). In general, expressing the dynamics of the
NTK in terms of its $1/\text{width}$ corrections from the static limit
requires convoluted schemes (Huang & Yau, 2020; Dyer & Gur-Ari, 2020) that are
intractable and difficult to interpret.

We circumvent this issue by exploiting the that fact to study kernel
alignment, we are only interested in the structural anisotropy of the kernel
in a specific direction (that of the target function). Our approach to track
the dynamics of kernel alignment is the following. We track the bilinear form
$B(\bm{z})=\bm{z}^{T}\bm{K}\bm{z}$ with $\bm{z}\in S^{P-1}$, a unit vector.
Kernel alignment would manifest as $B\left(\bm{y}/|\bm{y}|\right)$ growing
faster than
$\left<B({\bm{z}})\right>_{\bm{z}\sim\text{Unif}(S^{P-1})}\sim\frac{1}{P}\text{Tr}(\bm{K})$,
where the average is over all unit vectors with a uniform measure for large
$P$.

To make kernel alignment dynamics in such nonlinear NNs tractable, we make two
heuristic assumptions motivated by empirical observations. First, we assume
that the network is sufficiently wide that the sign of preactivations do not
change: $\text{sign}\left(\bm{w_{i}}\cdot\bm{x^{\mu}}\right)$ is static
throughout GD dynamics. We also assume that $\forall t>0:$
$\bm{y}^{T}\bm{f}(t)>0$. These assumptions will be justified with simulations.

###  6.1 Networks with Scalar Outputs

For brevity, we write $\phi_{i}^{\mu}\equiv\phi(\bm{w_{i}}\cdot\bm{x^{\mu}})$.
These are the hidden layer features which are dynamic. Then the output of two-
layer ReLU network with a scalar output can be written as

$f\left(\bm{x^{\mu}}\right)=\sum_{i=1}^{M}V_{i}\phi_{i}^{\mu},$ (13)

where $\phi_{i}^{\mu}=\text{ReLU}(\bm{w}_{i}\cdot\bm{x}^{\mu})$. For
simplicity, we consider the task of random binary classification where
$\\{\bm{x^{\mu}}\\}$ are drawn i.i.d. from the unit sphere and the target
function is a random binary label $y^{\mu}\in\\{-1,1\\}$. The effect of more
complex data structures is considered below in Sec.6.2. In this setting, we
would like to show that kernel alignment occurs and understand it
analytically.

A detailed derivation is provided in Appendix D.2 and we provide a sketch
here. For this network, $\bm{K}$ has two components, given by

$\displaystyle K_{\mu,\nu}$
$\displaystyle=\nabla_{\bm{V}}f^{\mu}\cdot\nabla_{\bm{V}}f^{\nu}+\nabla_{\bm{W}}f^{\mu}\cdot\nabla_{\bm{W}}f^{\nu}$
(14)
$\displaystyle\equiv(\bm{K}_{\bm{V}})_{\mu,\nu}+(\bm{K}_{\bm{W}})_{\mu,\nu},$
(15)

The first component is contributed by $\nabla_{\bm{V}}f$ and the second by
$\nabla_{\bm{W}}f$. Denote their respective contribution to $B(\bm{z})$ as
$\alpha(\bm{z})=\bm{z}^{T}\bm{K_{V}}\bm{z}$ and
$\beta(\bm{z})=\bm{z}^{T}\bm{K_{W}}\bm{z}$. We show in Appendix D that
gradient descent generates the following dynamics

$\displaystyle\frac{d\alpha}{dt}\left(\bm{y}/|\bm{y}|\right)$
$\displaystyle\approx\frac{1}{P}\eta\left[\left(\bm{y}-\bm{f}\right)^{T}\bm{y}\right]\left[\bm{y}^{T}\bm{f}\right]$
$\displaystyle\left\langle\frac{d\alpha}{dt}\left(\bm{z}\right)\right\rangle_{\bm{z}}$
$\displaystyle\approx\frac{1}{P}2\eta\left(\bm{y}-\bm{f}\right)^{T}\bm{f}$
(16) $\displaystyle\forall\bm{z}:\frac{d\beta}{dt}\left(\bm{z}\right)$
$\displaystyle\approx\frac{1}{2}\eta(\bm{y}-\bm{f})^{T}\bm{f}$

While these equations are not closed (they depend on $\bm{f}(t)$), it is clear
that the dynamics of $\beta(\bm{z})$, and by extension $\bm{K_{W}}$, are
independent of $\bm{z}$ and therefore isotropic. On the other hand, we can
consider the anisotropy of $\alpha$ by considering early stages of learning
where $\bm{y}-\bm{f}\approx\bm{y}$, yielding
$\frac{d\alpha}{dt}\left(\bm{y}/|\bm{y}|\right)\approx\eta\bm{y}^{T}\bm{f}$
and
$\left\langle\frac{d\alpha}{dt}\left(\bm{z}\right)\right\rangle_{\bm{z}}\approx\frac{1}{P}2\eta\bm{y}^{T}\bm{f}$.
Note that $|\bm{y}|^{2}=P$ since $y_{\mu}\in\\{\pm 1\\}$. These results
indicate that $\bm{K_{V}}$ grows $O(P)$ times faster in the direction of
$\bm{y}$ than in the other directions. The NTK, which is the sum of
$\bm{K_{V}},\bm{K_{W}}$, therefore develops an anisotropy in the form of
kernel alignment with $\bm{y}$.

To quantitatively test our theory, we trained (see details in Appendix E.4)
two-layer ReLU networks on the binary classification task using gradient
descent and found that Eq.16 nicely capture the dynamics and overall strength
of kernel alignment in the networks(Fig.6). First, as predicted, kernel
alignment is driven by alignment in the $\bm{K_{V}}$ component, as
predicted(Fig.6B). In addition, $\bm{K_{V}}$ grows at a much faster rate in
the direction of $\bm{y}$(Fig.6C) than in other directions(Fig.6E). On the
other hand, $\bm{K_{W}}$ grows at the same rate in all directions(Fig.6D,F).

Figure 6: Kernel alignment in two-layer ReLU network is well characterized by
theory.
A Training loss over time. B Alignment between the NTK
($\bm{K_{V}}+\bm{K_{W}}$) and the target functions increases before
saturating. As predicted by the theory, alignment is driven by anisotropy in
$\bm{K_{V}}$ while changes to $\bm{K_{W}}$ are isotropic. C Dynamics of the
projection of $\bm{K_{V}}$ on the target functions. D Dynamics of the
projection of $\bm{K_{W}}$ on the target functions. E Dynamics of the trace of
$\bm{K_{V}}$ (same as $\langle\bm{z}^{T}\bm{K_{V}}\bm{z}\rangle_{\bm{z}}$.
Notice that the growth is much weaker than the growth in the direction in
$\bm{y}$ (panel C). F Dynamics of the trace of $\bm{K_{W}}$. Notice that the
growth is of the same magnitude than the growth in the direction in $\bm{y}$
(panel D).

###  6.2 Dynamics of Kernel Specialization in a Classification Setting

To analytically understand kernel specialization, we now extend the above
analysis to NNs with multiple output heads. For concreteness, we consider a
two-layer ReLU network with $C$ output heads. Output from the $c$th head is
given by

$f_{c}^{\mu}=f_{c}(\bm{x}^{\mu})=\sum_{i=1}^{M}V_{i}^{c}\phi_{i}^{\mu}.$ (17)

For the task, we assume the scenario of training the network to separate a
mixture of $C$ Gaussian distributions,
$\\{\mathcal{N}(\bm{\mu_{c}},\sigma^{2}\mathbf{I})\\}_{c=1,...,C}$. From each
distribution, $P/C$ samples are drawn into the training set. We assume
$\\{\bm{\mu}_{c}\\}$ to be unit vectors that are pairwise perpendicular. Each
head has a separate target function, $\bm{y}_{c}$. We consider the case where
the target functions encode the correct class of each input. I.e.
$y^{\mu}_{c}=1$ if $\bm{x}^{\mu}$ is in class $c$ and $0$ otherwise. Here we
present the simplistic but illuminating limit where $\sigma^{2}\rightarrow 0$
and $\bm{x}^{\mu\in\text{class c}}=\bm{\mu_{c}}$ (the case of finite but small
$\sigma^{2}$ is discussed in Appendix D.3, but the derivation has the same
flavor as the one here). At this limit, the input correlation structure is
simply $\bm{x}^{\mu}\cdot\bm{x}^{\nu}=1$ if they are in the same class and
zero otherwise (this is possible for $C<N$).

Importantly, this structure in input space induces an analogous structure in
$\\{\bm{D}^{\mu}\\}_{\mu=1,...,P}$, where
$D^{\mu}_{i}=d_{x}\phi(x)|_{x=\bm{w_{i}}\cdot\bm{x^{\mu}}}$. We note that the
covariance of $D^{\mu},D^{\nu}$ over random hidden weight vectors $\bm{w}$
gives

$\displaystyle\text{Cov}(D^{\mu},D^{\nu})\approx\delta_{y_{\mu},y_{\nu}}\frac{3}{4}.$
(18)

Crucially, as we proceed to show, this class-dependent structure in
$D^{\mu}_{i}$ is necessary for kernel specialization in this setting. This is
also consistent with our finding that linear networks cannot specialize since
$D^{\mu}_{i}=1$ for all $i,\mu$ for linear nets.

To detect kernel specialization, we consider the structure of the $c$-th
subkernel by tracking the bilinear form
$B^{c}(\bm{y}_{d})=\bm{y}_{d}^{T}\bm{K}^{c,c}\bm{y}_{d}$. In order to show
kernel specialization, $B^{c}(\bm{y}_{d})$ must grow faster for $d=c$ then
$d\neq c$. We first observe that the subkernel has two components, similar to
the scalar-output network case, given by
$K^{c,c}_{\mu,\nu}=\left(K^{c,c}_{\bm{V}}\right)_{\mu,\nu}+\left(K^{c,c}_{\bm{W}}\right)_{\mu,\nu}$.
Note that $\bm{K}^{c,c}_{\bm{V}}$ is the same for all $c$ and therefore cannot
contribute to kernel specialization. By exclusion, kernel specialization must
arise from $\bm{K}^{c,c}_{\bm{W}}$. We can thus simply our analysis of
anisotropic dynamics of $\bm{K^{c,c}}$ by by focusing on its component coming
from $\bm{K}^{c,c}_{\bm{W}}$, given by
$B^{c}_{\bm{W}}(\bm{y}_{d})\equiv\bm{y}_{d}^{T}\bm{K}^{c,c}_{\bm{W}}\bm{y}_{d}$.
At the $\sigma^{2}\rightarrow 0$ limit, points from each class trivially
collapse to a single point, allowing us to write $\forall\mu\in\text{class
c}:D^{\mu}_{i}=D^{c}_{i},\phi^{\mu}_{i}=\phi^{c}_{i}$. One can show that the
dynamics of it follow

$\displaystyle\eta^{-1}\frac{dB^{c}_{\bm{W}}(\bm{y}_{d})}{dt}$
$\displaystyle=\sum_{i=1}^{M}\frac{dB^{c}_{\bm{W}}(\bm{y}_{d})}{d(\bm{V}_{c})_{i}}\frac{d(\bm{V}_{c})_{i}}{dt}$
$\displaystyle=$
$\displaystyle\left(\frac{P}{C}\right)^{2}\sum_{i=1}^{M}D^{d}_{i}({\bm{V_{c}}})_{i}\phi^{c}_{i}.$
(19)

The sum over neuron indices $i$ highlights how the correlation structure
between $\\{\bm{D}^{c}_{i}\\}_{i=1,...,M}$ and
$\\{\bm{D}^{d}_{i}\\}_{i=1,...,M}$ lead to kernel specialization. If $d=c$,
then $B_{c}(\bm{y}_{d})$ increases as

$\eta^{-1}\frac{dB^{c}_{\bm{W}}(\bm{y}_{d})}{dt}=\left(\frac{P}{C}\right)^{2}f_{c}(\bm{x}^{\mu\in\text{class
c}})>0.$ (20)

It follows from our assumption of $\sum_{\mu}y^{\mu}_{c}f(\bm{x}^{\mu})>0$
that $f_{c}(\bm{x}^{\mu\in\text{class c}})>0$ (since all data points from each
class collapse to one point, this is the same regardless of $\mu$). Now we
consider $d\neq c$. The sum in equation 19 will always contain fewer terms
than the full sum over $i\in\\{1,...,M\\}$ since $D_{i}^{\mu}\in\\{0,1\\}$. We
approximate $D_{i}^{\mu}$ as independent of $\phi_{i}^{\nu}$, giving
approximately half of the terms
$\\{(\bm{V}_{c})_{i}\phi^{c}_{i}\\}_{i=1,...,M}$ (see further discussion in
Appendix.59). This motivates the following approximation

$\frac{dB^{c}_{\bm{W}}(\bm{y}_{d})}{dt}\approx\frac{1}{2}\frac{dB^{c}_{\bm{W}}(\bm{y}_{c})}{dt}.$
(21)

Figure 7: Total kernel alignment is driven by both indiscriminate alignment
and kernel specialization.
A. In a ReLU network trained to classify a mixture of 10 Gaussians, the
$\bm{K_{V}}$ component of the NTK grows in the direction of all 10 target
functions at approximately the same rate. This component is the same for all
subkernels. Each trace is the projection on a different target function, with
its initial value subtracted. B. In the same setting as A, the
$\bm{K_{W}}^{c,c}$ component of the NTK grows in the direction of its
corresponding target function (“specialized”) at around twice the rate of that
in the direction of other target functions (“others”), as predicted by our
theory. Green dashed line is half the specialized rate. In all cases the
initial value is subtracted from the entire trace. The error bars are standard
deviations over different $c,d$. C In a ReLU network trained on the MNIST
task, $\bm{K_{V}}$ does not specialize as its alignment with all target
functions grows at a similar rate. D. In the same setting as C, the
$\bm{K_{W}}^{c,c}$ component of the NTK becomes highly specialized.
Average of 5 runs; standard error shown as color contour.

In other words, $K^{c,c}_{\bm{W}}$ grows approximately twice as fast in the
direction of $\bm{y}_{c}$ compared to $\bm{y}_{d}$, indicating kernel
specialization. Since, $K^{c,c}_{\bm{V}_{c}}$ does not differentiate between
different target functions, the full subkernel
$K^{c,c}_{\bm{W}}+K^{c,c}_{\bm{V}}$ becomes specialized over time. To test our
theory quantitatively, we trained two-layer ReLU networks with 10 output heads
to classify a mixture of 10 Gaussians (with $\sigma^{2}=0.01$, see details in
AppendixE.5) and measured how the two components of each subkernel
($\bm{K^{c,c}}$), $\bm{K_{V}^{c,c}}$ and $\bm{K_{W}^{c,c}}$ grow in the
directions of the 10 target functions. As predicted by our theory, the
$\bm{K_{V}^{c,c}}$ component grows at the same rate in the direction of all
target functions(Fig.7 A) and therefore does not contribute to kernel
specialization (it still contributes to alignment). On the other hand, the
$\bm{K_{W}^{c,c}}$ component grows at approximately twice the rate in the
direction of its corresponding target function ($\bm{y_{c}}$) compared to in
the directions of other target functions ($\bm{y_{c}\neq d}$, Fig.7 B)), as
predicted by Eq.21.

We also verified these qualitative predictions in more practical settings by
training a two-layer ReLU MLP on the 10-class MNIST classification task and
plotting how well $K^{c,c}_{\bm{W}},K^{c,c}_{\bm{V}}$ align to the 10 target
functions, respectively (Fig.7, $c=3$ in the figures). As predicted,
$K^{3,3}_{\bm{V}}$ becomes aligned with 10 target functions to a similar
extent (Fig.7C) while $K^{3,3}_{\bm{W}}$ preferentially aligns with
$\bm{y}^{3}$(Fig.7D).

##  7 Conclusions

This work demonstrated how kernel alignment emerges during NN learning
dynamics and in turn accelerates learning, both through experiments and
theory. Kernel alignment induces an anisotropic structure in the NTK over
time. We first demonstrated empirically that learning is accelerated by this
anisotropy in a way that cannot be simply accounted for by an increase in the
scale of the kernel. We analytically studied the optimal feature evolution
model to show that aligning the NTK with the target function is beneficial.
For the first time to our best knowledge, we demonstrated empirically that in
NNs with multiclass outputs, alignment manifests in the form of kernel
specialization, where the subkernel corresponding to each output head aligns
to its corresponding target function, but not others.

We then developed theoretical analyses of how the NTK acquires an anisotropic
structure that aligns with the target function during NN training. By
analyzing deep linear networks and two-layer ReLU networks, we give tractable
analytical descriptions of the alignment dynamics of the NTK. Our analyses
suggested that alignment occurs faster and more strongly in deeper networks
and that specialization is a phenomenon that requires nonlinear activation
functions and specific structures in the data.

### Limitations and Future Directions

A primary limitation of our work is the set of simplifying assumptions that we
have taken: focusing on the mean-squared error loss and assuming the simple
data structures for the theory in ReLU networks. Furthermore, our derivations
rely on heuristic assumptions and ansatzes, which should be made rigorous.
This may be possible through performing quenched averages over initial weights
and data points using techniques from statistical physics.

Another limitation is that we have focused on the kernel evaluated on the
training set through the $P\times P$ Gram matrix. It will be interesting to
extend some of the analysis to the test set to evaluate how kernel alignment
affects generalization.

## References

  * Advani et al. (2020) Advani, M. S., Saxe, A. M., and Sompolinsky, H.  High-dimensional dynamics of generalization error in neural networks.  _Neural Networks_ , 132:428–446, 2020.
  * Aitken & Gur-Ari (2020) Aitken, K. and Gur-Ari, G.  On the asymptotics of wide networks with polynomial activations, 2020\.
  * Arora et al. (2018) Arora, S., Cohen, N., and Hazan, E.  On the optimization of deep networks: Implicit acceleration by overparameterization.  In Dy, J. and Krause, A. (eds.), _Proceedings of the 35th International Conference on Machine Learning_ , volume 80 of _Proceedings of Machine Learning Research_ , pp. 244–253. PMLR, 10–15 Jul 2018.  URL https://proceedings.mlr.press/v80/arora18a.html.
  * Atanasov et al. (2021) Atanasov, A., Bordelon, B., and Pehlevan, C.  Neural networks as kernel learners: The silent alignment effect.  _arXiv preprint arXiv:2111.00034_ , 2021.
  * Bahri et al. (2021) Bahri, Y., Dyer, E., Kaplan, J., Lee, J., and Sharma, U.  Explaining neural scaling laws.  _arXiv preprint arXiv:2102.06701_ , 2021.
  * Baratin et al. (2021) Baratin, A., George, T., Laurent, C., Devon Hjelm, R., Lajoie, G., Vincent, P., and Lacoste-Julien, S.  Implicit regularization via neural feature alignment.  In Banerjee, A. and Fukumizu, K. (eds.), _Proceedings of The 24th International Conference on Artificial Intelligence and Statistics_ , volume 130 of _Proceedings of Machine Learning Research_ , pp. 2269–2277. PMLR, 13–15 Apr 2021.  URL http://proceedings.mlr.press/v130/baratin21a.html.
  * Bietti & Mairal (2019) Bietti, A. and Mairal, J.  On the inductive bias of neural tangent kernels, 2019.
  * Bordelon et al. (2020) Bordelon, B., Canatar, A., and Pehlevan, C.  Spectrum dependent learning curves in kernel regression and wide neural networks.  In _International Conference on Machine Learning_ , pp. 1024–1034. PMLR, 2020.
  * Chizat et al. (2019) Chizat, L., Oyallon, E., and Bach, F.  On lazy training in differentiable programming.  In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems_ , volume 32. Curran Associates, Inc., 2019.  URL https://proceedings.neurips.cc/paper/2019/file/ae614c557843b1df326cb29c57225459-Paper.pdf.
  * Cortes et al. (2012) Cortes, C., Mohri, M., and Rostamizadeh, A.  Algorithms for learning kernels based on centered alignment.  _J. Mach. Learn. Res._ , 13(1):795–828, March 2012.  ISSN 1532-4435.
  * Du et al. (2018) Du, S. S., Hu, W., and Lee, J. D.  Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced, 2018.
  * Dyer & Gur-Ari (2020) Dyer, E. and Gur-Ari, G.  Asymptotics of wide networks from feynman diagrams.  In _International Conference on Learning Representations_ , 2020.  URL https://openreview.net/forum?id=S1gFvANKDS.
  * Ergen & Pilanci (2020) Ergen, T. and Pilanci, M.  Convex geometry of two-layer relu networks: Implicit autoencoding and interpretable models.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 4024–4033. PMLR, 2020.
  * Fort et al. (2020) Fort, S., Dziugaite, G. K., Paul, M., Kharaghani, S., Roy, D. M., and Ganguli, S.  Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel, 2020.
  * Geiger et al. (2020) Geiger, M., Spigler, S., Jacot, A., and Wyart, M.  Disentangling feature and lazy training in deep neural networks.  _Journal of Statistical Mechanics: Theory and Experiment_ , 2020(11):113301, nov 2020.  doi: 10.1088/1742-5468/abc4de.  URL https://doi.org/10.1088/1742-5468/abc4de.
  * Ghorbani et al. (2021) Ghorbani, B., Mei, S., Misiakiewicz, T., and Montanari, A.  When do neural networks outperform kernel methods?*.  _Journal of Statistical Mechanics: Theory and Experiment_ , 2021(12):124009, Dec 2021.  ISSN 1742-5468.  doi: 10.1088/1742-5468/ac3a81.  URL http://dx.doi.org/10.1088/1742-5468/ac3a81.
  * Huang & Yau (2020) Huang, J. and Yau, H.-T.  Dynamics of deep neural networks and neural tangent hierarchy.  In III, H. D. and Singh, A. (eds.), _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pp. 4542–4551. PMLR, 13–18 Jul 2020.  URL http://proceedings.mlr.press/v119/huang20l.html.
  * Jacot et al. (2020a) Jacot, A., Gabriel, F., and Hongler, C.  Neural tangent kernel: Convergence and generalization in neural networks, 2020a.
  * Jacot et al. (2020b) Jacot, A., Şimşek, B., Spadaro, F., Hongler, C., and Gabriel, F.  Kernel alignment risk estimator: Risk prediction from training data, 2020b.
  * LeCun et al. (2015) LeCun, Y., Bengio, Y., and Hinton, G.  Deep learning.  _nature_ , 521(7553):436–444, 2015.
  * Lee et al. (2020a) Lee, J., Schoenholz, S. S., Pennington, J., Adlam, B., Xiao, L., Novak, R., and Sohl-Dickstein, J.  Finite versus infinite neural networks: an empirical study, 2020a.
  * Lee et al. (2020b) Lee, J., Xiao, L., Schoenholz, S. S., Bahri, Y., Novak, R., Sohl-Dickstein, J., and Pennington, J.  Wide neural networks of any depth evolve as linear models under gradient descent.  _Journal of Statistical Mechanics: Theory and Experiment_ , 2020(12):124002, Dec 2020b.  ISSN 1742-5468.  doi: 10.1088/1742-5468/abc62b.  URL http://dx.doi.org/10.1088/1742-5468/abc62b.
  * Luo et al. (2021) Luo, T., Xu, Z.-Q. J., Ma, Z., and Zhang, Y.  Phase diagram for two-layer relu neural networks at infinite-width limit.  _Journal of Machine Learning Research_ , 22(71):1–47, 2021.
  * Mei et al. (2019) Mei, S., Misiakiewicz, T., and Montanari, A.  Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit.  In _COLT_ , 2019.
  * Nguyen (2019) Nguyen, P.-M.  Mean field limit of the learning dynamics of multilayer neural networks, 2019.
  * Nguyen & Pham (2020) Nguyen, P.-M. and Pham, H. T.  A rigorous framework for the mean field limit of multilayer neural networks.  _CoRR_ , abs/2001.11443, 2020.  URL https://arxiv.org/abs/2001.11443.
  * Novak et al. (2020) Novak, R., Xiao, L., Hron, J., Lee, J., Alemi, A. A., Sohl-Dickstein, J., and Schoenholz, S. S.  Neural tangents: Fast and easy infinite neural networks in python.  In _International Conference on Learning Representations_ , 2020.  URL https://github.com/google/neural-tangents.
  * Paccolat et al. (2021) Paccolat, J., Petrini, L., Geiger, M., Tyloo, K., and Wyart, M.  Geometric compression of invariant manifolds in neural networks.  _Journal of Statistical Mechanics: Theory and Experiment_ , 2021(4):044001, 2021.
  * Roberts et al. (2021) Roberts, D. A., Yaida, S., and Hanin, B.  The principles of deep learning theory.  _arXiv preprint arXiv:2106.10165_ , 2021.
  * Safran & Shamir (2018) Safran, I. and Shamir, O.  Spurious local minima are common in two-layer relu neural networks.  In _International Conference on Machine Learning_ , pp. 4433–4441. PMLR, 2018.
  * Saxe et al. (2014) Saxe, A. M., McClelland, J. L., and Ganguli, S.  Exact solutions to the nonlinear dynamics of learning in deep linear neural networks.  In Bengio, Y. and LeCun, Y. (eds.), _2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings_ , 2014.  URL http://arxiv.org/abs/1312.6120.
  * Sirignano & Spiliopoulos (2020) Sirignano, J. and Spiliopoulos, K.  Mean field analysis of neural networks: A law of large numbers.  _SIAM Journal on Applied Mathematics_ , 80(2):725–752, 2020.
  * Yang & Hu (2021) Yang, G. and Hu, E. J.  Feature learning in infinite-width neural networks, 2021.
  * Yang & Salman (2020) Yang, G. and Salman, H.  A fine-grained spectral perspective on neural networks, 2020.
  * Zagoruyko & Komodakis (2017) Zagoruyko, S. and Komodakis, N.  Wide residual networks, 2017.

##  Appendix A Derivation of NTK for Linear Network with Small Initialization

The neural tangent kernel for the deep linear network
$f(\bm{x})=\bm{w}^{L+1\top}\bm{W}^{L}...\bm{W}^{1}\bm{x}$ is

$\displaystyle
K(\bm{x},\bm{x}^{\prime})=\sum_{\ell=1}^{L+1}\left<\frac{\partial
f(\bm{x})}{\partial\bm{W}^{\ell}},\frac{\partial
f(\bm{x}^{\prime})}{\partial\bm{W}^{\ell}}\right>_{F}$ (22)

Using the fact that, from small initialization,
$\bm{W}^{\ell}=u(t)\bm{r}_{\ell+1}(t)\bm{r}_{\ell}(t)^{\top}$, we find for
$\ell>1$

$\displaystyle\frac{\partial f(\bm{x})}{\partial\bm{W}^{\ell}}$
$\displaystyle=\left(\prod_{\ell^{\prime}>\ell}\bm{W}^{\ell^{\prime}}\right)^{\top}\bm{x}^{\top}\left(\prod_{\ell^{\prime}<\ell}\bm{W}^{\ell^{\prime}}\right)^{\top}=u(t)^{L-1}\bm{x}^{\top}\bm{r}_{1}(t)\bm{r}_{\ell+1}(t)\bm{r}_{\ell}(t)^{\top}$
$\displaystyle\left<\frac{\partial
f(\bm{x})}{\partial\bm{W}^{\ell}},\frac{\partial
f(\bm{x}^{\prime})}{\partial\bm{W}^{\ell}}\right>$
$\displaystyle=u(t)^{2L-2}\bm{x}^{\top}\bm{r}_{1}(t)\bm{r}_{1}(t)^{\top}\bm{x}^{\prime}\left<\bm{r}_{\ell+1}(t)\bm{r}_{\ell}(t)^{\top},\bm{r}_{\ell+1}(t)\bm{r}_{\ell}(t)^{\top}\right>_{F}$
$\displaystyle=u(t)^{2L-2}\bm{x}^{\top}\bm{r}_{1}(t)\bm{r}_{1}(t)^{\top}\bm{x}^{\prime}$
(23)

For $\ell=1$ we have

$\displaystyle\frac{\partial
f(\bm{x})}{\partial\bm{W}^{1}}=\left(\prod_{\ell^{\prime}>1}\bm{W}^{\ell^{\prime}}\right)^{\top}\bm{x}^{\top}\implies\left<\frac{\partial
f(\bm{x})}{\partial\bm{W}^{1}},\frac{\partial
f(\bm{x}^{\prime})}{\partial\bm{W}^{1}}\right>=u(t)^{2L-2}\bm{x}\cdot\bm{x}^{\prime}$
(24)

Thus, adding the contributions from each of the $L+1$ layers, we find

$\displaystyle
K(\bm{x},\bm{x}^{\prime})=u(t)^{2L-2}\bm{x}^{\top}\left[L\bm{r}_{1}(t)\bm{r}_{1}(t)+\bm{I}\right]\bm{x}^{\prime}$
(25)

This shows that the relative size of the rank-one spike will be controlled by
the depth of the network, $L$.

##  Appendix B Optimal Feature Evolution Induces Kernel Alignment

Let $\bm{\Delta}=\bm{f}-\bm{y}$ and let $\bm{\Psi}\in\mathbb{R}^{N\times P}$
represent the feature matrix whose inner product gives the kernel
$\bm{K}=\bm{\Psi}^{\top}\bm{\Psi}$. We will first discuss a discrete time
dynamical system before taking a gradient flow limit. Thus, we index
$\bm{\Delta}_{t}$ as the error at time $t$ and $\bm{\Psi}_{t}$ as the features
at time $t$. We consider the following simulataneous updates to
$\bm{\Psi}_{t}$ and $\bm{\Delta}_{t}$

$\displaystyle\bm{\Delta}_{t+1}(\bm{\Psi}_{t})=\bm{\Delta}_{t}-\eta\bm{\Psi}_{t}^{\top}\bm{\Psi}_{t}\bm{\Delta}$
$\displaystyle\bm{\Psi}_{t+1}$
$\displaystyle=\bm{\Psi}_{t}-\eta\gamma\frac{\partial}{\partial\bm{\Psi}_{t}}||\bm{\Delta}_{t+1}(\bm{\Psi}_{t})||^{2}$
(26)
$\displaystyle=\bm{\Psi}_{t}-\frac{1}{2}\eta\gamma\frac{\partial}{\partial\bm{\Psi}_{t}}||\bm{\Delta}_{t}-\eta\bm{\Psi}_{t}^{\top}\bm{\Psi}_{t}\bm{\Delta}_{t}||^{2}$
(27)

Expanding the last term and computing the derivative gives

$\displaystyle\frac{1}{2}\frac{\partial}{\partial\bm{\Psi}_{t}}||\bm{\Delta}_{t}-\eta\bm{\Psi}_{t}^{\top}\bm{\Psi}_{t}\bm{\Delta}_{t}||^{2}$
(28)
$\displaystyle=\frac{1}{2}\frac{\partial}{\partial\bm{\Psi}_{t}}\left[||\bm{\Delta}_{t}||^{2}-2\eta\bm{\Delta}_{t}^{\top}\bm{\Psi}_{t}^{\top}\bm{\Psi}_{t}\bm{\Delta}_{t}+\eta^{2}\bm{\Delta}^{\top}\bm{\Psi}_{t}^{\top}\bm{\Psi}_{t}\bm{\Psi}_{t}^{\top}\bm{\Psi}_{t}\bm{\Delta}\right]$
$\displaystyle=-\eta\bm{\Psi}_{t}\bm{\Delta}_{t}\bm{\Delta}_{t}^{\top}+\eta^{2}\bm{\Psi}_{t}\bm{\Psi}_{t}^{\top}\bm{\Psi}_{t}\bm{\Delta}_{t}\bm{\Delta}_{t}^{\top}+\eta^{2}\bm{\Psi}_{t}\bm{\Delta}_{t}\bm{\Delta}_{t}^{\top}\bm{\Psi}_{t}^{\top}\bm{\Psi}_{t}$
$\displaystyle=-\eta\bm{\Psi}_{t}\bm{\Delta}_{t}\bm{\Delta}_{t}^{\top}+O(\eta^{2})$
(29)

Taking the $\eta\to 0$ limit while taking the distance in time between
adjacent steps to zero, we find the following gradient flow dynamics

$\displaystyle\dot{\bm{\Delta}}(t)=-\eta\bm{\Psi}(t)^{\top}\bm{\Psi}(t)\bm{\Delta}(t)\
\quad\
\dot{\bm{\Psi}}(t)=\gamma\eta\bm{\Psi}(t)\bm{\Delta}(t)\bm{\Delta}(t)^{\top}$
(30)

This is a collection of coupled nonlinear ODEs. The $\gamma\to 0$ limit
recovers lazy learning where the features do not evolve. Increasing $\gamma$
increases the rate at which features evolve, thus we name it the feature
learning rate. Despite the nonlinearity, we will attempt to solve these
equations to gain insight into such optimal feature updates. The key trick is
that the equations can be decoupled through the use of a conservation law. To
motivate the conservation law, consider the scalar version of these
differential equations

$\displaystyle\dot{\Delta}=-\psi^{2}\Delta\quad\dot{\psi}=\gamma\psi\Delta^{2}$
(31)

Note that $\frac{1}{2}\frac{d}{dt}\Delta^{2}=-\Delta^{2}\psi^{2}$ while
$\frac{1}{2}\frac{d}{dt}\psi^{2}=\gamma\psi^{2}\Delta^{2}$. A particular
linear combination of these terms reveals a conservation law

$\frac{\gamma}{2}\frac{d}{dt}\Delta^{2}+\frac{1}{2}\psi^{2}=\frac{1}{2}\frac{d}{dt}\left[\gamma\Delta^{2}+\psi^{2}\right]=-\gamma\Delta^{2}\psi^{2}+\gamma\Delta^{2}\psi^{2}=0$
(32)

Thus $\gamma\Delta^{2}+\psi^{2}$ is a conserved quantity throughout the
dynamics. This indicates that, in $(\Delta,\psi)$ space, the trajectory can
only move along an ellipse, with the ratio of axis lengths determined by
$\sqrt{\gamma}$. Using the conservation law, we can introduce a constant
$C=\gamma\Delta^{2}+\psi^{2}=\gamma\Delta_{0}^{2}+\psi_{0}^{2}$, where
$\Delta_{0}$ and $\psi_{0}$ are the initial values. We note the similarity
between these elliptical differential equations and the hyperbolic geometry of
gradient descent in deep linear neural networks (Saxe et al., 2014), where the
conservation laws have the form $a^{2}-b^{2}=C$. Using our discovered
elliptical conservation law, the differential equations can now be decoupled

$\dot{\Delta}=-(C-\gamma\Delta^{2})\Delta\quad\dot{\psi}=\psi(C-\psi^{2})$
(33)

Letting $u=\frac{1}{2}\Delta^{2}$ and $v=\frac{1}{2}\gamma^{2}$, we find
$\dot{u}=-u(C-\gamma u)$ and $\dot{v}=v(C-v)$, which give solutions of the
form

$u=\frac{CA}{A+e^{2Ct}}\quad v=\frac{CB}{B+e^{-2Ct}}$ (34)

for constants $A$ and $B$ determined by the initial conditions. This indicates
that the loss and the kernel power increases as logistic functions. Since $u$
represents the loss, this indicates that at large times, a scaling of
$u\sim\exp(-2(\psi_{0}^{2}+\gamma\Delta_{0}^{2})t)$ is obtained, which
improves with increasing $\gamma$.

The scalar case was illuminating since it allowed us to identify a
conservation law and solve the differential equation. We now aim to extend
this argument to arbitrary dimensional matrices
$\bm{\Psi}(t)\in\mathbb{R}^{N\times P}$ and vectors
$\bm{\Delta}(t)\in\mathbb{R}^{P}$. Inspired by the elliptical geometry in the
scalar case, we make the following ansatz that
$\bm{C}=\gamma\bm{\Delta}\bm{\Delta}^{\top}+\bm{\Psi}^{\top}\bm{\Psi}$ is
conserved. Indeed, explicit differentiation reveals this to be the case.

$\displaystyle\frac{d}{dt}$
$\displaystyle\left[\gamma\bm{\Delta}\bm{\Delta}^{\top}+\bm{\Psi}^{\top}\bm{\Psi}\right]$
$\displaystyle=-\eta\gamma\bm{\Psi}^{\top}\bm{\Psi}\bm{\Delta}\bm{\Delta}^{\top}-\eta\gamma\bm{\Delta}\bm{\Delta}^{\top}\bm{\Psi}^{\top}\bm{\Psi}+\eta\gamma\bm{\Delta}\bm{\Delta}^{\top}\bm{\Psi}^{\top}\bm{\Psi}+\eta\gamma\bm{\Psi}^{\top}\bm{\Psi}\bm{\Delta}\bm{\Delta}^{\top}=\bm{0}$
(35)

Thus $\bm{C}=\gamma\bm{\Delta}\bm{\Delta}^{\top}+\bm{\Psi}^{\top}\bm{\Psi}$ is
a conserved matrix. We can use this fact to again decouple the dynamics

$\dot{\bm{\Delta}}=-\eta(\bm{C}-\gamma\bm{\Delta}\bm{\Delta}^{\top})\bm{\Delta}\quad\dot{\bm{\Psi}}=\eta\bm{\Psi}(\bm{C}-\bm{\Psi}^{\top}\bm{\Psi})$
(36)

Positive $\gamma$ has the effect of accelerating convergence of $\bm{\Delta}$
to zero, while the initial condition and final conditions can be explicitly
related
$\bm{C}=\bm{K}_{\infty}=\gamma\bm{y}\bm{y}^{\top}+\bm{\Psi}_{0}^{\top}\bm{\Psi}_{0}$,
demonstrating that the kernel will align more with the labels after training.
The dynamics of the loss and the kernel can be examined in the eigenbasis of
$\bm{C}$. Let
$\bm{C}=\bm{K}_{\infty}=\sum_{k}c_{k}\bm{v}_{k}\bm{v}_{k}^{\top}$ and let
$\bm{\Delta}=\sum_{k}\delta_{k}\bm{v}_{k}$ and
$\bm{K}=\bm{\Psi}^{\top}\bm{\Psi}=\sum_{k,\ell}A_{k,\ell}\bm{v}_{k}\bm{v}_{\ell}^{\top}$
for symmetric matrix $\bm{A}$. This generates the following differential
equations

$\displaystyle\frac{d}{dt}\ln\delta_{k}$ $\displaystyle=-\eta
c_{k}+\eta\gamma\sum_{\ell}\delta_{\ell}^{2}$ $\displaystyle\dot{A}_{k,\ell}$
$\displaystyle=\eta A_{k,\ell}(c_{k}+c_{\ell})-2\eta\sum_{j}A_{k,j}A_{j,\ell}$
(37)

To zero-th order in $\gamma$, the loss scales like
$L_{t}=\sum_{k}\left(\bm{v}_{k}^{\top}\bm{y}\right)^{2}\exp(-2\eta c_{k}t)$,
which in general will decay more quickly than the loss for the frozen kernel,
since $\bm{K}^{\infty}$ is more aligned with $\bm{y}$ than $\bm{K}_{0}$. When
$\gamma$ is small but non-negligible, we expect
$(\bm{v}_{1}^{\top}\bm{y})^{2}\gg\left(\bm{v}_{k}^{\top}\bm{y}\right)^{2}$ for
$k>1$. We thus get a loss that looks like $L_{t}=De^{-2\eta
c_{1}t}+\sum_{k>1}\left(\bm{v}_{k}^{\top}\bm{y}\right)^{2}\exp(-2\eta
c_{k}t)$. The first term dominates at small times since it has a large
prefactor constant $D=(\bm{v}_{1}^{\top}\bm{y})^{2}$, however once $t\approx
1/c_{1}$, the tail sum dominates and the loss falls as a power law, with a
possibly improved exponent due to better alignment.

Now, let’s consider the kernel’s dynamics. First, at small times $A_{k,\ell}$
is non-diagonal since $\bm{K}_{0}$ and $\bm{K}_{\infty}$ do not necessarily
commute. These off diagonal terms will eventually decay due to the
$-\sum_{j}A_{k,j}A_{j,\ell}$ term. Once $A$ is approximately diagonal, the
dynamics for the diagonal terms are

$\dot{A}_{k,k}=2\eta A_{k,k}c_{k}-2\eta A_{k,k}^{2}$ (38)

This is identical to the scalar equations studied above which we can solve
exactly

$A_{k,k}(t)=\frac{B_{k}c_{k}}{B_{k}+e^{-2\eta c_{k}t}}$ (39)

for some constants $B_{k}$ determined by the initial values $A_{kk}(0)$. Thus,
$A_{k,k}(t)$ increase as logistic functions with a time constant given by
$c_{k}$. Thus, the kernel is approximately

$\bm{K}(t)\sim\sum_{k}A_{kk}(t)\bm{v}_{k}\bm{v}_{k}^{\top}$ (40)

which gives an alignment of

$\left<\bm{y}\bm{y}^{\top},\bm{K}(t)\right>=\frac{1}{\gamma}\left<\bm{K}_{\infty}-\bm{K}_{0},\bm{K}(t)\right>_{F}=\frac{1}{\gamma}\sum_{k}\left(c_{k}-A_{k,k}(0)\right)A_{kk}(t)$
(41)

which increases as a weighted sum of logistic functions. The norm of the
kernel grows as $||\bm{K}(t)||_{F}^{2}=\sum_{k}A_{k,k}(t)^{2}$ so the
alignment curve has the form

$A(t)=\frac{\sum_{k}\left(c_{k}-A_{k,k}(0)\right)A_{kk}(t)}{\sqrt{\sum_{k}A_{k,k}(t)^{2}}\sqrt{||\bm{K}_{0}||^{2}-2\sum_{k}c_{k}A_{k,k}(t)+\sum_{k}A_{k,k}(t)^{2}}}.$
(42)

##  Appendix C Rescaling Alters Feature Learning Rate

In the paper, we consider the following rescaling of the output function
$g(\bm{x})=\frac{1}{\gamma}f(\bm{x})$ and let the learning rate be
$\eta=\eta_{0}\gamma^{2}$. With this choice, gradient flow on the loss
$L=\sum_{\mu}\ell(\frac{1}{\gamma}f(\bm{x}^{\mu}),y^{\mu})$ gives

$\displaystyle\frac{d\bm{\theta}}{dt}$
$\displaystyle=-\frac{\eta}{\gamma}\sum_{\mu}\frac{\partial\ell_{\mu}}{\partial
g_{\mu}}\frac{\partial
f_{\mu}}{\partial\bm{\theta}}=-\eta_{0}\gamma\sum_{\mu}\frac{\partial\ell_{\mu}}{\partial
g_{\mu}}\frac{\partial f_{\mu}}{\partial\bm{\theta}}=O_{\gamma}(\gamma)$ (43)
$\displaystyle\implies\frac{d}{dt}L$
$\displaystyle=\frac{1}{\gamma}\sum_{\mu}\frac{\partial\ell_{\mu}}{\partial
g_{\mu}}\frac{\partial
f_{\mu}}{\partial\bm{\theta}}\frac{d\bm{\theta}}{dt}=-\eta_{0}\sum_{\mu}\frac{\partial\ell_{\mu}}{\partial
g_{\mu}}\frac{\partial\ell_{\nu}}{\partial g_{\nu}}\frac{\partial
f_{\mu}}{\partial\bm{\theta}}\cdot\frac{\partial
f_{\nu}}{\partial\bm{\theta}}=O_{\gamma}(1)$ (44)
$\displaystyle\frac{d}{dt}\frac{\partial f(\bm{x})}{\partial\bm{\theta}}$
$\displaystyle=\frac{\partial
f^{2}}{\partial\bm{\theta}\partial\bm{\theta}}\cdot\frac{d\bm{\theta}}{dt}=O_{\gamma}(\gamma)$
(45)

We thus see an $O_{\gamma}(\gamma)$ increase in the relative rate of evolution
of the parameter gradients compared to the loss $L$. This is very similar to
the style of analysis in (Chizat et al., 2019; Geiger et al., 2020).

###  C.1 Linear networks with multiple outputs cannot specialize

We derived a general expression for
$K^{c,c^{\prime}}(\bm{x},\bm{x^{\prime}})$, defined in Eq.8, for networks of
any depth and show that it cannot show kernel specialization under an
assumption of symmetry between target functions. Define
$\bm{f}_{l}(\bm{x})=\bm{W^{l}}\bm{W^{l-1}}...\bm{W^{1}}\bm{x}$. Then

$\displaystyle
N^{L+1}K^{c,c^{\prime}}(\bm{x},\bm{x^{\prime}})=\nabla_{\Theta}{r^{(c)}(\bm{x}})^{T}\nabla_{\Theta}{r^{(c^{\prime})}(\bm{x^{\prime}}})$
(46) $\displaystyle=$
$\displaystyle\delta(c-c^{\prime})\bm{f}_{L}(\bm{x})^{T}\bm{f}_{L}(\bm{x^{\prime}})$
(47)
$\displaystyle+{\bm{V^{c}}}^{T}\bm{W^{L}}{\bm{W^{L}}}^{T}\bm{V^{c^{\prime}}}\bm{f}_{L-1}(\bm{x})^{T}\bm{f}_{L-1}({\bm{x^{\prime}}})$
(48)
$\displaystyle+{\bm{V^{c}}}^{T}\bm{W^{L}}\bm{W^{L-1}}{\bm{W^{L-1}}}^{T}{\bm{W^{L}}}^{T}\bm{V^{c^{\prime}}}\bm{f}_{L-2}(\bm{x})^{T}\bm{f}_{L-2}(\bm{x^{\prime}})$
(49)
$\displaystyle+{\bm{V^{c}}}^{T}\bm{W^{L}}\bm{W^{L-1}}\bm{W^{L-2}}{\bm{W^{L-2}}}^{T}{\bm{W^{L-1}}}^{T}{\bm{W^{L}}}^{T}\bm{V^{c^{\prime}}}\bm{f}_{L-3}(\bm{x})^{T}\bm{f}_{L-3}(\bm{x^{\prime}})$
(50) $\displaystyle+...$ (51)

Defining scalar functions for $l<L$

$\alpha_{l}(c,c^{\prime})\equiv{\bm{V^{c}}}^{T}\bm{W^{L}}\bm{W^{L-1}}...\bm{W^{l+1}}{\bm{W^{l+1}}}^{T}...{\bm{W^{L-1}}}^{T}{\bm{W^{L}}}^{T}\bm{V^{c^{\prime}}},$
(52)

one has

$N^{L+1}K^{c,c^{\prime}}(\bm{x},\bm{x^{\prime}})=\delta(c-c^{\prime})\bm{f}_{L}(\bm{x})^{T}\bm{f}_{L}(\bm{x^{\prime}})+\sum_{l=0}^{L-1}\alpha_{l}(c,c^{\prime})\bm{f}_{l}(\bm{x})^{T}\bm{f}_{l}(\bm{x^{\prime}}).$
(53)

It is thus a weighted sum of covariance of activations in all layers and the
input. We make a class-symmetry ansatz that

$\forall
l,c,c^{\prime}:\alpha_{l}(c,c)=\alpha_{l}(c^{\prime},c^{\prime})=\tilde{\alpha_{l}}.$
(54)

To see why this ansatz is reasonable, define
$\tilde{\bm{V}_{l}^{c}}\equiv{\bm{W^{l+1}}}^{T}...{\bm{W^{L-1}}}^{T}{\bm{W^{L}}}^{T}\bm{V^{c}}$.
$\alpha_{l}(c,c)=\lVert\tilde{\bm{V}_{l}^{c}}\rVert^{2}$; after learning,
$y^{(c)}(\bm{x})=r^{(c)}(\bm{x})=N^{-(L+1)/2}({\bm{W^{l+1}}}^{T}...{\bm{W^{L-1}}}^{T}{\bm{W^{L}}}^{T}\bm{V^{c}})^{T}\bm{f}_{l}(\bm{x})$.
We then assume the covariance of $\bm{f}_{l}(\bm{x})$ projected along the
direction of $\bm{V^{c}}$ to be approximately the same across $c$ and that
$r^{(c)}(\bm{x})$ to have approximately the same variance. This would suggest
that $\tilde{\bm{V}_{l}^{c}}$ should have the same norm across $c$.

Under the class-symmetry ansatz,

$N^{L+1}K^{c,c}(\bm{x},\bm{x^{\prime}})=\bm{f}_{L}(\bm{x})^{T}\bm{f}_{L}(\bm{x^{\prime}})+\sum_{l=0}^{L-1}\tilde{\alpha_{l}}\bm{f}_{l}(\bm{x})^{T}\bm{f}_{l}(\bm{x^{\prime}})$
(55)

does not have $c$ dependence and thus cannot specialize.

##  Appendix D Theory of Kernel Alignment in Two-layer ReLU Networks

As in the main text, we use the following notation

$\phi_{i}^{\mu}\equiv\phi\left(\bm{w_{i}}\cdot\bm{x^{\mu}}\right)$ (56)
$D_{i}^{\mu}\equiv\frac{d}{dx}\phi\left(x\right)\Bigg{|}_{x=\bm{w_{i}}\cdot\bm{x^{\mu}}}.$
(57) $\bm{f}\in\mathbb{R}^{P}\quad f^{\mu}=f\left(\bm{x^{\mu}}\right)$ (58)

###  D.1 A General Ansatz

Throughout the derivations for the training dynamics of two-layer ReLU
networks, we make use of the following ansatz:

$\forall\bm{x^{\mu}}\perp\bm{x^{\nu}},\forall
t:\sum_{i}D_{i}^{\nu}\left(t\right)V_{i}\left(t\right)\phi_{i}^{\mu}\left(t\right)\approx\frac{1}{2}\sum_{i}V_{i}\left(t\right)\phi_{i}^{\mu}\left(t\right)=\frac{1}{2}f^{\mu}.$
(59)

While we do not prove this result rigorously, we provide a heuristic argument
for this. Since we assume $D_{i}^{\mu}\left(t\right)$ to be static in time,
they are determined by their initial values. Since the weights are initialized
as i.i.d. Gaussian, $\\{D_{i}^{\mu}\\}_{i=1,...,M}$ are random Bernoulli
variables with mean of $1/2$.

Given this, our ansatz will be true if $\\{D_{i}^{\nu}\\}_{i=1,...,M}$ and
$\\{V_{i}\left(t\right)\phi_{i}^{\mu}\left(t\right)\\}_{i=1,...,M}$ are
uncorrelated or weakly correlated. At initialization, this is clearly true for
$\bm{x^{\mu}}\perp\bm{x^{\nu}}$. However, the dynamics of
$V_{i}\left(t\right)$ depends on $\phi_{i}^{\nu}$, which has the same sign as
$D_{i}^{\nu}$. Therefore overtime, $\\{V_{i}\\}_{i=1,...,M}$ may generate
correlation with $\\{D_{i}^{\nu}\\}_{i=1,...,M}$. However, we speculate that
when the dataset is isotropic ($\\{\bm{x^{\nu}}\\}_{\nu=1,..,P}$ are pairwise
perpendicular and $P$ is large), this coupling is weak. The dynamics are given
by

$\frac{dV_{i}}{dt}=\eta\sum_{\nu}\left(f^{\nu}-y^{\nu}\right)\phi_{i}^{\nu}.$
(60)

At initialization,
$\\{\left(f^{\nu}-y^{\nu}\right)\phi_{i}^{\nu}\\}_{\nu=1,...P}$ are pairwise
uncorrelated. Thus, $\\{\phi_{i}^{\nu}\\}$ has only a $1/P$ effect on the
dynamics of $V_{i}\left(t\right)$, which is small for large $P$. If changes to
the parameters are small (due to large network width), this should hold
approximately during training.

We numerically tested Eq.59 in the various settings and found excellent
agreement. However, proving our conjecture rigorously is left for future work.

###  D.2 Kernel Alignment in Two-layer ReLU Networks

Assuming we are training a two layer ReLU network defined by

$f^{\mu}=f\left(\bm{x^{\mu}}\right)=\sum_{i}V_{i}\phi\left(\bm{w_{i}}\cdot\bm{x^{\mu}}\right)$
(61)

on a training set with $P$ examples,
$\\{\bm{x^{\mu}},y^{\mu}\\}_{\mu=1,...,P}$ with a mean squared error loss.

We assume $\\{\bm{x^{\mu}}\\}$ to be all unit vectors that are pairwise
perpendicular and $y^{\mu}\in\\{-1,1\\}$. As mentioned in the main text, we
assume $\\{D_{i}^{\mu}\\}$ to be static for all $\mu,i$ and
$\bm{y}^{T}\bm{f}>0$ at all times.

Batch gradient descent gives the following dynamics

$\frac{d\bm{w_{i}}}{dt}=-\eta\sum_{\nu}\left(f^{\nu}-y^{\nu}\right)V_{i}D_{i}^{\nu}\bm{x^{\nu}}=-\eta\sum_{\nu}\left(f^{\nu}-y^{\nu}\right)V_{i}D_{i}^{\nu}\bm{x^{\nu}}$
(62)

$\frac{dV_{i}}{dt}=-\eta\sum_{\nu}\left(f^{\nu}-y^{\nu}\right)\phi{}_{i}^{\nu}$
(63)

$\displaystyle\frac{d\phi_{i}^{\mu}}{dt}$
$\displaystyle=\nabla_{\bm{w_{i}}}\phi_{i}^{\mu}\cdot\frac{d\bm{w_{i}}}{dt}=D_{i}^{\mu}\bm{x^{\mu}}\cdot\left[-\eta\sum_{\nu}\left(f^{\nu}-y^{\nu}\right)V_{i}D_{i}^{\nu}\bm{x^{\nu}}\right]$
(64)
$\displaystyle=D_{i}^{\mu}\bm{x^{\mu}}\cdot\left[-\eta\left(f^{\mu}-y^{\mu}\right)V_{i}D_{i}^{\mu}\bm{x^{\mu}}\right]$
(65) $\displaystyle=\eta D_{i}^{\mu}\left(y^{\mu}-f^{\mu}\right)V_{i}.$ (66)

As we did in the main text, we write the NTK as a sum of its two components

$K_{\mu,\nu}=\nabla_{\bm{V}}f^{\mu}\cdot\nabla_{\bm{V}}f^{\nu}+\nabla_{\bm{W}}f^{\mu}\cdot\nabla_{\bm{W}}f^{\nu}\equiv(\bm{K}_{\bm{V}})_{\mu,\nu}+(\bm{K}_{\bm{W}})_{\mu,\nu}$
(67)

Define (for a unit vector $\bm{z}\in\mathbb{R}^{P}$) bilinear forms

$\displaystyle\alpha\left(\bm{z}\right)$
$\displaystyle=\bm{z}^{T}\bm{K}_{\bm{V}}\bm{z}=\sum_{\mu,\nu}z^{\mu}\nabla_{\bm{V}}f^{\mu}\cdot\nabla_{\bm{V}}f^{\nu}z^{\nu}=\sum_{i}\left(\sum_{\mu}z^{\mu}\phi_{i}^{\mu}\right)^{2}$
(68) $\displaystyle\beta\left(\bm{z}\right)$
$\displaystyle=\bm{z}^{T}\bm{K}_{\bm{W}}\bm{z}=\sum_{\mu,\nu}z^{\mu}\nabla_{\bm{W}}f^{\mu}\cdot\nabla_{\bm{W}}f^{\nu}z^{\nu}$
(69)
$\displaystyle=\sum_{\mu,\nu}z^{\mu}\left[\sum_{i}V_{i}^{2}D_{i}^{\mu}D_{i}^{\nu}\left(\bm{x^{\mu}}\cdot\bm{x^{\nu}}\right)\right]z^{\nu}$
(70) $\displaystyle=\sum_{\mu}v^{\mu
2}\left[\sum_{i}V_{i}^{2}D_{i}^{\mu}\right]$ (71)

Then

$\displaystyle\frac{d\alpha}{dt}\left(\bm{z}\right)$
$\displaystyle=\sum_{i}\sum_{\mu}\frac{d\alpha}{d\phi_{i}^{\mu}}\frac{d\phi_{i}^{\mu}}{dt}$
(72)
$\displaystyle=2\sum_{i}\left(\sum_{\nu}z^{\nu}\phi_{i}^{\nu}\right)\sum_{\mu}z^{\mu}\eta
D_{i}^{\mu}\left(y^{\mu}-f^{\mu}\right)V_{i}$ (73)
$\displaystyle=2\eta\sum_{\mu}\left(y^{\mu}-f^{\mu}\right)z^{\mu}\sum_{\nu}z^{\nu}\left(\sum_{i}\phi_{i}^{\nu}D_{i}^{\mu}V_{i}\right)$
(74)

We now consider the object $\sum_{i}\phi_{i}^{\nu}D_{i}^{\mu}V_{i}.$ For
$\mu=\nu$, $\sum_{i}\phi_{i}^{\nu}D_{i}^{\nu}V_{i}=f^{\nu}.$ For $\mu\neq\nu$,
we use the ansatz in Eq.59. Hence

$\displaystyle\frac{d\alpha}{dt}\left(\bm{z}\right)$
$\displaystyle\approx\eta\sum_{\mu}\left(y^{\mu}-f^{\mu}\right)z^{\mu}z^{\mu}f^{\mu}+\eta\sum_{\mu}\left(y^{\mu}-f^{\mu}\right)z^{\mu}\sum_{\nu\neq\mu}z^{\nu}f^{\nu}$
(75)
$\displaystyle=\eta\sum_{\mu}\left(y^{\mu}-f^{\mu}\right)z^{\mu}z^{\mu}f^{\mu}+\eta\left[\left(\bm{y}-\bm{f}\right)^{T}\bm{z}\right]\left[\bm{z}^{T}\bm{f}\right]$
(76)

To track the anisotropy of $\bm{K}_{\bm{V}}$, we first consider its average
over all unit vectors with a uniform measure

$\displaystyle\left\langle\frac{d\alpha}{dt}\left(\bm{z}\right)\right\rangle_{\bm{z}}$
$\displaystyle=\frac{1}{P}\eta\left(\bm{y}-\bm{f}\right)^{T}\bm{f}+\frac{1}{P}\eta\left[\left(\bm{y}-\bm{f}\right)^{T}\bm{f}\right]$
(77) $\displaystyle=\frac{1}{P}2\eta\left(\bm{y}-\bm{f}\right)^{T}\bm{f}$ (78)
$\displaystyle\approx\frac{1}{P}2\eta\bm{y}^{T}\bm{f}$ (79)

we then compute it for a unit vector in the direction of $\bm{y}$,
$\bm{y}/\sqrt{P}$

$\displaystyle\frac{d\alpha}{dt}\left(\bm{y}/\sqrt{P}\right)$
$\displaystyle=\frac{1}{P}\eta\sum_{\mu}\left(y^{\mu}-f^{\mu}\right)f^{\mu}+\frac{1}{P}\eta\left[\left(\bm{y}-\bm{f}\right)^{T}\bm{y}\right]\left[\bm{y}^{T}\bm{f}\right]$
(80)
$\displaystyle\approx\frac{1}{P}\eta\bm{y}^{T}\bm{f}+\eta\bm{y}^{T}\bm{f}\approx\eta\bm{y}^{T}\bm{f}.$
(81)

Hence, for $\bm{y}^{T}\bm{f}>0$ and large $P$, we have

$\frac{d\alpha}{dt}\left(\bm{y}/\sqrt{P}\right)\gg\left\langle\frac{d\alpha}{dt}\left(\bm{z}\right)\right\rangle_{\bm{z}}.$
(82)

This show that one part of the NTK, $\bm{K}_{\bm{V}}$, is growing in the
direction of $\bm{y}$ and static or contracting in directions perpendicular to
$\bm{y}$, thus aligning to $\bm{y}$ over time. We now analyze the anisotropy
of $\bm{K_{W}}$.

Now for $\beta$, since we assume $\\{\bm{D}^{\mu}\\}$ to be static, its
dynamics are entirely driven by $\bm{V}$.

$\frac{d\beta}{dV_{i}}\left(\bm{z}\right)=2\sum_{\mu}z_{\mu}^{2}D_{i}^{\mu}V_{i}$
(83)

$\displaystyle\frac{d\beta}{dt}\left(\bm{z}\right)$
$\displaystyle=\sum_{i}\frac{d\beta}{dV_{i}}\frac{dV_{i}}{dt}=2\eta\sum_{i}\sum_{\mu}z_{\mu}^{2}D_{i}^{\mu}V_{i}\sum_{\nu}\left(y^{\nu}-f^{\nu}\right)\phi{}_{i}^{\nu}$
(84)
$\displaystyle=2\eta\sum_{\mu}z_{\mu}^{2}\sum_{\nu}\left(y^{\nu}-f^{\nu}\right)\sum_{i}D_{i}^{\mu}V_{i}\phi{}_{i}^{\nu}$
(85)

Analysis of the object $\sum_{i}D_{i}^{\mu}V_{i}\phi{}_{i}^{\nu}$ is the same
as above, resulting in

$\frac{d\beta}{dt}\left(\bm{z}\right)\approx\eta\left\|\bm{z}\right\|^{2}\left(\bm{y}-\bm{f}\right)^{T}\bm{f}.$
(86)

###  D.3 Kernel Specialization in Two-layer ReLU Networks

We consider a two-layer ReLU network with $C$ output heads. The output of the
$c$th head is

$f_{c}^{\mu}\equiv
f_{c}\left(\bm{x^{\mu}}\right)=\sum_{i}V_{i}^{c}\phi\left(\bm{w_{i}}\cdot\bm{x^{\mu}}\right)$
(87)

Assume we are classifying a mixture of $C$ Gaussians,
$\\{\mathcal{N}(\bm{\mu}_{c},\sigma^{2}\mathbb{I}\\}_{c=1,...,C}$. $P/C$
points are drawn from each Gaussian. Centers of the Gaussians,
$\\{\bm{\mu}_{c}\\}$, are all unit vectors and pairwise perpendicular. Then
data points from class $c$ can be written as

$\bm{x}^{\mu\in\text{class c}}=\bm{\mu}_{c}+\delta\bm{x}^{\mu}.$ (88)

This leads to the following input correlation structure

$\bm{x}^{\mu}\cdot\bm{x}^{\nu}=\begin{cases}1&\text{if}\mu=\nu\\\
m&\text{if}\mu\neq\nu\in\text{same class}\\\ 0&o.w.\end{cases}$ (89)

Assume $\sigma^{2}$ to be sufficiently small that

$\phi\left(\bm{w}_{i}\cdot\bm{\mu}_{c}+\delta\bm{x}^{\mu}\right)\approx\phi\left(\bm{w}_{i}\cdot\bm{\mu}_{c}\right)+\left(\bm{w}_{i}\cdot\delta\bm{x}^{\mu}\right)\phi^{\prime}\left(\bm{w}_{i}\cdot\bm{\mu}_{c}\right)$
(90)

$\phi^{\prime}\left(\bm{w}_{i}\cdot\bm{\mu}_{c}+\delta\bm{x}^{\mu}\right)\approx\phi^{\prime}\left(\bm{w}_{i}\cdot\bm{\mu}_{c}\right)+\left(\bm{w}_{i}\cdot\delta\bm{x}^{\mu}\right)\phi^{\prime\prime}\left(\bm{w}_{i}\cdot\bm{\mu}_{c}\right)$
(91)

If $P/C$ is sufficiently large, then

$\sum_{\mu\in\text{class c}}\phi_{i}^{\mu}\approx
P/C\phi_{i}^{c},\quad\phi_{i}^{c}\equiv\phi\left(\bm{w}_{i}\cdot\bm{\mu}_{c}\right)$
(92)

$\sum_{\mu\in\text{class c}}D_{i}^{\mu}\approx P/CD_{i}^{c},\quad
D_{i}^{c}\equiv\phi^{\prime}\left(\bm{w}_{i}\cdot\bm{\mu}_{c}\right).$ (93)

It follows that $f_{c}\left(\bm{\mu}_{c}\right)\approx\sum_{\mu\in\text{class
c}}f_{c}^{\mu}.$ Since we assume
$\sum_{\mu}f_{c}\left(\bm{x^{\mu}}\right)y_{c}^{\mu}>0$ and
$y_{c}^{\mu\in\text{class c}}=1$, $f_{c}\left(\bm{\mu}_{c}\right)>0$.

We would like to analyze the anisotropy in the dynamics of $c$th subkernel,
defined as

$\displaystyle\bm{K^{cc}}$
$\displaystyle=\sum_{d}\nabla_{\bm{V^{d}}}f_{c}^{\mu}\cdot\nabla_{\bm{V^{d}}}f_{c}^{\nu}+\nabla_{\bm{W}}f_{c}^{\mu}\cdot\nabla_{\bm{W}}f_{c}^{\nu}=\nabla_{\bm{V^{c}}}f_{c}^{\mu}\cdot\nabla_{\bm{V^{c}}}f_{c}^{\nu}+\nabla_{\bm{W}}f_{c}^{\mu}\cdot\nabla_{\bm{W}}f_{c}^{\nu}$
(94) $\displaystyle\equiv\bm{K_{v}^{cc}}+\bm{K_{W}^{cc}}.$ (95)

In particular, we want to show that it grows in the direction of $\bm{y_{c}}$
at a rate faster than in the directions of $\bm{y_{d\neq c}}$. As argued in
the main text, $\bm{K_{v}^{cc}}$ cannot be anisotropic for different
$\bm{y_{d}}$. Thus, we only track the anisotropic dynamics of
$\bm{K_{W}^{cc}}$ with the bilinear form (all $\\{\bm{y_{d}}\\}_{d=1,...,C}$
have the same norm)

$B_{\bm{W}}^{c}$($\bm{y_{d}}$)
$\displaystyle=\bm{y_{d}}^{T}\bm{K_{W}^{cc}}\bm{y_{d}}=\sum_{\mu,\nu}^{P}y_{d}^{\mu}y_{d}^{\nu}\bm{x}_{\mu}^{T}\bm{x}_{\nu}\sum_{i}D_{i}^{\mu}D_{i}^{\nu}(\bm{V}_{c})_{i}^{2}.$
(96) $\displaystyle=\sum_{\mu\in\text{class
d}}\sum_{i}\left(D_{i}^{\mu}\right)^{2}(\bm{V}_{c})_{i}^{2}+m\sum_{k}\sum_{\mu\neq\nu\text{$\in\text{class
d}$}}^{P/C}\sum_{i}D_{i}^{\mu}D_{i}^{\nu}(\bm{V}_{c})_{i}^{2}$ (97)
$\displaystyle=\sum_{\mu\in\text{class
d}}\sum_{i}D_{i}^{\mu}(\bm{V}_{c})_{i}^{2}+m\sum_{k}\sum_{\mu\neq\nu\text{$\in\text{class
d}$}}^{P/C}\sum_{i}D_{i}^{\mu}D_{i}^{\nu}(\bm{V}_{c})_{i}^{2}\quad\text{because
}\text{$D_{i}^{\mu}$$\in\\{0,1\\}$}$ (98)

The gradient descent dynamics of $\bm{V_{c}}$ is

$\frac{d(\bm{V}_{c})_{i}}{dt}=-\eta\sum_{\mu}\left(f_{c}^{\mu}-y_{c}^{\mu}\right)\phi_{i}^{\mu}\approx\eta\sum_{\mu\in\text{class
c}}\phi_{i}^{\mu}.$ (99)

Then dynamics of $B_{\bm{W}}^{c}$($\bm{y_{d}}$) follow

$\displaystyle\eta^{-1}\frac{dB_{\bm{W}}^{c}(\bm{y_{d}})}{dt}$
$\displaystyle=\eta^{-1}\sum_{i}\frac{dB_{\bm{W}}^{c}(\bm{y_{d}})}{d(\bm{V}_{c})_{i}}\frac{d(\bm{V}_{c})_{i}}{dt}=\sum_{i}(\bm{V}_{c})_{i}\left[\sum_{\mu\in\text{class
d}}D_{i}^{\mu}+m\sum_{\mu\neq\nu\text{$\in\text{class
d}$}}^{P/C}D_{i}^{\mu}D_{i}^{\nu}\right]\sum_{\omega\in\text{class
c}}\phi_{i}^{\omega}$ (100)
$\displaystyle\approx\frac{P}{C}\sum_{i}\left[\frac{P}{C}D_{i}^{d}(\bm{V}_{c})_{i}+m\sum_{\mu\neq\nu\text{$\in\text{class
d}$}}^{P/C}D_{i}^{d}(\bm{V}_{c})_{i}\right]\phi_{i}^{c}$ (101)
$\displaystyle=\frac{P}{C}\left[\frac{P}{C}+mP^{2}/C^{2}\right]\sum_{i}D_{i}^{d}(\bm{V}_{c})_{i}\phi_{i}^{c}$
(102)

For $B_{\bm{W}}^{c}(\bm{y_{c}})$, this simplifies to
$\frac{P}{C}\left[\frac{P}{C}+mP^{2}/C^{2}\right]f_{c}\left(\bm{\mu}_{c}\right)>0.$
On the other hand, for $B_{\bm{W}}^{c}(\bm{y_{d\neq c}}),$ one has
$\sum_{i}D_{i}^{d}(\bm{V}_{c})_{i}\phi_{i}^{c}\approx\frac{1}{2}f_{c}\left(\bm{\mu}_{c}\right)$(following
Eq.59) and thus

$\eta^{-1}\frac{dB_{\bm{W}}^{c}(\bm{y_{d}})}{dt}\approx\frac{1}{2}\frac{P}{C}\left[\frac{P}{C}+mP^{2}/C^{2}\right]f_{c}\left(\bm{\mu}_{c}\right)>0$
(103)

##  Appendix E Experimental Details

We train our models on a Google Colab GPU and include code to reproduce all
experimental results in the supplementary materials. To match our theory, we
use fixed learning rate SGD. Both evaluation of the infinite width kernels and
training were performed with the Neural Tangents API (Novak et al., 2020).

###  E.1 Wide Res-Net with CIFAR-10

We used the Neural Tangents API implementation of the Wide ResNet model which
can be found on the ReadME of the github https://github.com/google/neural-
tangents#infinitely-wideresnet. We used a width factor of $k=3$ and a
blocksize of $b=2$. The only change between our experiments was rescaling the
output of the network $g(\bm{x})=\frac{1}{\gamma}f(\bm{x})$. We used $100$
CIFAR-10 images taken at random from the first two classes and used binary
labels $y_{\mu}\in\\{\pm 1\\}$.

###  E.2 ReLU networks trained on MNIST

The network has 2 hidden layers with hidden dimension of 500. The network was
trained with full-batch gradient descent on a subset of the MNIST dataset
(1000 examples). The target functions are ”10-hot” vectors encoding the
classes. The learning rate was fixed at 20 (the loss function is averaged over
the batch and the 10 classes).

These simulations were performed in Google Colab using GPU acceleration and
jax.

###  E.3 ReLU convolutional networks trained on CIFAR-10

The network has two convolutional layers, each with 33 channels and each
filter is $7\times 7$. The target functions are ”10-hot” vectors encoding the
classes. The network was trained on 9000 random samples from the CIFAR dataset
with a batchsize of 100 and the mean squared error loss function. The learning
rate was fixed at 20 (the loss function is averaged over the batch and the 10
classes). To save computational resources, the kernel was estimated using Gram
matrices computed on 300 examples.

These simulations were performed in Google Colab using GPU acceleration and
jax.

###  E.4 Two-layer ReLU networks trained on random binary classification.

We used input dimension ($N$) 1000, hidden dimension ($M$) 2000 and the number
of examples ($P$) 100. The network has one hidden layer. The network did not
have biases. The first layer weights $\bm{W}$ were initialized as i.i.d.
Gaussian with variance $1/N$, and the second layer weights $\bm{V}$ where
initialized as i.i.d. Gaussian with variance $1/M$.

The input vectors were sampled from an $N$-dimensional Gaussian distribution
$\mathcal{N}(\bm{0},1/\sqrt{N}\mathbb{I})$ and the labels were random binary
labels $y^{\mu}\in{-1,1}$ with equal probability.

Full-batch gradient descent with a fixed learning rate (0.1) was performed on
$\bm{W,V}$ using the mean squared error loss function
$L=P^{-1}\sum_{\mu}(y^{\mu}-f^{\mu})^{2}$. These simulations were performed on
a personal computer using pytorch and no GPU.

###  E.5 Two-layer ReLU networks trained on classifying a mixture of
Gaussians.

We used input dimension ($N$) 1000, hidden dimension ($M$) 2000 and the number
of examples ($P$) 500. The network has one hidden layer. There are 10 classes
($C$). The target function for each output node is an one-hot encoding vector.
The network did not have biases. The first layer weights $\bm{W}$ were
initialized as i.i.d. Gaussian with variance $1/N$, and the second layer
weights $\bm{V}$ where initialized as i.i.d. Gaussian with variance $1/M$.

The centers of the Gaussians are randomly drawn from the $N$-dimensional unit
sphere and the variance of all Gaussians ($\sigma^{2}$) is 0.01. The network
has 10 output nodes. The classes are balanced such that for each class there
are $P/C=50$ data points in the training set.

Full-batch gradient descent with a fixed learning rate (0.2) was performed on
$\bm{W,V}$ using the mean squared error loss function
$L=P^{-1}\sum_{\mu}\sum_{c}^{C}(y^{\mu}_{c}-f^{\mu}_{c})^{2}$. These
simulations were performed on a personal computer using pytorch and no GPU.