What can linearized neural networks actually say about generalization?
Authors: Guillermo Ortiz-Jimenez
Seyed-Mohsen Moosavi-Dezfooli
Pascal Frossard
Venue: NeurIPS 2021 Poster (2021)
Keywords: Deep learning theory
NTK
generalization
inductive bias
linearized networks
neural tangent kernel

Abstract: For certain infinitely-wide neural networks, the neural tangent kernel (NTK) theory fully characterizes generalization, but for the networks used in practice, the empirical NTK only provides a rough first-order approximation. Still, a growing body of work keeps leveraging this approximation to successfully analyze important deep learning phenomena and design algorithms for new applications. In our work, we provide strong empirical evidence to determine the practical validity of such approximation by conducting a systematic comparison of the behavior of different neural networks and their linear approximations on different tasks. We show that the linear approximations can indeed rank the learning complexity of certain tasks for neural networks, even when they achieve very different performances. However, in contrast to what was previously reported, we discover that neural networks do not always perform better than their kernel approximations, and reveal that the performance gap heavily depends on architecture, dataset size and training task. We discover that networks overfit to these tasks mostly due to the evolution of their kernel during training, thus, revealing a new type of implicit bias.

What can linearized neural networks actually say about generalization?
Guillermo Ortiz-Jiménez EPFL, Lausanne, Switzerland guillermo .et.errmmoooiiiiini... ort.zzjmeeezzeppllll... ch
Seyed-Mohsen Moosavi-Dezfooli ETH Zurich, Zurich, Switzerland seyed ethz.
Pascal Frossard
EPFL, Lausanne, Switzerland pascal frossardeepfl.. ch
Abstract
For certain infinitely-wide neural networks, the neural tangent kernel (NTK) theory fully characterizes generalization, but for the networks used in practice, the empiri- cal NTK only provides a rough first-order approximation. Still, a growing body of work keeps leveraging this approximation to successfully analyze important deep learning phenomena and design algorithms for new applications. In our work, we provide strong empirical evidence determine the practical validity of such approximation by conducting a systematic comparison of the behavior of different neural networks and their linear approximations on different tasks. We show that the linear approximations can indeed rank the learning complexity of certain tasks for neural networks, even when they achieve very different performances. However, in contrast to what was previously reported, we discover that neural networks do not always perform better than their kernel approximations, and reveal that the performance gap heavily depends on architecture, dataset size and training task. We discover that networks overfit to these tasks mostly due the evolution of their kernel during training, thus, revealing a new type of implicit bias.
Introduction
Due to their excellent practical performance, deep learning based methods are today the de facto standard for many visual and linguistic learning tasks. Nevertheless, despite this practical success, our theoretical understanding of how, and what, can neural networks learn is still in its infancy (T]. Recently, a growing body work has started explore the use of linear approximations analyze deep networks, leading to the neural tangent kernel (NTK) framework (2).
The NTK framework is based the observation that for certain initialization schemes, the infinite- width limit of many neural architectures can be exactly characterized using kernel tools (2). This reduces key questions in deep learning theory the study of linear methods and convex functional analysis, for which a rich set theories exist (3). This intuitive approach has been proved very fertile, leading to important results in generalization and optimization of very wide networks (4).
The NTK theory, however, can only fully describe certain infinitely wide neural networks, and for the narrow architectures used in practice, it only provides first-order approximation of their training dynamics (see Fig. [ Despite these limitations, the intuitiveness of the NTK, which allows to use a powerful set of theoretical tools exploit it, has led to rapid increase in the amount of research that successfully leverages the NTK in applications, such as predicting generalization [10] and training
speed [TI]. explaining certain inductive biases [2 13
designing new classifiers [16] 17
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

So. + (0NTK
0b) Vo)'VoSa
V f
So
00o
H
F
Figure 1: Conceptual illustration of the NTK approximation (see See 2 The empirical NTK defines a linear function space tangent H to the non-linear function space F defined by the network. In the limit of infinite width, the neural network space loses its curvature and coincides with the tangent space. Training a linearized network restricts the optimization trajectory to lie in the tangent space.
Recent reports, however, have started questioning the effectiveness of this approach, as one can find multiple examples in which kernel methods are provably outperformed by neural networks [18 20) Most importantly, it has been observed empirically that linearized models computed using a first- order Taylor expansion around the initialization of neural network perform much worse than the networks they approximate on standard image recognition benchmarks [21]: phenomenon that has been coined the non-linear advantage. However, has also been observed that, if one linearizes the network at a later stage of training, the non-linear advantage is greatly reduced. The reasons behind this phenomenon are poorly understood, yet they are key explain the success of deep learning.
(18)20)
Building from these observations, we delve deeper into the source of the non-linear advantage, trying to understand why previous work could successfully leverage the NTK in some applications. In particular, we shed new light on the question: When can the NTK approximation be used to predict generalization, and what does it actually say about it? We propose, for the first time, to empirically study this problem from the perspective of the characteristics of the training labels. To that end, we conduct systematic analysis comparing the performance of different neural network architectures with their kernelized versions, on several problems with the same data support, but different labels. Doing so, identify the alignment of the target function with the NTK key quantity governing important aspects of generalization in deep learning. Namely, one can rank the learning complexity of solving certain tasks with deep networks according to their kernel alignment.
We, then, study the evolution of the aligmment during training to see its influence on the inductive bias. Prior work had shown that, during optimization, deep networks significantly increase their alignment with the target function, and this had been strongly conjectured to be positive for generalization [22]- 24 In contrast, in this work, offer more nuanced view of this phenomenon, and show that it does not always have a positive effect for generalization. In fact, we provide multiple concrete examples where deep networks exhibit a non-linear disadvanlage compared their kernel approximations.
The main contributions of our work are:
We show that the alignment with the empirical NTK at initialization can provide good measure of relative learning complexity in deep learning for a diverse set of tasks.
We use this fact to shed new light on the directional inductive bias of most convolutional neural networks, as this alignment can be used to identify neural anisotropy directions (25)
Moreover, we identify a set of non-trivial tasks in which neural networks perform worse than their linearized approximations, and show this is due their non-linen lon-linear dynamics.
We, hence, provide fine-grained analysis of the evolution of the kernel during training, and show that the NTK rotates mostly in a single axis. This mechanism is responsible for the rapid convergence of neural networks to the training labels, but interestingly we find that it can sometimes hurt generalization, depending on the target task.
We see our empirical findings an important step forward in our understanding of deep learning. Our work paves the way for new research avenues based on our newly observed phenomena, but also provides a fresh perspective to understand how to use the NTK approximation in several applications.

---

2 Preliminaries
Let fe : E denote a neural network parameterized by set of weights 0 E R". Without loss of generality, but to reduce the computational load, in this work, we only consider the binary classification setting, where the objective is to learn an underlying target function f R - {+1} by training the weights minimize the empirical risk R(fo) m Σx= Liign( fa ff over finite set of i.i.d. training data S = { (=:)) == from underlying distribution D. We broadly say that a model generalizes whenever it achieves a low expected risk R(fo) = Eg sign(fo(=))
In small neighborhood around the weight initialization Oo. neural network can be approximated using a first-order Taylor expansion (see Fig. [
fx( fo (z: 0u) = fo, (x) (e- 0o) Vofo,
where VoJa, (ax) E R.T denotes the Jacobian of the network with respect to the parameters evaluated at dg. Here, the model fo represents linearized network which maps weight vectors to functions living in a reproducible kernel Hilbert space (RKHS) 7 C (2(",,, determined by the empirical NTK 4 at On. ao (æt, x²) (a (z)), VaJa, (x') )). Unless stated otherwise, we will generally drop the dependency on On and use θ to refer to the NTK at initialization.
In most contexts, the NTK evolves during training by following the trajectory of the network Jacobian fe, computed at checkpoint O, (see Fig. T.. Remarkably, however, was recently discovered that for certain types of initialization, and in the limit of infinite network-width, the approximation in M is exact, and the NTK constant throughout training 2]. In this regime, one can, then, provide generalization guarantees for neural networks using generalization bounds from kernel methods, and show [26), for instance, that with high probability
1 F. (xr, x))
S ( +0
(2)
where /* arg minNeR Rh) + 1// with >>>>>>0> O denoting a regularization constant, and I being the RKHS norm of the target function, which for positive definite kernels can be computed as
20 1
| = Σ j=1
X (BaDD (x)
(3)
Here, the couples Of) denote the cigenvalue-rigenfuuntion pairs, in order of decreasing
eigenvalues, of the Mercer's decomposition of the kernel, i.e., O(x, x') Σj= メyoر (x) (z) (æ'). This means that, in kernel regimes, the difference between the empirical and the expected risk is smaller when training on target functions with lower RKHS norm. That is, whose projection on the cigenfunctions of the kernel is mostly concentrated along its largest eigenvalues. One can then see the RKHS norm as a metric that ranks the complexity of learning different targets using the kernel O
Estimating (3) in practice is challenging as it requires access to the smallest eigenvalues of the kernel. However, one use the following lemma to compute more tractable bound of the RKHS norm, which shows that a high target-kernel alignment is a good proxy for a small RKHS norm
Lemma 1. Let af)) = Ex x' DD (f f²)f((()) denote the alignment of the target f E H with the kernel @ Then |If1l& llfl)//(f)) Moreover, for the NTK, a(f) ||E [f (f (x))
Proof See Appendix.
At this point, it is important to highlight that in most practical applications we do not deal with infinitely-wide networks, and hence 2 can only be regarded as learning guarantee for linearized models, i.c. fo. Furthermore, from now on, will interchangeably use the terms NTK and empirical NTK to simply refer to the finite-width kernels derived from (T). In our experiments, we compute those using the neural_ neural_tangento library [27] built on top of the JAX framework 28 which we also use to generate the linearized models.
Similarly, as it is commonly done in the kernel literature, we will use the eigenvectors of the Gram matrix to approximate the values of the eigenfunctions Φby (a:) over finite dataset. We will also use the terms eigenvector and eigenfunction interchangeably. In particular, we will use & E Rm to denote the matrix containing the jth Gram eigenvector Φ; E IRm its jth row, where the rows are ordered according to the vector of decreasing eigenvalues À R
3

---

3 Linearized models can predict relative task complexity for deep networks
A growing body of work is using the linear approximation of neural networks as kernel methods to analyze and build novel algorithms. Meanwhile, recent reports, both theoretical and empirical, have started to question if the NTK approximation can really tell anything useful about generalization for finite-width networks. In this section, we try to demystify some of these confusions and aim shed light on the question: What can the empirical NTK actually predict about generalization?
To that end, we conduct systematic study with different neural networks and their linearized approximations given by [ which we train to solve a structured array of predictive tasks with different complexity. Our results indicate that for many problems the linear models and the deep networks do agree in the way they order the complexity of learning certain tasks, even if their performance on the same problems can greatly differ. This explains why the NTK approximation can be used in applications where the main goal is to just predict the relative difficulty of different tasks.
3.1 Learning NTK eigenfunctions
In kernel theory, the sample and optimization complexity required to learn given function is normally bounded by its kernel norm (3) which intuitively measures the alignment the target function with the eigenfunctions of the kemnel. The eigenfunctions themselves, thus, represent natural set of target functions with increasingly high learning complexity according to the increasing value of their associated eigenvalues- for kernel methods. Since our goal is to find if the kernel approximation can indeed predict generalization for neural networks, we evaluate the performance of these networks when learning the eigenfunctions of their NTKs.
In particular, we generate a sequence of datasets constructed using the standard CIFARIO [29] samples, which we label using different binarized versions of the NTK eigenfunctions. That is, to every sample æ in CIFARIO we assign the label sign(() eigenfunction of the NTK at initialization (see Sec. ( In this construction, the choice of CIFARI 10 as supporting distribution makes our experiments close to real settings which might be conditioned by low dimensional structures in the data manifold [19 23); while the choice of eigenfunctions as targets guarantees a progressive increase in complexity, at least, for the linearized networks. Specifically, for 3j the alignment is given by a ((r) X (see Appendix).
where Φ, represents the jth
Accurey 1000
MLP
Ascunay
LoNot
Accurecy 100.0%
RosNod188
100.0%
90.0%
90.0%%
50.0%
80.0%
8000
70.0%%
60.0%%
0
250
500
750
1000
0
250
500
750
1000
0
250
500
750
1000
Eigenffmeting tinn incies
Eigen/inction index
Eigenfinction index
Noural oetwork
Limstiaed model
Accuracy of NN on lunnry CIFAR1U
Figure 2: Validation accuracy of different neural network architectures and their linearizations when trained binarized eigenfunctions of the NTK at initialization, i.e., - sógn( (z)). As baseline, we also provide the accuracies on CIFAR2 (see Sec. 4
We train different neural network architectures selected cover the spectrum of small large models (30 and their linearized models given by [ Unless stated otherwise, we always use the same standard training procedure consisting of the use of stochastic gradient descent (SGD) to optimize a logistic loss, with decaying learning rate starting at 0.05 and momentum set to 0.9. The values of our metrics are reported after 100 epochs of training
Fig. 2 summarizes the main results of our experiments Here, we can clearly see how the validation accuracy of networks trained to predict targets aligned with {Ø, progressively drops with decreasing eigenvalues for both linearized models as predicted by the theory as well as for neural networks. Similarly, Fig. 3 shows how the training dynamics of these networks also correlate with
I 'Our code can be found at htthtttp:///giiiii
searrtiiii/innarzeed-roon Resslls with equivalent findings for other training schemes and datasets be found in the Appendix.

---

Epoch
Convergence
6g dist.
Params 6 dist.
cos dist
Parame dist.
60-
0.15
0.10
10-
01
0-
0.00
250
750
250
750
Eigenfimnciion inckee
ElgenFunction index
Elyonfuncciin index
Figure 3: Correlation of different training metrics with the index of the eigenfunction the network is
trained on. Plots show the number of training iterations taken by the network to achieve a 0.01 training loss, and the & and cosine distances between initialization and final parameters for 8 trained on the binarized eigenfunctions of the NTK at initialization.
cigenfunction index. Specifically, we see that networks take more time to fit eigenfunctions associated smaller eigenvalues, and need travel a larger distance in the weight space to do so.
Overall, our observations reveal that sorting tasks based on their alignment with the NTK is a good predictor of learning complexity both for linear models and neural networks. Interestingly, however, we can also observe large performance gaps between the models. Indeed, even if the networks and the kernels agree which eigenfunctions are harder to learn, the kernel methods perform comparatively much better. This clearly differs from what was previously observed for other tasks [18 and hence, highlights that the existence of non-linear advantage is not always certain.
3.2 Learning linear predictors
The NTK eigenfunctions are one example of a canonical set of tasks with increasing hardness for kernel methods, whose learning complexity for neural networks follows the same order. However, could there be more examples? And, are the previously observed correlations useful to predict other generalization phenomena? In order, to answer these questions, we propose analyze another set of problems, but this time using sequence of targets of increasing complexity for neural networks.
In this sense, it has recently been observed that, for convolutional neural networks (CNNs), the set of linear predictors - i.e., hyperplanes separating two distributions - represents a function class with a wide range of learning complexities among its elements (25). In particular, it has been confirmed empirically that possible to rank the complexity for a neural network to learn different linearly separable tasks based only on its neural anisotropy directions (NADs).
Definition (Neural anisotropy directions). The NADs of a neural network are the ordered sequence of orthonormal vectors { dd which form a full basis of the imput space and whose order is determined by the sample complexity required to learn the linear predictors {x + sign(v} x)}{=
In 25 the authors provided several heuristics to compute the NADs of neural network. However, we now provide a new, more principled, interpretation of the NADs, showing one can also obtain this sequence using kernel approximation. To that end, will make use of the following theorem.
Theorem 1. Let E gd-1 be unitary vector that parameterizes linear predictor gu gu(x) u xx and let x ~ NOO I). The alignment of gu with O is given by
a(gu)
E
[V Vo
u
where V2 (2) E IRaxd denotes the derivative of fe with respect to the weights and the input.
Proof See Appendix.
Theorem 1 gives alternative method to compute NADs. Indeed, in the kernel regime, the NADs are simply the right singular vectors of the matrix of mixed-derivatives E,, V2 V2.o.o, (x) of the network Note however. that this interpretation is just based on an approximation, and hence there is no explicit guarantee that these NADs will capture the direcional inductive bias of deep networks. Our experiments show otherwise, as they reveal that CNNs actually rank the learning complexity of different linear predictors in a way compatible with Theorem
All predictors have the same Le norm. Hence, their alignment is inversely proportional to their kernel norm.
5

---

Anstieyy
LeNet
Accusscy
ResNet18
Alignment
.007
10a
LeNet
ResNer18
80%
90%
600
0056
10-1
250
750
NAD index
N.D Inder
NAD Index
Linesriod modié
Figure 4: (Left) Performance comparison of different neural network architectures with their lin- earizations when learning linear target functions aligned with increasing NADs, i.e., 2 ++ sign x). (Right) Predicted value of the aligmment of the predictors with θ i.c., a(ge,) (see Theorem
Indeed, as shown in Fig. 4 when trained to classify a set of linearly separable dataset aligned with the NTK-ppeddcctdd NADs, CNNs perform better on those predictors with higher kernel alignment (i.e., corresponding to the first NADs) than on those with a lower one (i.e., later NADs). The fact that NADs can be explained using kernel theory constitutes another clear example that theory derived from naïve linear expansion of a neural network can sometimes capture important trends in the inductive bias of deep networks; even when we observe clear performance gap between linear and non-linear models. Surprisingly, neural networks exhibit a strong non-linear advantage on these tasks, even though these NADs were explicitly constructed be well-aligned with the linear models.
On the other hand, the fact that the empirical NTK of standard networks presents such strong direc- tional bias is remarkable its own, as it reveals important structure of the underlying architectures. Interestingly, recent theoretical studies 33 have found that the standard rotational invariance assump- tion in kernel theory (3) might be too restrictive explain generalization in many settings. Hence, showing that the kernels of neural networks have a strong rotational variance, clearly strengthens the link between the study of these kernels and deep learning theory.
Overall, our results explain why previous heuristics that used the NTK to rank the complexity were successful in doing so. Specifically, by systematically evaluating the performance of neural networks on tasks of increasing complexity for their linearized approximations, we have observed that the non-linear dynamics these networks do not change the way in which they sort the complexity of these problems. However, we should not forget that the differences in performance between the neural networks and their approximation are very significant and whether they favor or not neural networks depends on the task.
of learning certain tasks [10
4 Sources of the non-linear (dis)advantage
In this section, we study in more detail the mechanisms that separate neural networks from their linear approximations and that lead to their significant performance differences. Specifically, we will show that there are important nuances involved in the comparison of linear and non-linear models, which depend on number of training samples, architecture and target task.
To shed more light on these complex relations, we will conduct fine-grained analysis of the dynamics of neural networks and study the evolution of their empirical NTK. We will show that the kernel dynamics can explain why neural networks converge much faster than kernel methods, even though this rapid adaptation can sometimes be imperfect and lead the networks overfit.
4.1 The non-linear advantage depends on the sample size
As we have seen, there exist multiple problems in which neural networks perform significantly better than their linearized approximations (see Sec. 3.2). but also others where they do not (see Sec. S.c. We now show, however, that the magnitude of these differences is influenced by the training set size.
We can illustrate this phenomenon by training several neural networks to predict some semantically- meaningful labels of an image dataset. In particular, and for consistency with the previous two-class
Fuull details of the experiment be found in the Appendix.

---

Aucuracy
MLP
Accuracy
LeNet
Avcuneyy
ResNetll
80.0%
05.056
70.0%
70.0%
80.0%
60.075
60.0%
10P
10
62 diss.
10'
10
6g die.
10P
10
4000
4000
2000
2000
10P
10
10
10
Number training somples
Number traaning samples
Number training somples
Neural network
Linsarized model
Figure 5: Comparison of test accuracy (top) and parameter distance to initialization (bottom) between neural networks and their linear approximations trained on CIFAR2 with different training set sizes. Plots show average over five different random seeds and when the test accuracy has saturated.
examples, deal with binary version of CIFARIO and assign label +1 all samples from the first five classes in the dataset, and label 1 to the rest. We will refer to this dataset as CIFAR2. Indeed, as seen in Fig. some neural networks exhibit a large non-linear advantage this task, but this advantage mostly appears when training on larger datasets. This phenomenon suggests that the inductive bias that boosts neural networks' performance involves important element of scale.
One can intuitively understand this behavior by analyzing the distance traveled by the parameters during optimization (see bottom row of Fig. 5) Indeed, for smaller training set sizes, the networks can find solutions that fit the training data closer to their initialization more easily As a result, the error incurred by the linear approximation in these cases is smaller. This explain why there significant performance gaps between NTK-based models and neural networks for small-data (16). and it also highlights the strength of the linear approximation in this regime.
4.2
The kernel rotates in a single axis
So far have mostly analyzed results dealing with linear expansions around the weight initialization 0.. However, recent empirical studies have argued that linearizing at later stages of training induces smaller approximation errors (21]. suggesting that the NTK dynamies better explain the final training behavior of a neural network. To the best of our knowledge, this phenomenon is still poorly understood, mostly because hinges on understanding the non-linear dynamics of deep networks. We now show, however, that understanding the way the spectrum of the NTK evolves during training can provide important insights into these dynamics.
To that end, we first analyze the evolution of the principal components of the empirical NTK in relation the target function. Specifically, let C denote the matrix of first K eigenvectors of the Gram matrix θ, obtained by linearizing the network after epochs of training, and let E R be the vector of training labels. In Sec. 3.1 we have seen that both neural networks and their linear approximations perform better on targets aligned with the first eigenfunctions of Oo- We propose, vis.ylle therefore, track the energy concentration ||<b.v || /////|l/ ///l of the labels the first eigenfunctions
with the aim to identify a significant transformation of the principal eigenspace (4)).
Fig. shows the result of this procedure applied to CNN trained to classify CIFAR2. Strikingly, the amount of energy of the target function that is concentrated on the K = 50 first cigenfunctions the NTK significantly grows during training. This is heavily mon-linear phenomenon by definition
'Note that linear and non-linear models achieve their maximum test accuracy in roughly the same number of epochs regardless the training set size. This contrasts with dynamics the training loss, which linear models take significantly more iterations to perfectly fit the training data than non-lineat ones. In this sense, the results of Fig. present a snapshot taken when both networks approximately achieve their maximum generalization performance.

---

Energy cone.
80%
(c))
Luttislianionn
Crpu
Bpach
60%
10
10*
10'
40%
10a
10
50
150
200
500
1000
1000
Epoch:
Eigenfunction index
Eigenfunction index
Figure 6: Energy concentration of 10, 000 CIFAR2 training labels on the first K = 50 eigenvectors of the ker- nel Gram matrices 4 of a LeNet.
Figure 7: Alignment of the eigenfunctions of the NTK initialization for a LeNet with the NTKs computed at the beginning (left) and at the end of 100 epochs of training (right) to predict the 500th eigenfunction at initialization.
the linearized models have fixed kernel and hinges on a dynamical realignment of O during training. That is, training a neural network rotates θ in a way that increases | 12/ly/12.
Prior work has also observed similar phenomenon, albeit in more restricted experimental setup: These observations have been confirmed only in a few datasets, and required the use of minibatches to approximate the alignment of the NTK to track the evolution of small portion of the eigenspace (22 24j. However, we now show that that the kernel rotation is prevalent across training setups, and that can also be observed when training solve other problems. In fact, more fine-grained inspection reveals that the kernel rotation mostly happens in a single functional axis. It maximizes the alignment of ① with the target function i.e., as || [f (=) fe. (x)] l1.. but does not greatly affect the rest of the spectrum. Indeed, we can see how during training a grows significantly more for the target than for any other function.
This is clearly illustrated in Fig. / where we compare ag and 0ruu for the first 1, 000 eigenfunctions of $000 (x) (z)). Strikingly, Oo {oy ) 11,000 when training to predict and arbitrary cigenfunction x + sign( we can see that after training to predict $500- ((0500) increases much more than (())) ) for any other eigenfunction Φj. In fact, the relative alignment between all other eigenfunctions does not change much. Note, however, that the absolute values of all alignments have also grown, a phenomenon which is due general increase in the Jacobian norm E. V oooo (a:) ll2 during training.
These results show that there is great potential in using linear approximations to investigate important deep learning phenomena involving pretrained networks as in [10 17). Indeed, (f()) is higher at the end of training. can be expected that a linear expansion at this late stage will be able to capture the inductive bias needed to fine-tune geometry of the NTK RKHS H. Note that in H. fl 0 Onl12 || [3]: and recall that, indicated by Lemma [ target functions with small I also have high (f)) This means that in H those tasks with a high af)) are indeed represented by weights closer the origin of the linearization, which thus makes the approximation error of the linearization smaller for these targets as observed in [21]
targets similar to f. The intuition behind this lies in the
4.3 Kernel rotation improves speed of convergence, but can hurt generalization
The rotation of O during training is an important mechanism that explains why the NTK dynamics can better capture the behavior of neural networks at the end of training- However, is also fundamental in explaining the ability of neural networks to quickly fit the training data. Specifically, important to highlight the stark contrast, at a dynamical level, between linear models and neural networks. Indeed, we have consistently observed across our experiments that neural networks converge much faster to solutions with a near-zero training loss than their linear approximations.
We can explain the influence of the rotation of the NTK on this phenomenon through a simple experiment. In particular, we train three different models to predict another arbitrary eigenfunction + sign( sign(daooof (=)0(=))) (=)): i) ResNet18, ii) its linearization around initialization, and iii) an unbiased linearization around the solution of the ResNet18 0", i.e., fo(x: (0 0u) Ve (æ:).
Fig. g compares the dynamics of these models, revealing that the neural network indeed converges much faster than its linear approximation. We see, however, that the kernelized model constructed
Recall that (0)))) Ay- Similar results for other networks and tasks can be found in the Appendix

---

Train Iress
1.00
0.75
Linearized model (pretrriine)) Acc. 57.8%
0.50
Neural network- Acc. 64.7%
Linearized model (init) Acc. 92.8%
0.25
0.00
20
60
80
100
Epochs
Figure 8: (Left) Evolution of training loss while learning sign()100(z)) for ResNet18 and two linearized models based on O and 100 of the ResNet18. (Right) Test accuracy on + (igo (2)) for the three models.
using the pretrained NTK of the network has also a faster convergence, We conclude, therefore, that, since the difference between the two linear models only lies on the kernel they use, it is indeed the transformation of the kernel through the non-linear dynamics of the neural network that makes these models converge so quickly. Note, however, that the rapid adaptation of θ to the training labels can have heavy toll generalization, i.e., the model based on the pretrained kernel converges much faster than the randomly initialized one, but has much lower test accuracy (comparable to the one of the neural network).
The dynamics of the kernel rotation are very fast, and we observe that the NTK overfits to the training labels in just few iterations. Fig. iillustrates this process. where we see that the performance of the linearized models with kernels extracted after few epochs of non-linear pretraining decays very rapidly. This observation is analogous the one presented in (2T). although in the opposite direction. Indeed, instead of showing that pretraining can greatly improve the performance of the linarized networks, Fig. shows that when the training task does not abide by the inductive bias of the network, pretraining can rapidly degrade the linear network performance. On the other hand, on CIFAR2 (see Sec. 4.T the kernel rotation does greatly improve test accuracy for some models.
The fact that the non-linear dynamics can both boost and hurt generalization highlights that the kernel rotation is subject its own form of inductive bias. In this sense, we believe that explaining the non-trivial coupling between the kernel rotation, aligmment, and training dynamics is an important avenue for future research, which will allow to better understand deep networks.
MLP
LeNer
BesNet18
90.0%
90.0%
90.0%
80.055
85.0%l
70.0%5
80.0%
80.0%
60.0
201
Pretzaining epoch
Pretraining epoci
Pretmnining upoch
Linsarised model
Figure 9: Performance of linearized networks with kernels extracted after different epochs of pretraining of a non-linear network learning æ + sign (c400(1))) The dashed line represents the performance of fully non-linear training on the same task, and the value epoch 0 corresponds the linearized model at initialization.
5 Final remarks
Explaining generalization in deep learning [ has been the subject of extensive research in recent years [34-38 The NTK theory is part of this trend, and it has been used to prove convergence and generalization very wide networks 4-). Most these efforts, however, still cannot explain the behavior of the models used in practice. Therefore, multiple authors have proposed to study neural networks empirically, with the aim to identify novel deep learning phenomena which can be later
explained theoretically 21 41.

---

In this work, we have followed a similar approach, and presented a systematic study comparing the behavior of neural networks and their linear approximations on different tasks. Previous studies had shown there exist tasks that neural networks can solve but kernels cannot [18 Our work complements those results, and provides examples of tasks where kernels perform better than neural networks (see I. We see this result as an important milestone for deep learning theory, as it shifts the focus of our research from asking "why are non-linear networks better than kernel methods?" to "what is special about our standard tasks which makes neural networks adapt so well to them?". Moving forward, knowing which tasks a neural network can and cannot solve efficiently will be fundamental to explain their inductive bias [22]
Our findings complement the work in (21 where the authors also empirically compared neural networks and their linear approximations, but focused on a single training task. In this work, we have built from these observations and studied models on a more diverse set of problems. Doing so, we have shown that the alignment with the empirical NTTK can rank the learning complexity of certain tasks for neural networks, despite being agnostic to the non-linear (dis)advantages. In this sense, we have revealed that important factors such as sample size, architecture, target task can greatly influence the gap between kernelized models and neural networks.
Finally, our dynamical study of the kernel rotation complements the work of 22 and provides new important insights for future research. For example, the fact that the kernel rotates in a single axis, and that its tendency to overfit is accurately predicted by the NTK eigenfunctions can aid in the development of new models of training. Moreover, it opens the door to new algorithmic developments that could slow down the kernel rotation and potentially reduce overfitting on many tasks.
Overall, our work paves the way for new research avenues based deep learning theory, while it also provides actionable insights to use the NTK approximation in many applications.
Acknowledgements
We thank Emmanuel Abbé, Alessandro Favero, Apostolos Modas and Clément Vignac for their fruitful discussions and feedback.
Funding disclosure
This work has been partially supported by Google via a Postdoctoral Fellowship and GCP Research Credit Award.
References
[1] C. Zhang, S. Bengio, M. Hardt, B. Recht, and O. Vinyals, "Understanding deep learning requires rethinking generalization," in International Conference on Learning Representations (ICLR), 2017.
[2] A. Jacot, C. Hongler, and F. Gabriel, "Neural tangent kernel: Convergence and generalization in neural networks," in Advances in Neural Information Processing Systems (NeurIPS), 2018.
[3] B. Schölkopf and A. J. Smola, Learning with Kernels: support vector machines, regularization, optimization, and beyond. Adaptive computation and machine learning series, MIT Press, 2002.
[4] J. Lee, L. Xiao, S. S. Schoenholz, Y. Bahri, R. Novak, J. Sohl-Dickstein, and J. Pennington, "Wide neural networks of any depth evolve as linear models under gradient descent," in Advances in Neural Information Processing Systems (NeurIPS), 2019.
[5] S. S. Du, J. D. Lee, H. Li, L. Wang, and X. Zhai, "Gradient descent finds global minima of deep neural networks," " in International Conferenece on Machine learning (ICML), 2019.
[6] S. Arora, S. Du, W. Hu, Z. Li, R. Salakhutdinov, and R. Wang. "On exact computation with an infinitely wide neural net," in Advances in Neural Information Processing Systems (NeurIPS), 2019.
[7] A. Bietti and Mairal, "On the inductive bias of neural tangent kernels," Advances in Neural Information Processing Systems (NeurIPS), 2019.
10

---

[8] D. Zou and Q. Gu, "An improved analysis of training over- parammeetteee deep neural networks." in Advances in Neural Information Processing Systems (NeurIPS). 2019.
[9] C. Liu, Zhu, and M. Belkin, "On the linearity of large non-linear models: when and why the tangent kernel is constant," in Advances in Neural Information Processing Systems (NeurIPS), 2020.
[10] A. Deshpande, A. Achille, A. Ravichandran, H. Li, L. Zancato, C. Fowlkes, R. Bhotika, S. Soatto, and P. Perona, "A linearized framework and new benchmark for model selection for fine-tuning," arXiv:2102.0008: [cs], 2021.
[11] L. Zancato, A. Achille, A. Ravichandran, R. Bhotika, and S. Soatto, "Predicting training time without training." " in Advances Neural Information Processing Systems (NeurIPS), 2020.
[12] H. Mobahi, M. Farajtabar, and P. L. Bartlett, Sesfl-Disiillatinn Amplifies Regularization in Hilbert Space," in Advances in Neural Information Processing Systems (NeurIPS), 2020.
[13] M. Tancik, P. P. Srinivasan, B. Mildenhall, S. Fridovich-Keil, N. Raghavan, U. Singhal, R. Ra- mamoorthi, J. T. Barron, and R. Ng. "Fourier features let networks learn high frequency functions in low dimensional domains," in Advances in Neural Information Processing Systems (NeurIPS), 2020.
[14) T. Gebhart, U. Saxena, and P. Schrater, "A Unified Paths Perspective for Pruning at Initialization," 10255 2021.
[15] G. Bachmann, S. Moosavi-Dezfooli, and T. Hofmann, "Uniform convergence, adversarial spheres and a simple remedy," arxiv:2105.0349) 2021.
[16] S. Arora, S. S. Du, Z. Li, R. Salakhutdinov, R. Wang, and D. Yu, "Harnessing the power of infinitely wide deep nets on small-data tasks," in International Conference on Learning Representations (ICLR), 2020.
[17] W. Maddox, S. Tang, P. G. Moreno, A. G. Wilson, and A. Damianou, "Fast adaptation with linearized neural networks," in International Conference on Artificial Intelligence and Statistics (AISTATS). 2021.
[18] 7. Allen-Zhu and Y. Li, "What can resnet leam efficiently, going beyond kernels?," in Advances in Neural Information Processing Systems (NeurIPS), 2019.
[19] B. Ghorbani, S. Mei, T. Misiakiewicz, and A. Montanari, "When do neural networks outperform kernel methods?," in Advances in Neural Information Processing Systems (NeurIPS), 2020.
[20] E. Malach, P. Kamath, E. Abbe, and N. Srebro, "Quantifying the Benefit of Using Differentiable Learning over Tangent Kernels," arXiv:2103.00 002010 [cs, stat], 2021.
[21] S. Fort, G. K. Dziugaite, M. Paul, S. Kharaghani, D. M. Roy, and S. Ganguli, "Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel," in Advances in Neural Information Processing Systems (NeurIPS), 2020.
[22] D. Kopitkov and V. Indelman, "Neural spectrum aligmment: Empirical study," in International Conference on Artificial Neural Networks (ICANN), 2020.
[23] Paccolat, I.. Petrini, M. Geiger, K. Tyloo, and M. Wyart, "Geometric compression of invariant manifolds in neural nets," Journal of Statistical Mechanics: Theory and Experiment, no. 4, 2021.
[24] A. Baratin, T. George, C. Laurent, R. D. Hjelm, G. Lajoie, P. Vincent, and S. Lacoste-Julic "Implicit regularization via neural feature alignment," in International Conference Artificial Intelligence and Statistics (AISTATS), 2021.
[25] G. Ortiz-Jimenez, A. Modas, S.-M. Moosavi-Dezfooli, and P. Frossard, "Neural Anisotropy Directions," in Advances in Neural Information Processing Systems 34 (NeurIPS), 2020.
[26] L. Bartlett and S. Mendelson, "Rademacher and gaussian complexities: Risk bounds and structural results," in Computational Learning Theory (COLT), 2001.
[27] R. Novak, L. Xiao, J. Hron, J. Lee, A. A. Alemi, J. "Neural tangents: Fast and easy infinite neural networks in python," International Conference on Learning Representations (ICLR), 2020.
and S. S. Schoenholz,

---

[28] J. Bradbury, R. Frostig, P. Hawkins, M. Johnson, C. Leary, Maaclurrin, G. Necula, A. Paszke, J. VanderPlas, Wanderman-Milnn and Q. Zhang, "JAX: composable transformations of Python+Nummyy programs,"" 2018.
[29] A. Krizhevsky, "Learning Multiple Layers of Features from Tiny Images," tech. rep., University of Toronto, 2009.
[30] F. Rosenblatt, "The Perceptron: A Probabilistic Model for Information Storage and Organization in The Brain," " Psychological Review, no. 6, 1958.
[31] Y. Lecun, L.. Bottou, Y. Bengio, and Haffner, "Gradient-based learning applied to document recognition," Proceedings of the IEEE, no. 11, 1998.
[32] K. He, X. Zhang, S. Ren, and J. Sun, "Deep residual learning for image recognition." in 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, 2016.
[33] K. Donhauser, M. Wu, and Yang, "How rotational invariance of common kernels prevents generalization in high dimensions," in International Conference Machine Learning (ICML), 2021.
[34] B. Neyshabur, Z. Li, S. Bhojanapalli, LeCun, and N. Srebro, "The role of over-parametrization in generalization of neural networks," in International Conference on Learning Representations (ICLR), 2019.
[35] D. Soudry, Hoffer, M. S. Nacson, and N. Srebro, "The implicit bias of gradient descent on separable data," in International Conference on Learning Representations (ICLR), 2018.
[36] S. Gunasekar, J. D. Lee, D. Soudry, and N. Srebro, "Characteriving implicit bias in terms of optimization geometry," in International Conferenece on Machine learning (ICML), 2018.
[37] P. I.. Bartlett, N. Harvey, C. Liaw, and A. Mehrabian, "Nearly-tight VC-dimension and pseudodi- mension bounds for piecewise linear neural networks," Journal Machine Learning Research, vol. 20, pp. 63:1-63: 2019.
[38] E. Abbe and C. Sandon, "On the universality of deep learning," in Advances in Neural Informa- tion Processing Systems (NeurIPS), 2020.
[39] V. Nagarajan and Z. Kolter, "Uniform convergence may be unable explain generalization in deep learning," in Advances in Neural Information Processing Systems (NeurIPS). 2019.
[40] Nakkiran, B. Neyshabur, and H. Sedghi, "The deep bootstrap framework: Good online learners are good offline generalizers," in International Conference Learning Representations (ICLR), 2021.
[41] H. Maennel, M. Alabdulmohsin, I. O. Tolstikhin, R. J. Baldock, O. Bousquet, S. Gelly, and D. Keysers, "What do neural networks learn when trained with random labels?," in Advances in Neural Information Processing Systems (NeurIPS), 2020.
[42] D. H. Wolpert, "The lack of A priori distinctions between learning algorithms," Neural Compu- tation, vol. 8, no. 7, 1996.
12

---

Paper checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See discussion at the end of
Sec. and Sec. as well as our Final remarks.
(c) Did you discuss any potential negative societal impacts of your work? [N/A] Our paper deals with fundamental questions regarding our understanding of deep networks. In this sense, subject to the same ethical concerns as the machine learning field as a whole, which makes it hard to identify potential direct risks or benefits associated to our empirical findings.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes] See Appendix.
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main ex- perimental results (either in the supplemental material or URL)? [Yes] The code to reproduce our experiments can be found at lhttps:/// //gi thub. com wwt./ooooooo.oe gortizii/
1 inearized-netwwrkss
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] Full details of all our experiments can be found in the Appendix.
(c) Did you report error bars (c.g.. with respect the random seed after running experi- ments multiple times)? [No] Computing the kernel Gram matrix for different random seeds is computationally expensive procedure (see Appendix) which with limited resources would have taken too much time to compute for all the settings in our experi- ments. However, because none of our claims depend precise metrics, and rather supported by clear trends which we validated on different settings, we are confident that our results reproducible.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] Seep Appendix.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you deseribe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
13


Reviews:
[{"review_id": "vLV_0rvm27O", "reviewer": "Reviewer_NLbM", "summary": "Disclaimer: I'm only superficially familiar with the related literature, so this is rather low-confidence review.\n\nThe paper investigates whether the study of linearized neural networks via the neural tangent kernel is insightful for understanding generalization in deep learning. It argues that this is indeed the case based primarily on empirical evidence. Overall I am not fully convinced by the significance of the results in this paper, so that I am leaning towards a reject, but I would expect the other reviews to be more informative.", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, I think this is a polished and well-written paper. Considering the growing body of recent works on the NTK the topic of the submission is clearly of interest to the community. In particular the paper attempts to address recent concerns over the NTK not being informative regarding the generalization capabilities of finite-width neural networks.\n\nTo this end, the paper presents the following results:\n- Based on labels generated from the eigenfunctions of the NTK, both linearized and finite-width NNs decrease in performance with increasing complexity of the learning taks. Further it observes that similarly on the recently proposed neural anisotropy directions performance decreases concurrently with increasing task complexity. These results complement previous works which have used the NTK to predict task complexities and gives them a slightly stronger theoretical foundation. However, even though the experiments are based on carefully chosen synthetic labels, the support is primarily empirical, which seems somewhat incremental and not nearly as strong as any theoretical guarantees.\n- Next, the paper finds that neural networks primarily outperform their linearized counterparts in the large data regimes and argues that this is due to the distance of the optimum from initialization, which impacts the accuracy of the linearization approximation. It also notes that the alignment of the labels with the top eigenfunctions of the NTK increases through training.\n\nWhile the results in this paper seem to fill some gaps in the literature, they feel somewhat incremental and mostly appear to add to previous observations rather than making novel ones of their own. Given this incremental nature I find that the predominantly empirical analysis lacks significance. If the authors think that this understates their contribution, I would encourage them to clarify what avenues of future research their work opens up specifically. The discussion at the end remains rather vague on this and (to my taste) focuses too much on summarizing the paper.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "n/a", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2.5"}, "extra_text": {}}, {"review_id": "HnnBi-MPq1Q", "reviewer": "Reviewer_rfX8", "summary": "This paper studies the training dynamics of both neural networks and their linear approximations, and show that the finite-width NTK aligns to the target function during training, which in turn induces rapid convergence to a particular solution.", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "While this submission is well-written and covers an interesting avenue of research, most of its content has previously been published in a AISTAT2021 paper: Baratin et al., namely:\n\n - the dynamical bias of linear models towards functions that are aligned with the kernel (which is a standard result)\n - the increase in alignment of the NTK to the target during training of a neural network\n - the consequence of this increase in alignment on (linear) generalization guarantees\n\nTherefore, depending on the NeurIPS policy regarding prior work, it looks to me that this work cannot be considered an original contribution. I am putting a low score only for that reason.\n\nA comment regarding the experiment in section 4.1/fig 5: As you increase the training set size while keeping the minibatch size fixed, you also increase the number of minibatches in an epoch, thus there are more training iterations for the same number of epochs. Given that, I am not so surprised by the increasing l2 distance to initialization. Do these different behaviours between linear and non-linear dynamics occur because of the increasing number of examples, or because of the increasing number of training iterations?\n\n\nReference:\nBaratin, George, Laurent, Hjelm, Lajoie, Vincent, and Lacoste-Julien, Implicit regularization via Neural Feature Alignment, AISTATS2021 https://arxiv.org/abs/2008.00938", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1.5"}, "extra_text": {}}, {"review_id": "FHulDm7Gkz", "reviewer": "Reviewer_4dYZ", "summary": "The submitted paper studies the differences in the behavior of typical neural networks and the predictions made by the Neural Tangent Kernel (NTK) framework by showing that the NTK evolves and adapts to the current task during training. The experiments show that the predictions of NTK framework correctly rank tasks in terms of hardness but over- and under-estimate their hardness in different cases. The paper then studies the role of sample size and the directions where NTK evolves.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The main finding of the paper, that NTK evolves through the training to align its top eigenvectors to the task, has already been shown in recent work. There are still new and interesting results about how NTK evolves and its potential for overfitting and I'm leaning towards acceptance.\n\nAdaptation of NTK to the task has been recently shown by [1] (also see [2] for an empirical study). Those papers also talk about benefits of alignment for convergence and generalization. I have not read those two papers in depth and I'm happy to discuss if the authors believe the submitted paper is showing a different phenomenon.\n\nThe most interesting result to me was Figure 8 that shows NTK alignment can hurt generalization. While some previous works describe NTK alignment as a sort of implicit regularizer and also the generalization bound in Eq (2) suggests that alignment helps generalization, Figure 8 shows that the story is more complicated and allowing the network to evolve its kernel can hurt generalization. Do the authors have an explanation for this apparent paradox? I suspect that in small data regime the NTK itself overfits to the training data so that its alignment increases on train data but stays constant (or maybe even decreases) on test data. A plot like Figure 6 but evaluated on both train and test data and in a setting where NTK evolution hurts generalization can clear this up.\n\nAnother interesting finding is Figure 7 (and the replications in the appendix) where NTK alignment happens mostly in a single direction, although I'm not sure how this would extend to more general cases. For example, if the target function was aligned with two eigenvectors of the initial NTK, would the evolution be in those two directions or still in the direction of one eigenvector?\n\nThe experiments highlight the role of alignment on convergence while the theory in the introduction only motivates alignment with a generalization bound. Bringing an alignment-dependent convergence analysis beside the generalization bound can help with motivating alignment. An example is Theorem 4.1 in [3] (that paper is about ranking of tasks in terms of hardness and does not study evolution of NTK).\n\nMinor comments:\n\nIn some experiments like section 3.1 the eigenvector has to be binarized to be used as training labels. Any idea how much this binarization changes the direction of this vector?\n\nIt looks to me that the curve in Figure 7 (left) becomes equal to lambda^2. Mentioning this in the text can help with understanding the plot. For the curve on the right, it is good to mention that phi_i still refers to eigenvectors of the initial NTK if this is the case.\n\nThe choice of K=50 for Figure 6 looks arbitrary and the result might be different for another K. Can one use the scalar measure in Lemma 1 to get rid of K?\n\n\n[1] https://arxiv.org/abs/2008.00938\n\n[2] https://arxiv.org/abs/1910.08720\n\n[3] https://arxiv.org/abs/1901.08584\n\nUpdate: increased the score as the comment about related work is addressed.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors address limitations and societal impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "10"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper explores what we can learn about neural network generalization through linearization, such as the empirical NTK. There were a number of supportive sentiments, but also concerns, discussed at length in the rebuttal period. One of the key concerns was novelty with respect to Baratin et. al. After discussion, it was understood that there are still novel contributions, such as the relative complexity of tasks being predicted with the empirical NTK, some experiments demonstrating the superiority of linearity model, and how alignment can hurt generalization. The final version of this paper should _carefully_ discuss Baratin et. al, and other related work, and address all of the reviewer questions, including experiments that were part of the rebuttal in the paper.

Decision: Accept (Poster)