Neural Networks as Kernel Learners: The Silent Alignment Effect
Authors: Alexander Atanasov
Blake Bordelon
Cengiz Pehlevan
Venue: ICLR 2022 Poster (2022)
Keywords: Neural Tangent Kernel
Feature Learning
Inductive Bias of Neural Networks

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 the depth. We also demonstrate that non-whitened data can weaken the silent alignment effect.

Published as a conference paper at ICLR 2022
NEURAL NETWORKS AS KERNEL LEARNERS: THE SILENT ALIGNMENT EFFECT
Alexander Atanasov", Blake Bordelon" & Cengiz Pehlevan Harvard University
Cambridge, MA 02138, USA
{atanasov, bl ake bordelor on,,
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 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 empirically 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 fully connected 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 's tangent kernel. The early spectral learning of the kernel depends on the depth. We also demonstrate that non-whitened data can weaken the silent alignment effect.
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 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 certain parameterization gives rise to linearized dynamics in parameter space (Lee et al.) 2019 Liu iet al. 2020) and consequently a kernel regression solution with kernel known as the neural tangent kernel (NTK) in function space (Jacot et al. 2018 Arora 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)) SSchölkoff Smola 2002). Moreover, they share many of the phenomena that overparameterized neural networks exhibit, such 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, phenomenon widely believed to be crucial to their success.
In this present study, we pursue alternative perspective the NTK, and ask whether neural network with an NTK that changes significantly during training can ever be a kernel machine for a data-dependen kernel: i.e. does there exist kernel function K for which the final neural network function f is f(z) Ep= a" Em=ra"K (æ, x²) with coefficients a that depend only on the training data? We answer in the affirmative: that large class of neural networks 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 kemel machines and deep learning which, unlike the infinite width limit, allows for the kernel to be shaped by the data.
TThese authors contributed equally.

---

Published as a conference paper at ICLR 2022
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 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 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 provided in Appendix B
we demonstrate the silent alignment effect by considering a simplified model where
2. In Section 3]
we provide an analysis of the NTK evolution of two layer linear MLPs 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
3. In Section 4 extend our analysis to deep MLPs by showing that the time required for align- ment scales with initialization the same way as the time for the loss to decrease appreciably. Still, these time scales can be sufficiently separated to lead to the silent alignment effect for which 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]
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 non- linear networks with multiple output classes. For anisotropic data, 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 kernel machine that uses the final NTK.
1.1 RELATED WORKS
Jacot et al. (2018) demonstrated that infinitely wide neural networks with an appropriate parame- terization trained on mean square error loss evolve their predictions as a linear dynamical system limitation of this kernel regime is that the neural network inter- nal 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 a that depend on the test point æ. This dependence on makes the construction not kernel method in the traditional sense that we pursue here (see Remark in (Domingos) 2020)).
with the NTK at initalization.
Phenomenological studies and models of kernel evolution have been recently invoked 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 al.] 2021 SSan Shaa & Bordelon 202T, Woodworth al., 2020 Chen al. |020 Geiger et al. 2021 Bai et al. 2020). Experiments have shown that lazy learning can be faster but less robust than feature learning (Flesch al.) 202T) and that the generalization advantage that feature learning
provides to the final predictor is heavily task and architecture dependent (Lee et al. 2020p. Fort et al. (2000) found that networks can undergo a rapid change of kernel early on in training after which the network's 's output function is well-approximated by kernel method with a data-dependent NTK. Our findings are consistent with these results.
Stöger& & Soltanolkotabi| (2021) recently obtained a similar multiple-phase training dynamics involy- ing 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 al. 2021), where it is referred to as the aadle-to-saadll regime.
2

---

Published as a conference paper at ICLR 2022
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 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 controlled by the time evolution of the NTK. As a concrete example, we will consider learning scalar target function with mean square error loss through gradient flow. Let æ E II represent an arbitrary input to the network f(x) and let {x*, y | be a supervised learning training set. Under gradient flow the parameters 0 of the neural network will evolve, so the
output function is dimedeepndentt and we write this as f(x, t). The evolution for the predictions of the network on test point can be written in terms of the NTK K( .)
af(m.t). 0f(w' .t] 08
88
as
f fxx, ([x,,)=] = K(x x." t)(y* - f((x" tt),
where 1) is the learning rate. If one had access to the dynamics of K( æ², 1) throughout all t, one could solve for the final learned function f" with integrating factors under conditions discussed in Appendix A
ELe dt (x) cxp (-1 Kr at' ) (y" "-ooff("))
(xxx)= fu(x))
(2)
Here, ke(a))" = K(x,, xx" t). [K]]ar K( xx" f), and - is the initial error on point x' We see that the final function has contributions throughout the full training interval e (0, (0). The seminal work by Jacot et al. (2018) considers infinite-width limit of neural networks, where the kernel function K, (æ, x') stays constant throughout training time. In this setting where the kernel is constant and fo(x") O, then obtain a true kernel regression solution f(x) En, k(z))* K yd for a kernel (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 data-dependent manner. In this paper, we consider a setting where the initial kernel small in scale, aligns its eigenfunctions early on during gradient descent, and then increases only in scale monotonically. As concrete phenomenological model, consider depth L networks with homogenous activation functions with weights initialized with variance o². At initialization Ko (a(ar,'') a²) O(o2L-2), [a())) O(a²) (see Appendix B). We further assume that after time t, the kernel only evolves in scale in a constant direction
K(x, x', (
[021-2 ((t))
t t>t
(3)
where K xr, !) evolves from an initial kernel at time increases monotonically from 2--2 to 1. In this model, one also obtains a kernel regression solution in the limit where a with the final, rather than the initial kernel: f(z) = (x) K y +
K (x, x') by = and g(t)
")a²). We provide 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 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 a universal and predictable way.
We show empirically that our results carry over to nonlinear networks with ReLU and tanh activa- tions under the condition that the data is whitened. For example, see Figure [ where we show the silent alignment effect on ReLU networks with whitened MNIST and CIFAR-IO images. We define
alignment as the overlap between the kernel and the target function
y Ky where y E Rº is
| <
3

---

Published as a conference paper at ICLR 2022
Kernel and Loss
1.0
Test Prediction NTK
ReLU MLP on Whitened MINIST
Test Prediction NTK
ReLU MLP on Whitened CIFAR
Alggmmen
0.6
0.2
500
-55 55.0 -2.5
0.0
25
5.0
7.5
Test Pred ction NN
Test Predi Prediction NN
(a) Whitened Data MLP Dynamics
(b) Prediction MNIST
(c) Prediction CIFAR-10
Loss Alignment
1.
Test Prodiction NTK
Wide Res-Net on Whitened CIFAR
0.8
Algrement
9.-
0.4
0.0-
250
5d0
750
1000
1250
1500
0
2
Test Prediction NN
(d) Wide Res-Net Dynamies
(e) Prediction Res-Net
Figure 1: A demonstration of the Silent Alignment effect. (a) We trained 2-layer ReLU MLP on P = 1000 MNIST images of handwritten O's and I's which were whitened. Early in training, around & = 50, the NTK aligns to the target function and stay fixed (green). The kernel's overall scale (orange) and the loss (blue) begin to move at around maximal final alignment value linear networks is overlayed (dashed green), see Appendix [E.2] (b) We compare the predictions of the NTK and the trained network on MNIST test points. Due to silent alignment, the final learned function is well described as a kernel regression solution with the final NTK K However, regression with the initial NTK is not good model of the network's predictions. (c) The same experiment on classes. Here use MSE loss on width 100 network with initialization scale
300. The analytic solution for the
1000 whitened CIFAR- 10 images from the first two
0.1. (d) Wide-ResNet with width multiplier = 4 and blocksize of h = trained with = 100 training points from the first two classes of CIFAR-10. The dashed orange line marks when the kernel starts growing significantly, by which point the alignment has already finished. (c) Predictions of the final NTK are strongly correlated with the final NN function.
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 well. In Figure 1 (d) and (e) show learning dynamics and network predictions for Wide-ResNet (Zagoruyko & Komodakis 2017) on whitened CIFAR-10 trained with the Adam op- timizer (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, we discuss in Section 5 and Appendix ! Our results suggest that the final NTK may be useful for analyzing generalization and transfer as we discuss for the linear case in Appendix
KERNEL EVOLUTION IN 2 LAYER LINEAR NETWORKS
We will first study shallow linear networks trained with small initialization before providing analysis for deeper networks in Section 4 We will focus 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 demonstrate that our analytic solutions match empirical simulations in Appendix C.5
We assume the data points xn E RRD nu 1,..., P of zero mean with correlation matrix Σ = ΣP x TT Further, we assume that the target values are generated by a linear teacher function yd = sBr x for a unit vector By. The scalar s merely quantifies the size of the su- pervised learning signal: the variance of [y) BB We define the two-layer linear neu-
4

---

Published as a conference paper at ICLR 2022
87
P1
81 (c) Phase
(a) Initialization
(b) Phase
Figure 2: The evolution of the kernel's eigenfunctions happens during the early alignment phase for
t R x. but significant evolution in the network predictions happens for t > t2 log(sa-"), (a) Contour plot of kernel's norm for linear funetions (z) 3 æ. The black line represents the space of weights which interpolate the training set, ie X 3 = At initialization, the kernel is isotropic, resulting in spherically symmetric level sets of RKHS norm. The network function represented as blue dot. (b) During Phase the kernel's eigenfunctions have evolved, enhancing power in the direction of the min-norm interpolator, but the network function has not moved far from the origin. (c) In Phase II, the network function W moves from the origin to the final solution.
ral network with N hidden units as f(x) = a Wa. Concretely, we initialize the weights with standard parameterization ag (O, /N), Wiy (O DD/ Understanding the role of in the dynamics will be crucial to our study. We analyze gradient flow dynamics on MSE cost 2P En (f(x") )"
= 1
Under gradient flow with learning rate 1, the weight matrices in each layer evolve as
dt a = da
aL = WE (sBr w a), d
W = ow a (s3r W a)
aL
T
(4)
dt
The NTK takes the following form throughout training.
(( = W Wa' + |a|2>> x'
(5)
Note that while the second term, simple isotropic linear kernel, does not reflect the nature of the learning task, the first term W W x' can evolve yield anisotropic kernel that has learned representation from the data.
3.1 PHASES OF TRAINING IN TWO LAYER LINEAR NETWORK
We next show that there are essentially two phases of training when training two-layer linear network from small initialization on whitened-input data.
Phase alignment phase which occurs for J. In this phase the weights align to their low rank structure and the kernel picks up a rank-one term of the form x 333 x'. In this setting, since 0, which is saddle point of the loss function, the gradient
the network is initialized near W. the loss is small. Consequently, the magnitudes of the weights and kernel evolve slowly.
Phase II: A data fitting phase which begins around escapes the initial saddle point W. 0 and loss decreases to zero. In this setting both the kernel's overall scale and the scale of the function fxx t) increase substantially.
Log(so 2). In this phase, the system
If Phase I and Phase II are well separated in time, which can be guaranteed by making small, then the final function solves 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
3.1.1 PHASE I: EARLY ALIGNMENT FOR SMALL INITIALIZATION
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 matrix X has all of its nonzero singular values equal. We let
5

---

Published as a conference paper at ICLR 2022
B represent the normalized component of By 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 a - (a==w3+0(00)
dt w saB + +(o").
(6)
dt
Truncating terms order 0² and higher, we can solve for the kernel's dynamics early on in training
(xx, tt) = go cosh(2nst) x' [33 + +1
t t log(s/o²)
(7)
where do is an initialization dependent quantity, see Appendix CT The bound on the enror is ob- tained in Appendix C.2 We see that the kernel picks up a rank one-correction 33 which points in the direction of the task vector B. indicating that the kernel evolves in direction sensitive to the target function = 8Br x. This term grows exponentially during the early stages of train- ing, and overwhelms the original kernel Ko 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
3.1.2 PHASE II: SPECTRAL LEARNING
We now assume that the weights have approached their low rank structure, as predicted from the previous analysis of Phase 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 al. 2018 Du et al. 2018), we note the following conser- vation law 28 (at)a(tt) W (t))W(t) 0 which holds for all time. If we assume small initial weight variance g², aa WWT O 0 at initialization, and stays that way during the training due to the conservation law. This condition is surprisingly informative, since it indicates that W is rank-one up ((r) corrections. From the analysis of the alignment phase, we also have that W T W x /3 These two observations uniquely determine the rank one structure of W to be a3 + ((r). Thus, from equation 5 it follows that in Phase II, the kernel evolution takes the form
K(x, x'; = u(t)²xx [33 + I] x' + O(r).
(8)
where (t))² = = 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 uut)? analytically, u(1)² 20.2st (c_at - 1 + s/u8)=1, using the results of (Fukumizu 1998 Saxe et al.)
2014) as we discuss in C.4 This is sigmoidal curve which starts at un and approaches 8. The transition time where active learning begins occurs when ea Rs 8/u6 = t R s- log(s/o"). This analysis demonstrates that the kernel only evolves in scale during this second phase in training from the small initial value un O(a²) to its asymptote.
Hence, kernel evolution in this scenario is equivalent to the assumptions discussed in Section 2 with gtt) = u(t)?, showing that the final solution is well approximated by kernel regression with the final NTK. We stress that the timescale for the first phase f1 ~ 1/s, where eigenvectors evolve, is independent of the scale of the initialization o2, whereas the second phase occurs around l2 s l log(s/0* This separation of timescales i < l2 for small a guarantees the silent alignment effect. We illustrate these learning curves and for varying a in Figure C.2
3.2 UNWHITENED DATA
When data is unwhitened, the right singular vector of W aligns with ΣB early in training, as we show in Appendix C.3 This happens since, early on, the dynamics for the first layer are d W a(t) E. Thus the early time kernel will have a rank-one spike in the Σ direction. However, this configuration
di
not stable as the network outputs grow. In fact, at late time W must realign to converge to W x a3 since the network function converges to the optimum and f= a Wx = sB x, which is the minimum la norm solution (Appendix G.I.. Thus, the final kernel will always look like (x,
sæ |3 + +I However, since the realignment of W's s singular vectors happens during the Phase 1l 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 we explain in Appendix

---

Published as a conference paper at ICLR 2022
4 EXTENSION TO DEEP LINEAR NETWORKS
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: RRD + R takes the form f(x) = wL WL-1... l.. The gradient flow dynamics under mean squared error (MSE) loss become
ØL d W = 11 OW?
II F'
(+B T '
II w PCc
(9)
di
where i- W W w..wb RD is shorthand for the effective one-layer linear network weights.
Inspired by observations made in prior works (Fukumizu) [998] 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 d 2i [W'WeT W W
O. This condition indicates 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.
4.1 NTK UNDER SMALL INIIIALIZATION
We now consider the effects of small initialization. When the initial weight variance r² is sufficiently small, W e+1+TT W O(a²) s 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 w (t) W u(t)re++(()rr) T where (t( are unit vectors. Using only this balancing condition and expanding to leading order in a, we find that the NTK's dynamics look like
K(x, x' = 1(t)2(t-1) x T [(L ())(((()))))) + +1/2's =1()(0(())) +O()).
(10)
We derive this formula in the Appendix E We observe that the NTK consists of rank-1 correction to the isotropic linear kemel x at' with the rank-one spike pointing along the rl) direction. This true dynamically throughout training under the assumption of small σ. At convergence r(t) + B, which is the unique fixed point reachable through gradient descent. We discuss evolution of u() below. The alignment of the NTK with the direction ,3 increases with depth L.
4.1.1 WHITENED DATA VS ANISOTROPIC DATA
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 ri (1) quickly aligns to /3 since the first layer weight matrix evolves in the rank-one direction d W = v(1),3 throughout training for a time dependent vector function vtt). As consequence, early in training the top eigenvector of the NTK aligns to B. Due to gradient descent dynamics, W W grows only in the 33 direction. Since the T1 quickly aligns 33 due W growing only along the ,3 direction, then the global scalar function e(t) = satisfies the dynamics ē(l) = c(1)2-2/2 [s - c(() in the whitened data case, which is consistent with the dynamics obtained when starting from the orthogonal initialization scheme of [axe al. (2014) We show in the Appendix E.I] that spectral learning occurs over a L+2 where f is the time required to reach half the
dt
timescale on the order of f1/2 x -(L-2] 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 a for L > 2, empirically alignment in deep networks occurs faster than spectral learning. Hence, the silent alignment conditions of Section are satisfied. In the case where the data is unwhitened, the r (t) vector aligns with Σ early in training. This happens since, early on, the dynamics for the first layer are d W ~ vtt) Σ for time dependent vector vtt). However, for the same reasons discussed in Section 3.2 the kernel must realign at late times, violating the conditions for silent alignment.
'Toough we focus on neglecting the O(o²) initial weight matrices in the main text. an approximate anal- ysis for wide networks at finite o and large width is provided in Appendix H.2 which reveals additional dependence relative layer widths.
7

---

Published as a conference paper at ICLR 2022
Loss Alignment
1.
16h
0.8
100
103
lage
g 165
06
40 100
=100
-10
10"
16
80
10
107
10
10
10-
(a) ODE Time to Learn
(b) Dynamies
(c) Time To Learn
Figure 3: (a) Time half loss scales in power law with for networks with L > 3:
f1/2 ~ [I-2] L+2 (black dashed) is compared with numerically integrating the dynamics o(() = (2-2/ c) (solid). The power law scaling of 11/2 with a is qualitatively different than what happens for L = 2, where we identified logarithmic scaling f1/2 log(o-3). (b) Linear networks with D 30 inputs and N 50 hidden units trained on synthetic whitened data with | = 1. We show for a L = 3 linear network the cosine similarity of W W with /33 (dashed) and the loss (solid) for different initialization scales. (c) The time to get 1/2 the initial loss and the time for the cosine similarity of W'''' with 33 to reach 1/2 both scale as a L+2 however
one can see that alignment occurs before half loss achieved.
4.2 MULTIPLE OUTPUT CHANNELS
We next discuss the where the network has multiple C' output channels. Each network output,
we denote as fe(x') resulting in C2 kernel sub-blocks Ke (zz, x²) = Vfe(x) Vfe In this
context, the balanced condition W'W' x W W implies that each of the weight matrices rank-C, implying rank-C' kemnl.. We give an explicit formula for this kernel in Appendix [H] For concreteness, consider whitened input data Σ = I and teacher with weights 3 E RCD The singular value decomposition of the teacher weights f Ea 8, S, determines the evolution of
each mode (Sax et al., 2014h. Each singular mode begins to be learned at ta To guarantee silent alignment, we need all of the Phase time constants to be smaller than all of the Phase II time constants, In the case of two layer network, this is equivalent to the condition Smez log (sman"0)) so that the kernel aligmment timescales are well separated from the timescales spectral learning. We see that alignment precedes learning Figure H. (a). For deeper networks, discussed in 4.1.1 alignment scales in the same way as the time for learning.
log (s. un
5 SILENT ALIGNMENT ON REAL DATA AND RELU NETS
In this section, we empirically demonstrate that many of the phenomena described in the previous sections carry over 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 pre- dictions. This argument is based on phenomenon discussed by Chizat et al. (2119). Consider an initial scaling of the parameters by a. We find that the relative change the loss compared to the
V C 0(a-4) which becomes very large for C
relative change in the features has the form
small initialization a as we show in Appendix D 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. How- 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 In even more realistic settings, like ResNet in Figure [ (d), we also signatures of the silent alignment effect since the kemnel does not grow in magnitude until the alignment has stabilized.
We 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 = USVT and construct partially whitened dataset X,, USV where 7 (0,1). As 7 0
8

---

Published as a conference paper at ICLR 2022
the dataset becomes closer 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 y + 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 consequence, the NN function deviates from a kernel machine with the final NTK.
Partial Whitened Spectrum
Classification Tist Error
109
1.2
0.32
0.5
0.32
3.23 10
2.00
0.5
y-
y-
0.2
0.24
133
200
402
600
803
0.22
200
430
600
910
(a) Input Spectra
(b) Train Loss
(c) Test Error
Kemail Norm
Kernel Aligmment
Precicttor NN
000
0.5
[K(t)]
0.1
402
200-
0.0
600
0.2
1.0
(d) Kernel Norm
(c) Phase I alignment
(1) Predictor Comparison
Figure 4: Anisotropy in the data introduces multiple timescales which can interfere with the silent alignment effect in ReLU network. Here we train an MLP to do two-class regression using Adam at learning rate 5 x 10 3 (a) We consider the partial whitening transformation on the 1000 CIFAR- 10 images A X for > E (0, for covariance eigenvalues Σvg- = X Uk. (b) The loss dynamics for unwhitened data have a multitude of timescales rather than a single sigmoidal learning curve. As a consequence, kernel alignment does not happen all once before the loss decreases and the final solution is not a kernel machine with the final NTK. (c) The network's ' test error on classification. (d) Anisotropic data gives a slower evolution in the kernel's Frobenius norm. (e) The kernel alignment very rapidly approaches an asymptote for whitened data but exhibits longer timescale for the anisotropic data. (f) The final NTK predictor gives a better predictor for the neural network when the data is whitened, but still substantially outperforms the initial kernel even in the anisotropic case.
CONCLUSION
We provided an example of case where neural networks can learn 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 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 at odds with the learned neural network function being a kernel regression solution; i.e. superposition of kernel function on the training data. While we provide one mechanism for a richly trained neural network to learn a kernel regres- sion 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 kernel machines for some data- dependent kernel, one can hopefully shed light on their generalization and transfer learning capabilities (Bordelon et al.
2020 Canatar al. 2021 Loureiro al. 2021 Geiger et al. 2021) and see Appendix F

---

Towards Published as a conference paper at ICLR 2022
ACKNOWLEDGMENTS
CP acknowledges support from the Harvard Data Science Initiative. AA acknowledges support from an NDSEG Fellowship and Hertz Fellowship. BB acknowledges the support of the NSF-Simons Center for Mathematical and Statistical Analysis of Biology at Harvard (award # #1764269) and the Harvard Q-Bio Initiative. We thank Jacob Zavatone- Veth and Abdul Canatar for helpful discussions and feedback.
REFERENCES
Kyle Aitken and Guy Gur-Ari. On the asymptotics of wide networks with polynomial activations. ArXiv, Abs/2006.06687, 2020.
Guillaume Alain and Yoshua Bengio. Understanding intermediate layers using linear classifier probes. arXiv preprint arXiv:1610.01644, 2016.
Anders Johan Andreassen and Ethan Dyer. Asymptotics of wide convolutional neural networks. ArXiv, 0asp220000667755 2020.
Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 244-253. PMLR, 10-15 Jul 2018. URL httpps://prrceest m1 mlr. www.press/////aate
Ba. html
Sanjeev Arora, Simon Shaolei Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In NeurlPS, 2019.
Michael Baake and Ulrike Schlaegel. The peano-baker series. Proceedings of the Steklov Institute of Mathematics, 275(1):155-159, 20.1.
Yu Bai, Ben Krause, Huan Wang, Caiming Xiong, and Richard Socher. Taylorized training: better approximation of neural network training at finite width, 2020.
Aristide Baratin, Thomas George, César Laurent, R. Devon Hjelm, Guillaume Lajoie, Pascal Vin- cent, and Simon Juclste-Julien. Implicit regularization via neural feature aligmment. In AISTATS, 2021.
Mikhail Belkin, Siyuan Ma, and Soumik Mandal. To understand deep learning we need to un- derstand kernel learning. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning. volume 80 of Proceedings of Machine Learn- ing Research, pp. 541-549 PMLR, 10-15 Jul 2018. URL https: ptttps://proceedings mir.
press /aaea/v80///0 w/b/ 1 kin n 8a. html
Blake Bordelon, Abdulkadir Canatar, and Cengiz Pehlevan. Spectrum dependent learning curves in kernel regression and wide neural networks. In Hal Daumé III and Aarti Singh (eds.), Pro- ceedings of the 37th International Conference Machine Learning, volume 119 of Proceed- ings of Muchine Learning Rexearch, pp. 1024-1034. PMLR, 13-18 Jul 2020. URL https:
//proceedings.. wttp::/swwwmpppps.somooemmse wttpee/s//l11//breeeooo/22 www.eloood...........o.......... hhmll
Roger W Brockett. Finite dimensional linear systems. SIAM, 2015.
Abdulkadir Canatar, Blake Bordelon, and Cengiz Pehlevan. Spectral bias and task-model alignment explain generalization in kernel regression and infinitely wide neural networks. Nature Commu- nications, 12, 2021.
Shuxiao Chen, Hangfeng He, and Weijie J. Su. Label-aware neural tangent kemne:: Toward better generalization and local elasticity, 2020.
Lénaïc Chizat, Edouard Oyallon, and Francis R. Bach. On lazy training in differentiable program- ming. In NeurlPS, 2019.
Uri Cohen, Sue Yeon Chung, Daniel D Lee, and Haim Sompolinsky. Separability and geometry of object manifolds in deep neural networks. Nature communications, 11(1):1-13, 2020.
10

---

Published as a conference paper at ICLR 2022
Corinna Cortes, Mehryar Mohri, and Afshin Rostamizadeh. Algorithms for learning kemnlls based on centered alignment. The Journal of Machine Learning Research, 13(1):795- 828, 2012.
Pedro Domingos. Every model learned by gradient descent is approximately a kernel machine. arXiv preprint arXiv:2012.00152, 2020.
Simon Shaolei Du, Wei Hu, and J. Lee. Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced. In NeurIPS, 2018.
Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from feymman diagrams. Interna- (ttienrreieeeeet/
tional Conference on Learming Representations, 2020. URL
Sorum?id-SlgFvANe
Timo Flesch, Keno Jucchems, Tsvetomira Dumbalska, Andrew Saxe, and Christopher Summerfield. Rich and lazy learning of task representations in brains and neural networks. bioRxiv, 2021. doi: 10.1 10.1101/2021 04. 24.23.4444.... 4412234441112
biorxiv..vvvv........i..ii 18. URL hnttps://ww/ http://////wbbbooi...oooooo wttp://.wwr..ootee/eeaaaea
2021/04/ 000000////22222 04. 23. 44 128
Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M Roy, and Surya Ganguli. Deep learning versus kernel learning: an empirical study of loss landscape geom- etry and the time evolution of the neural tangent kernel. H. Larochelle, M. Ranzato, Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, wtw.:roceeedinng.......i..e..e neurips. Papereapppppp.
pp. 5850-5861. Curran Associates, Inc., 2020. URL
Cc /cc/paper/ae seaper/2020/2 2tt///220022//2 / ottpee//555000000001551111122220
Kenji Fukumizu. Effect of batch learning in multilayer neural networks. Gen, 1(04):1E-03, 1998.
Mario Geiger, Leonardo Petrini, and Matthieu Wyart. Landscape and training regimes in deep learning. Physics Reports, 924:1-18, 2021. ISSN 0370-1573. doi: htpp://doi.org/10.1000 physrep. 2021 1. 00.000000000 040001...... URL http:: http://ww...eeeeeoeeeeeeoeceeees wwwsscceiiiiiieeeeeeoooeeoeieeee rect. com/ science/article Landscape and training regimes in deep learning.
pii 031/8037015732100120
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: convergence and gen- eralization in neural networks (invited paper). Proceedings of the 53rd Annual ACM SIGACT Symposium on Theory of Computing, 2018.
Arthur Jacot, François Ged, Franck Gabriel, Berfin Şimşek, and Clément Hongler. Deep linear networks dynamics: Low-rank biases induced by initialization scale and 12 regularization. arXiv preprint arXiv:2106. 15933, 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1412. 6980, 2014.
Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jascha Sohl-Dickstein. Wide neural networks of any depth evolve as linear models under gradicnt descent. ArXiv, abs/ (bbb//990.00700, 2019.
Jaehoon Lee, Samuel S. Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study, 2020.
Tengyuan Liang and Alexander Rakhlin. Just interpolate: Kernel "ridgeless" regression can gener- alize. CoRR, abs/1808.00 003.00387, 2018. htt/ttp:://// http://arxvv.v/.. http///awx......a..../. 180. 00387
Chaoyue Liu, Libin Zhu, and Misha Belkin. On the linearity of large non-linear models: when and why the tangent kernel is constant. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Bal- can, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 18954-15964 Curran Associates, Inc., 2020. URL URL_https:://prrceeddnng hww.iiiinnnnnnnnnissi...om.. pdf
cc 6c72pp//2000111000eeeee222222333
Bruno Loureiro, Cédric Gerbelot, Hugo Cui, Sebastian Goldt, Florent Krzakala, Mare Mézard, and Lenka Zdeborová. Capturing the learning curves of generic features maps for realistic data sets with a teacher-student model. CoRR, abs/2102.081277, 2021. URL https: :ttppii.////.....//////////.mm.m orq/
abs 22 02.00..4 08127

---

Published as a conference paper at ICLR 2022
Roman Novak, Lechao Xiao, Jiri Hron, Jachoon Lee, Alexander A. Alemi, Jascha Sohl-Dickstein, and Samuel S. Schoenholz. Neural tangents: Fast and easy infinite neural networks in python. In International Conference Learning Representations, 2020. httpttt//ppeeeeee/////
net./forum? http:/orrummiidxll1y3yyy
Carl Edward Rasmussen and Christopher K. I. Williams. Gaussian processes for machine learning. Adaptive computation and machine learning. MIT Press, 2006. ISBN 026218253X.
Andrew M. Saxe, James L. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dy- namics of learning in deep linear neural network. In In International Conference on Learning Representations, 2014.
Bernhard Schölkopf and Alexander Smola. Learning with kernels: support vector machines, reg- ularization, optimization, and beyond. Adaptive computation and machine learning. MIT Press, 2002. BRLLttp:: /tt:///wwww... wor wow ldcat.
Haozhe Shan and Blake Bordelon. Rapid feature evolution accelerates learning in neural networks, 2021.
Dominik Stöger and Mahdi Soltanolkotabi. Small random initialization akin to spectral learning: Optimization and generalization guarantees for overparameterized low-rank matrix reconstruc- tion. arXiv preprint 2tt.::2206..10013. 2021.
Blake Woodworth, Suriya Gunasekar, Jason D. Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Jacob Abernethy and Shivani Agarwal (eds.), Proceedings of Thirty Third Conference on Learning The- ory, volume 125 of Proceedings of Machine Learning Research, pp. 3635-3673. PMLR, 09-12 Jul 2020. httptttss:/pprcceeeeee
mlr press /tt//vv12//////// hw11////wwwott/00aa html
Chulhee Yun, Shankar Krishnan, and Hossein Mobahi. A unifying view on implicit bias in training linear neural networks. arXiv preprint arXi:22010002511, 2020.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks, 2017.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. In 5th International Conference on Learning Representationss ICLR 2017, Toulon, France, April 24-26. 2017, Conference Track Proceedings. OpenReview.net, 2017. URL https: https::/openzeeier. net / /t//ooounnit Syt.:
12

---

Published as a conference paper at ICLR 2022
Appendix
A DERIVATION OF EQUATION 2
A.1 TRAINING POINT PREDICTIONS WITH TIME VARYING KERNEL
Given a training set of P data points {(x", ) the dynamics of the network training er- rors [A,] = f(xx", yd close terms of a time-varying neural tangent kernel [Ki]ar K [(x", .t.))
dl A: -KA
(11)
We introduce the transition matrix D, E RPPP which has the property that A, 4; and Po I, we obtain the matrix evolution equation = - This equation can be solved formally in terms of the Beano-Baker series (Baake Schlaegel 2011 Brockett 2015)
+ =I [ Ks + ddssk... /" ds2K,2
(12)
|0- dsp Kei 1-" dsyKes lr dsyk,, +...
02
(13)
which can easily be verified to solve de (t) - ((t) with initial condition +(0) = I. Under
the condition that L" (t) commutes with (t),, which 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
[%/ dskk
1 [[ ds ds&g
1 1 [. dsk,
3 -..
(-1) k!
[ dsk,
Φ =I
6
= exp (-/-" K,ds)
(14)
Thus, under the condition that K, commutes with S K,ds we can exactly solve for the training error dynamics in terms of integrating factors
A: 4 = exp (-/, K,ds (KKKK Ao.
(15)
We expect this formula to hold approximately whenever the eigenvectors of K are approximately equal the eigenvectors of fo K,ds.
A.2 TEST POINT PREDICTIONS WITH TIME VARYING KERNEL
Given access 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 test point f(x is given by
d _e 1((0) = k.
(16)
where (x))² = K(x, .t,t). t). This gives the final value of f to be
(z) = (z) fo(x) + / 00
dt (x))" eep (-n (nf Kvdt'
]
(y fo).
(17)
This is exactly equation
13

---

Published as a conference paper at ICLR 2022
B KERNEL EVOLUTION IN SCALE ONLY
We consider the model of kernel evolution introduced in Section scale for t 7()) and is of small overall size for < 7(1),
2 where the kernel evolves only in
[eKo(z,0',0) K (xx, x', t) (g( (2,
1 (()
(18)
t> T(e)
This model allows for alignment of the kernel while small in the time window E (0, r). followed by scale growth only for t (T, oc). The time threshold will generally depend on the initial kernel scale c. For example, in depth linear MLPs, e ~ o 2f.-2 and T T~ 1.+2 with initialization scale as we show in Figure E.3 We will define differentiable function h(t) = S g(t')dt' so that (1()) g(t). l((T) = O. and lime h(t) = oc. This last condition follows from g's continuity and the assumption that Do g(l) = 1. We will first show that the final neural network has the form f(x)= = fo (x) + (z) K (y - fo) + O(eT(e)). First, we need to calculate the errors
A (t) y ft) E Rº made on the training examples. These satisfy the dynamics
d A()
(EKo(t)A(t) (x
T
(19)
di
(1)
1>T*
where Ko E PxP are P x P gram matrices; e.g. [Ko(t) Ko( w²,t). t). The vector k(z) has entries given by [μ = K(xxx"). For t E (0, r), the error vector follows the dynamics
d (t) - ((t)
(20)
dt
Introducing operator norm of a matrix, c op = max E2l2, we will now bound the operator norm of the change the transition matrix (() introduced section A.
Lemma 1. Let Kn (ax, (K) op represent the maximum operator norm of Ko achieved on the interval (0, T) Let (t) E PxP be the transition matrix for the linear dynamics of equation 20 so that (t) (t) (t) and (0)) I. Then,
+ (T) (0) < (e)
(21)
Proof. We begin by noting that, due to the triangle inequality,
+ (r) {0))op = c I (o(t)
op $ef, [Ko(t) (t) dt
(22)
S E ( (o((t) (()) dt S eko l(¹ [0())) dt.
We will now establish that [ op S 1. Note that for any vector E R that
1 d (t) = v + (t)u -e (t) K(t) S0.
(23)
where the final inequality follows from the fact that K(I) is positive semidefinite for all I. Therefore we have shown that ot $ (0))opp Topp Using this inequality, we find that
(T) (0))op S cko l (t))opp dt S ckor(c).
(24)
With the above Lemmal I we can bound the discrepancy A (7) and A A()), namely
A (0))2 = ( (7) a
(25)
< (T E (0) op A (0) S (e)
This inequality must therefore hold entry well, so that
( (0) + (e)k)
(26)
We will now establish how the training predictions A (1) evolve for the second interval E (7, 20).
14

---

Published as a conference paper at ICLR 2022
Lemma 2. Suppose from t (T, oo) that (t) obeys the dynamics di (t) <x At)
where A (T) is as in equation 26 Then, for all E (T, 00),
A (t) = exp (-((t)((kKK))) [ fo) + O(cr(e))a))
(27)
Proof. The differential equation can be solved through cigendecomposition and integrating factors. Let (t()) represent the k-th component of A (t) in the eigenbasis of K which is static for
(r, oo). Let the corresponding eigenvalue of Koc be λς. The scalar variable (x(t) obeys the
dynamics
dt 4x(1) = (ah()=>>h(())Ail)))
(28)
This can be solved with integrating factors, noting that d [EAsA(H) AR((]] 0. This implies that # (x(t) e htx (T( Written vector, (t) exp (- (th(t)K>>) (T)) Since by Lemma we have A (T) = (0) + (E)koT(E)), we obtain the desired result.
We will now combine the results of the previous two lemmas which analyze the evolution of the network predictions on the training set give our main silent alignment result, which specifies what the neural network function predicts for an arbitrary test point æ.
Theorem 1. Let the kernel have dynamics of Equation 18] where g(t) is continuous, integrable function with lim g(t) 1. The function learned by the neural network is
f(x) fo(x) (x)) K y + Oc(er(e)).
(29)
Proof. Using Lemma and 2 we know the full dynamics for training predictions A (4). Using (1), we can solve for the final predictor f(x:) by integrating dynamics f(x, k( () (t).
f(æ) - fo(x)= = L/ ko(x, A (t)ddt kz)()) {
h1(() exp (-h(t) A (T) dt (30)
We will now bound the first term. Taking Ko max |ko( t) we get that
[ ko(o, (11) S ( (M)
''" [ko t)
S (1 (0) (e(e(()))
We can now integrate the matrix exponential in the second term, using the fact that
N'(t) exp (- (-h(t) = exp (-p(- dh = K=1
(31)
Using the fact that A(7) (0) + +(((e))) from Lemma |
we arrive at the desired result
f(x) fo(x) A (x) K Ao (e(r(e))
(32)
We have now established that, given the kernel dynamics in Equation 18 f(x) - fo(x)) converges to the kernel regression solution with final NTK as - 0 provided Tim CT(C) = 0. This is generic in the settings we consider in this paper for networks with small initialization. In this small initialization setting, fo is also negligible that f(æ) itself is kernel regression solution. For example, in a linear depth neural network with initial weight scale a,, the initial scale of the kernel is o 2r.-2 while the time to alignment scales
o2 thus et ~ o can be made arbitrarily small by taking a + 0. Lastly, the initial network outputs fo(x) ~ I² can also be made arbitrarily small.
15

---

Published as a conference paper at ICLR 2022
C PHASES OF LEARNING AT SMALL INITIALIZATION
C.I PHASE I: TWO LAYER NETWORK AND KERNEL ALIGNMENT
We now present an analysis distinet 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 æ M with M = W W + la I. Our goal is to determine the eigendecomposition of M. Introduce the variables (t) 3 MB [| + 3 W W and r(t) W 3. These dynamics form closed two dimensional linear system early in training
d (())
0
1 0
(())
= 2njs
+O(0"), 1, 0
dt ((())
[
r(t)
(33)
(()) r())
[]
2 + 2 (20 - ro) [-1] 2nst + 0
=
(go + ro)
-o)
The variable q(t) represents the alignment of the NTK with the optimal direction B while r(t) defines the alignment of the network with the teacher. We see that this aligmment increases expo- nentially with timescale & ~ --- While the above equations hold for early time and small initialization for any initial condition qu- To. we can further estimate these initial values under random initialization provided the input dimension is large. We stress that this limit is not nec- essary for the silent alignment, but allows for a nice simplification. For Gaussian initialization a, N/ /NN). Wiy NOO o² DD) with large D, we have
(go) = 2 211 N (
o²
(ro) 0, (r0) = 04
(34)
D
In the large D limit, we have with high probability > ro and thus r(t) = 4u sinh(2nst) and q = qo c2sh()n)t).. Note that this gives the quamtities (t) " (tt = 3 W a early
training. Now consider a unit vector v which is orthogonal to the solution 3 v = We find that the projection of M along this direction evolves dynamically as:
v M(t)u= = =2v [Ba T W + T all
dt
(35)
= 2r((t).
We can conclude that v Me is equal g(t) up to an additive initialization constant. We see that this is evolving half as quickly as B M3 = (t( Since v Mov ~ O(o²) is small compared to the exponentially growing (t), the only matrix that satisfies these two conditions must necessarily take the form
M (t) go cosh(2nst) [33T + = Mo.
(36)
The first term, which growing exponentially t will eventually overwhelm the randomly initial- ized kernel Mo, which is O
C.2 PHASE I: ERROR IN THE LEADING ORDER APPROXIMATION
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
d w - a (sB W WT)
w (s33 (sB- dt
(37)
dt
One can use these equations to solve for the dynamics of the r, variables:
d
dt gt) = 2sr - TT [W aa W + a TWW a] B = 2sr - 12,2 +
Σ (a T w
(38)
1.
d
di (1) = 2sq a TWWW + aa WB 2sg 2/a 2r +a Taa"
WW W3
(39)
16

---

Published as a conference paper at ICLR 2022
The second equality in equation 38 comes from inserting complete basis of states including 3 and v; 1 3 into the last term of the left-hand side. The second equality in equation 39 comes from writing WW WWT + aa 2aa + (WW) T aa ). Note that the final term in brackets the right hand side a conserved quantity for linear networks, and so is always of order O(a²).
Assuming the solutions for q. are valid to order o², we get that dt a We; = O
d
dt We; (B 3) W We (B B) T)Du;
(40)
WW T O(o²) is also constant in Further noting that because of the conservation law aa time. This gives us that
aa [aa ww W3 < o² W = o
(41)
We now note that r; |a both grow as a (a, s-independent) constant times times d2 The correc- tion to the dynamics of both equations is then bounded by constant times o Ast This will be less than o² as long as t satisfies
t<
log
(42)
43
For o² < s, the alignment time t 1/s falls within this range and we guaranteed alignment to the Ganguli-Saxe configuration.
The error of the full solution time be bounded by the integral of this error bound from to t, namely constant times o <s/s. As long as s > ot, we are guaranteed that the error of the kernel
(aa²) as given in equation 7
C.3 PIASE I: TWO LAYER ANALYSIS WITH UNWIITENED DATA
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. this case,
a sWE3,
(43)
dl
W sa."
(44)
dt
which holds asymptotically as t/ log closed linear dynamical system
0. We introduce the following variables which form
(1( = 3 W a, r2(t) = ja|3,
rs(t) = 3 W WEB,
(45)
r1(t) 3 ΣWT. a,
ro(() =3 EW T
Introduce the constants to show that 33 ΣT
B ΣaB. Using the weight dynamics, it straightforward
0
0
0
0
0
2
0
+(t)
b
0
r(t), log(o 2).
(46)
0
0
0
2b
0
This matrix has eigenvalues À (0,-vb,
2VD). Since there are only two positive cigenvalues Vb, 2v6, suffices to consider evolution along those two cigendirections, where the kernel and neural network function will be amplified. Evolution along these direcions give
0
vb avb b bvb
r(t) ~ aht
+ C2e
2o
(47)
0
17

---

Published as a conference paper at ICLR 2022
t= 0
t=t 1
t =
62
82
52
B1
B1
B1
(a) Initialization
(b) Transition Time
(c) Min l norm solution
Figure C.I: Kernel evolution on anisotropic data consists of two alignment phases. (a) At ini- tialization the level curves of (3 M exhibit spherical symmetry. (b) After the initial phase I alignment, the matrix M exhibits spike in the Σ direction. (c) At long times, the network func- tion and kernel's spiked direction need converge to the minimum & norm solution as we explain in G.I This requires realignment of the kernel at late times, eliminating the preconditions for the silent alignment effect.
where C1, C2 are constants determined by intialization. At large time, the large eigenvalue mode = 2vb will dominate. Decomposing W Wo + at)) (t) + (,) we find that the only self consistent solution is (t) = 0, (t() = b 1/2, This implies that the kernel evolution will take the form
K(x, x .t) K( x', 00) + (a(t)) Mæ',
1
(48)
M
VB ΣマB ΣØ3 Σ+
We see that the kernel evolves along the directions Σ early in training for unwhitened data. We visualize the two stages of learning unwhitened data in Figure C.1
C.4 PHASE WHITENED DATA
Consider two layer network / a W æ where balance has been achieved W u()) and att) = u(t)â. Once this balance condition is stable for fixed à we can calculate the time derivative of (t)
d a(t) ú(t) unt) - u() à
(49)
dt
Letting c(t) u find that a(t) 2t is u(t 20(t) (s c(t)), which the two layer dynamics derived in Saxe al. (2014). This dynamics has solution c(t) 20t-1tt/ca
C.5 SOLUTIONS TO THE FULL TRAINING DYNAMICS OF LINEAR NETWORKS AT SMALL INITIALIZATION
By combining the analyses of the subsection with the exact solutions discussed in we can match both solutions obtain formulas for r(t) and q(t) for the entire network's training path that are exact up to ( corrections. Up to ( ) we then have that
2 sinh(2st)
r(t) s 2et 1) + 28/90*
(50)
2 cosh(22t)
((t) (c2st 1) + 2s/90
(51)
This yields that the initialization constant /2/2 plays the effective role of cn in the Ganguli-Saxc solution for phase II. Equation 34 yields the expected value of this initialization constant. We have
18

---

Published as a conference paper at ICLR 2022
100
1250
Experierr
10
-4 In
0.8
1070
0.8
0.6
850
0.4
650
02-
420
0.0
500
1006
1500
2000
1000
1938
2302
(a) Loss Dynamics
(b) Time to Le = 0.5
(c) Alignment = 10
Figure C.2: Initialization scale controls the time spent in phase I, where the network escapes the saddle point near W, 0 and the kernel aligns to the task. (a) The loss curves for two-layer linear networks with small initialization follow sigmoidal trajectories as in Saxe et al. (2014) which transition from their maximum to minimum at a time which decreases with initialization scale. Theory is shown in black dashed lines. (b) Verification of the Phase I time 1/2- measured as the time for the loss to reach one half its original value. This scales logarithmically with 0². (c) The alignment the kernel eigenfunctions happens before the loss appreciably decreases for = 10 10-4, evidenced by the kernel alignment curve. The analytically obtained maximum alignment value is overlayed in dashed green.
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
Exact Analytic Solutions
Exact vs Analytic Solutions
10a
10
10
<10 empirical
ati analytk
110 empirical
atdl analytic
10
analytic
aralytic
200
4000
6:00
20
40
60
(a) Synthetic
(b) Whitened MNIST
Figure C.3: Overlay of empirical and exact solutions for r()) in two layer linear feedforward networks for synthetic and two-class MNIST whitened datasets. (a) We take = 25, N = 10. (b) We take D 784, 100.
D BALANCING OF WEIGHTS IN DEEP LINEAR NETWORKS
The balance condition discussed in the main text holds for deep linear trained with any loss function of the form C = En e(S", ((",)))) (not just MSE) since, was shown by Arora et al. (2018): Du et al. 2018
n dt (W'WeT)
1 d
Σ Ot
05h WET + We 0fa T OW
0ff
OW
di 0FF
0.n
W + WI+T 0F OW
(52)
SWH+I
1 d
(WL+IT w(t))
77 dt
19

---

Published as a conference paper at ICLR 2022
The second line follows from the first since the following quantities identical
W WL+1 WL-1T wt (x* (TT ...TTT W
(53)
OW
W 0
WIT OW W (WeT ...w)) TTT W
(54)
By inspection these two quantities are equal. Thus we have, for any loss function, a deep linear network has the following conservation laws
di [w' - W W O.
(55)
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 where the network has single output, we inductively prove that each layer's 's weight matrix is W = u(t)r (()r((t) We will assume this formula true for layer & + 1 and prove it must hold for layer l since
W W w vet))
(56)
This implies that W is rank one with right singular vector equal to rt-1(t). Thus, the decomposi- tion for each layer the form W ((t))=++(())r)) for some unit vector rz((). Similar analysis can be performed for the multi-class setting.
NTK FORMULA FOR DEEP LINEAR NETWORKS E
The neural tangent kernel for a linear network f(x) = W.L...... L is defined as an inner- product over all gradients
0f(x) 0f(x') OW ow?
K(x Σ
(57)
II
w x x' +
Σ t=2
II
w
x WT w w .'.
Under the balanced assumption that W = u(t))e+ - O(o), expanding the kernel to leading order yields the following form:
xxx,x') u(t)2L-2g x [LL ())=)))++((0)()))))+ O(Go,
(58)
1.0
1.0
Vernel Aligmment
0.6
Ws Alignment
0.8
0.6
0.4
0.6
0.4
0.2
0.2
0.0
10
102
10
109
10a
162
10>
10-
102
Figure E.1: The depth dependence of loss dynamics, kernel alignment and the alignment first layer weights in linear network on synthetic whitened data in D 30 dimensions. (a) The loss reaches half its imitial value after t ~ o 1.+2 steps for > The decay rate of the loss becomes sharper with depth. (b) Final kernel alignment increases monotonically with depth T. and approaches 1 as + oo. (c) The alignment of the first layer weights W, with the optimal direction B approaches for all models.
20

---

Published as a conference paper at ICLR 2022
Alignmont
LD55
04
10 ReLU MLP Depth
45
49
(a) Loss
(b) Alignment
(c) Test Comparison
Figure E.2: ReLU networks across depths trained 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.
E.1 DEEP LINEAR NETWORK DYNAMICS UNDER BALANCE
In this section, will consider the dynamics of the variable once the balance condition is satisfied. Let wt = Tt)(t). Then the dynamics for u(t) under the balancing assumption is
w1 r d d dt dt
1-1 II w (s- x(t)
1-1
u(t)
(s-u(t)4
II w' u(t)2l- n(t)
(59)
which implies the fact i(t) u(t)² 1--111166 -s c) Changing variables to c(t) =(t)² we obtain
?(t) = u(t) = u(1)21- (ss - u(()²) = 2/ 2/ c(t)).
(60)
When of is initialized to very small value compared to we can
d c(t) c(t))-2/// =
+ (o(t) -1+2/L GD
-1+2/L
L
st
dt
(61)
- 1ي2
(L - 2)
1-22
= c(t =
Co
st
I
This implies a timescale to learn of s-1
2 L+2
L-2
We can approximate timescale for first layer's singular vector (1) to align to as well. Let be vector orthogonal to B.
d ( 3 2 W W.. 0
(62)
dt W 2
W
This suggests that alignment in deep network should also occur on a timescale of t o L+2 While there is no strict separation of timescales in terms of the scaling of alignment and learning with a, find that alignment tends to precede a significant drop the loss show in Figure
E.2 FINAL NTK FOR DEEP LINEAR NETWORKS
Independent the structure of the data, the first vector r -> and the final NTK has the following form:
K(x, x') 2 [L 1) + +1] x' O(o).
(63)
This formula is merely consequence of the balancing condition and convergence to the optimum u(t) J. + s. We provide empirical support that final kernel alignment increases with depth in Figure
E3
21

---

Published as a conference paper at ICLR 2022
Kernel Alignment
0.3
k-
0.1
2000
(a) Alignment Dynamics and Depth
(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 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 form K x (L-1)yy) 1)yy" Ko, illustrating why alignment of the final NTK to class structure increases with depth L.
partially whitened ( 0.25
E.3 FORMULA FOR THE INTERMEDIATE NNGP KERNELS
Let (3 represent the unit vector pointing in the optimal direction for the scalar output To calculate this, all that needs to be assumed that the network has converged to its optimum and that the weights satisfy the balance property that W u'ri+1"²- By the convergence assumption u s 1//h and ri = B.. We stress that this holds for arbitrary input correlation structure E, the final NNGP kernel for layer l can also be computed. Expanding the weights and collecting terms at leading order yields:
Ke(x,x') 221 33 2'+ +0().
(64)
Evaluating on the training set gives K = (yy E grx.
F GENERALIZATION ERROR IN TRANSFER LEARNING TASK
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 compares with the inductive bias of the original isotropic kernel x'. In particular, we study the offline generalization performance of kernels of the form
K (x, x'; A) = x AA33 + x'
(65)
In Section we showed that could be altered by changing the network depth. Concretely, our transfer learning problem consists of training a linear probe 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 (x. x. A) with A which depends on the chosen layer and the depth of the network. For simplicity. we assume that the data are generated according to a simple Gaussian distribution æ ~ (O, I) and that the target values are generated with linear function y w.
We decompose the new task vector w a, + /(1 - a²]w where W 3 = O. 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 (a 1). If w B. increasing the alignment A strictly decreases the generalization error. This is illustrated in Figure
F.1 DERIVATION OF LEARNING CURVES
We will discuss the average case generalization error in the transfer learning setting. Prior work has shown that the generalization performance of kernel regression be calculated through kernll
22

---

Published as a conference paper at ICLR 2022
Kernel Alignment Strength
Transfer Task Alignment
Eg with P= 10
1.0
14
0.8
0.75
41
0.7
0.5
A-0
2-0.0
A-5
A-1
2-0. 0.5
2-022
0.25
A-10
05
(a) a = 0.75
(b) =
(c) Fixed P
Figure F.1: The offline generalization error a transfer task with the learned linear kernel K [BT + +I] w' and Vnew uB v (1-o0)wy
(a) For new transfer task which is correlated with the learned function 3, 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 B 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 a are both large.
eigenvalue problem (Bordelon al. 2020 Canatar al. 2021 Loureiro al., 2021
(x, x' (x)())) (=))
(66)
Once this integral eigenvalue problem is solved for eigenvalues Aς and orthonormal eigenfunctions Φk. the average case generalization error at P training examples
Eg 1-7 Σ
(x))) k)? = A Σ
EXP+R*
X P
A
X
(67)
(NpP+
In our case, interested in the generalization performance of the linear kernel
K(x, x') xx [33 + I]a'. x'.
(68)
Since K is a linear kernel, its eigenfunctions should be linear functions (x(x)) = Φi Assuming that the data distribution has identity covariance, we find
(x T [A3 + +1] x' x. & [A33 + +1]x =
(69)
This implies that the Φ; vectors are eigenvectors of M. The first eigenvector is Φ G with
eigenvalue ê1 A + 1. The other D eigenvectors can be chosen as any frame the D
dimensional subspace orthogonal to B. Each of these D - eigenvectors has eigenvalue A 1. Using these results, and the fact that w aΦi V1 Ldw we can calculate the expected generalization error.
a
1 1 a²
Eg
((1 A)P x)²
(P k)2
1 A
D 1
(1+ D-1 (1+
7
(1+A)P+*
P+R
(P+k)*
By the result proven in Canatar et al. (202I). the lowest possible eπror for fixed A occurs by maxi- mizing the fraction of variance along the large eigenvalue direction, corresponding to a = +l.
G LINEAR NTKS DURING GD LEARN THE SAME FUNCTION
In the overparameterized setting where D > P. all linear networks discussed in the Section converge the minimum norm interpolator when the data is whitened. Specifically, letting the learned neural network function be written (xx 3 at, and the data matrix X E RDx and
23

---

Published as a conference paper at ICLR 2022
target labels y R represent the training data.
The solution vector & solves the constrained
optimization problem
min |1811 s.t. X 3 y, B
(70)
which is the kernel regression solution for the initial kernel Ko(x,x²) x x'. This un- surprising due to a symmetry argument: when Bo R O, the only privileged point on the affine space X B y is the point closest to the origin, which is precisely the solution above. Surpris- ingly, the final -icudo-inverse solution sB 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 ((x, t)
[A(1)33T r] x' (see Section 4.1) which would induce the following kemnll
interpolation problems
min 3 [A(1)33T
' s.t. s.t. x --
(71)
The solution to this optimization problem is indeed the kernel regression solution with kemnel ((t) since the learned function takes the form fx) En a ((x, x" .t) with -(t)-ly. Using
the Sherman therman-Morrrson rule, 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+ A(t) (8.3)²
BT [A(t) + +1] B 10
(72)
Now, we let 3 ===00 =-8 B1. where B B = =0. This is the general decomposition for the set of interpolators which have the property X [B + ++1]]= ===yy
min |31
(73)
91
The solution is merely to _1 0. Thus the optimal solution is therefore the same for any fi-
nite value of A. However, the final RKHS norm the learned function , [A(t) + r] B decreases with time, indicating that the kernel becomes more aligned with the pseudo-inverse direc- tion as A increases.
G.1 DEEP LINEAR NETWORKS FROM SMALL INITIALIZATION LEARN PSEUDO-INV INVERSE
In this subsection of the appendix, we will use our theoretical technology for balanced linear net- works to demonstrate the universal learned function for any data, not just whitened input, pro- viding an alternative derivation to the result proven in Theorem of Yun al. (2020). This analysis is performed in the - 0 limit, where We u(t)r (t)r(() as showed in Section Under this condition the learned function f(x)
A is defined through weights 3 W1 W₂ ...WL, (rt))rr((t)))) (t). We see that the direction of the learned function is controlled entirely by rı (t). It suffices to prove that (f) E span (ai.,
..., xp) for all f to show that the network
learns the pseudo-inverse solution 3 X(X X)- y, where X R. Dx P and y E R.P are the training data and targets respectively. Note that by gradient descent, we have
dt Wt(t) = WTJ W.,W, Σ (in-fu(t))a [.( = ((t)--1r2(0 Σ 1( (on-fa(t)) ft(t)a=T.
(74)
1
From the balance condition Wt(t) u (t)) (T) we also have
d (t) [t())==(t) (t)) (()))))t))))
(75)
dt
Equating the two above expressions for de left gives the following
Wi(t) and taking an inner product with r2(() from the
(t) = u(t)* 1-1 Σ (oan-fa(t):) (t))))²" - [ū(t) + (()))=()) r())).
(76)
24

---

Published as a conference paper at ICLR 2022
Thus, if Tı (t) span(22,..., xp) then (1()) E span(21,..,.PP} that the full dynamics of rı (t) lie in the subspace spanned by the training data. At initialization, have Wi(t) x =(1) - , + +(²²) so the initial T1 vector will indeed align with the span of the training data in the + 0 limit.
By the fact that ß = u (ra) the learned linear coefficients & must also be in span ..... pp} so B En ay These must also interpolate the data provided P. giving the following condition
3.x"
Σ Rv Bp y/ a
x
(77)
This is exactly the minimum l norm interpolating solution which solves
min B
B s.t. X T y.
(78)
While anisotropy of the data makes no impact on what function is ultimately learned in the linear network case, the anisotropy alignment are satisfied in a nonlinear network, which can prevent the final function from being NTK regressor with final NTK.
have a signficant influence on whether the preconditions for silent
H FINAL KERNEL IN MULTI-CLASS NETWORKS
For network with C' output channel, balancing and alignment guarantee that the configuration of the network is orthogonal and balanced as in the setting Saxe et al. (2014). One can then integrate each mode separately to obtain the final kernel
W
un (072(1)
Ke.o' (x, = x Moow'
(79)
L-1
+ Mc.e = ded Σ aa (t
Σ
2(e 2(L 2(L- (t)e e Σ Σng 3
up (t)rffT
where the Cartesian unit vectors E RC are one-hot class output This shows that the contributions to kernel depend on how well the class output channels align with unit vectors rf. Further, the singular values Ua can evolve at different timescales depending the structure of the data.
Specifically, for depth L network, both the the alignment time {(( and the time to learn given singular value Sa scale as s 1=10=22 2--L as shown in appendix E. The differences alignment times
Aoo == ((T)) - -((" for modes Sa. $g therefores scales At (L) = a3-
H.1 FINAL NNGP IN MULTI-OUTPUT CASE
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. Let B e RCDD D represent the optimal weight matrix which has the property (in the over-parameterrize case) 3x = y At the optimum, the neural network must learn , = WL !. Computing the SVD of B En Saz.v reveals that " g1/L and r Va and r Zc. Using these facts, it is easy to derive the final NNGP kernel for layer
Kg(0,x)) = x IT ...
.....TTT t.....W'z'
= Σ (w rdrd IT x' ( E/L z.
(80)
Evaluating on the training set X DxP gives K₁ X 13T 1 8 x. which interpolates between X X at layer e = and YT layer L.
25

---

Published as a conference paper at ICLR 2022
Then Dapth
Kare) und
62-
(a) Linear depth
Linear depth
(c) Tanh depth
Figure H.l: Demonstration of 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 o² 10 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 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.
H.2 MORE REFINED BALANCING ANALYSIS
We can derive corrections to our 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) (f) W (t) 1+1 (t) W.w, W
W 1 W
(81)
Note that Wo E X has entries with zero mean and variance o²/Ne in the standard pa- rameterization. The products of initial matrices therefore Wishart distributed. For sufficiently large widths, we can approximate the initial weight matrix products with their expectation over the random initialization
Ne+2
W W1+11 W o²
I.
(82)
No+1
This concentration becomes more accurate as the widths No, Ne+1 >> oc. For the last layer the number of classes C Nr is sufficiently large, we also obtain similar concentration for the last layer W WL xo 0² Nu I. Repeating the backward induction on the conservation law, we find the following recursively defined singular value decompositions
W' = Σ (000)
(1)² o (t)² + 0² (1
Net2 (t)² -0
N41
(83)
1-1 Nx-1
9t = L
N
We note that the corrections ge vanish if all layers A > é have the same width. Let ua (I.) (t).
The condition for convergence is
II
II (u,) 2 00 =
(84)
In the 2² - 0 limit, we can solve that ua ~ g/LL as before. However, we can now obtain leading order corrections (in 0²) which take the form of the form
go²
I.
L(L -
Nk+1
[ua]² 1 2²
2Lsa
1/L
0² +0,
98
N
(85)
t=1
26

---

Published as a conference paper at ICLR 2022
which reveals that the size of the correction depends not only but also the depth and network widths. Suppose all network widths were equal N₂ = Mr. then the term = and there is no contribution from the first moment of the random weights.
1 LAZINESS IN HOMOGENOUS NETWORKS
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 change compared to the loss early in training.
Let f R represent the function outputs the training set. We will compute the time derivatives of the loss and the network gradients.
d
d0
di Vof(x)
V2 f(x)
f2)(x) Σ(O P") -")
(86)
di
d
d0 2
C
IV (y f)
(87)
dt
dt
Here, | lop denotes the operator norm of a matrix. We are interested in the ratio of the loss' time derivative to the gradient's 's time derivative. With an initialization scale of f ~ ((a4) we find
[
C
V2 f2x( En f" fl2
(2)(2) E,, V f'll"
Vfl
V (y f)
VfIVF 3/2
(88)
where in the last step we approximated yd - f zs y for small initialization scale since yd 0(1)
and fur ~ ( Now we will estimate the scale of each of the terms above. For a homogenous model f ~ (()--1) and V ~ 0(a--2). Counting powers of in numerator and denominator,
we find that this quantity of interest scales as
f())
L 0 de
(89)
Vf
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.
J RESNET EXPERIMENTAL DETAILS
Below, we provide the alignment and loss dynamics for wide resnet for CIFAR- 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 al. (220)) and is based
the original architecture of Zagoruyko & Komodakis (017)) with a widening factor of k = 4 and single block per ResNet group b 1, giving final network with 8 trainable layers. For both Figure | (d) and (e) well Figure J.1 use Adam with a learning rate of = 10-5 and initial weight scale of a = 0.3 in standard parameterization for all intermediate blocks. For the first conv layer, we used a 6.0. We find that small initial weight variance in the first layer gives rise to less stable learning and worse kernel alignment.
we provide comprehensive results for different depths which we control by increasing the number of blocks per group corresponding to WRNs with 6b + 1 trainable conv layers.
Below, in Figure J.2
J.1 ADAPTIVE OPTIMIZERS AND THE RELEVANT KERNEL
Many adaptive gradient methods compute updates to parameters 6, according to
ac
(,()) (t) 00
(90)
27

---

Published as a conference paper at ICLR 2022
Wide Res-Net Unwhitened CIFAR
Loss and Alignment
0.8
(100 Alignment
Test Prediction NTK
0.6
0.4
0
500
1000
1500
2000
2500
-2
-1
2
Test Prediction NN
(a) Dynamics
(b) Predictions
Figure J.1:
The dynamics and predictions of a Wide-Resnet with & 4 and b on 100 unwhited CIFAR-10 images from the first two classes.
1.0
0.8
Kernel Norm
4
Loss
0.4
2
0.2
0.0
1000
1000
2000
3000
(a) Training Loss
(b) Kernel Norm
0.5
FeTx foul![
0.8
Alignment
0.4
0.6
0.3
1000
2000
3000
Wide Res Net Blocks
(c) Alignment Dynamics
(d) Predictor Comparison
Figure J.2: The silent alignment effect is preserved across a large range of depths in WideResNet trained on whitened 10 images. The number of blocks per group 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.
where (t) are time-varying functions which are computed in terms the history of gradient mo- ments for parameter f, or in terms of its instantaneous gradient. The relevant kernel at time f which governs instantancous evolution of network predictions is
Σ (t) 0f(z, t) 0f(x',t)
K(x, ((. =
(91)
00,
00,
since /(x) En K(x, x0. 1)v" - f(x ()). Though we do not calculate this kernel which is relevant to the adaptive learning rate scheme since it not supported Neural Tangents API, this could be worthy future investigation.
28

---

Published as a conference paper at ICLR 2022
Kernel Norm
Loss
(a) Train Loss Dynamies
(b) Kernel Norm
Alignment
0.2
0.2
0.0
500
Wide Res-Net Width
(c) Alignment
(d) Predictor Comparison
Figure J.3: Varying the ResNet widening parameter also alters the kernel and loss dynamics. (a)
The loss curve for b WRNs with widening factor k. Wider networks train
quickly. (b)
The kernel norm increases more rapidly for wider networks but changes by a smaller amount. (c)
Alignment reaches close its asymptote by the time the kernel (dashed). (d) The final kernel is much better predictor of the NN function than the initial kernel.
grows to 10% its final value
29


Reviews:
[{"review_id": "nkhJeFRBPy", "reviewer": "Reviewer_puin", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n-\tThe authors identify an interesting silent alignment effect, where the trained network acts as kernel predictor, and provide a thorough analysis for 2-layers linear networks and deep linear networks.\n-\tThe authors demonstrate the effect empirically.\n\nWeaknesses:\n-\tI think the authors should contrast their results with recent results on implicit bias of linear networks with small initialization. Specifically, it is shown in [1] (theorem 7 therein) and [2] (theorem 6.3 therein, for the case of 2 layers) that linear fully connected networks with small initialization converge to minimum $\\ell_2$-norm solution (i.e. kernel solution with linear kernel). Note that this is correct also for anisotropic data.\n-\tThe effect of depth is not so clear from the simulations. It can be useful to show on the same plot the kernel alignment for a fixed initialization scale but different depths.\n\nQuestion:\n-\tPlease explain why in Eq. (8) the error is $O(\\sigma^2)$.\n-\tPlease explain the transition in Eq. (15), why the integral starts from 0 and not from $\\tau$ ?\n\nMinor:\n-\tThere is a typo in the definition of $K_{c,c'}$.\n-\tIn figure 5, only one $\\sigma$ is used.\n\n\n[1] Yun et. al. A unifying view on implicit bias in training linear neural networks. ICLR 2021\n\n[2] Azulay et. al. On the Implicit Bias of Initialization Shape: Beyond Infinitesimal Mirror Descent. ICML 2021\n", "summary_of_the_paper": "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).", "summary_of_the_review": "Overall the paper is interesting but the relation to implicit bias results is not discussed. I will be happy to adjust my score if the authors can address my questions, and also compare the results to known implicit bias results.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "MU2P-iPWhFL", "reviewer": "Reviewer_dMqU", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Summary\n\nKernel 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. \n\nThis 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.\n\n## Strengths\n\n* From the claims and the cited references this seems to be the first work that figures out the so-called \"silent alignment\" effect. The effect is analyzed on (deep) linear networks and is a bit surprising that how well it aligned with empirical evaluations. \n* Combining the conservation law shown in Du et al. (2018) with early dynamics from small initialization seems to be the main \"magic\" that enables the analysis. The proof technique can itself be interesting to the community.\n* I checked the derivation of major claims made in the paper and they looked correct to me.\n\n## Weaknesses\n\n* The analysis was done on (deep) linear networks. The justification on ReLU nets is through a combination of the observation of Chizat et al. (2019) and empirical evidence. It is unclear from the text if this assumption is widely tested on a variety of ReLU nets or just the single example of 1K whitened CIFAR images used in the paper. It would be nice to see if the silent alignment behavior is consistent across shallow/deep, narrow/wide nonlinear nets.  \n\n* The set-up of the simulation is a bit unclear to me. From the figures and captions I cannot see what the time span is for values like (t=1000). Does the results hold if the training spans as long as what people do in practice? The behavior might also change if the dataset contains way more patterns than the 1000 image experiments presented in the paper.\n\n## Overall suggestions\n\nI vote for acceptance of the paper. I should note that the submission should be best reviewed by deep learning theory people, since I am not particularly familiar with the latest literature that focuses on learning dynamics beyond the lazy regime. My assessment is based on the technical correctness and proper justification of the claims but may have missed some important related work that lowers the significance of contribution. \n\n## Questions \n\n* Figure 1b/1c: does the individual points correspond to training data or test data?\n* Section 3.1: I don't understand what you mean by \"In this setting the overall scale of the weights and kernel, evolves slowly since the network is initialized close to the saddle point at W, a=0\".\n* I don't see how the left hand side of (11) is approximated by the right.\n\n## minor points\n\n* Section 2: \"However, we analytically show in sections 3 and 4 analytically\"\n* Missing period on page 5, last paragraph.", "summary_of_the_paper": "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. \n\nThis 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.", "summary_of_the_review": "I vote for acceptance of the paper. It provides a new perspective of the \"rich regime\" of neural network training. The claims are backed by analysis on (deep) linear networks and align well with empirical evaluation. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "1tB1HsB0Ta-", "reviewer": "Reviewer_XmeY", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n1. Understanding the behavior of neural networks outside the NTK regime is a major question in the theory of deep learning. This paper gives a precise characterization of the kernel evolution in certain settings, in which the learning can be cleanly divided into two phases. The end result is that the final predictor is equivalent to a kernel regressor wrt a data-dependent kernel.\n2. The paper also investigates how varying certain design choices (e.g. depth, relative learning rates, covariance of data) affects the silent alignment effect.\n\nWeaknesses:\n1. There appears to be a small mistake in the treatment of deep linear networks. Since different layers have different learning rates, the corresponding NTK should also be a weighted sum of the contributions from all layers, instead of an unweighted sum (Eq. (36)).\n2. The analysis in the paper is not completely rigorous because of all the approximations used, and many times equations are used for approximate equalities. It would be much better if the approximations can be controlled explicitly and rigorously.\n3. Figure 5b, 5c: why does silent alignment not clearly happen for deeper networks (i.e. alignment keeps improving while the loss is decreasing)? Does the theory predict this?\n4. Overall speaking, the results are quite restricted since they are only for linear networks and whitened data. It would be very interesting if some characterizations can be given for the final NTK in non-linear networks.\n\nMinor points:\n1. page 5 line 5: \"achieved low loss\" --> \"achieved high loss\"?\n2. Eq. (10): $K_0(x, x')$ missing? Approximation error missing?\n3. Figure 5 caption: $10^{-2}, 10^{-1}, 10^{-2}$?", "summary_of_the_paper": "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.", "summary_of_the_review": "This paper makes some interesting and novel contributions towards understanding the learning dynamics of neural networks through the lens of NTK evolvement. On the other hand, the setting where a complete answer can be provided is restricted, and the analysis in the paper is not rigorous.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors make a case for a phenomenon of deep network training
that they call the "silent alignment effect": that, while the training
error is still large, the NTK associated with the network aligns its eigenvectors
with key directions in "feature space".  They support this with non-rigorous theoretical
analysis of linear networks, and extensive experiments with real networks on real data.
The consensus view was that this paper provides novel and useful insight into training
dynamics, in particular regarding feature learning.
Decision: Accept (Poster)