What can linearized neural networks actually say about generalization?

1 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 [ 1 ] . Recently , a growing body of work has started to explore the use of linear approximations to analyze deep networks , leading to the neural tangent kernel ( NTK ) framework [ 2 ] . The NTK framework is based on the observation that for certain initialization schemes , the infinitewidth limit of many neural architectures can be exactly characterized using kernel tools [ 2 ] . This reduces key questions in deep learning theory to the study of linear methods and convex functional analysis , for which a rich set of theories exist [ 3 ] . This intuitive approach has been proved very fertile , leading to important results in generalization and optimization of very wide networks [ 4–9 ] . The NTK theory , however , can only fully describe certain infinitely wide neural networks , and for the narrow architectures used in practice , it only provides a first-order approximation of their training dynamics ( see Fig . 1 ) . Despite these limitations , the intuitiveness of the NTK , which allows to use a powerful set of theoretical tools to exploit it , has led to a rapid increase in the amount of research that successfully leverages the NTK in applications , such as predicting generalization [ 10 ] and training speed [ 11 ] , explaining certain inductive biases [ 12–15 ] or designing new classifiers [ 16 , 17 ] . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 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 firstorder Taylor expansion around the initialization of a neural network – perform much worse than the networks they approximate on standard image recognition benchmarks [ 21 ] ; a phenomenon that has been coined as the non-linear advantage . However , it 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 to explain the success of deep learning . 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 a 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 , we identify the alignment of the target function with the NTK as a 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 alignment 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 , we offer a 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 disadvantage compared to their kernel approximations . The main contributions of our work are : • We show that the alignment with the empirical NTK at initialization can provide a 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 to their non-linear dynamics . • We , hence , provide a 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 as 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 it also provides a fresh perspective to understand how to use the NTK approximation in several applications . 2 Preliminaries . Let fθ : Rd → R denote a neural network parameterized by a set of weights θ ∈ Rn . 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 : Rd → { ±1 } by training the weights to minimize the empirical risk R̂ ( fθ ) = 1m ∑m i=1 1sign ( fθ ( xi ) 6=f ( xi ) over a finite set of i.i.d . training data S = { ( xi , f ( xi ) ) } mi=1 from an underlying distribution D. We broadly say that a model generalizes whenever it achieves a low expected riskR ( fθ ) = Ex∼D [ 1sign ( fθ ( x ) 6=f ( x ) ] . In a small neighborhood around the weight initialization θ0 , a neural network can be approximated using a first-order Taylor expansion ( see Fig . 1 ) fθ ( x ) ≈ f̂θ ( x ; θ0 ) = fθ0 ( x ) + ( θ − θ0 ) > ∇θfθ0 ( x ) , ( 1 ) where∇θfθ0 ( x ) ∈ Rn denotes the Jacobian of the network with respect to the parameters evaluated at θ0 . Here , the model f̂θ represents a linearized network which maps weight vectors to functions living in a reproducible kernel Hilbert space ( RKHS ) H ⊆ L2 ( Rd ) , determined by the empirical NTK [ 4 ] at θ0 , Θθ0 ( x , x ′ ) = 〈∇θfθ0 ( x ) , ∇θfθ0 ( x′ ) 〉 . Unless stated otherwise , we will generally drop the dependency on θ0 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 ∇θfθt computed at a checkpoint θt ( see Fig . 1 ) . Remarkably , however , it was recently discovered that for certain types of initialization , and in the limit of infinite network-width , the approximation in ( 1 ) is exact , and the NTK is 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 R ( f ? ) ≤ R̂ ( f ? ) +O √‖f‖2ΘEx [ Θ ( x , x ) ] m , ( 2 ) where f ? = argminh∈H R̂ ( h ) + µ‖h‖2Θ , with µ > 0 denoting a regularization constant , and ‖f‖2Θ being the RKHS norm of the target function , which for positive definite kernels can be computed as ‖f‖2Θ = ∞∑ j=1 1 λj ( Ex∼D [ φj ( x ) f ( x ) ] ) 2 . ( 3 ) Here , the couples { ( λj , φj ) } ∞j=1 denote the eigenvalue-eigenfunction pairs , in order of decreasing eigenvalues , of the Mercer ’ s decomposition of the kernel , i.e. , Θ ( x , x′ ) = ∑∞ j=1 λjφj ( x ) φj ( x ′ ) . This means that , in kernel regimes , the difference between the empirical and the expected risk is smaller when training on target functions with a lower RKHS norm . That is , whose projection on the eigenfunctions 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 Θ. Estimating ( 3 ) in practice is challenging as it requires access to the smallest eigenvalues of the kernel . However , one can use the following lemma to compute a 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 α ( f ) = Ex , x′∼D [ f ( x ) Θ ( x , x′ ) f ( x′ ) ] denote the alignment of the target f ∈ H with the kernel Θ . Then ‖f‖2Θ ≥ ‖f‖42/α ( f ) . Moreover , for the NTK , α ( f ) = ‖Ex [ f ( x ) ∇θfθ0 ( x ) ] ‖ 2 2 . 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 a learning guarantee for linearized models , i.e . f̂θ . Furthermore , from now on , we will interchangeably use the terms NTK and empirical NTK to simply refer to the finite-width kernels derived from ( 1 ) . In our experiments , we compute those using the neural_tangents 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 φj ( x ) over a finite dataset . We will also use the terms eigenvector and eigenfunction interchangeably . In particular , we will use Φ ∈ Rm×m to denote the matrix containing the jth Gram eigenvector φj ∈ Rm in its jth row , where the rows are ordered according to the vector of decreasing eigenvalues λ ∈ Rm+ . 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 to shed light on the question : What can the empirical NTK actually predict about generalization ? To that end , we conduct a systematic study with different neural networks and their linearized approximations given by ( 1 ) , 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 .

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.