Neural Networks as Kernel Learners: The Silent Alignment Effect

Neural networks in the lazy training regime converge to kernel machines . Can neural networks in the rich feature learning regime learn a kernel machine with a data-dependent kernel ? We demonstrate that this can indeed happen due to a phenomenon we term silent alignment , which requires that the tangent kernel of a network evolves in eigenstructure while small and before the loss appreciably decreases , and grows only in overall scale afterwards . We show that such an effect takes place in homogenous neural networks with small initialization and whitened data . We provide an analytical treatment of this effect in the linear network case . In general , we find that the kernel develops a low-rank contribution in the early phase of training , and then evolves in overall scale , yielding a function equivalent to a kernel regression solution with the final network ’ s tangent kernel . The early spectral learning of the kernel depends on both depth and on relative learning rates in each layer . We also demonstrate that non-whitened data can weaken the silent alignment effect . 1 INTRODUCTION . Despite the numerous empirical successes of deep learning , much of the underlying theory remains poorly understood . One promising direction forward to an interpretable account of deep learning is in the study of the relationship between deep neural networks and kernel machines . Several studies in recent years have shown that gradient flow on infinitely wide neural networks with a certain parameterization gives rise to linearized dynamics in parameter space ( Lee et al. , 2019 ; Liu et al. , 2020 ) and consequently a kernel regression solution with a kernel known as the neural tangent kernel ( NTK ) in function space ( Jacot et al. , 2018 ; Arora et al. , 2019 ) . Kernel machines enjoy firmer theoretical footing than deep neural networks , which allows one to accurately study their training and generalization ( Rasmussen & Williams , 2006 ; Schölkopf & Smola , 2002 ) . Moreover , they share many of the phenomena that overparameterized neural networks exhibit , such as interpolating the training data ( Zhang et al. , 2017 ; Liang & Rakhlin , 2018 ; Belkin et al. , 2018 ) . However , the exact equivalence between neural networks and kernel machines breaks for finite width networks . Further , the regime with approximately static kernel , also referred to as the lazy training regime ( Chizat et al. , 2019 ) , can not account for the ability of deep networks to adapt their internal representations to the structure of the data , a phenomenon widely believed to be crucial to their success . In this present study , we pursue an alternative perspective on the NTK , and ask whether a neural network with an NTK that changes significantly during training can ever be a kernel machine for a data-dependent kernel : i.e . does there exist a kernel function K for which the final neural network function f is f ( x ) ≈ ∑P µ=1 α µK ( x , xµ ) with coefficients αµ that depend only on the training data ? We answer in the affirmative : that a large class of neural networks at small initialization trained on approximately whitened data are accurately approximated as kernel regression solutions with their final , data-dependent NTKs up to an error dependent on initialization scale . Hence , our results provide a further concrete link between kernel machines and deep learning which , unlike the infinite width limit , allows for the kernel to be shaped by the data . The phenomenon we study consists of two training phases . In the first phase , the kernel starts off small in overall scale and quickly aligns its eigenvectors toward task-relevant directions . In the second phase , the kernel increases in overall scale , causing the network to learn a kernel regression solution with the final NTK . We call this phenomenon the silent alignment effect because the feature learning happens before the loss appreciably decreases . Our contributions are the following 1 . In Section 2 , we demonstrate the silent alignment effect by considering a simplified model where the kernel evolves while small and then subsequently increases only in scale . We theoretically show that if these conditions are met , the final neural network is a kernel machine that uses the final , data-dependent NTK . A proof is provided in Appendix B . 2 . In Section 3 , we provide an analysis of the NTK evolution of two layer linear neural networks with scalar target function with small initialization . If the input training data is whitened , the kernel aligns its eigenvectors towards the direction of the optimal linear function early on during training while the loss does not decrease appreciably . After this , the kernel changes in scale only , showing this setup satisfies the requirements for silent alignment discussed in Section 2 . 3 . In Section 4 , we extend our analysis to deep linear networks by showing that the time required for alignment scales with initialization the same way as the time for the loss to decrease appreciably . Still , these time scales may be sufficiently separated to lead to the silent alignment effect for which we provide empirical evidence . We further present an explicit formula for the final kernel in linear networks of any depth and width when trained from small initialization , showing that the final NTK aligns to task-relevant directions . 4 . In Section 5 , we show empirically that the silent alignment phenomenon carries over to nonlinear networks trained with ReLU and Tanh activations on isotropic data , as well as linear and nonlinear networks with multiple output classes . For anisotropic data , we show that the NTK must necessarily change its eigenvectors when the loss is significantly decreasing , destroying the silent alignment phenomenon . In these cases , the final neural network output deviates from a kernel machine that uses the final NTK . 1.1 RELATED WORKS . Jacot et al . ( 2018 ) demonstrated that infinitely wide neural networks initialized with an appropriate parameterization and trained on mean square error loss evolve their predictions as a linear dynamical system with the NTK at initalization . A limitation of this kernel regime is that the neural network internal representations and the kernel function do not evolve during training . Conditions under which such lazy training can happen is studied further in ( Chizat et al. , 2019 ; Liu et al. , 2020 ) . Domingos ( 2020 ) recently showed that every model , including neural networks , trained with gradient descent leads to a kernel model with a path kernel and coefficients αµ that depend on the test point x . This dependence on x makes the construction not a kernel method in the traditional sense that we pursue here ( see Remark 1 in ( Domingos , 2020 ) ) . Phenomenological studies and models of kernel evolution have been recently invoked to gain insight into the difference between lazy and feature learning regimes of neural networks . These include analysis of NTK dynamics which revealed that the NTK in the feature learning regime aligns its eigenvectors to the labels throughout training , causing non-linear prediction dynamics ( Fort et al. , 2020 ; Baratin et al. , 2021 ; Shan & Bordelon , 2021 ; Woodworth et al. , 2020 ; Chen et al. , 2020 ; Geiger et al. , 2021 ; Bai et al. , 2020 ) . Experiments have shown that lazy learning can be faster but less robust than feature learning ( Flesch et al. , 2021 ) and that the generalization advantage that feature learning provides to the final predictor is heavily task and architecture dependent ( Lee et al. , 2020 ) . Fort et al . ( 2020 ) found that networks can undergo a rapid change of kernel early on in training after which the network ’ s output function is well-approximated by a kernel method with a data-dependent NTK . Our findings are consistent with these results . Stöger & Soltanolkotabi ( 2021 ) recently obtained a similar multiple-phase training dynamics involving an early alignment phase followed by spectral learning and refinement phases in the setting of low-rank matrix recovery . Their results share qualitative similarities with our analysis of deep linear networks . The second phase after alignment , where the kernel ’ s eigenspectrum grows , was studied in linear networks in ( Jacot et al. , 2021 ) , where it is referred to as the saddle-to-saddle regime . Unlike prior works ( Dyer & Gur-Ari , 2020 ; Aitken & Gur-Ari , 2020 ; Andreassen & Dyer , 2020 ) , our results do not rely on perturbative expansions in network width . Also unlike the work of Saxe et al . ( 2014 ) , our solutions for the evolution of the kernel do not depend on choosing a specific set of initial conditions , but rather follow only from assumptions of small initialization and whitened data . 2 THE SILENT ALIGNMENT EFFECT AND APPROXIMATE KERNEL SOLUTION . Neural networks in the overparameterized regime can find many interpolators : the precise function that the network converges to is controlled by the time evolution of the NTK . As a concrete example , we will consider learning a scalar target function with mean square error loss through gradient flow . Let x ∈ RD represent an arbitrary input to the network f ( x ) and let { xµ , yµ } Pµ=1 be a supervised learning training set . Under gradient flow the parameters θ of the neural network will evolve , so the output function is time-dependent and we write this as f ( x , t ) . The evolution for the predictions of the network on a test point can be written in terms of the NTK K ( x , x′ , t ) = ∂f ( x , t ) ∂θ · ∂f ( x′ , t ) ∂θ as d dt f ( x , t ) = η ∑ µ K ( x , xµ , t ) ( yµ − f ( xµ , t ) ) , ( 1 ) where η is the learning rate . If one had access to the dynamics of K ( x , xµ , t ) throughout all t , one could solve for the final learned function f∗ with integrating factors under conditions discussed in Appendix A f∗ ( x ) = f0 ( x ) + ∑ µν ∫ ∞ 0 dt kt ( x ) µ [ exp ( −η ∫ t 0 Kt′ dt ′ ) ] µν ( yν − f0 ( xν ) ) . ( 2 ) Here , kt ( x ) µ = K ( x , xµ , t ) , [ Kt ] µ , ν = K ( xµ , xν , t ) , and yµ − f0 ( xµ ) is the initial error on point xµ . We see that the final function has contributions throughout the full training interval t ∈ ( 0 , ∞ ) . The seminal work by Jacot et al . ( 2018 ) considers an infinite-width limit of neural networks , where the kernel function Kt ( x , x′ ) stays constant throughout training time . In this setting where the kernel is constant and f0 ( xµ ) ≈ 0 , then we obtain a true kernel regression solution f ( x ) =∑ µ , ν k ( x ) µK−1µν y ν for a kernel K ( x , x′ ) which does not depend on the training data . Much less is known about what happens in the rich , feature learning regime of neural networks , where the kernel evolves significantly during time in a data-dependent manner . In this paper , we consider a setting where the initial kernel is small in scale , aligns its eigenfunctions early on during gradient descent , and then increases only in scale monotonically . As a concrete phenomenological model , consider depth L networks with homogenous activation functions with weights initialized with variance σ2 . At initialization K0 ( x , x′ ) ∼ O ( σ2L−2 ) , f0 ( x ) ∼ O ( σL ) ( see Appendix B ) . We further assume that after time τ , the kernel only evolves in scale in a constant direction K ( x , x′ , t ) = { σ2L−2K̃ ( x , x′ , t ) t ≤ τ g ( t ) K∞ ( x , x ′ ) t > τ , ( 3 ) where K̃ ( x , x′ , t ) evolves from an initial kernel at time t = 0 to K∞ ( x , x′ ) by t = τ and g ( t ) increases monotonically from σ2L−2 to 1 . In this model , one also obtains a kernel regression solution in the limit where σ → 0 with the final , rather than the initial kernel : f ( x ) = k∞ ( x ) ·K−1∞ y + O ( σL ) . We provide a proof of this in the Appendix B . The assumption that the kernel evolves early on in gradient descent before increasing only in scale may seem overly strict as a model of kernel evolution . However , we analytically show in Sections 3 and 4 that this can happen in deep linear networks initialized with small weights , and consequently that the final learned function is a kernel regression with the final NTK . Moreover , we show that for a linear network with small weight initialization , the final NTK depends on the training data in a universal and predictable way . We show empirically that our results carry over to nonlinear networks with ReLU and tanh activations under the condition that the data is whitened . For example , see Figure 1 , where we show the silent alignment effect on ReLU networks with whitened MNIST and CIFAR-10 images . We define alignment as the overlap between the kernel and the target function y > Ky ‖K‖F |y|2 , where y ∈ R P is a vector of the target values , quantifying the projection of the labels onto the kernel , as discussed in ( Cortes et al. , 2012 ) . This quantity increases early in training but quickly stabilizes around its asymptotic value before the loss decreases . Though Equation 2 was derived under assumption of gradient flow with constant learning rate , the underlying conclusions can hold in more realistic settings as well . In Figure 1 ( d ) and ( e ) we show learning dynamics and network predictions for Wide-ResNet ( Zagoruyko & Komodakis , 2017 ) on whitened CIFAR-10 trained with the Adam optimizer ( Kingma & Ba , 2014 ) with learning rate 10−5 , which exhibits silent alignment and strong correlation with the final NTK predictor . In the unwhitened setting , this effect is partially degraded , as we discuss in Section 5 and Appendix J . Our results suggest that the final NTK may be useful for analyzing generalization and transfer as we discuss for the linear case in Appendix F .

Summary: Kernel regimes have been a major workhorse in studying the training dynamics of deep neural networks. Previous work has shown that networks with large initialization scale and trained with gradient flow will converge to the kernel regression solution under a data-independent kernel, known as the NTK. However, it is now known that the practical training of neural networks often operates in a different regime called the "rich" regime, which makes it more difficult to analyze.   This paper bridges the gap between the two by pointing out that the rich regime can be split into two phases. The first phases is aligning the kernel spectrum with the training labels, which constructs a data-dependent kernel. The second phase only increases the scale of this kernel. This behavior is analytically shown on (deep) linear networks by writing out the dynamics of the NTK. Therefore, after the first phase, the neural network training will converge to the kernel regression solution of the adapted NTK. The same behavior, although not theoretically shown, are empirically verified on nonlinear ReLU nets.