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: This paper identifies a phenomenon called the "silent alignment effect," which happens in linear networks as well as ReLU networks on whitened data, with sufficiently small initialization of the network weights. In such a regime, the NTK of the network changes its eigenvectors to align with the problem structure at the beginning while remaining small in scale; then, the kernel essentially stops changing its eigenstructure and just grows in scale, eventually leading to a kernel regression solution wrt the NTK after the initial phase. The paper also demonstrates that non-whitened data can weaken this effect, for which the NTK needs to evolve later in training.