Year: 2021

Abstract:
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.

text: 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.

heading: 1 INTRODUCTION
text: 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), cannot 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.

heading: 1.1 RELATED WORKS
text: 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.

heading: 2 THE SILENT ALIGNMENT EFFECT AND APPROXIMATE KERNEL SOLUTION
text: 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.

heading: 3 KERNEL EVOLUTION IN 2 LAYER LINEAR NETWORKS
text: We will first study shallow linear networks trained with small initialization before providing analysis for deeper networks in Section 4. We will focus our discussion in this section on the scalar output case but we will provide similar analysis in the multiple output channel case in a subsequent section. We stress, though we have systematically tracked the error incurred at each step, we have focused on transparency over rigor in the following derivation. We demonstrate that our analytic solutions match empirical simulations in Appendix C.5.
We assume the P data points xµ ∈ RD, µ = 1, . . . , P of zero mean with correlation matrix Σ = 1P ∑P µ=1 x
µxµ>. Further, we assume that the target values are generated by a linear teacher function yµ = sβT · xµ for a unit vector βT . The scalar s merely quantifies the size of the supervised learning signal: the variance of |y|2 = s2β>T ΣβT . We define the two-layer linear neural network with N hidden units as f(x) = a>Wx. Concretely, we initialize the weights with standard parameterization ai ∼ N (0, σ2/N),Wij ∼ N (0, σ2/D). Understanding the role of σ in the dynamics will be crucial to our study. We analyze gradient flow dynamics on MSE cost L = 12P ∑ µ (f(x µ)− yµ)2.
Under gradient flow with learning rate η = 1, the weight matrices in each layer evolve as d
dt a = −∂L ∂a = WΣ
( sβT −W>a ) , d
dt W = − ∂L ∂W = a
( sβT −W>a )> Σ. (4)
The NTK takes the following form throughout training.
K(x,x′; t) = x>W>Wx′ + |a|2x>x′. (5) Note that while the second term, a simple isotropic linear kernel, does not reflect the nature of the learning task, the first term x>W>Wx′ can evolve to yield an anisotropic kernel that has learned a representation from the data.

heading: 3.1 PHASES OF TRAINING IN TWO LAYER LINEAR NETWORK
text: We next show that there are essentially two phases of training when training a two-layer linear network from small initialization on whitened-input data.
• Phase I: An alignment phase which occurs for t ∼ 1s . In this phase the weights align to their low rank structure and the kernel picks up a rank-one term of the form x>ββ>x′. In this setting, since the network is initialized nearW ,a = 0, which is a saddle point of the loss function, the gradient of the loss is small. Consequently, the magnitudes of the weights and kernel evolve slowly.
• Phase II: A data fitting phase which begins around t ∼ 1s log(sσ −2). In this phase, the system
escapes the initial saddle point W ,a = 0 and loss decreases to zero. In this setting both the kernel’s overall scale and the scale of the function f(x, t) increase substantially.
If Phase I and Phase II are well separated in time, which can be guaranteed by making σ small, then the final function solves a kernel interpolation problem for the NTK which is only sensitive to the geometry of gradients in the final basin of attraction. In fact, in the linear case, the kernel interpolation at every point along the gradient descent trajectory would give the final solution as we show in Appendix G. A visual summary of these phases is provided in Figure 2.

heading: 3.1.1 PHASE I: EARLY ALIGNMENT FOR SMALL INITIALIZATION
text: In this section we show how the kernel aligns to the correct eigenspace early in training. We focus on the whitened setting, where the data matrixX has all of its nonzero singular values equal. We let β represent the normalized component of βT in the span of the training data {xµ}. We will discuss general Σ in section 3.2. We approximate the dynamics early in training by recognizing that the network output is small due to the small initialization. Early on, the dynamics are given by:
d dt a = sWβ +O(σ3) , d dt W = saβ> +O(σ3). (6)
Truncating terms order σ3 and higher, we can solve for the kernel’s dynamics early on in training K(x,x′; t) = q0 cosh(2ηst) x > [ββ> + I]x′ +O(σ2), t s−1 log(s/σ2). (7)
where q0 is an initialization dependent quantity, see Appendix C.1. The bound on the error is obtained in Appendix C.2. We see that the kernel picks up a rank one-correction ββ> which points in the direction of the task vector β, indicating that the kernel evolves in a direction sensitive to the target function y = sβT · x. This term grows exponentially during the early stages of training, and overwhelms the original kernel K0 with timescale 1/s. Though the neural network has not yet achieved low loss in this phase, the alignment of the kernel and learned representation has consequences for the transfer ability of the network on correlated tasks as we show in Appendix F.

heading: 3.1.2 PHASE II: SPECTRAL LEARNING
text: We now assume that the weights have approached their low rank structure, as predicted from the previous analysis of Phase I dynamics, and study the subsequent NTK evolution. We will show that, under the assumption of whitening, the kernel only evolves in overall scale.
First, following (Fukumizu, 1998; Arora et al., 2018; Du et al., 2018), we note the following conservation law ddt [ a(t)a(t)> −W (t)W (t)> ] = 0 which holds for all time. If we assume small initial weight variance σ2, aa> −WW> = O(σ2) ≈ 0 at initialization, and stays that way during the training due to the conservation law. This condition is surprisingly informative, since it indicates thatW is rank-one up to O(σ) corrections. From the analysis of the alignment phase, we also have thatW>W ∝ ββ>. These two observations uniquely determine the rank one structure ofW to be aβ> +O(σ). Thus, from equation 5 it follows that in Phase II, the kernel evolution takes the form
K(x,x′; t) = u(t)2x> [ ββ> + I ] x′ +O(σ), (8)
where u(t)2 = |a|2. This demonstrates that the kernel only changes in overall scale during Phase II. Once the weights are aligned with this scheme, we can get an expression for the evolution of u(t)2 analytically, u(t)2 = se2st(e2st − 1 + s/u20)−1, using the results of (Fukumizu, 1998; Saxe et al., 2014) as we discuss in C.4. This is a sigmoidal curve which starts at u20 and approaches s. The transition time where active learning begins occurs when est ≈ s/u20 =⇒ t ≈ s−1 log(s/σ2). This analysis demonstrates that the kernel only evolves in scale during this second phase in training from the small initial value u20 ∼ O(σ2) to its asymptote. Hence, kernel evolution in this scenario is equivalent to the assumptions discussed in Section 2, with g(t) = u(t)2, showing that the final solution is well approximated by kernel regression with the final NTK. We stress that the timescale for the first phase t1 ∼ 1/s, where eigenvectors evolve, is independent of the scale of the initialization σ2, whereas the second phase occurs around t2 ≈ t1 log(s/σ
2). This separation of timescales t1 t2 for small σ guarantees the silent alignment effect. We illustrate these learning curves and for varying σ in Figure C.2.

heading: 3.2 UNWHITENED DATA
text: When data is unwhitened, the right singular vector of W aligns with Σβ early in training, as we show in Appendix C.3. This happens since, early on, the dynamics for the first layer are d dtW ∼ a(t)β
>Σ. Thus the early time kernel will have a rank-one spike in the Σβ direction. However, this configuration is not stable as the network outputs grow. In fact, at late time W must realign to converge to W ∝ aβ> since the network function converges to the optimum and f = a>Wx = sβ · x, which is the minimum `2 norm solution (Appendix G.1). Thus, the final kernel will always look like K∞(x,x′) = sx> [ ββ> + I ] x′. However, since the realignment of W ’s singular vectors happens during the Phase II spectral learning, the kernel is not constant up to overall scale, violating the conditions for silent alignment. We note that the learned function still is a kernel regression solution of the final NTK, which is a peculiarity of the linear network case, but this is not achieved through the silent alignment phenomenon as we explain in Appendix C.3.

heading: 4 EXTENSION TO DEEP LINEAR NETWORKS
text: We next consider scalar target functions approximated by deep linear neural networks and show that many of the insights from the two layer network carry over. The neural network function
f : RD → R takes the form f(x) = wL>WL−1...W 1x. The gradient flow dynamics under mean squared error (MSE) loss become
d dt W ` = −η ∂L ∂W ` = η (∏ `′>` W ` ′ )> (sβ − w̃)>Σ (∏ `′<` W ` ′ )> , (9)
where w̃ = W 1>W 2>...wL ∈ RD is shorthand for the effective one-layer linear network weights. Inspired by observations made in prior works (Fukumizu, 1998; Arora et al., 2018; Du et al., 2018), we again note that the following set of conservation laws hold during the dynamics of gradient descent ddt [ W `W `> −W `+1>W `+1 ] = 0. This condition indicates a balance in the size of weight updates in adjacent layers and simplifies the analysis of linear networks. This balancing condition between weights of adjacent layers is not specific to MSE loss, but will also hold for any loss function, see Appendix D. We will use this condition to characterize the NTK’s evolution.

heading: 4.1 NTK UNDER SMALL INITIALIZATION
text: We now consider the effects of small initialization. When the initial weight variance σ2 is sufficiently small,W `W `>−W `+1>W `+1 = O(σ2) ≈ 0 at initialization.1 This conservation law implies that these matrices remain approximately equal throughout training. Performing an SVD on each matrix and inductively using the above formula from the last layer to the first, we find that all matrices will be approximately rank-one wL = u(t)rL(t) , W ` = u(t)r`+1(t)r`(t)>, where r`(t) are unit vectors. Using only this balancing condition and expanding to leading order in σ, we find that the NTK’s dynamics look like
K(x,x′, t) = u(t)2(L−1)x> [ (L− 1)r1(t)r1(t)> + I ] x′ +O(σ). (10)
We derive this formula in the Appendix E. We observe that the NTK consists of a rank-1 correction to the isotropic linear kernel x · x′ with the rank-one spike pointing along the r1(t) direction. This is true dynamically throughout training under the assumption of small σ. At convergence r(t)→ β, which is the unique fixed point reachable through gradient descent. We discuss evolution of u(t) below. The alignment of the NTK with the direction β increases with depth L.

heading: 4.1.1 WHITENED DATA VS ANISOTROPIC DATA
text: We now argue that in the case where the input data is whitened, the trained network function is again a kernel machine that uses the final NTK. The unit vector r1(t) quickly aligns to β since the first layer weight matrix evolves in the rank-one direction ddtW
1 = v(t)β> throughout training for a time dependent vector function v(t). As a consequence, early in training the top eigenvector of the NTK aligns to β. Due to gradient descent dynamics, W 1>W 1 grows only in the ββ> direction. Since the r1 quickly aligns to β due to W 1 growing only along the β direction, then the global scalar function c(t) = u(t)L satisfies the dynamics ċ(t) = c(t)2−2/L [s− c(t)] in the whitened data case, which is consistent with the dynamics obtained when starting from the orthogonal initialization scheme of Saxe et al. (2014). We show in the Appendix E.1 that spectral learning occurs over a timescale on the order of t1/2 ≈ Ls(L−2)σ
−L+2, where t1/2 is the time required to reach half the value of the initial loss. We discuss this scaling in detail in Figure 3, showing that although the timescale of alignment shares the same scaling with σ for L > 2, empirically alignment in deep networks occurs faster than spectral learning. Hence, the silent alignment conditions of Section 2 are satisfied. In the case where the data is unwhitened, the r1(t) vector aligns with Σβ early in training. This happens since, early on, the dynamics for the first layer are ddtW
1 ∼ v(t)β>Σ for time dependent vector v(t). However, for the same reasons we discussed in Section 3.2 the kernel must realign at late times, violating the conditions for silent alignment.

heading: 4.2 MULTIPLE OUTPUT CHANNELS
text: We next discuss the case where the network has multiple C output channels. Each network output, we denote as fc(x′) resulting in C2 kernel sub-blocks Kc,c′(x,x′) = ∇fc(x) · ∇fc′(x′). In this
1Though we focus on neglecting the O(σ2) initial weight matrices in the main text, a more refined analysis for wide networks at finite σ2 is provided in Appendix H.2, which reveals additional dependence on relative layer widths.
context, the balanced conditionW `W `> ≈W `+1>W `+1 implies that each of the weight matrices is rank-C, implying a rank-C kernel. We give an explicit formula for this kernel in Appendix H. For concreteness, consider whitened input data Σ = I and a teacher with weights β ∈ RC×D. The singular value decomposition of the teacher weights β = ∑ α sαzαv > α determines the evolution of
each mode (Saxe et al., 2014). Each singular mode begins to be learned at tα = 1sα log ( sαu −2 0 ) . To guarantee silent alignment, we need all of the Phase I time constants to be smaller than all of the Phase II time constants. In the case of a two layer network, this is equivalent to the condition
1 smin
1smax log ( smaxu −2 0 ) so that the kernel alignment timescales are well separated from the
timescales of spectral learning. We see that alignment precedes learning in Figure H.1 (a). For deeper networks, as discussed in 4.1.1, alignment scales in the same way as the time for learning.

heading: 5 SILENT ALIGNMENT ON REAL DATA AND RELU NETS
text: In this section, we empirically demonstrate that many of the phenomena described in the previous sections carry over to the nonlinear homogenous networks with small initialization provided that the data is not highly anisotropic. A similar separation in timescales is expected in the nonlinear L-homogenous case since, early in training, the kernel evolves more quickly than the network predictions. This argument is based on a phenomenon discussed by Chizat et al. (2019). Consider an initial scaling of the parameters by σ. We find that the relative change in the loss compared to the relative change in the features has the form |
d dt∇f | |∇f | L | ddtL| ≈ O(σ−L) which becomes very large for small initialization σ as we show in Appendix I. This indicates, that from small initialization, the parameter gradients and NTK evolve much more quickly than the loss. This is a necessary, but not sufficient condition for the silent alignment effect. To guarantee the silent alignment, the gradients must be finished evolving except for overall scale by the time the loss appreciably decreases. However, we showed that for whitened data that nonlinear ReLU networks do in fact enjoy the separation of timescales necessary for the silent alignment effect in Figure 1. In even more realistic settings, like ResNet in Figure 1 (d), we also see signatures of the silent alignment effect since the kernel does not grow in magnitude until the alignment has stabilized.
We now explore how anisotropic data can interfere with silent alignment. We consider the partial whitening transformation: let the singular value decomposition of the data matrix be X = USV > and construct a new partially whitened dataset Xγ = USγV >, where γ ∈ (0, 1). As γ → 0 the dataset becomes closer to perfectly whitened. We compute loss and kernel aligment for depth 2 ReLU MLPs on a subset of CIFAR-10 and show results in Figure 4. As γ → 0 the agreement between the final NTK and the learned neural network function becomes much closer, since the
kernel alignment curve is stable after a smaller number of training steps. As the data becomes more anisotropic, the kernel’s dynamics become less trivial at later time: rather than evolving only in scale, the alignment with the target function varies in a non-trivial way while the loss is decreasing. As a consequence, the NN function deviates from a kernel machine with the final NTK.

heading: 6 CONCLUSION
text: We provided an example of a case where neural networks can learn a kernel regression solution while in the rich regime. Our silent alignment phenomenon requires a separation of timescales between the evolution of the NTK’s eigenfunctions and relative eigenvalues and a separate phase where the NTK grows only in scale. We demonstrate that, if these conditions are satisfied, then the final neural network function satisfies a representer theorem for the final NTK. We show analytically that these assumptions are realized in linear neural networks with small initialization trained on approximately whitened data and observe that the results hold for nonlinear networks and networks with multiple outputs. We demonstrate that silent alignment is highly sensitive to anisotropy in the input data.
Our results demonstrate that representation learning is not necessarily at odds with the learned neural network function being a kernel regression solution; i.e. a superposition of a kernel function on the training data. While we provide one mechanism for a richly trained neural network to learn a kernel regression solution through the silent alignment effect, perhaps other temporal dynamics of the NTK could also give rise to the neural network learning a kernel machine for a data-dependent kernel. Further, by asking whether neural networks behave as kernel machine for some data-dependent kernel function, one can hopefully shed light on their generalization properties and transfer learning capabilities (Bordelon et al., 2020; Canatar et al., 2021; Loureiro et al., 2021; Geiger et al., 2021) and see Appendix F.

heading: A DERIVATION OF EQUATION 2

heading: A.1 TRAINING POINT PREDICTIONS WITH TIME VARYING KERNEL
text: Given a training set of P data points {(xµ, yµ)}Pµ=1, the dynamics of the network training errors [∆t]µ := f(xµ, t) − yµ close in terms of a time-varying neural tangent kernel [Kt]µν = K(xµ,xν , t)
d dt ∆t = −Kt∆t. (11)
We introduce the transition matrix Φt ∈ RP×P which has the property that ∆t = Φt∆0 and Φ0 = I , we obtain the matrix evolution equation Φ̇t = −KtΦt. This equation can be solved formally in terms of the Peano-Baker series (Baake & Schlaegel, 2011; Brockett, 2015)
Φt =I − ∫ t 0 ds1Ks1 + ∫ t 0 ds1Ks1 ∫ s1 0 ds2Ks2 (12)
− ∫ t 0 ds1Ks1 ∫ s1 0 ds2Ks2 ∫ s2 0 ds3Ks3 + ... (13)
which can easily be verified to solve ddtΦ(t) = −K(t)Φ(t) with initial condition Φ(0) = I . Under the condition that ∫ t 0 K(t) commutes with K(t), which is true in the settings of interest in this paper, specifically the setting discussed in Appendix B, we can simplify the Peano-Baker series into a simple matrix exponential
Φt = I − [∫ t
0
dsKs
] + 1
2 [∫ t 0 dsKs ]2 − 1 6 [∫ t 0 dsKs ]3 ...+ (−1)k k! [∫ t 0 dsKs ]k + ...
= exp ( − ∫ t 0 Ksds ) . (14)
Thus, under the condition that Kt commutes with ∫ t 0 Ksds we can exactly solve for the training error dynamics in terms of integrating factors
∆t = Φt∆0 = exp ( − ∫ t 0 Ksds ) ∆0. (15)
We expect this formula to hold approximately whenever the eigenvectors of K are approximately equal to the eigenvectors of ∫ t 0 Ksds.

heading: A.2 TEST POINT PREDICTIONS WITH TIME VARYING KERNEL
text: Given access to the value of the function on training points, one can evaluate the function on test points. We have that the evolution of the function on a test point f(x) is given by
d dt ft(x) = −kt(x)∆t, (16)
where [kt(x)]µ := K(x,xµ, t). This gives the final value of f to be
f∗(x) := f∞(x) = f0(x) + ∫ ∞ 0 dtkt(x) µ [ exp ( −η ∫ t 0 Kt′dt ′ )] (y − f0). (17)
This is exactly equation 2.

heading: B KERNEL EVOLUTION IN SCALE ONLY
text: We consider the model of kernel evolution introduced in Section 2 where the kernel evolves only in scale for t > τ and is of small overall size for t < τ ,
K(x, x′, t) =
{ K0(x, x
′, t) t ≤ τ g(t)K∞(x, x ′) t > τ .
where = σ2L−2 for depth L homogenous networks with initialization scale σ. We will define a differentiable function h(t) = ∫ t τ dt′g(t′) so that h′(t) = g(t), h(τ) = 0 and limt→∞ h(t) = ∞. This last condition follows from g’s continuity and the assumption that limt→∞ g(t) = 1. We will first show that the final neural network has the form f(x) = f0(x) +k∞(x) ·K−1∞ (y−f0) +O( ). First, we need to calculate the errors ∆(t) = y−f(t) made on the training examples. These satisfy the dynamics
d dt ∆ = { K0(t)∆(t) t ≤ τ h′(t)K∞∆ t > τ ,
where K0(t),K∞ ∈ RP×P are gram matrices. After the first interval of time, where the kernel evolves significantly t ∈ (0, τ), the error vector has the form
∆(τ) = ∆0 = y − f0 +O( ). (18) In the second interval t ∈ (τ,∞),
d dt ∆(t) = −h′(t)K∞∆(t)
=⇒ ∆(t) = exp (−h(t)K∞) ∆(τ) = exp (−h(t)K∞) (y − f0) +O( ). (19)
We utilized the assumption that g(τ) = 0. Now that we know ∆(t), we can solve for the final predictor. Integrating the dynamical system ḟ(x, t) = k(x, t) ·∆(t), we find
f(x)− f0(x) = ∫ τ 0 k0(x, t) ·∆(t)dt+ k∞(x) · ∫ ∞ τ dt h′(t) exp (−h(t)K∞) ∆0
= k∞(x) · ∫ ∞ 0 dh exp (−hK∞) ∆0 +O( ) = k∞(x) ·K−1∞ ∆0 +O( ). (20)
Up to this point, we have established that f(x) = f0(x) + k∞(x) ·K−1∞ (y − f0) + O( ). This establishes that the predictor f(x) − f0(x) is arbitrarily close to a kernel regression solution on labels y− f0 provided that the kernel evolves only while small in scale during the first phase t < τ . For depth L neural networks with homogenous activations and small initialization, = σ2L−2, and f0 ∼ O(σL). Using this assumption of homogeneity, we have f(x) = f0(x) + k∞(x) ·K−1∞ y +O(σ2L−2) = k∞(x) ·K−1∞ y +O(σL). (21) This argument establishes that, if the kernel and initial network function start out small in overall scale, and the kernel obeys the dynamics discussed above, then the final learned network function f(x) is a kernel regression solution with the final NTK.

heading: C PHASES OF LEARNING AT SMALL INITIALIZATION

heading: C.1 PHASE I: TWO LAYER NETWORK AND KERNEL ALIGNMENT
text: We now present an analysis distinct from that of the previous subsection to go beyond the first step of gradient descent. The NTK for the two layer linear network has the form K = x>Mx′ with M = W>W + |a|2I . Our goal is to determine the eigendecomposition of M . Introduce the variables q(t) = 12β >Mβ = 12 [ |a|2 + β>W>Wβ ] and r(t) = a>Wβ. These dynamics form a closed two dimensional linear system early in training d
dt
[ q(t) r(t) ] = 2ηs [ 0 1 1 0 ] [ q(t) r(t) ] +O(σ3) , t, σ → 0
=⇒ [ q(t) r(t) ] = 1 2 (q0 + r0) [ 1 1 ] e2ηst + 1 2 (q0 − r0) [ 1 −1 ] e−2ηst +O(σ3)
(22)
The variable q(t) represents the alignment of the NTK with the optimal direction β while r(t) defines the alignment of the network with the teacher. We see that this alignment increases exponentially with timescale t ∼ η−1s−1. For Gaussian initialization ai ∼ N (0, σ2/N),Wij ∼ N (0, σ2/D) with large D, we have
〈q0〉 = σ2
2
( 1 + N
D
) , 〈r0〉 = 0 , 〈 r20 〉 = σ4
D . (23)
In the large D limit, we have with high probability q0 r0 and thus r(t) = q0 sinh(2ηst) and q = q0 cosh(2ηst). Note that this gives the quantities q(t) = 12β
>M(t)β, r(t) = β>W>a early in training. Now consider a unit vector v which is orthogonal to the solution β>v = 0. We find that the projection of M along this direction evolves dynamically as:
d dt v>M(t)v = 2v>
[ βa>W + β>W>aI ] v
= 2r(t). (24)
We can conclude that v>Mv is equal to q(t) up to an additive initialization constant. We see that this is evolving half as quickly as β>Mβ = 2q(t). Since v>M0v ∼ O(σ2) is small compared to the exponentially growingM(t), the only matrix that satisfies these two conditions must necessarily take the form
M(t) = q0 cosh(2ηst) [ ββ> + I ] +M0. (25)
The first term, which is growing exponentially in t will eventually overwhelm the randomly initialized kernelM0, which is O(σ2).

heading: C.2 PHASE I: ERROR IN THE LEADING ORDER APPROXIMATION
text: In solving the equations of the previous section, we truncated the full gradient descent equations at order σ3. It is important to confirm that the error generated by this truncation remains bounded. We will argue by self-consistency. The full equations are
d dt a = W
( sβ −W>a ) , d
dt W = a
( sβ −W>a )> . (26)
One can use these equations to solve for the dynamics of the r, q variables:
d dt q(t) = 2sr − β>
[ W>aa>W + a>WW>a ] β = 2sr − [2r2 + ∑ vi⊥β (a>Wvi) 2] (27)
d dt r(t) = 2sq − a>[WW> + aa>]Wβ = 2sq − 2|a|2r + a>[aa> −WW>]Wβ (28)
The second equality in equation 27 comes from inserting a complete basis of states including β and vi ⊥ β into the last term of the left-hand side. The second equality in equation 28 comes from writing WW> + aa> = 2aa> + (WW> − aa>). Note that the final term in brackets on the right hand side is a conserved quantity for linear networks, and so is always of order O(σ2).
Assuming the solutions for q, r are valid to order σ2, we get that ddta >Wvi = O(σ 4).
d dt a>Wvi = (β − β̂)>W>Wvi + |a|2(β − β̂)>vi (29)
Further noting that because of the conservation law aa> −WW> = O(σ2) is also constant in time. This gives us that
|a>[aa> −WW>]Wβ| ≤ σ2|a>Wβ| = σ2r. (30)
We now note that r, |a|2 both grow as a (σ, s-independent) constant times times σ2e2st. The correction to the dynamics of both equations is then bounded by a constant times σ4e4st. This will be less than σ2 as long as t satisfies
t 1 4s log ( s σ2 ) . (31)
For σ2 s, the alignment time t = 1/s falls within this range and we are guaranteed alignment to the Ganguli-Saxe configuration.
The error of the full solution at time t can be bounded by the integral of this error bound from 0 to t, namely a constant times σ4/s. As long as s σ4, we are guaranteed that the error of the kernel is O(σ2) as given in equation 7.

heading: C.3 PHASE I: TWO LAYER ANALYSIS WITH UNWHITENED DATA
text: We now study the same linearization around the initial fixed point used in the main text but for the two layer network with unwhitened data. In this case,
d dt a ∼ sWΣβ, (32)
d dt W ∼ saβ>Σ. (33)
which holds asymptotically as t/ log(σ−1) → 0. We introduce the following variables which form a closed linear dynamical system
r1(t) = β >W>a,
r2(t) = |a|2, r3(t) = β >W>WΣβ,
r4(t) = β >ΣW>a,
r5(t) = β >ΣW>WΣβ.
(34)
Introduce the constants a = β>Σβ , b = β>Σ2β. Using the weight dynamics, it is straightforward to show that
ṙ(t) ∼ s  0 a 1 0 0 0 0 0 2 0 b 0 0 a 0 0 b 0 0 1 0 0 0 2b 0  r(t) , t log(σ−2). (35)
This matrix has eigenvalues λ ∈ {0,− √ b, √ b,−2 √ b, 2 √ b}. Since there are only two positive eigenvalues √ b, 2 √ b, it suffices to consider evolution along those two eigendirections, where the kernel and neural network function will be amplified. Evolution along these direcions give
r(t) ∼ c1es √ bt  1 0√ b
0 0
+ c2e2s√bt  a√ b a √ b
b
b √ b  , (36) where c1, c2 are constants determined by intialization. At large time, the large eigenvalue mode λ = 2 √ b will dominate. Decomposing W = W0 + a(t) [v1(t)β + v2(t)Σβ]
> we find that the only self consistent solution is v1(t) = 0, v2(t) = b−1/2. This implies that the kernel evolution will take the form
K(x,x′, t) ∼ K(x,x′, 0) + |a(t)|2x>Mx′,
M = [ 1√
β>Σ2β Σββ>Σ + I
] .
(37)
We see that the kernel evolves along the directions Σβ early in training for unwhitened data. We visualize the two stages of learning for unwhitened data in Figure C.1.

heading: C.4 PHASE II: WHITENED DATA
text: Consider a two layer network f = a>Wx where balance has been achieved W = u(t)âβ> and a(t) = u(t)â. Once this balance condition is stable for fixed â, we can calculate the time derivative of u(t)
d dt a(t) = u̇(t)â = u(t)
[ s− u(t)2 ] â. (38)
Letting c(t) = u(t)2, we find that ċ(t) = 2u(t)2 [ s− u(t)2 ] = 2c(t) [s− c(t)], which is the two layer dynamics derived in Saxe et al. (2014). This dynamics has solution c(t) = se 2st
e2st−1+s/c0 .

heading: C.5 SOLUTIONS TO THE FULL TRAINING DYNAMICS OF LINEAR NETWORKS AT SMALL INITIALIZATION
text: By combining the analyses of the subsection C.1 with the exact solutions discussed in C.4 we can match both solutions to obtain formulas for r(t) and q(t) for the entire network’s training path that are exact up to O(σ2) corrections. Up to O(σ2) we then have that
r(t) = s 2 sinh(2st)
(e2st − 1) + 2s/q0 , (39)
q(t) = s 2 cosh(2st)
(e2st − 1) + 2s/q0 . (40)
This yields that the initialization constant q0/2 plays the effective role of c0 in the Ganguli-Saxe solution for phase II. Equation 23 yields the expected value of this initialization constant. We have empirically verified that these exact equations hold to high accuracy across a variety network sizes, initialization scales, and whitened datasets. We illustrate some of these in figure C.3.

heading: D BALANCING OF WEIGHTS IN DEEP LINEAR NETWORKS
text: The balance condition discussed in the main text holds for deep linear trained with any loss function of the form L = ∑ µ `(f µ, yµ) (not just MSE) since, as was shown by Arora et al. (2018); Du et al.
(2018) 1
η
d
dt
( W `W `> ) = ∑ µ ∂` ∂fµ
[ ∂fµ
∂W ` W `> +W `
∂fµ
∂W `
> ]
= ∑ µ ∂` ∂fµ
[ ∂fµ
∂W `+1
> W `+1 +W `+1> ∂fµ
∂W `+1
]
= 1
η
d
dt
( W `+1>W `+1 ) .
(41)
The second line follows from the first since the following quantities are identical ∂fµ
∂W ` W `> = W `+1>...WL−1>wL
( xµ>W 1>...W `−1> ) W `>, (42)
W `+1> ∂fµ
∂W `+1 = W `+1>
( W `>...wL ) xµ>W 1>...W `>. (43)
By inspection these two quantities are equal. Thus we have, for any loss function, a deep linear network has the following conservation laws
d
dt
[ W `W `> −W `+1>W `+1 ] = 0. (44)
We show in the next section that this balancing condition is very helpful in identifying the time evolution of the neural tangent kernel. In the case where the network has a single output, we can
inductively prove that each layer’s weight matrix isW ` = u(t)r`+1(t)r`(t)>. We will assume this formula is true for layer `+ 1 and prove it must hold for layer ` since
W `W `> = W `+1>W `+1 = u(t)2r`+1(t)r`+1(t) >. (45)
This implies that W ` is rank one with right singular vector equal to r`+1(t). Thus, the decomposition for each layer the form W ` = u(t)r`+1(t)r`(t)> for some unit vector r`(t). Similar analysis can be performed for the multi-class setting.

heading: E NTK FORMULA FOR DEEP LINEAR NETWORKS
text: The neural tangent kernel for a linear network f(x) = wL>WL−1...W 1x is defined as an innerproduct over all gradients
K(x,x′) = L∑ `=1 〈 ∂f(x) ∂W ` , ∂f(x′) ∂W ` 〉
= ∣∣∣∣∣ L∏ `=2 W ` ∣∣∣∣∣ 2 x · x′ + L∑ `=2 ∣∣∣∣∣∏ `′>` W ` ′ ∣∣∣∣∣ 2 x>W 1>...W `−1>W `−1...W 1x′.
(46)
Under the balanced assumption that W ` = u(t)r`+1(t)r`(t)> + O(σ), expanding the kernel to leading order in σ yields the following form:
K(x,x′) = u(t)2L−2x> [ (L− 1)I + r1(t)r1(t)> ] x′ +O(σ), (47)
100 101 102 103 t
0.0
0.2
0.4
0.6
0.8
1.0
Lo ss
L = 2 L = 3 L = 4 L = 5
100 101 102 103 t
0.2
0.4
0.6
Ke rn
el A
lig nm
en t L = 2
L = 3 L = 4 L = 5
100 101 102 103 t
0.2
0.4
0.6
0.8 1.0 W 1 A lig nm en t L = 2L = 3 L = 4 L = 5
Figure E.1: The depth dependence of loss dynamics, kernel alignment and the alignment of first layer weights in a linear network on synthetic whitened data in D = 30 dimensions. (a) The loss reaches half its initial value after t ∼ σ−L+2 steps for L > 2. The decay rate of the loss becomes sharper with depth. (b) Final kernel alignment increases monotonically with depth L and approaches 1 as L→∞. (c) The alignment of the first layer weightsW1 with the optimal direction β approaches 1 for all models.

heading: E.1 DEEP LINEAR NETWORK DYNAMICS UNDER BALANCE
text: In this section, we will consider the dynamics of the variable u once the balance condition is satisfied. Let wL = ŵu(t). Then the dynamics for u(t) under the balancing assumption is
d dt wL = ŵ d dt u(t) = ( L−1∏ `=1 W ` ) (s− u(t)L) L−1∏ `=1 W `β = u(t)2L−2(s− u(t)L)ŵ, (48)
which implies the fact u̇(t) = u(t)L−1(s− u(t)L). Changing variables to c(t) = u(t)L we obtain
ċ(t) = u(t)L−1u̇ = u(t)2L−2(s− u(t)L) = c(t)2−2/L(s− c(t)). (49)
When cL0 is initialized to a very small value compared to s we can
d dt c(t) ∼ c(t)2−2/Ls =⇒
[ c(t)−1+2/L − c−1+2/L0 ] = −L− 2
L st =⇒ c(t) = [ c −L−2L 0 −
(L− 2) L st
]− LL−2 .
(50)
100 101 102 103
t
0.0
0.2
0.4
0.6
0.8
1.0
Lo ss
L = 2 L = 3 L = 4 L = 5
(a) Loss
102 103
t
0.00
0.02
0.04
0.06
0.08
Al ig
nm en
t
(b) Alignment
1.5 2.0 2.5 3.0 3.5 4.0 4.5 5.0 5.5 ReLU MLP Depth 0.00
0.05
0.10
0.15
0.20
0.25
|f N TK
f N N |/|
f N N |
K0 K
(c) Test Comparison
Figure E.2: ReLU networks across depths trained on two-class whitened CIFAR. The hidden widths are all of size 100. We use 3k train points and 7k test points. (a) The loss exhibits the same scaling as σ−L+2 as in the linear setting. (b) Deep networks with nonlinearities are seen to undergo the silent alignment effect early on in training. The dashed lines indicate when the kernel has grown to 10% of its final value. (c) The trained network outputs on test data match closely the kernel regression with the final learned kernel, but do not match regression the initial kernel.
This implies a timescale to learn of t ∼ s−1 LL−2σ −L+2.
We can approximate the timescale for the first layer’s singular vector r1(t) to align to β as well. Let v be a vector orthogonal to β.
d
dt
|W (1)β|2 |W (1)v|2 = 2β>W (1)>...wL |W (1)v|2 ∼ O(σL−2). (51)
This suggests that alignment in a deep network should also occur on a timescale of t ∼ σ−L+2. While there is no strict separation of timescales in terms of the scaling of alignment and learning with σ, we find that alignment tends to precede a significant drop in the loss as we show in Figure 3.

heading: E.2 FINAL NTK FOR DEEP LINEAR NETWORKS
text: Independent of the structure of the data, the first vector r1 → β and the final NTK has the following form:
K(x,x′) = s2L−2x> [ (L− 1)ββ> + I ] x′ +O(σ). (52)
This formula is merely a consequence of the balancing condition and convergence to the optimum u(t)L → s. We provide empirical support that final kernel alignment increases with depth in Figure E.3.

heading: E.3 FORMULA FOR THE INTERMEDIATE NNGP KERNELS
text: Let β represent the unit vector pointing in the optimal direction for the scalar output case. To calculate this, all that needs to be assumed is that the network has converged to its optimum and that the weights satisfy the balance property so that W ` = u∗r∗`+1r ∗ ` . By the convergence assumption u∗ = s1/L and r∗1 = β. We stress that this holds for arbitrary input correlation structure Σ, the final NNGP kernel for layer ` can also be computed. Expanding the weights and collecting terms at leading order in σ yields:
K`(x,x ′) ∼ s2`/Lx>ββ>x′ +O(σ). (53)
Evaluating on the training set givesK` = ( yy> )`/L ∈ RP×P .

heading: F GENERALIZATION ERROR IN TRANSFER LEARNING TASK
text: The structure of the final kernel can alter the ability of the network to flexibly transfer to new tasks with a small amount of data. In this section, we examine how learned intermediate representations
0 500 1000 1500 2000 2500 t
0.1
0.2
0.3
0.4
Ke rn
el A
lig nm
en t L = 2
L = 4 L = 6
(a) Alignment Dynamics and Depth
K 0
L = 2
K
L = 4 L = 6
(b) Initial and Final NTKs
Figure E.3: The final NTK for a deep linear network aligns with the class specific directions with strength that depends on network depth L. This experiment is a partially whitened (γ = 0.25 see Section 5) subset of 100 CIFAR-10 images. (a) The dynamics of alignment for different depth models. (b) The final gram matrix has the formK∞ ∝ (L−1)yy>+K0, illustrating why alignment of the final NTK to class structure increases with depth L.
compares with the inductive bias of the original isotropic kernel x · x′. In particular, we study the offline generalization performance of kernels of the form
K(x,x′;A) = x> [ Aββ> + I ] x′. (54)
In Section 4 we showed that A could be altered by changing the network depth or the relative learning rates in the first layer. Concretely, our transfer learning problem consists of training a linear probe on one of the intermediate layers of the network (Alain & Bengio (2016); Cohen et al. (2020)). This would also produce a kernel regression solution for kernel K(x,x′, A) with A which depends on the chosen layer, the learning rates and the depth of the network. For simplicity, we assume that the data are generated according to a simple Gaussian distribution x ∼ N (0, I) and that the target values are generated with a linear function y(x) = w · x. We decompose the new task vector w = αβ+ √ (1− α2)w⊥ wherew⊥ ·β = 0. The expected generalization error after training with P samples can be computed with methods from the physics of disordered systems (Bordelon et al., 2020; Canatar et al., 2021; Loureiro et al., 2021). For any A > 0, the easiest transfer task is w = β (α = 1). If w = β, increasing the alignment A strictly decreases the generalization error. This is illustrated in Figure F.1.
0 5 10 15 20 25 P
0.4
0.5
0.6
0.7
0.8
0.9
1.0
E g
Kernel Alignment Strength
A = 0 A = 1 A = 5 A = 10
(a) α = 0.75
0 5 10 15 20 25 P
0.5
0.6
0.7
0.8
0.9
1.0
E g
Transfer Task Alignment
= 0.0 = 0.2 = 0.5 = 0.7
(b) A = 1
0.0 2.5 5.0 7.5 10.0 A
1.0
0.75
0.5
0.25
0.0
Eg with P = 10
0.2
0.4
0.6
0.8
(c) Fixed P
Figure F.1: The offline generalization error in a transfer task with the learned linear kernel K = x> [ Aβ>β> + I ] x′ and ynew = [ αβ + √ (1− α2)w⊥ ] · x. (a) For a new transfer task which
is correlated with the learned function β, neural networks with large feature learning A give lower generalization error at small sample sizes. (b) For fixed A = 1, the tasks w which are strongly correlated with β are easier to learn during transfer. (c) The lowest generalization error in a transfer learning setup occurs when feature learning strength A and correlation between tasks α are both large.

heading: F.1 DERIVATION OF LEARNING CURVES
text: We will discuss the average case generalization error in the transfer learning setting. Prior work has shown that the generalization performance of kernel regression can be calculated through a kernel eigenvalue problem (Bordelon et al., 2020; Canatar et al., 2021; Loureiro et al., 2021)
〈K(x,x′)φk(x)〉x = λkφk(x). (55)
Once this integral eigenvalue problem is solved for eigenvalues λk and orthonormal eigenfunctions φk, the average case generalization error at P training examples is
Eg = 1 1− γ ∑ k 〈y(x)φk(x)〉2 (λkP + κ)2 , κ = λ+ κ ∑ k λk λkP + κ , γ = P ∑ k λ2k (λkP + κ)2 . (56)
In our case, we are interested in the generalization performance of the linear kernel K(x,x′) = x> [ ββ> + I ] x′. (57)
Since K is a linear kernel, its eigenfunctions should be linear functions φk(x) = φk · x. Assuming that the data distribution has identity covariance, we find〈
x> [ Aββ> + I ] x′φ>k x 〉 = φ>k [ Aββ> + I ] x′ = λkφ > k x ′. (58)
This implies that the φk vectors are eigenvectors of M . The first eigenvector is φ1 = β with eigenvalue λ1 = A + 1. The other D − 1 eigenvectors can be chosen as any frame in the D − 1 dimensional subspace orthogonal to β. Each of these D − 1 eigenvectors has eigenvalue λk = 1. Using these results, and the fact that w = αφ1 + √ 1− α2w⊥, we can calculate the expected generalization error.
Eg = 1
1− γ
[ α2
((1 +A)P + κ)2 +
1− α2
(P + κ)2 ] κ = λ+ κ [ 1 +A
(1 +A)P + κ + D − 1 P + κ
] , γ = P [ (1 +A)2
((1 +A)P + κ)2 + D − 1 (P + κ)2
] .
By the result proven in Canatar et al. (2021), the lowest possible error for fixed A occurs by maximizing the fraction of variance along the large eigenvalue direction, corresponding to α = ±1.

heading: G LINEAR NTKS DURING GD LEARN THE SAME FUNCTION
text: In the overparameterized setting where D > P , all linear networks discussed in the Section 4 converge to the minimum norm interpolator when the data is whitened. Specifically, letting the learned neural network function be written as f(x) = β̂>x, and the data matrix X ∈ RD×P and target labels y ∈ RP represent the training data. The solution vector β solves the constrained optimization problem
min β̂ ||β̂||2 , s.t.X>β̂ = y, (59)
which is the kernel regression solution for the initial kernel K0(x, x′) = x · x′. This is unsurprising due to a symmetry argument: when β0 ≈ 0, the only privileged point on the affine space X>β = y is the point closest to the origin, which is precisely the solution above. Surprisingly, the final pseudo-inverse solution sβ also minimizes the RKHS norms for any of the kernels throughout gradient descent. Up to an overall scale, the kernels throughout evolution take the form K(x,x′; t) = x> [ A(t)ββ> + I ] x′ (see Section 4.1) which would induce the following kernel interpolation problems
min β̂
β̂> [ A(t)ββ> + I ]−1 β̂ , s.t.X>β̂ = y. (60)
The solution to this optimization problem is indeed the kernel regression solution with kernel K(t) since the learned function takes the form f(x) = ∑ µ α µK(x,xµ, t) with α = K(t)−1y. Using
the Sherman-Morrison rule, we show that the solution to each of these problems t ≥ 0 gives the same result, namely the pseudo-inverse solution. This can be seen from the following
β̂> [ A(t)ββ> + I ]−1 β̂ = |β̂|2 − A(t)
1 +A(t)
( β · β̂ )2 . (61)
Now, we let β̂ = sβ + β⊥, where β · β⊥ = 0. This is the general decomposition for the set of interpolators which have the propertyX> [β + β⊥] = y.
min β⊥ |β⊥|2. (62)
The solution is merely to set β⊥ = 0. Thus the optimal solution is therefore the same for any finite value of A. However, the final RKHS norm of the learned function β> [ A(t)ββ> + I ]−1 β decreases with time, indicating that the kernel becomes more aligned with the pseudo-inverse direction as A increases.

heading: G.1 DEEP LINEAR NETWORKS FROM SMALL INITIALIZATION LEARN PSEUDO-INVERSE
text: In this subsection of the appendix, we will use our theoretical technology for balanced linear networks to demonstrate the universal learned function for any data, not just whitened input, providing an alternative derivation to the result proven in Theorem 7 of Yun et al. (2020). This analysis is performed in the σ → 0 limit, where W` = u(t)r`+1(t)r`(t)> as we showed in Section D. Under this condition the learned function f(x) = β̂ · x is defined through weights β̂ = W>1 W > 2 ...wL = u(t)
Lr1(t). We see that the direction of the learned function is controlled entirely by r1(t). It suffices to prove that r1(t) ∈ span{x1, ...,xP } for all t to show that the network learns the pseudo-inverse solution β∗ = X(X>X)−1y, where X ∈ RD×P and y ∈ RP are the training data and targets respectively. Note that by gradient descent, we have
d dt W1(t) = W > 2 ...wL ∑ µ (yµ − fµ(t))x>µ = u(t)L−1r2(t) ∑ µ (yµ − fµ(t))x>µ . (63)
From the balance conditionW1(t) = u(t)r2(t)r1(t)>, we also have d
dt W1(t) = [u̇(t)r2(t) + u(t)ṙ2(t)] r1(t)
> + u(t)r2(t)ṙ1(t) >. (64)
Equating the two above expressions for ddtW1(t) and taking an inner product with r2(t) from the left gives the following
u(t)ṙ1(t) = u(t) L−1 ∑ µ (yµ − fµ(t))xµ − [u̇(t) + u(t)r2(t) · ṙ2(t)] r1(t). (65)
Thus, if r1(t) ∈ span{x1, ...,xP } then ṙ1(t) ∈ span{x1, ...,xP } so that the full dynamics of r1(t) lie in the subspace spanned by the training data. At initialization, we have Ẇ1(t) ∝ z(t) ∑ µ y µx>µ + O(σ 3) so the initial r1 vector will indeed align with the span of the training data in the σ → 0 limit. By the fact that β̂ = u(t)Lr1(t), the learned linear coefficients β̂ must also be in span{x1, ...,xP } so β̂ = ∑ µ αµxµ. These must also interpolate the data provided D ≥ P , giving the following condition
β̂ · xν = ∑ µ xν · xµαµ = yν =⇒ α = (X>X)−1y =⇒ β̂ = X(X>X)−1y. (66)
This is exactly the minimum `2 norm interpolating solution which solves
min β̂ |β̂|22 , s.t. , X>β̂ = y. (67)
While anisotropy of the data makes no impact on what function is ultimately learned in the linear network case, the anisotropy can have a signficant influence on whether the preconditions for silent alignment are satisfied in a nonlinear network, which can prevent the final function from being a NTK regressor with final NTK.

heading: H FINAL KERNEL IN MULTI-CLASS NETWORKS
text: For a network with C output channel, balancing and alignment guarantee that the configuration of the network is orthogonal and balanced as in the setting of Saxe et al. (2014). One can then integrate each mode separately to obtain the final kernel as
W ` = C∑ α=1 uα(t)r α `+1(t)r α ` (t) > , Kc,c′(x,x ′) = x>Mc,c′x ′,
Mc,c′ = δc,c′ ∑ α u2(L−1)α (t)r α 1 r α 1 + L−1∑ `=1 ∑ α u2(L−`)α (t)e > c r α Lr α> L ec′ ∑ β u 2(`−1) β (t)r β 1 r β> 1 , (68)
where the Cartesian unit vectors ec ∈ RC are one-hot on class output c. This shows that the contributions to the kernel depend on how well the class output channels align with the unit vectors rαL. Further, the singular values uα can evolve at different timescales depending on the structure of the data.
Specifically, for a depth L network, both the the alignment time t(L)α and the time to learn a given singular value sα scale as s−1α σ
2−L, as shown in appendix E.1. The differences in alignment times ∆tαβ := t (L) α − t(L)β for modes sα, sβ therefores scales as ∆t (L) αβ = σ 2−L∆t (2) αβ .
0 50 100 150 200 250 300 t
0.0
0.2
0.4
0.6
0.8
1.0
Ke rn
el a
nd L
os s
Linear Depth 2
L |K| Kernel Alignment
(a) Linear depth 2
0 1000 2000 3000 4000 5000 6000 t
0.0
0.2
0.4
0.6
0.8
1.0
Ke rn
el a
nd L
os s
Linear Depth 4
(b) Linear depth 4
0 200 400 600 800 1000 1200 1400 t
0.0
0.2
0.4
0.6
0.8
1.0
Ke rn
el a
nd L
os s
Tanh Depth 3
(c) Tanh depth 3
Figure H.1: Demonstration of a separation between alignment and spectral learning phases across networks trained on multi-class data. Here we train on whitened MNIST. Each network has σ so that σL = 10−4 for L the depth of the network. (a) Depth two multi-class dynamics are very similar to single class. The analytically predicted final alignment is in dashed green. (b) For deeper multi-class networks, each singular values is learned far apart in time, and alignment does not as clearly precede loss decay. (c) Similar dynamics are obtained for tanh networks. Note for deeper networks there is a stronger separation of the times to learn each singular value, resulting in the ten separated drops in the loss.

heading: H.1 FINAL NNGP IN MULTI-OUTPUT CASE
text: We can also gain intuition about the learned representations in each layer by looking at the NNGP kernels, which merely take inner-products between layer activations for different inputs. For simplicity, we will assume all learning rates are equal η` = η. Let β ∈ RC×D represent the optimal weight matrix which has the property (in the over-parameterized case) βxµ = yµ. At the optimum, the neural network must learn β = WL...W 1. Computing the SVD of β = ∑ α βαzαv > α reveals that u∗α = β 1/L α and r1α = vα and r L α = zα. Using these facts, it is easy to derive the final NNGP kernel for layer `.
K`(x,x ′) = x>W 1>...W `>W `...W 1x′
= x> [∑ α (u∗α) 2` r1αr 1> α ] x′ = x> [ β>β ]`/L x′.
(69)
Evaluating on the training set X ∈ RD×P gives K` = X> [ β>β ]`/L X , which interpolates betweenX>X at layer ` = 0 and Y >Y at layer L.

heading: H.2 MORE REFINED BALANCING ANALYSIS
text: We can derive a more refined decomposition of the kernel by including the initial conditions in our derived conservation laws. In particular, we will consider balancing for large width networks. Note that the network does not need to be in the lazy regime. The balancing condition is
W `(t)W `>(t)−W `+1>(t)W `+1(t) = W `0W `>0 −W `+1>0 W `+1 0 . (70)
Note that W `0 ∈ RN`+1×N` has entries with zero mean and variance σ2/N` in the standard parameterization. The products of initial matrices are therefore Wishart distributed. For sufficiently large widths, we can approximate the initial weight matrix products with their expectation over the random initialization
W `0W `> 0 −W `+1>0 W `+1 0 ≈ σ2
( 1− N`+2
N`+1
) I. (71)
This concentration becomes more accurate as the widths N`, N`+1 → ∞. For the last layer if the number of classes C = NL is sufficiently large, we also obtain similar concentration for the last layer WL>WL ≈ σ2 NLNL−1 I . Repeating the backward induction on the conservation law, we find the following recursively defined singular value decompositions
W ` = ∑ α u`α(t)r α `+1(t)r α ` (t) >,
u`α(t) 2 = u`+1α (t)
2 + σ2 (
1− N`+2 N`+1
) = uLα(t) 2 + σ2g`,
g` = L− `− L−1∑ k=` Nk+1 Nk .
(72)
We note that the corrections g` vanish if all layers k > ` have the same width. Let uα = u (L) α (t). The condition for convergence is∏ ` [ u(`)α ]2 = ∏ ` [ (uα) 2 + σ2g` ] = s2α. (73)
In the σ2 → 0 limit, we can solve that uα ∼ s1/L as before. However, we can now obtain leading order corrections (in σ2) which take the form of the form
[uα] 2 ∼ s2α −
gσ2
2Ls 1/L α , σ2 → 0 , g = L∑ `=1 g` = L(L− 1) 2 − L∑ `=1 L∑ k=` Nk+1 Nk , (74)
which reveals that the size of the correction depends not only on σ but also on the depth and network widths. Suppose all network widths were equal N` = Nk, then the term g = 0 and there is no contribution from the first moment of the random weights.

heading: I LAZINESS IN HOMOGENOUS NETWORKS
text: In this section we recapitulate the argument found in Chizat et al. (2019). The goal is to estimate how rapidly the gradient features on a test point ∇f(x) change compared to the loss early in training. Let f ∈ RP represent the function outputs on the training set. We will compute the time derivatives of the loss and the network gradients.∣∣∣∣ ddt∇θf(x) ∣∣∣∣ = ∣∣∣∣∇2f(x) · dθdt ∣∣∣∣ = ∣∣∣∣∣∇2f(x) ·∑ µ
(yµ − fµ)∇fµ ∣∣∣∣∣ , (75)∣∣∣∣ ddtL ∣∣∣∣ = ∣∣∣∣dθdt ∣∣∣∣2 = |∇f · (y − f)|2 . (76)
Here, | · |op denotes the operator norm of a matrix. We are interested in the ratio of the loss’ time derivative to the gradient’s time derivative. With an initialization scale of f ∼ O(σL) we find∣∣ d dt∇f(x) ∣∣
|∇f | L∣∣ d dtL ∣∣ = |∇2f(x) ·
∑ µ(y µ − fµ)∇fµ||y − f |2
|∇f ||∇f · (y − f)|2 ≈ |∇2f(x) ·
∑ µ y µ∇fµ||y|2
|∇f ||∇f · y|2 ,
(77)
where in the last step we approximated yµ− fµ ≈ yµ for small initialization scale since yµ ∼ O(1) and fµ ∼ O(σL). Now we will estimate the scale of each of the terms above. For a homogenous model∇f ∼ O(σL−1) and∇2f ∼ O(σL−2). Counting powers of σ in numerator and denominator, we find that this quantity of interest scales as∣∣ d dt∇f(x) ∣∣
|∇f | L∣∣ d dtL ∣∣ = O(σ−L). (78)
This result indicates that, from small initialization, the gradient NTK features and thus the kernel itself will evolve much more rapidly than the loss. This effect can be amplified by increasing depth and decreasing initialization scale.

heading: J RESNET EXPERIMENTAL DETAILS
text: Below, we provide the alignment and loss dynamics for wide resnet for CIFAR-10 with 100 training points. Because the loss decreases significantly before the kernel reaches its final alignment value, the final NTK is not perfectly correlated with the final neural network function. The wide ResNet model is taken from Novak et al. (2020) and is based on the original architecture of Zagoruyko & Komodakis (2017) with a widening factor of k = 4 and a single block per ResNet group b = 1, giving a final network with 8 trainable layers. For both Figure 1 (d) and (e) as well as Figure J.1 use Adam with a learning rate of η = 10−5 and initial weight scale of σ = 0.3 in standard parameterization for all intermediate blocks. For the first conv layer, we used σ = 6.0. We find that small initial weight variance in the first layer gives rise to less stable learning and worse kernel alignment.
0 500 1000 1500 2000 2500 t
0.0
0.2
0.4
0.6
0.8
1.0
Lo ss
a nd
A lig
nm en
t
Lt |K(t)| Alignment
(a) Dynamics
2 1 0 1 2 Test Prediction NN
2
1
0
1
2
Te st
P re
di ct
io n
NT K
Wide Res-Net on Unwhitened CIFAR K0 K
(b) Predictions
Figure J.1: The dynamics and predictions of a Wide-Resnet with k = 4 and b = 1 on P = 100 unwhited CIFAR-10 images from the first two classes.
Below, in Figure J.2, we provide comprehensive results for different depths which we control by increasing the number of blocks per group b, corresponding to WRNs with 6b + 1 trainable conv layers.

heading: J.1 ADAPTIVE OPTIMIZERS AND THE RELEVANT KERNEL
text: Many adaptive gradient methods compute updates to parameters θj according to
θ̇j(t) = −ηj(t) ∂L ∂θj,t
(79)
0 1000 2000 3000 t
0.0
0.2
0.4
0.6
0.8
1.0
Lo ss
b = 1 b = 2 b = 3 b = 4
(a) Training Loss
0 1000 2000 3000 t
0
1
2
3
4
5
Ke rn
el N
or m
(b) Kernel Norm
0 1000 2000 3000 t
0.0
0.1
0.2
0.3
0.4
0.5
Al ig
nm en
t
(c) Alignment Dynamics
1 2 3 4 Wide Res-Net Blocks
0.0
0.2
0.4
0.6
0.8
|f N TK
f N N |/|
f N N | K0 K
(d) Predictor Comparison
Figure J.2: The silent alignment effect is preserved across a large range of depths in WideResNet trained on whitened CIFAR-10 images. The number of blocks per group b alters the total number of conv layers (6b + 1 total conv layers). (a) The deeper models train faster with Adam. (b) The final NTK norm increases with depth. (c) The alignment achieves close to its final value by the time the kernel norm reaches 10% of its final value (dashed line), indicating successful silent alignment. (d) The neural network predictions are very close to the predictions of the final NTK but is not accurately predicted by the initial NTK.
0 1000 2000 3000 4000 5000 6000 7000 t
0.0
0.2
0.4
0.6
0.8
1.0
Lo ss
k = 1 k = 2 k = 3 k = 4
(a) Train Loss Dynamics
0 500 1000 1500 2000 2500 3000 3500 t
0
1
2
3
4
5
6
7
Ke rn
el N
or m
(b) Kernel Norm
0 500 1000 1500 2000 2500 3000 3500 t
0.0
0.1
0.2
0.3
0.4
0.5
Al ig
nm en
t
(c) Alignment
0.5 1.0 1.5 2.0 2.5 3.0 3.5 4.0 4.5 Wide Res-Net Width 0.0
0.2
0.4
0.6
0.8
1.0
|f N TK
f N N |/|
f N N |
K0 K
(d) Predictor Comparison
Figure J.3: Varying the ResNet widening parameter k also alters the kernel and loss dynamics. (a) The loss curve for b = 2 WRNs with widening factor k. Wider networks train more quickly. (b) The kernel norm increases more rapidly for wider networks but changes by a smaller amount. (c) Alignment reaches close to its asymptote by the time the kernel norm grows to 10% its final value (dashed). (d) The final kernel is a much better predictor of the NN function than the initial kernel.
where ηj(t) are time-varying functions which are computed in terms of the history of gradient moments for parameter θj or in terms of its instantaneous gradient. The relevant kernel at time t which governs instantaneous evolution of network predictions is
K(x,x′, t) = ∑ j ηj(t) ∂f(x, t) ∂θj ∂f(x′, t) ∂θj (80)
since ḟ(x) = ∑ µK(x,x
µ, t)(yµ − f(xµ, t)). Though we do not calculate this kernel which is relevant to the adaptive learning rate scheme since it is not supported in Neural Tangents API, this could be a worthy future investigation.

Review summaries:
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.

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.

In this paper the authors show that in some cases, a neural network, trained in the rich regime, can behave like a kernel predictor where the kernel is the neural tangent kernel at the end of training. Specifically, this happens when the kernel changes before the loss significantly decreases, but then stays fixed up to changes in scale. The authors call it “silent alignment”, and show that this happens for deep linear networks with small initialization and whitened data. Some experiments show that silent alignment can happen also for non-linear networks (e.g. with ReLU activations).