Feature learning as alignment: a structural property of gradient descent in non-linear neural networks
Authors: Daniel Beaglehole
Ioannis Mitliagkas
Atish Agarwala
Venue: Accepted by TMLR ()

Abstract: Understanding the mechanisms through which neural networks extract statistics from input-label pairs through feature learning is one of the most important unsolved problems in supervised learning. Prior works demonstrated that the gram matrices of the weights (the neural feature matrices, NFM) and the  average gradient outer products (AGOP) become correlated during training, in a statement known as the neural feature ansatz (NFA). Through the NFA, the authors introduce mapping with the AGOP as a general mechanism for neural feature learning. However, these works do not provide a theoretical explanation for this correlation or its origins. In this work, we further clarify the nature of this correlation, and explain its emergence. We show that this correlation is equivalent to alignment between the left singular structure of the weight matrices and the newly defined pre-activation tangent features at each layer. We further establish that the alignment is driven by the interaction of weight changes induced by SGD with the pre-activation features, and analyze the resulting dynamics analytically at early times in terms of simple statistics of the inputs and labels. We prove the derivative alignment occurs with high probability in specific high dimensional settings. Finally, motivated by the observation that the NFA is driven by this centered correlation, we introduce a simple optimization rule that dramatically increases the NFA correlations at any given layer and improves the quality of features learned.

Published in Transactions on Machine Learning Research (11/2024)
Feature learning as alignment: structural property of gra- dient descent in non-linear neural networks
Daniel Beaglehole*.
Sheaglehole@ucsd.cc edu
UC San Diego
Ioannis Mitliagkas
koonmismr
Mila, Université de Montréal
Google Mind
Atish Agarwala
thetish @agogll..ggoooooooo com
Google DeepMind
Reviewed OpenReview: https: openreview. net/ forum? id= JXCe22XZrr
Abstract
Understanding the mechanisms through which neural networks extract statistics from input- label pairs through fcature learning is one of the most important unsolved problems in supervised learning. Prior works demonstrated that the gram matrices of the weights (the neural feature matrices, NFM) and the average gradient outer products (AGOP) become correlated during training, in statement known as the neural feature ansatz (NFA). Through the NFA, the authors introduce mapping with the AGOP as a general mechanism for neural feature learning. However, these works do not provide a theoretical explanation for this correlation or its origins. In this work, we further clarify the nature of this correlation, and explain its emergence. We show that this correlation is equivalent aligmment between the left singular structure of the weight matrices and the newly defined pre-activation langent features at each layer. We further establish that the alignment is driven by the interaction of weight changes induced by SGD with the pre-activation features, and analyze the resulting dynamics analytically at early times in terms of simple statistics of the inputs and labels. We prove the derivative alignment occurs almost surely in specific high dimensional set tings.. Finally, we introduce a simple optimization rule motivated by our analysis of the centered correlation which dramatically increases the NFA correlations al any given layer and improves the quality of features learned.
1 Introduction
Neural networks have emerged as the state-of-the-art machine learning methods for scemingly complex tasks, such as language generation (Brown et al., 2020), image classification (Krizhevsky et al., 2012), and visual rendering (Mildenhall et al., 2021). The precise reasons why neural networks generalize well have been the subject intensive exploration, beginning with the observation that standard generalization bounds from statistical lcarning theory fall short of explaining their performance (Zhang al., 2021).
A promising line of work emerged in the form of the neural langent kernel, connecting neural networks to kernels in the wide limit (Jacot et al., 2018; Chizat et al., 2019). However, subsequent research showed that the success of neural networks relies critically on aspects of learning which are absont in kernel approximations (Ghorbani et al., 2019; Allen- Zhu and Li, 2019; Yehudai and Shamir, 2019; Li et al., 2020; Refinetti et al., 2021). Other
"Corresponding author. Work partially done as intern Google DeepMind.

---

Published in Transactions on Machine Learning Research (11/2024)
work showed that low width suffices for gradient descent to achieve arbitrarily small test error (Ji and Zhu, 2020), further refuting the idea that extremely wide networks are necessary.
Subsequently, the success of neural networks has been largely attributed to feature learning - the ability of neural networks to learn representations of data which are useful for downstream tasks. However, the specific mechanism through which features are learned is an important unsolved problem deep learning theory. A number of works have studied the abilities of neural networks to learn features in structured settings (Abbe et al., 2022; Ba et al., 2022; Nichani et al., 2023; Barak et al., 2022; Damian et al., 2022; Moniri et al., 2023; Parkinson et al., 2023). Some of that work proves strict separation in terms of sample complexity between neural networks trained with stochastic gradient descent and kernels (Mousavi-Hosscini et al., 2022).
The work above studies simple structure, such as learning from low-rank data or functions that are hierarchical compositions of simple elements. Recent work makes a big step towards generalizing these assumptions by proposing the neural feature ansatz (NFA) (Radhakrishnan et al., 2024a; Beaglchole et al., 2023), a general structure that emerges in the weights of trained neural networks. The NFA stales that the gram matrix of the weights at a given layer (known as the neural feature matriz NFM) is aligned with the average gradient outer product (AGOP) of the network with respect to the imput to that layer. In particular, the NFM and AGOP are highly correlated in all layers of trained neural networks of general architectures, including practical models such as VGG (Simonyan and Zisserman, 2014), vision transformers ( (Dosovitskiy et al., 2021), and GPT--fmily models (Brown et al., 2020).
A major missing element is an explanation for how and why the AGOP and NFM become correlated through training with gradient descent. In this paper, we precisely describe the emergence of this correlation. We establish that the NFA is equivalent to alignment between the left singular structure of the weight matrices and the uncentered covariance of the pre-aclivation langent kernel (PTK) (Section 2) features. We then introduce the centered newral feature correlation (C-NFC) which isolates this alignment process. We show empirically that the C-NFC is close to its maximum value of 1 at early times, and fully captures the NFA at late times for a variety of architectures (fully-connected, convolutional, and attention layers) over a diverse collection of datasets (Section 3). Our experiments suggest that the C-NFC drives the development of the NFA. Through this centering, we show that the dynamics of the C-NFC can be understood analytically in show
terms of the statistics of the data and labels at carly times (Section 4). Using this docomposition, that the NEA emerges as a structural property of gradient descent (analytical result in the commonly studied setting of uniform data on the sphere, under certain assumptions on the activation and target functions). In particular, the first non-zero derivatives of the centered NFM and AGOP will be asymptotically identical. We further characterize how the NEA depends on the data distribution, and explore this analytically and experimentally.
Finally, motivated by our theory, we design an intervention increase the influence of the C-NFC and make the NFA more robust: Speed Limited Optimization, layerwise gradient normalization scheme (Section 5). The effectiveness of the latter update rule suggests path towards rational design of architectures and training procedures that maximize the NEA notion of feature learning by promoting alignment dynamics.
2 Alignment between the weight matrices and the pre-activation tangents
In this section, we decompose the AGOP into the weight matrices and the feature covariance of the pre- activation tangent kernel (PTK), and demonstrate that the NEA is equivalent to alignment between the weights and these PTK features. We include glossary of terms for ease of reference in Appendix A.
2.1 Preliminaries
We consider fully-connected neural networks with a single output of depth I. 2 1, where I is the number of hidden layers, written f I&d R. We write the input to layer E E (o....,,, L} as Ie, where To x is the original datapoint, and the pro-activation he(x). Then,
he(z) = (00 xz,
Je+1 o(he(z))
(1)
2

---

Published in Transactions on Machine Learning Research (11/2024)
where @ is an element-wise nonlinearity, W(O E R is weight matrix, and ke is the hidden dimension at layer f. We restrict 1 to be the number of output logits, and set ko = d, where d is the input dimension of the data. Note that f(x) = (x) and we ignore LL-1- We train f by gradient descent on a loss function C(, X), where X is an imput dataset, and I is the collection of weights.
We consider a supervised learning setup where we are provided n input-label pairs (xe(1, ,,()))((1))) ,y()), .... (x(h), ((()) e E IR. We denote the network inputs (datapoints) X [ and the labels y E Rxx1, For a given network, the inputs to intermediate layers E {O....,LI L are written Xe RTx&, where Xo = X. We train a fully-connected neural network to learn the mapping from network inputs to labels by minimizing a standard loss function, such as mean-squared-crror or cross-cntropy, on the datasct.
EGOP
W W (NFM)
W KW (AGOP)
W QW
Ek= Th mod r ((++1) mod r with (first plot) captures the low-rank structure of the task. The NFM (WTW) (second plot) and AGOP (W T KW) (third plot) of a fully-connected network are similar to each other and the EGOP. Replacing K with a symmetric matrix Q with the same spectrum but independent eigenvectors obscures the low rank structure (fourth plot), and reduces the correlation from GF. = 0.93 to p (F, NF,,TQW) QW) 0.53.
Figure 1: Various feature learning measures for target function y(x)
r and inputs drawn from standard normal. The EGOP Exzp
Dy dv
On Bar
One can define two objects associated with neural networks that capture learned structure. For a given layer f, the neural feature matrir (NFM) F& is the gram matrix of the columns of the weight matrix W(),, i.c. Fe = (W(1)) TWO, Fe depends on the right singular vectors (and corresponding singular values) of W(O. The second fundamental object we consider is the average gradient outer product (AGOP) Ge, defined as )f(a(") 8f(a!")) T
G = A Σ.
duc
Since both these matrices have the same dimensions (ke+1 x ke41), we can consider their cosine similarity. We define the neural feature corvelation (NFC) by:
1/2 p (Fe, = tr Gc) tr (FT -1/2 tr (G, Gt)
G1)-1/2 1/2
(NFC)
This takes on values in [0, 1] since F, and Ge are both PSD.
Prior work has shown that in trained neural networks, the NFC will generally be close to 1 to varying degree of approximation. This notion is formalized in the Neural Feature Ansalz (NFA):
Ansatz 1 (Neural Feature Ansatz (Radhakrishnan et al., 2024a)). The Newral Feature Ansatz slates that, for all layers& E E [4] of a fully-connected neural network with L hidden layers trained on input data x(1) ,.... the NFC will satisfy p (Ge, Ft) 11
(.),...,!!!) (n))
The inputs to the covariance are the NFM and AGOP with respect to the input of layer f. Here, af(x) dmg E k1
denotes the gradient of the function f with respect to the intermediate representation zt. For simplicity, we may concatenate these gradients across the input dataset into a single matrix a((x) dec E Baxkr Note we consider scalar outputs in this work, though the NFA relation is identical when there are c 2 outputs, where in this case df(z)
.2t E xe is the input-output. Jacobian of the model f-
We note that while the ansatz states exact proportionality, in practice, the NFM and AGOP are highly correlated with correlation less than 1, where the final correlation depends on many aspects of training and architecture choice (discussed further in Section 3).
3

---

Published in Transactions on Machine Learning Research (11/2024)
The relation between the NFM and the AGOP is significant, in part, because for neural network that has learned enough information about the target function, the AGOP of this model with respect to the first-layer inputs will approximate the expected gradient outer product (EGOP) of the target function (Yuan et al., 2023). In particular, the EGOP of the target function contains task-specific structure that is completely independent of the model used to estimate it. Where the labels are generated from a particular larget function v(x) Ied + on data sampled from distribution μ4, the EGOP is defined as
By Oy Ox Ox
EGOP 11) BEnn
If the neural feature ansatz (Ansatz 1) holds, the correlation of the EGOP and the AGOP at the end of training also implies high correlation between the NFM of the first layer and the EGOP, so that the NFM has encoded this task-specific structure.
To demonstrate the significance of high correlation between the NFM and the AGOP in successfully trained networks, we consider the following choin-monomial low-rank task:
y(x)= Σ 2h mod r I(k+1) mod
(2)
where the data imputs are sampled from isotropic Gaussian distribution μ == =N, T). In this case, the entries EGOP μ)) will be 0 for rows and columns outside of the X sub-matrix corresponding to 17 (Figure 1), as y does not vary with coordinates I,+1....,22 Za. Within this sub-matrix, the diagonal entries will have value 2, while the off-diagonal entries will be either 1 or O. Therefore, EGOP(y, μ) will be rank r, where r is much less than the ambient dinension.
We verify for this task that the AGOP of the trained model resembles the EGOP (first and third panels of Figure 1). Here the NFA holds and the NFM (second panel) resembles the AGOP and therefore the EGOP as well. Therefore, the neural network has Icarned the model-independent and task-spocific structure of the chain-monomial lask in the right singular values and vectors of the first layer weight malrix, as these are determined by the NFM. In fact, previous works have demonstrated that the NFM of the first layer of a well-trained neural network is highly correlated with the AGOP of a fixed kernel method trained on the same dataset (Radhakrishnan et al., 2024a).
This insight has inspired iterative kernel methods which can match the performance of fully-connected networks (Radhakrishnan et al., 2024a;b; Aristoff et al., 2024) and improve over fixed convolutional kernels (Beaglehole et al,, 2023). Additional prior works demonstrate the benefit of including the AGOP features to improve feature-less predictors (Hristache et al., 2001; Trivedi et al., 2014; Kpotufe et al., 2016). Additionally, because the NFM is correlated with the AGOP, the AGOP can be used to recover the features from feature-less methods, such as kernel machines. The AGOP has also been show to caplure surprising phenomena of neural networks beyond low-rank feature learning including deep neural collapse (Beaglehole et al., 2024).
2.2 Alignment decomposition
In order to understand Ansatz 1, it is useful to decompose the AGOP. Doing so will allow us to show that the neural fcature corrclation (NFC) can be interpretod as an alignment metric between weight matrices and the pre-aclivation langent kernel (PTK). The PTK xCt)) 2) is defined with respect to a layer of neural network and two inputs I, 2, The kernel evaluates to:
(x, af(x) 0f(z)
(PTK)
Thg Thg
When the arguments x and z are omitted, K(O) Rax consists of the matrix of kernel evaluations on all pairs of datapoints. We may also consider the covariance ((0) E x of the features associated with the PTK, ax(x) E Bunke In particular, we define the feature covariance of the PTK as,
She
( 0f(x)
0f(x)
(PTK feature covariance)
The
Th
4

---

Published in Transactions on Machine Learning Research (11/2024)
Crucially, for any layer L, we re-write the AGOP in terms of this feature covariance as,
Gt = = k(9
This gives us the following proposition:
Proposition 2 (Alignment decomposition of NFC).
p GF) = ((wl0) w( (w(o)) (w0,) k(0))))
This alignment holds trivially and exactly K(() is the identity. However, the correlation can be high in trained networks where KT0)) non-trivial. For example, in the chain monomial task (Figure 1), K(0) is far from identity (standard deviation of its eigenvalues is 5.9 times its average eigenvalue), but the NFA correlation is 0.93 at the end of training. We also note that if K() is independent of W the alignment is lower than in trained networks; in the same example, replacing K(O) with a matrix Q with equal spectrum but random cigenvectors greatly reduces the correlation to 0.53 and qualitatively disrupts the structure relative to the NFM (Figure 1, rightmost column). We show the same result for the CelebA dataset (see Appendix M). Therefore, the NFA is a consequence of alignment between the left eigenvectors of W(()) and K() in addition to spectral considerations.
Note that K1 itself is the AGOP of the neural network f with respect to the pre-activations at layer e, i.e.
K-O0- af(a) (a) ) The
(3)
Ohe
For convolutional and attention layors, the K(()) are computed by additionally averaged over all patches and token positions in the input, respectively (see Radhakrishnan et al. (2024a); Beaglehole el. al. (2023) for the formulation of the NFA in these architectures).
3 Centering the NFC isolates weight-PTK alignment
We showed that the ncural feature ansatz is oquivalent to PTK-woight aligument (Proposition 2). We now ask: is the increase in the NFC due alignment of the weight matrices to the current PTK, or the aligument of the PTK to the current weights? In practice, both effects matter, but numerical evidence suggests that changes the PTK do not drive the early dynamics of the NFC (Appendix D). Instead, we observe that the alignment between the weights and the PTK feature covariance is driven by centered NFC that captures alignment between the parameler changes and the PTK. We then show this centered NFC can hold robustly in settings where the NFC holds with correlation less than 1, such as early in training and/ /ororr with large initialization. Finally, we establish analytically through the centered NFC how the NEA emerges as a structural property of gradient descent.
We bogin by considering a docomposition of the NFM and AGOP into parts that depend on initialization and a part that depends on the changes in the weight matrix. Let wi(f) and k:(t) be the weight matrix and PTK feature covariance for layer at time t. We can write the NFM and AGOP in terms of the initial weights (((") and the change in weights W-(0) = w(()) - ( as follows:
W. = WT W; + WT W, + W. Wo + W. Wo: F W,T,,
(4)
G-W, W. K,W, = W. K,W, + WT K,W, + W. K.W, + W. K.Wo
where we omitted é from all terms for ease of notation. The first term in each decomposition isolates the gram matrix of the changes in the weight matrix. We call W W the centered NFM and W KW the centered AGOP. The centered AGOP in particular measures the alignment of weight updates with the current PTK feature covariance. Both terms are 0 at initialization, and if the weight matrices change significantly (that is, if ||l|| > [Wol)), both the NPM and AGOP dominated by the centered terms. (Note: in the limit that.
5

---

Published in Transactions on Machine Learning Research (11/2024)
A. (MLP SVHN)
B.
NFA across MLP layers
Layer
0.9
0.8
0.7
0.
0.6
0.0
0.5
0.4
0.4
UC-
Mean
0.3
C NFC
0.2
0.2
09
93
50
120
150
Spochs
3
Layer
UC-NFC, epoch
UC-NFC, full training C-NFC, full training
C-NFC, 1 epoch
(VGG-11 on CIFAR-10)
Lager:
Lese2
UC NFC
Loyer
UC- NFC
C-NFC
Spouts
Earch
D. (GPT Shakespeare characters)
query
1.0
0.9
C-NFC
02
0.8
UC NFC C-NEC
1.0
m
M
W
0.6
UC NFC
820
1020
650
660 1.00
600
82o
1050
Epochs
Epochs
Epochs
Figure 2: Uncentered and centered neural feature correlations across (A,B) fully-connectedd, (C) convolutional, and (D) attention layers with large initialization scale. (A,C,D) show trajectories of C/UC-NFC over training. (B) shows NFC values across all layers of an MLP with five hidden layers, averaged over CIFAR-10, CIFAR-100, SVHN, MNIST, GTSRB, and STL-10 datasets. (A-C) are additionally averaged over three random seeds. Each row of (D) is an attention block (ordered from first to last in the GPT model), while the columns show correlations for query, key, and value layers, respectively.

---

Published in Transactions on Machine Learning Research (11/2024)
initialization scale vanishes to 0 (i.e small initialization), then we expect the centered and uncentered NFM and AGOP to coincide, as W = W,- - Wo + W,, where & is a step at the end of training.)
The decomposition in Equation 4 suggests that updates which drive correlation of the centered NFM and centered AGOP can also drive the overall value of the original NFC. Inspired by this observation, we define the centered NFC (C-NFC) as,
p ( Wi (lo, (((
For the remainder of the paper, will refer to the original NFC as the uncentered NFC (UC-NFC) to avoid ambiguity.
The centered NFC is consistently higher than the uncentered NFC across training times, architectures, and datasets (Figure 2). This especially holds at early times and in deeper layers of a network (Figure 2C and D, VGG-11 on CIFAR-10 10 and the GPT model on Shakespeare respectively). For MLPs, we conduct broader set of experiments and verify that the trends hold a wide range of vision datasets. We also note that the C-NFC at carly times is relatively robust to the initialization statistics of the weight matrix W unlike the UC-NFC (for additional experiments, see Appendix H). We will return to this point in Section 5. These experiments suggest the C-NFC is responsible for improvements in the uncentered correlation.
High C-NFC values can drive increases in the UC-NFC as long as the contered NFM and contered AGOP are increasing in magnitude during training. We can confirm that the weights move significantly from initialization, which drives the contribution from the centered NFM and centered AGOP (Figures 14, and 15 in Appendix M). The increased importance of the C-NFC leads the UC-NFC to converge to the C-NFC at late times in most of our experimental settings, the contribution from the weight changes dominates the contribution from the initialization. These findings validate that the C-NFC is an important contributor to increases in the UC-NFC.
Our experiments suggest that studying the C-NFC is a useful first step to understanding the neural feature ansatz. In Section 4, we develop a theoretical analysis to understand why the C-NFC is generically large early times. We then use this analysis to predict the value of the C-NFC and motivate interventions which can keep the C-NFC high and promote the NFA carlier in training (Section 5).
4 Theoretical analysis of the C-NFC at early times
4.1 Gradient flow dynamics
We now theoretically identify why the centered AGOP and NFM become correlated structural property of gradient descent, for at least the carly training times. We consider the setup introduced in Equation (1). For theoretical convenience we focus on the case of training under gradient flow, where the dynamics of a weight matrix W(O) trained on loss L is given by
dW(( dt
W
-V C
(5)
where W²(() is the time derivative of ((), We expect our results to apply to gradient descent with small and intermediate learning rates as well. For simplicity of notation, we omit the layer index l.
We note that W = 0. This gives us the following proposition, which shows that early time dynamics of the centered NFM and AGOP are dominated by the second time derivatives:
Proposition 3 (Centered NFC dynamics). Let W be the weights of a fully-connected layer of a neural network at initialization and be the inputs to that layer. Then, when the neural network trained by gradient flow on loss function E, we have W WTW WT KW =0, (WTW) di (W KW) 0, and the first non-zero time derivatives satisfy,
d² de
d² (W KW) = 2- X xxx
W WW)
W) =2 2- 2-xTxxx
(6)
de
7

---

Published in Transactions on Machine Learning Research (11/2024)
where È the n x n diagonal motriz of logit derivatives č diag Df
Proof of Proposition 3. At initialization, all terms containing at least one copy of W vanish, leading the zero and first time derivatives to be O. Further, we have W K 0f(x) T 0f(x) and K. = 0f(x) 0f(x) T
0f(x) ah LX. Therefore, using that
8h
ōh
5h
1 d²
(W W) W W X L of(x) (((x))(x))
2 2
Oh
Oh
and,
1 d²
(W (WTW) W) = kyT= X L 01()) 0S(X) T
K
LX
2 2
Oh
Oh
= x c of(x) 0f(x) 0f(x) 0f(x)*
Oh. Oh
0h
0h
XT
We immediately see how gradient-based training can drive the C-NFC towards 1 at early times; the first non-trivial time derivatives are often highly correlated as they differ by only single matrix power of K, and have the same dependence on the labels. If A. is proportional to a projection matrix, then the two derivatives will have perfect correlation; even if not, IC and X²² will have identical eigenvectors, and hence we might expect range of spectra will enable high correlation between them. In the case that KC and KC² are sufficiently correlated, the C-NFC is driven to high value immediately upon training, and the high value of the C-NFC eventually drives the UC-NFC and causes the NFA to hold.
In the remainder of this section, we study this correlation in two different high-dimenssioal limits. Our analysis suggests that for large models the C-NFC has a generic tendency to increase early in training.
4.2 Maximum C-NFC for data on the sphere
We first provide a general and well-studied setting where the first non-zero derivatives of the centered NFM and AGOP are perfectly correlated uniform data the sphere in high dimensions (Ghorbani et al., 2020; 2021; Misiakiewicz, 2022). This limit corresponds to an infinite width network which, combined with the data symmetry, induces the PTK matrix X. to act like a projection matrix.
Data setup We sample data ~ v Unif (S6-1) uniformly distributed on the sphere in d dimensions with radius Vd. We assume the labels are generated from a target function that maps f gd + [-d, d]. Le. the label for the point is equal to f*)x). We train the parameters a, W for one-hidden layer fully-connected neural network f(z) = f(1;a, W) a ((Wx) with element-wise activation function 6. For learning trajectory (az, Wt), the initial values ag and Wo are sampled i.i.d. standard Gaussian with variance at most O(1/k). Le. for all i, e, we have (i). Wo ~ NN(,,) for = O(1/k).
Notation We will use asymptotic notation Oz. ou, Oa. Qa. wa in the usual way, where the limits are taken with respect to the data dimension d. Oα, σ2, θa, Sa. We are equivalent to their previous definitions but hide dependencies of polylogarithmic functions of d. We will use || || to refer to & vector norm and the operator norm of a matrix, i.e. || = max,eR" ||Ai|l for a square matrix A E [Rdxd Recall we sample n datapoints [[o] and use hidden dimension & so that a, E R and W, E RRd We write the vector of all ones in dimensions as 1 E 1 We will omit the time index t for ease of notation.
We make the assumption that f has bounded by a polylogarithmic factor of dimension with superlinearly vanishing probability:
8

---

Published in Transactions on Machine Learning Research (11/2024)
Assumption 4 (Labels bounded). We have, for any positive 8 > 0, with probability 1 O2(d-1-9), (x*( = Oz(log for a constant inleger L.
We also assume that the target function has a non vanishnis linear component: Assumption 5 (Non-trivial lincar component).. We have, ||E, [f)()) || = Q1(
These assumptions hold for example when f is exactly lincar, i.c. f*(x 3 for || S In this case, gd-1
the inner product x 3 = Oa(1) with probability Oa(d-1-9) for x uniform on v and any positive d.
We make an additional assumption on the diffcrentiability of the PTK, inherited from the activation ø, according to the conditions from Misiakiewicz (2022). This assumption needed for our analysis, in which differentiability enables us to Taylor expand the PTK matrix (El Karoui, 2008). We restate this assumption. Assumption 6 (Diflerentiability of the PTK). Let K(r, = Eay wy [f x)) (w] x')] = ha((e, /dd be the PTK for the network f. where ha : [-1,1] + R a positive semi semi-defing definite kernel function, defined for each dimension d. We assume there erist finite h(O), h/ (0), (0) > 0 such that lima-+so ha(0) = h(0), lime h2) (0), and limg-+oo (0) (0), where the first and second derivatives of ha, h/2 and hlj, are assumed to erist on [-1,1] for all d.
Note this is a sub-case of Assumption 1 in Misiakiewicz (2022) for level é= 1, and is satisfied for ha that is twice differentiable everywhere. We now introduce our final simplifying assumption on the activation function and data:
We will consider f trained using mean-squared error (MSE) loss. With this loss, A corresponds to the diagonal matrix of the residuals y - fxx). If the outputs of the network is 0 on the training data, then L L=Y = diag (y), the labels themselves. We can guarantee this by either of the following methods:
Method 7 (Subtract copy). Subtracting off an identical (untrained) copy of the neural network from each oulput at inilialization. L.e. we train the parameters of the network f on the loss with respect to oulputs f((x) = fz) - fo(z), where fo(z) is copy of f at initialization and fa(z) = f(;; a,, W.) Method 8 (Small initialization). Initializing W= = 0 or ||l|| = c for e > 0 arbitrorily small.
We now state our main thcorem.
Theorem 9 (Maximum C-NFC). Suppose the data X are sampled uniformly at random from gå-1 and the labels are generated by I satisfying Assumptions 4 and Suppose we train f with MSE loss and initialize the network so that the initial outputs are O by either Method 7 or Method above. Assume the activation function & satisfies Assumplion 6. We consider the regime that wa (dlog"t-1 d) S n S oa (d=-)) for some 8 > O, and width k + co. Then, at t = 0 in training (in the setting of Proposition 3), almost surely over the random X as n, d + 20,
p
d ar (az ( WW, (WT (TTWW) +1.
d
de
The proof follows from the mean term of KC giving the leading order terms in the C-NFC derivatives calculation here, and 11 a rank-1 projector. The proof deferred to Appendix C.
We clarify that although we take infinite width, we are not necessarily in the NTK regime, as we allow for arbitrary scaling of the weights, including the yP parametrization (Yang and Hu, 2020). Hence, our setting allows for feature learning.
In the next section we will construct a dataset which interpolates between adversarial and aligned eigenstructure to demonstrate the range of possible values the derivative correlations can take, and theoretically predict their values.
4.3 Early time C-NFC dynamics in co-scaling regime
Another interesting limit is the lincar co-scaling regime, where n, d, + oo with n/d = wi and k/d = 1/2- Here, we show it possible theoretically prodict the correlation of the derivatives of the centered NFM and AGOP.
9

---

Published in Transactions on Machine Learning Research (11/2024)
In this regime, we expect the following two properties (Adlam and Pennington, 2020), which we state assumptions. First, our key quantities can be written as traces of products of large random matrices; as the dimensions of all such matrices increase, we expect these traces to converge to their average values. This property is known as relf-averaging.
Assumption 10 (Self-averaging.. We assume that the expected traces appearing in the NFC across inilial- ications are equal to the traces themselves.
The elements weight matrices W are drawn from independent Gaussians (Li.d. within each matrix). the data X were also standard Gaussian, we would be able to apply free probability the noncommutative analog of classical independence W and X in the limit of large dimensions (Mingo and Speicher, 2017). To apply free probability for more general X, require the following assumption:
to compute traces of matrix products involving analytic functions of
Assumption 11 (Asymptotic freedom initial parameters, and input-label pairs). We assume W and are asymplotically freely independent. Further, we assume the labels ((X) are asymplottoally free of W (bul not X).
For example, the labels in the student teacher sctups of Adlam and Pennington (2020); Adlam et al. (2019)
satisfy this condition.
From Proposition 3, we know that for MSE loss the correlation of NFM/AGOP time derivatives under gradient flow initialization can be written as:
P (WT W. W KW)= tr XXXXXX tr (XXTXXXX))) 1/2 tr ((XT)X)X))2 (xx"xx)))""""""xxxxvxxx")x) - 1/2 (7)
In our high-dimensional limit, we expect that, for example, the first term can be given as
lim tr (X TXTTYxxxTYXXX n.d.k-o
lim Eo [tr (XTYRXXXTXX'X)) n,d
(8)
We can decompose the average as
Eg [tr (XT XXXXXXXXXXXXXI =tr YE oo [K] [ YX) + tr (Cov (XXXXX X XxKYX))
(9)
with similar decompositions for the denominator torm.
Therefore if the statistics of can be understood as a function of X, we can compule the correlation this linear triple-scaling limit. We focus for now on a one-hidden layer quadratic network (similar to the previous section) to avoid the branching of terms that arises in more complicated networks. We provide some additional analysis of the first term in Equation 9 for more complicated architectures in Appendix E.
4.4 Exact predictions with one hidden layer and quadratic activations
Concretely, we study a neural network which can be written as f(x) = a T (Wa)², where the readout layer, and the square is element-wise, Let (44 "|1r- (xTYX xx), and M =x x, and Fa W WT diag (a²) W. Then, the numerator of Equation can be written as:
tr (XTTXXXX = (M (MXY FaM F F.M
We reiterate our assumptions and compute this trace (as well as those for the denominator terms) in Appendix F using standard results from random matrix theory. These calculations show us that the correlation is determined by traces of powers of M x and M () by the calculations in Section F, which are properties of the input-label pairs, and Fa, which specific to the architecture and initialization statistics.
10

---

Published in Transactions on Machine Learning Research (11/2024)
p(W W. W XW)
P(W W. w KW)
1.0
Msan
Mean
AWV sP" KW)
Corrolation
0.5
605
035
0.76
0.25
6.55
Balance
Balance
pWWTW W. WTTWWTTN KW)
Figure 3: Predicted versus observed correlation of the second derivatives of centered F and G on the alignment. reversing dataset. Different shaded color curves correspond to four different seeds for the dataset. The solid blue curve is the average over all data sceds. The rightmost sub-figure a scatter plot of the predicted versus observed correlations of these second derivatives, with one point for each balance value. We instantiate the dataset in the proportional regime where width, imput dimension, and dataset size are all equal to 1024.
Manipulating the C-NFC To numerically explore the validity of the random matrix theory calculations, we developed a method to generate datasets with different values of (W W, W M..WTWW, KW). We construct a random dalaset called the aligument reversing datasel, parameterized by a balance parameter E (0, 1] to adversarially disrupt the NFA near initialization in the regime that width k, input dimension d, and cataset size n are all equal (n k d 1024). By Proposition 16, for the aforementioned neural architecture, the expected second derivative of the centered NFM satisfies, IE [ W/ X YE [K] YX (XTYX)* while the expected second derivative of the centered AGOP, E [W KW] = YE [C]] YX, has an additional component X X X Xy Our construction exploits this difference in that X becomes adversarially unaligned to Y as the balance parameter decreases.
The construction exploits that can manipulate X x X freely of the NFM using certain choice of Y. We desigu the dataset such that this AGOP- term is close to identity, while the NFM second derivative has many large off-diagona entries, leading to low correlation between the second derivatives of the NFM and AGOP
In our experiment, we sample multiple random datasets with this construction and compute the predicted and observed correlation of the second derivatives of the centered NFC at initialization. For specific details the construction,
Appendix G.
We observe in Figure 3 that the centered NFC predicted with random matrix theory closely matches the observed values, across individual four random seeds and for the average of the correlation across them. Crucially, a single neural network is used across the datasets, confirming the validity of the self-averaging assumption. The variation the plot across seeds come from randomness in the sample of the data, which cause deviations from the adversarial construction.
5 Increasing the centered contribution to the NFC strengthens feature learning
Our theoretical and experimental work has established that gradient based training leads to alignment of the weight matrices to the PTK feature covariance. This process is driven by the C-NFC. There- fore, one path towards improving the neural feature correlations is to increase the contribution of the measured, for example, by the
C-NFC to the dynamics tr (WT WWT KW) tr TWTWWT KW)- (WT) When this ratio is large, the C-NFC contributes significantly in magnitude to the UC-NFC, indicating successful fcature learning as mcasured by the uncentered NFC. We will discuss how small initialization promotes feature learning, and design an optimization rule, Speed Limited Optimization, that increases the C/UC ratio and drives the value of the UC-NFC to 1.
or, C/UC, ratio
11

---

Published in Transactions on Machine Learning Research (11/2024)
5.1 Feature learning and initialization
One factor that modulates the level of feature learning is the scale of the initialization in each layer. In our experiments training networks with unmodified gradient descent, we observe that the centered NFC will increasingly dominate the uncentered NFC with decreasing initialization (third panel, first row, Figure 4). Further, the centered NFC increases in contribution to the uncentered quantity, the strength of the UC-NFC, and to lesser extent, the C-NFC increases (Figure 4). The decreases in correspondence between these quantities is also associated with a decrease in the feature quality of the NFM (Appendix I), for the chain monomial task.
For fully-connected networks with homogeneous activation functions, such as ReLU, and no normalization layers, decreasing initialization scale is equivalent to decreasing the scale of the outputs, since we can write f(Wx) a p =f(aWx) for any scalar for any homogenous activation f. This in turn is equivalent to increasing the scale of the labels. Therefore, decrcasing initialization forces the weights to change more in order to fit. the labels, leading to more change in F from its initialization. Conceptually, this aligns with the substantial line of empirical and theoretical evidence that increasing initialization scale or output scaling transitions training between the lazy and feature learning regimes (Chizat et al., 2019; Woodworth et al., 2020; Agarwala et al., 2020; Lyu et al., 2023).
This relationship suggests that small initialization can be broadly applied to increase the change in F, and the value of the UC-NFC. However, this may not be ideal; for example, if the activation function is differentiable at O, small initialization leads to network which is approximately linear. This may lead to low expressivity unless the learning dynamics can increase the weight magnitude.
5.2 Speed Limited Optimization
We can instead design an intervention which can increase feature learning without the need to decrease the initialization scale. We do by fixing the learning speed layerwise to constant values, which causes the C-NFC to dominate the UC-NFC dynamics. For weights at layer l and learning rate 17 O, we introduce Speed Limated Optimization (SLO), which is characterized by the following update rulc,
V WWnC C
(! w!(0) nn. -..Ce
IV Wite CI
where the hyperparameter C. 2 0 controls the amount of learning in layer E. We expect this rule to increase the strength of the UC-NFC in layers where C2 is large relative to Cm for m + f, as W(O) will be forced to change significantly from initialization. By forcing the weights in particular layer l to have fixed learning speeds, these weights will have fixed updates sizes at every epoch, regardless of the loss. We downscale the
speeds other layers # m to prevent training instability. As a result, [: |w!"' causing the centered and uncentered NFC to coincide for this layer.
for large 4,
We demonstrate the effects of SLO on the thain-monomias task (Equation (2)). We found that fixing the learning speed to be high in the first layer and low in the remaining layers causes the ratio of the unnormalized C-NFC to the UC-NFC to become clase to 1 across initialization scales (Figure 4). The same result holds for the SVHN dataset (see Appendix M). We note that this intervention can be applied to target underperforming layers and improve generalization in deeper networks (see Appendix K for details).
We observe that both the C-NFC and the UC-NFC become close to 1 after training with SLO, independent of initialization scale (Appendix I). Further, the quality of the features learned, measured by the similarity of F' and & to the true EGOP, significantly improve and are more similar to each other with SLO, even with large initialization. In contrast, in standard training the UC-NFA fails to develop with large initialization, as F' resembles identity (no feature learning), while G only slightly captures the relevant features.
Speed Limited Optimization is step toward the design of optimizers that improve generalization and training times through maximizing feature learning.
12

---

Published in Transactions on Machine Learning Research (11/2024)
UC-NFC
C-NFC
C/UC Ratio
Train Losses
Conslatio
Value
64
02-
B.
Steps
Steps
Stepe
Steps
33-
Conreitiion
Lose
0.2
310
5400
/80
Steps
Steps
Steps
Steps
Initialization
0.001
0.01
1.0
C
D.
NFM
AGOP
NFM
AGOP
Figure 4: The effect of SLO on C/UC neural feature correlations and feature learning on the chain monomial task. In the first two rows, we plot the uncentered and centered NFA for the first layer weight matrix as function of initialization, (A) with standard training and (B) with SLO. We consider two hidden layer network with ReLU activations, where we set Co 500, and Ci C2 0.002. The third column shows the ratio of the unnormalized C-NFC to the UC-NFC: tr (W WWT T KW) (W (WTWWT [WT The fourth column shows the training loss. In the third row, we plot the NFM and AGOP from trained network with (C) standard training and (D) with Speed Limited Optimization with fixed initialization scale of 1.0.
6 Discussion
Analyzing more general settings In principle our analysis can be extended to a larger sample size and activation functions with uncentered derivatives. Both of these settings require analyzing more terms in the Taylor expansion of the PTK matrix. This is complex calculations require understanding degree 8 polynomials of the inputs even in the simplest case of a one-hidden layer network (as opposed to degree 4 to understand the loss). We may also want to understand the case that n = de for integer E (i.e. without the logarithmic factor we consider). The appropriate Taylor expansions in these cases are discussed in Misiakiewicz (2022), and will likely require additional structure on the coefficients of the larget function f in the basis of spherical harmonics.
complicated than studying the loss,
most
Centered NFC through training Our theoretical analysis in this work shows that the PTK feature covariance at initialization has a relatively simple structure in terms the weights and the data. To predict the NFC later in training, we will likely need to account for the change in this matrix. One should be able to predict this development at short times by taking advantage of the fact that eigenvectors of the Hessian change slowly during training (Bao et al., 2023). An alternative approach would be to quadratic model
13

---

Published in Transactions on Machine Learning Research (11/2024)
for neural network dynamics (Agarwala et al., 2022; Zhu et al., 2022) which can capture dynamics of the empirical NTK (and therefore the PTK). One promising avenue to study the time dynamics of the PTK could be to adapt the results of Wang et al. (2024), which study how the effects of feature learning can propagate to the conjugate kernel for neural networks. It may also be useful to unify our notion of aligument with the observations of Arous et al. (2023), where they find that the weights align with the averaged outer products of the loss gradients.
Speed Limited Optimization There are prior works that implement differential learning rates across layers (Howard and Ruder, 2018; Singh et al., 2015), though these differ from differential learning speeds as in SLO, in which the L2-norm of the gradients are fixed, and not the scaling of the gradients. Our intervention is more similar to the LARS (You et al., 2017) and LAMB (You et al., 2019) optimizers that fix learning speeds to decrease training time in ResNet and BERT architectures. However, these optimizers fix learning speeds to the norm of the weight matrices, while our intervention sets learning speeds to free hyperparamelens in order to increase feature learning, as measured by the NFC. We also point out that as an additional explanation for the success of SLO in our experiments, the PTK feature covariance may change slowly (or just in a single direction) with this optimizer, because the learning rate is small for later layers, allowing the first layer weights to align with the PTK.
Other works consider interventions that vary the strength of feature learning across layers. The authors of Yang and Hu (2020) down-scale the final layer weights so that the weights in earlier layers must move significantly in lx norm to fit the data. The SLO intervention is more extreme than this approach in that SLO overrides the natural GD dynamics to exactly set the weight movement rates. In Chizat and Netrapalli (2024) the authors consider adjusting learning rates in each layer so as to force the activation vectors to substantially change across all layers. Instead, SLO sets learning speeds of the weights in each layer so that some weights move much faster than others, depending on the layer where we want to maximize feature learning.
Adaptive and stochastic optimizers We demonstrate in this work that the dataset (by our construction in Section 4) and optimizer (with SLO in Section 5) play a significant role in the strength of feature learning. Important future work would be to understand other scttings where significant empirical differences exist between optimization choices. In particular, analyzing the NFC may clarify the role of gradient batch size and adaptive gradient methods in generalization (Zhu et al., 2023).
7 Acknowledgements
We thank Lechao Xiao for detailed feedback on the manuscript. We also thank Jelfrey Pennington for helpful discussions. This work used the programs (1) XSEDE (Extreme science and engineering discovery environment) which is supported by NSF grant numbers ACI-1548562, and (2) ACCESS (Advanced cyberinfrastructure coordination ecosystem: services & support) which supported by NSF grants numbers #2138259, #2138286, #2138307, #2137603, and #2138296. Specifically, we used the resources from SDSC Expanse GPU compute nodes, and NCSA Delta system, via allocations TG-C1S220009.
References
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, al. Language models are few-shot learners. Advances in neural information processing systems, 13:18 2020.
Alex Krizhevsky, Πyy Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. Advances in newral information processing systems, 25, 2012.
Ben Mildenhall, Pratul P Srinivasan, Matthew Tancik, Jonathan T Barron, Ravi Ramamoorthi, and Ren Ng. Nerf: Representing scenes as neural radiance fields for view synthesis. Communications of the ACM, 65(1): 106, 2021.
14

---

Published in Transactions on Machine Learning Research (11/2024)
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communicctiion of the ACM, 2021.
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural Tangent Kernel: Convergence and Generalization in Neural Networks. In Advances in Neural Information Processing Systems, 2018.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019.
B. Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Limitations of Lazy Training of Two-layers Neural Networks. In Advances Neural Information Processing Systems, 2019.
Zeyuan Allen-Zhu and Yuanzhi Li. What Can ResNet Learn Efficiently, Going Beyond Kernels? In Advances in Neural Information Processing Systems, 2019.
Gilad Yehudai and Ohad Shamir. On the Power and Limitations of Random Features for Understanding Neural Networks. In Advances in Neural Information Processing Systems, 2019.
Yuanzhi Li, Tengyu Ma, and Hongyang R Zhang. Learning over-parametrized two-layer neural networks beyond NTK. In Conference on Learning Theory, 2020.
Maria Refinetti, Sebastian Goldt, Florent Krzakala, and Lenka Zdeborová. Classifying high-dimensional Gaussian mixtures: Where kernel methods fail and neural networks succeed. In International Conference on Machine Learning, 2021.
Guangda Ji and Zhanxing Zhu. Knowledge distillation in wide neural networks: Risk bound, data efficiency and imperfect teacher. In Advances in Newral Information Processing Systems, 2020.
Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-stairease property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory, 2022.
Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation. ar preprint arXiv:2205.0144 01445,, 2022.
Eshaan Nichani, Alex Damian, and Jason Lee. Provable guarantees for nonlinear feature learning in three-layer neural networks. arXiv preprint arXiv:2305.0680 06986, 2023.
Boaz Barak, Benjamin L Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress deep learning: Sgd learns parities near the computational limit. arxiv preprint ar xrXi:22207.08790 2022.
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on Learning Theory, pages 5413 5452. PMLR, 2022.
Behrad Moniri, Donghwan Lee, Hamed Hassani, and Edgar Dobriban. A theory of non-linear feature learning with one gradient step in two-layer noural networks. arXiv preprint arXiv:2810. 00.9991, 2023.
Suzanna Parkinson, Greg Ongie, and Rebecca Willett. Linear neural network layers promote learning single-and multiple-index models. arXiv preprint arXiv:2305.1580 2023.
Alireza Mousavi-Hosseini, Sejun Park, Manuela Girotti, Toannis Mitliagkas, and Murat A Erdogdu. Neural networks efficiently learn low-dimensional representations with sgd. arXiv preprint ar Xix: 2022.
Adityanarayanan Radhakrishnan, Daniel Beaglehole, Parthe Pandit, and Mikhail Belkin. Mechanism for feature learning in neural networks and backpropagation-free machine learning models. Science, 383(6690):
1461-1467. 2024a.
Daniel Beaglehole, Adityanarayanan Radhakrishnan, Parthe Pandit, and Mikhail Belkin. Mechanism of feature learning in convolutional neural networks. Xiu preprint or Xrxi:23309 005000000000.0000 2023.
15

---

Published in Transactions on Machine Learning Research (11/2024)
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arxiv preprinl Xrx::1449. 1556, 2014.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale, 2021.
Gan Yuan, Mingyue Xu, Samory Kpotufe, and Daniel Hsu. Efficient estimation of the central mean subspace via smoothed gradient outer products. arXiv preprint arXiu:2312.15469, 2023.
Adityanarayanan Radhakrishnan, Mikhail Belkin, and Dmitriy Drusvyatskiy. Linear recursive feature machines provably recover low-rank matrices, 2024b.
D Aristoff, M Johnson, G Simpson, and RJ Webber. The fast committor machine: Interpretable prediction with kernels. xrxiv preprinl 20Xi:22455.14410 2024.
Marian Hristache, Anatoli Juditsky, Jorg Polzchl, and Vladimir Spokoiny. Structure adaptive approach for dimension reduction. Annals of Slatislics, pages 1537-1566, 2001.
Shubhendu Trivedi, Jialei Wang, Samory Kpotufe, and Gregory Shakhnarovich. A consistent estimator of the expected gradient outerproduct. In UAT, pages
2014.
Samory Kpotufe, Abdeslam Boularias, Thomas Schultz, and Kyoungok Kim. Gradients weights improve regression and classification. Journal of Machine Learning Research, 2016.
Daniel Beaglehole, Peter Súkeník, Marco Mondelli, and Mikhail Belkin. Average gradient outer product as a mechanism for deep neural collapse, 2024.
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When Do Neural Networks Outperform Kernel Methods? In Advonces in Neural Information Processing Systems, 2020.
Behrooz Ghorbani, Song Mei, Theodor Misiakiewiez, and Andrea Montanari. Linearized two-layers neural networks in high dimension. The Annals of Statistics, 49(2):1029-1055) 2021.
Theodor Misiakiewicz. Spectrum of inner- product kernel matrices in the polynomial regime and multiple descent phenomenon kernel ridge regression. arXiv preprint orXiv:2204114455 2022.
Noureddine El Karoui. Spectrum estimation for large dimensional covariance matrices using random matrix theory. 2008
Greg Yang and Edward J Hu. Feature learning in infinite-width neural networks. arXiv preprint arXiv:2011.14222 2020.
Ben Adlam and Jeffrey Pennington. The neural tangent kernel in high dimensions: Triple descent and multi-scale theory of generalization. In International Conference on Machine Learning, pages 74-84. PMTR, 2020.
James A Mingo and Roland Speicher. Free probability and random matrices, volume 35. Springer, 2017.
Ben Adlam, Jake Levinson, and Jeffrey Pennington. random matrix perspective on mixtures of nonlinearities for deep learning. arXiv preprint Xrv: 00,000 2019.
Blake Woodworth, Suriya Gunasekar, Jason D Lee, Fdward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, pages 3635 3673. PMLR, 2020.
Atish Agarwala, Jeffrey Pennington, Yann Dauphin, and Sam Schocnholz. Temperature check: theory and practice for training models with softmax-ccoos-entroop losses. ar preprinl ar 0..20000.7744, 2020.
Kaifeng Lyu, Jikai Jin, Zhiyuan Li, Simon S Du, Jason D Lee, and Wei Hu. Dichotomy of early and late phase implicit biases provably induce grokking. arxiv preprint or Xiv:2311 11. 188817, 2023.
16

---

Published in Transactions on Machine Learning Research (11/2024)
Xuchan Bao, Alberto Bictti, Aaron Defazio, and Vivien Cabannes. Hessian inertia in neural networks. 1st Workshop on High-dimensional Learning Dynamies, 2023.
Atish Agarwala, Fabian Pedregosa, and Jeftrey Pennington. Second-order regression models exhibit progressive sharpening to the edge of stability. aix preprint orXi:2220004860, 2022.
Libin Zhu, Chaoyue Liu, Adityanarayanan Radhakrishnan, and Mikhail Belkin. Quadratic models for understanding neural network dynamics. arXiv preprint 2Xxiu:2055117857 2022.
Zhichao Wang, Denny Wu, and Zhou Fan. Nonlinear spiked covariance matrices and signal propagation in deep neural networks. arXiv preprint arXiv:2402.10127, 2024.
Gerard Ben Arous, Reza Gheissari, Jiaoyang Huang, and Aukosh Jagannath. High-dimensional sgd aligns with emerging outlier eigenspaces. arXiv preprint arXiv:2310.0300 2023.
Jeremy Howard and Sebastian Ruder. Universal language model fine-tuning for text classification. arXiv preprinl ar 180: 06146, 2018
Bharat Singh, Soham De, Yangmuzi Zhang, Thomas Goldstein, and Gavin Taylor. Layer-specific adaptive learning rates for deep networks. In 2015 IEEE 14th International Conference on Machine Learning and Applications (ICMLA), pages 364 368. IEEE, 2015.
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arviv preprint arviv: 038888 2017.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint arXic: 00962, 2019.
Lénaïc Chizat and Praneeth Netrapalli. The feature speed formula: a flexible approach to scale hyper- parameters of deep neural networks, 2024. URL https://arxas http:////xxxvvvogab/////ss// wtt.wrr///......aaaaogg
Libin Zhu, Chaoyue Liu, Adityanarayanan Radhakrishnan, and Mikhail Belkin. Catapults in sgd: spikes in the training loss and their impact on generalization through feature learning. or preprint Xrx:222006. 04815, 2023.
Hong Hu and Yue M Lu. Universality laws for high-dimensional learning with random features. IEEE Transactions on Information Theory, (33331993211 2022.
Jimmy Ba, Murat Erdogdu, Taiji Suzuki, Denny Wu, and Tianzong Zhang. Generalization of two-layer neural networks: An asymptotic viewpoint. In International Conference on Machine Learning, 2019.
Chi Jin, Praneeth Netrapalli, Rong Ge, Sham M Kakade, and Michael Jordan. short note on concentration inequalities for random vectors with subgaussian norm. Xie preprint or Xiv: 037603 2019.
17

---

Published in Transactions on Machine Learning Research (11/2024)
A Glossary
1. Expected gradient outer product (EGOP). Defined with respect to target function f" and an input data distribution over data in d dimensions.
EGOP(f .p) ERw (=f Rdxd
2. Average gradient outer product (AGOP). Defined with respect to predictor / and set of inputs (x())...,.._))
0f(z(a))
)f(x E ERddd
AGOP(S; (x()))))= = a=1 dr
dr
In the context of deep neural network, write the AGOP G2 with respect to the inputs of given layer as
G. = Ol(x!a)) a T
The
Ohe
3. Neural feature matrix (NFM). Given neural network f with weight matrix W(e at layer l, the NFM Fg is defined as,
Fi = (W(O)) w(f)
4. Neural feature correlation (NFC). Defined with respect to layer of a neural network, this is the correlation p Gr).
5. Neural feature ansatz (NFA). This refers to the statement that the NFC will have correlation approximately equal to 1.
6. Pre-activation tangent kernel (PTK). This is the kernel function corresponding, defined with respect 0f(a)
to neural network and layer f, that evaluates two imputs to 0f(z)
8he Ohe E IR.
7. PTK feature covariance. For given layer of a neural network and a dataset X, the PTK feature covariance (())) is defined as ('0)
0f(x) [f(xx] She E RR
The
8. Centered NFC (C-NFC). Consider layer W of a neural network with PTK feature covariance with respect to that layer. Let W W - Wo, where Wo weights at initialization. Then, the C-NFC is equivalent to the correlation (WT W, W KW).
9. Uncentered NFC (UC-NFC). This quantity identical the NFC defined above.
18

---

Published in Transactions on Machine Learning Research (11/2024)
B Connection of the PTK to the empirical NTK
We note the following connection between the PTK entries and the empirical NTK (ENTK). In particular, the PTK is significant component of the ENTK.
Proposition 12 (Pre-aclivation to neural langent identity). Consider a depth-L neural network f(x) with inputs X₃ E Tgnoke to weight matrices W(O E Be for & E [T.]. Consider the empirical NTK where grudients are taken only with respect to W(O evaluated between two points a and z denoted by (xo. = awier awar).
(a(()) af(se)). Then for all e, θe salisfies,
@e(xo, 20) = k( (IU, 20) x
Proof of Proposition 12. Note that
0f(a)) owNs
0f(ao) TT Ohe
30 E RRiixke Then,
Oe(xo, 20) = tr (05(z0) x/
0f(=0)
the
She
0f(zu) T Of((z)
2t 28
The
Ōhe
K(0 20) x 2t.
C
Additional proofs and statements
Theorem (Maximum C-NFC). Suppose the data X are sampled uniformly at random from Sd-1 and the labels are generated by f satisfying Assumptions 4 and Suppose we train f with MSE loss and initialize the network so that the initial outputs are O by either Method 7 Method above. Assume the activation function 6 satisfies Assumption 6. We consider the regime that wa (dlog" d) S n S Oa (d--) for some 8 > O, and width k + co. Then, at = 0 in training (in the setting of Proposition 3), almost surely over the rndom X as n, d + 2o,
W), d² (W WWTKW) +1.
ae W (WT)),
de
Proof of Theorem 9. Applying Proposition 3 and using that we train with MSE and zero initial oulpuls, have
1 d
1 d 2 (W : x
2 de (W x
Given that :: 00, 00, IC is equal to a deterministic kernel matrix conditioned on the inputs (i.e. does not depend on the sampled initial weights). We then use that we are in the proportional limit to approximate the kernel matrix K. by its Taylor expansion (Misiakiewicz, 2022; El Karoui, 2008; Hu and Lu, 2022; Ba et al., 2019). Namely, we have that with probability 1 1- oo((,,
x= + x + + o A.
(10)
d
where = ha(0))+ + (17(0) b170)) a= N²(0), م == h(1) - h(0) - hO),, and A with |||||2 = 1 is a matrix. In general, we will write standalone asymptotic variables to indicate matrices A of that order spectral norm.
By the linearization in equation (10) and the structure of XXT, we will show,
K k= 0(1) -1T + + AI, x² x²= = O(d) 11 +d A
19

---

Published in Transactions on Machine Learning Research (11/2024)
where A1, Ag have spectral norm 1. To prove this, first note that x X I + o(1) uniform data on the sphere are ((1) Ouubauubannniann and have identity covariance. Therefore, as n = (d), and the spectra of X X 0(1).
and X are identical, we have that | = 11 x X||
We now analyze the numerator in the correlation (using the normalized trace). To simplify notation we write (A, B) = tr (AT B) for matrices A, B. The numerator is then:
(X XX 6(d) tr ((X (X xx)x)²))
+ +(d) tr ((X Ty X))
+ +(dd) Ir (r(T (oTTyTT)(( x)(XX
+ +(d) tr ((XTYAXXX)XXX>>X>Y)
As the labels are bounded, - spectral norm and that ||X|||||X|| =||]||||]] and || YA2Y = 0(1) ((())-IXTXX ||XTX||- Finally, again applying ||X X| = O(d),
Therefore, by the sub-multiplicatiive property of the
XXTXX XI|, we have y Y X|| S | = 0 |X x]]
tr ((XTYA, vy x)) s 0 1 || 1Txl²= y11²
Ö(d) ||X T yl?
and,
tr ((((TT
( o ([d))-[XxT])
The fourth trace then satisfies,
tr ((XTYANY)) s Ō0( | x1I2 =
Combining the four traces and using that tr ((X yy x) tr ((y) xx (xm))²) || | yl1
(XTXXXX XXXX (dd |X w/1* + Ō(d-) ||xT y11 + (
Recall by Assumption 5, we have | (x)] 9(1). We apply Lemma 13 for our choice of n so that x y y [x)] = o(1) w.p. 1 - o(1). As consequence, we have ||XT yl = (d) w.h.p. Therefore, the leading term the numerator of the correlation that of ((XT "yy" x))) the term due the interaction of the mean terms of K. and X22
Meanwhile the first denominator term decomposes follows:
(x X = ē(1) tr ((X yr xx²²)
+0(1 -r ((XTYAVX))( yy x))
+((1) tr ((XTYA,))))
By similar argument to the numerator term, we observe that the leading sub-term for the first denominator term Õ(1) tr ((XT)) " X)²).
Further, the second denominator terms decomposes as follows:
(XTREY,, ((d) tr ((X Tv xx2) + +((f) ((XTYAXXX)(X) yy X))
+ Ō(P)) -r In ((XTXXXX)))
By similar argument, the second denominator term has leading sub-term O(d2) tr ((XT my x)2).
20

---

Published in Transactions on Machine Learning Research (11/2024)
Putting the numerator and denominator terms together, have the correlation has the following asymptotics:
( d² Te (WTW) W),
d² (WT KW))
((d) tr ((X yy xx)) + O(d")
de
(XT tr (5(1) yy xy² x 1/2 tt (ō(d)) (XT yy "yyyyy""" + 0(d4) 1/2
+1,
completing the proof.
Lemma 13. Suppose [f*()) S c. log (d) with probability 1 - o(1) for constant e > 0, and the data are sub- Gaussian with constant parameter K and some > 0. Then, , y - E [x)) = o(1) w.p. - o(1), provided
21+1 (d)).
Proof. As the data x are E-sub-Gaussiie for a universal constant K, and the labels are bounded by clog' (d), have that zy is sub- GabsSanssi with parameter M Ke Kolog Therefore, the empirical expectation of zy, I "E= FiYi is sub-Gaussian with parameter M/vn. Directly applying Lemma 1 in Jin et al. (2019), we that,
Pr
1
Σ
n2 E [ry] 21]
Fal:
2exp (-2cNPP)
fn
for a universal constant c 0. Therefore, provided n wdd (d)), we have with probability o(1, completing the proof.
Em= z,y;
[xy]]
= (1)
D Additional centerings of the NFC
Double-centered NFC One may additionally center the PTK feature map to understand the evolution of the PTK feature covariance and the weight matrices. We consider such centering that we refer to as the Couble-centedd NFC, where we measure p ((W()) W(O (W K(Ww()) where
(f(x)) Ols
(fo(X) [F(X) The
afo(X) and fo is the neural network at initialization. The
The
However, the double-centered NFC term corresponds contribute the centered and uncentered NFC (Figure 5) when initialization is large or for early periods of training. Note however this term becomes relevant over longer periods of training.
higher-order dynamics that do not significantly
Unnormalized DC/C NFA
Layer
Layer
1,0-1
Rato
1,8-2
1.844
Stupe
Staga
Intalization
0.001
0.01
0.1
1.0
Figure 5: Ratio of the unnormalized double-centered NFC the centered NFC throughout neural network training, In particular, we plot tr (W W RW) tr trW"" (WWTKW) throughout training for both layers of two-hidden layer MLP with ReLU activations.
21

---

Published in Transactions on Machine Learning Research (11/2024)
Isolating alignment of the PTK to the initial weight matrix One may also center just the PTK feature map; while substituting the initial weights for W to isolate how the PTK feature covariance aligns to the weight matrices. To measure this alignment, we consider the PTT-ceetered NFC, which is defined as the correlation p ( " (T) k(O! where W the initial weight matrix at layer L.
However, this correlation decreases through training, indicating that the correlation of these quantities does not drive alignment between the uncentered NFM and AGOP (Figure 6).
Layer PTK-NFC
Layer 2 PTK-NFC
Train Losses
Carre
0.5
Loos
0.0
327
687
840
100
300
200
600
180
220
400
640
Steps
Steps
Initialization scale
0.001
0.01
0.1
1.0
(a) Isotropic data.
Layer PTK-NFC
Layer PTK-NFC
Train Losses
1.0
Correletion
Lose
0.5
0.0
480
640
160
320
160
610
Staps
Steps
Stops
Initialization scale
0.001
0.1
1.0
(b) Data spectrum with decay 0g 1442
Figure 6: PTK-centered NFC throughout training for both layers of a two-hidden layer MLP with ReLU activations on Gaussian data with two different spectra.
E Extending our theoretical predictions to depth and general activations
Precise predictions of the C-NFC become more complicated with additional depth and general activation functions. However, we note that the deep C-NFC will remain sensitive to a first-order approximation which K is replaced by its expectation. We demonstrate that this term qualitatively captures the behavior of the C-NEC for 2 hidden layer architectures with quadratic and, to lesser extent, ReLU activation functions in Figure 7. In this experiment, eigenbasis. We parameterize the eigenvalue decay of the covariance matrix by a parameter a, called the data decay rate, so that the eigenvalues have values A 1 As a approaches or xo the data covariance approaches a projector matrix.
sample Gaussian data with mean 0 and covariance with a random
In this experiment, we see that the data covariance spectrum will also parameterize the eigenvalue decay of E. [K], allowing to vary how close the expected PTK matrix (and its dual, the PTK feature covariance) is
22

---

Published in Transactions on Machine Learning Research (11/2024)
from a projector, where the NFA holds exactly. We see that for intermediate values of a, both the observed and the predicted derivatives of the C-NFC decreases in value.
We plot the observed values in two settings corresponding to different asymptotic regimes. One setting is the proportional regime, where n = = = 128. The other is the NTK regime where n = d 128 and = 1024. For the quadratic case, as the network approaches infinite width, the prediction more closely matches the observed values. Additional terms corresponding the nonlinear part of ø' in ReLU networks, the derivative of the activation function, required to capture the correlation more accurately in this case.
Eigenvaauu decay of [[K]
C-NFC
1 predicted C-NFC
1.
162
atton
164
011
125
0.
85
06
C5
1.5
Decny
Decay
Data decay rane
Label distribttion
Label distribution
Randem
(a) Quadratic.
Eigenvalue decay of B[K]
C-NFC
Pl-prder predicted C-NFC
Conretston
164
Conree
06
C5
13
Dacny
Decay
Data decay rase
Label distribttion
Label distribution Randoin
Random
(b) ReLU.
Figure 7: Observed versus the first-order predicted C-NFC for the input to the first layer of two hidden layer MLP. The dashed line is neural network width k = n = d = 128, where n and d are the number of data point and data dimension, respectively, while the solid line uses n d = 128 and k = 1024.
F Free probability calculations of C-NFC
In order to understand the development of the NFC, we analyze the centered NFC in the limit that learning rate is much smaller than the initialization for one hidden layer MLP with quadratic activations. We write this particular network as,
where a E 1xk and W e gxxd where d is the input dimension and k is the width. In this case, the NFC has the following form,
tr (XTTXXXXXXXXXX
P (F.G)
(11)
ir (X
where K= XWT diag (a) 2 WXT.
23

---

Published in Transactions on Machine Learning Research (11/2024)
We assume two properties hold in the finite dimensional case consider, that will hold asymptotically in the infinite dimensional limit.
Assumption 14 (Self-averaging). We assume that computing the average of the NFC quantities across initializations is equal to the quantities themselves in the high-dimensional limit.
Assumption 15 (Asymptotic freeness). We assume that the collections {X,Y} and (W, a} are asymptotically free with respect to the operator IE [tr(-)], where tr[M] = - Ma-
We will compute the expected values of the centered NFC under these assumptions. In the remainder of the section we will drop the [+] in the trace for ease of notation.
F.1 Free probability identities
The following lemmas will be useful: let {C:) and (Ra) be freely independent of each other with respect to tr, with tr[C.] = 0. Alternating words have the following products:
Rr]
(12)
M =
(13)
tr ACCHAGGAACCCAA--AAAAAAAAAAACCCCC
(14)
(15)
where R = Rg - tr)
Applying these identities to the one hidden layer quadratic case, use the following definitions:
R W diag (a²) W, A (XTxx)², B= x
(16)
Crucially, R is freely independent of the set {A, B}. We will also use the notation M to indicated the centered version of M, M = M - tr(M].
F.2 Numerator term of NFC
The numerator in Equation (11)
tr wwW.VXXXTxXXXXX = ir (ARARBR)
(17)
Re-writing A = À tr(A] and B = B tr[B]
we have:
tr ((XTxrxxxxx<*3 tr
(r
(18)
This expands to Ir =
+ 24r (A) (ARB))))) (ĀRÄR²)
(19)
tr tr (BR)) + (A) (AR*) + tr tz) (B) (R)
Using the identitics we arrive at:
Ir (XTVXXXXXYXX> -/A (R²)
+ (tt (AB)
(20)
+ [tr[B]a (A2)
+ tr (A2B)
24

---

Published in Transactions on Machine Learning Research (11/2024)
F.3 First denominator term of NFC
The first denominator term in Equation (11) is
tr (x XXYXX (XXX) tr (ARAR)
(21)
This is a classic free probability product:
tr (x XXXXXXXXXXXXXXX -r(A)) + tr[A]2tr (R²) tr[
(22)
which can be derived from the lemmas.
F.4 Second denominator term of NFC
For the second denominator term of Equation (11) we have
tr (x'T(erxx' = to (ARBRAARBR)
(23)
Expanding the first have
tr (XYYXXX tr (ARBRARBR) Next we expand the first B:
tr(A]]r -ARBBA))))) + (RABRARB)
(24)
ir
(ARBRAARR + tr[B]rr (AR' ARBR) = -r (ARBRARBR) +br[A]ur[BJur (R² ARB) + ur[A]tr (R
(25)
The next A gives us
tr (XTVXXXXXVxX = =((ABRRĀRBR) + 2tr(([]trr (ARBBR +r([B (AR ARBR) +2tr[A]tr[BB B|tr RR (RARB) + (rA) (R*B) trr[A]]tr (R2
(26)
Expanding the final B we have
tr (x
(ARBRARBR) + 2tr(B)tr (ARBRAR*) + 2tr/A (ARBR*BR)
+ (Br (R ARB) + Bu [R A) + 2tr (B/tr (R
+ Br] + [rAA] B ur
AR"
(27)
Now all terms are in the form of alternating products from the lemma. This means we can factor out the non-zero traces of the other terms. Simplifying we have:
Lr (=xTyx3xxx
ur[ ((AB))) (117)
+ 2tr((R (tr Brrr (A² tr[A]tr ( (AB²)) + 4t]4A (AB) - u/444
(28)
+ (B + tr tr (B9
All terms of the NFC are now in terms of traces of the matrices A, B, and and functions on each term separately. The matrices A and B are determined by the data, while the moments of the eigenvalues of R are determined by the initialization distribution of the weights in the neural network, and neither training nor the data.
25

---

Published in Transactions on Machine Learning Research (11/2024)
G Alignment reversing dataset
The data consists of mixture of two distributions from which two subsets of the data X and X are sampled from, and is paramctrized by balance parameter 7 E (0, 1] and two variance parameters 61, 62 > 0. The subset (1) which has label y1 1 and constitutes a y fraction of the entire dataset, is sampled from a multivariate Gaussian distribution with mean 0 and covariance د 11' + I Then the second subset, - ((x())xx))++ and has labels y/2 Then, for balance parameter y sufficiently small, the AGOP second derivative approximately satisfies,
x is constructed such that (X(²)) x x²
E [ KW] YXX XX (X(I) x X ((X x (x T (x(e)) (x() X()) Tx()
11
In contrast, the NFM second derivative, E | W| = (X yx = ( = E², will be significantly lar from identity.
Specifically, we construct (2) by the following procedure:
1. Extract singular values S and right singular vectors U1 from a singular-valus decomposition (SVD) of (X(I))XX0 X
2. Extract the left singular vectors 1/2 from a sample X₂ that is sampled from the same distribution as x(),
3. Construct x(2) = UTT
4. Where X = (() @ x Set X c X + 62 where 7~ ~ NOO, T).
5. Set + y 10 1.
Note that U1S V25 "T
V = U1S- Tx therefore, we should set X(22 V2ST U, 2 Regarding the variance parameters, in practice we El 0.5 and
to get (X(2) ( = ((XI) X
62 1
Proposition 16 (Expected NFM and AGOP). For one hidden layer quadratic network, f(z) = a ((Wx), with a ~ NO, I) and W ~ Vh" NN(0,)) T),
B,ww [VTW] = [W
(xTx)
and,
B..w [WT KW] = tr (X X) YX
XXTXX +6XT XX
Proof of Proposition 16.
E W W| x YXE
W. diag (a) Wo] x
xxx
Further,
E W = x YE X K2| YX.
We note that,
k² WT diag (a)² WoX T XW. diag (a) Wo
(29)
b n
ΣΣΣ
(30)
81.82 P1.P2
L 2 W V...ppXo.pp Xa, W Wo, XW WTXT.
(31)
26

---

Published in Transactions on Machine Learning Research (11/2024)
Therefore, applying Wick's theorem, element i.j of K² satisfies,
72
k
n
E [a s.p. Xa, Xa W..py X, X.W.W W.W, x,
P1-P2
P1-P2-01-92
E [ [aW W. s.p
Wa s.p W s.q W w. Xapp Xapps {
Σ
P1.P2-01.92
E W SP
W., Wa
E (Ws
sn W. +
E WW p1
W s.g1
E W s2 W w.p +
E WWa, s. W s.ga
E (.., s2 W. W..r]
Xa Xa.p Xi.1 X,9
d
P1.
(pppp deres 22 + Sp Spega + 6p1g. Ops)))
Xapp Xapp Xi.gn Xi
3
Xa.p Xam Xi Xznn
P1:41
d
+
X Xa.p Xa.p2 Xim X Xj.p
d
+
Σ
Xa.p X Xa.p Xi.p
P1.P2
= 33 tr x) X, + > xy X
+3 Σ X X X
=3 (x X) X, +3X,X XXXXXXXX,+33 +x,+
Finally, we conclude,
E [K*] (tr (x (xTx) x) Xx + +2xX xxx)x
giving the second statement of the proposition.
Varying the data distribution
We verify that our observations for isotropic Gaussian data hold even when the data covariance has a significant spectral decay. (Figures 9 and 10). We again consider Gaussian data that is mean 0 and where the covariance is constructed from random eigenbasis. In Figure 9, we substitute the eignevalue decay 1 We plot the values of the UC-NFC, C-NFC, train loss,
ds 1 14k: while in Figure 10, we use A 1442- and Lest loss throughout training for the first and second layer of a two hidden layer network with ReLU activations, while additionally varying initialization scale. Similar to Figure &, we observe that the C-NFC is more robust to the initialization scale than the UC-NFC, and UC-NFC value become high through training, while being small at initialization. We see that the test loss improves for smaller initializations, where the value of the C-NFC and UC-NFC higher.
27

---

Published in Transactions on Machine Learning Research (11/2024)
C-NFC
Train
Steps
Stops
Staps
UC-NFC
C-NFC
Test
300
640
160
200
450
940
Steps
Siepa
Staps
Intiialization
0.10
0.01
10
Figure 8: Centered neural feature correlations. Data covariance decay rate A = 1. Top row is layer 1, bottom row is layer 2. Train (test) losses are scaled by the maximum train (test) loss achieved so that they are
between 0 and
UC-NFC
C-NFC
Train Loeses
C4
Steps
Stops
Stope
UC-NFC
C-NFC
Test Losses
280
640
180
100
60
000
64)
Steps
Sheps
Stape
Intiializelion
0.001
0.01
0.1
1.0
Figure 9: Centered neural feature correlations. Data covariance decay rate d T+E Top row is layer 1, bottom row is layer 2. Train (test) losses are scaled by the maximum train (test) loss achieved so that they are between and 1.
1
28

---

Published in Transactions on Machine Learning Research (11/2024)
C-NFC
C4-
640
190
300
Steps
Stops
Stops
C-NFC
080
640
190
000
64)
Steps
Staps
Staps
Intiialization
0,001
0.01
0,1
1.0
Figure 10: Centered neural feature correlations. Data covariance decay rate Ag ~ THKT. Top row is layer 1, bottom row layer Train (test) losses are scaled by the maximum train (test) loss achieved so that they are between and 1.
29

---

Published in Transactions on Machine Learning Research (11/2024)
Effect of initialization on feature learning
We see that when initialization is small, the C-NFC and UC-NFC are high at the end of training with and without fixing the learning speed (Figure 4). This is reflected by the quality of the features learned by the NFM and the qualitative similarity of the NFM and AGOP at small initialization scale (Figure 11). Further, we notice that as we increase initialization, without fixing speeds, the correspondence between the NFM and decreases and the quality of the NFM features decreases (at a faster rate than the AGOP). Strikingly, when learning speeds are fixed, the quality of the features in the AGOP and NFM becomes invariant to the initialization scale.
0.001
0.01
0.1
1.0
(a) NFM, spceds not fixed.
2.001
0.01
1.0
(b) AGOP, speeds not fixed.
0.001
0.01
1.0
(c) NFM, speeds fixed.
0.001
0.01
0.1
1.0
(d) AGOP, speeds fixed.
Figure 11: The NFM the end of training as a function of initialization scale in the first layer weights, with and without fixing learning speeds. The task again is the chain monomial with rank r = 5. The title of each plot is the initialization scale of the first layer sp-
30

---

Published in Transactions on Machine Learning Research (11/2024)
J Experimental details
We describe the neural network training and architectural hyperparameters in the experiments of this paper. Biases were not used for any networks. Further, in all polynomial tasks, we scaled the label vector to have standard deviation 1.
Corrupted AGOP For the experiments in Figure 1, we used n = 384 data points, d d= 32, k: = 128 as the width in all layers, isotropic Gaussian data, initialization scale 0.01 in the first layer and default scale in the second. We used ReLU activations and two hidden layers. For the experiments in Figure 8,9,10,5, and 6, we used two hidden layer network with ReLU activations, learning rate 0.05, 800 steps of gradient decent, and took vorrelation/covariann measurements every 5 steps.
C/UC-NFC calculations on real datasets We describe the experimental details for Figure 2. For (A,B) we trained a five layer MLP on the first 50,000 datapoints of Streetview House Numbers (SVHN), CIFAR-10, CIFAR-100, 1-,,, S0L-10, MNIST, and German Traffic Sign Recognition Benchmark (GTSRB) datasets. We used the default Py PyTcrc initialization (scale of 1) for all layers. We used width 256 in all layers, and trained with SGD with batch size 128. For SVHN and CIFAR, we trained for 150 epochs with learning rate 0.2. For STL-10 and GTSRB
used learning rate 0.1 for 150 epochs. For MNIST we trained for 50 epochs with learning rate 10. For the VGG-11 experiments on CIFAR-10, we used the default architecture from torchvision with batch-norm layers removed. We trained for 50 epochs and learning rate 1. For experiments with a GPT-family htops://gtthhn
model, we adapted the model and dataset from NanoGPT used all the default settings and the default Adam optimizer, but used no weight docay, learning rate 5c-3, and removed all dropout layers. We also reduced the number of alleution layers to from the original 6. We trained on the Shakespeare characters dataset.
We
Alignment reversing dataset For the experiments in Figure 3, we used k: k= n d 1024 for the width, dataset size, and input dimension, respectively. Further, the traces of powers of Fa are averaged over 30 neural net sceds to decouple these calculated values from the individual neural net sceds. The mean value plotted in the first two squares of figure is computed over 10 data seeds.
SLO experiments For the SLO figures (Figures 11, 4), we use isotropic Gaussian data, 600 steps of gradient descent. The learning rates are chosen based on initialization scale in the first layer. For initialization scales 1, 0,1, 0.01, and 0.001, we used learning rates 0.03, 0.1, 0.2, 0.4, respectively. We again used two hidden layers with ReLU activations. We chose 256, d = 32, and k 256 as the width. We divided the linear readout weights by 0.01 at initialization to promote feature learning, and modified SLO to scale gradients by (e + ||V rather than just the inverse of the norm of the gradient, for = 0.1. This technique smooths the training dynamics as the parameters approach a loss minimum, allowing the network to interpolate the labels.
Predictions with depth For the Deep C-NFC predictions (Figure 7), we used n = 128, d = 128, initialization scale of 1. The low rank task is just the chain monomial of rank r = 5. The high rank polynomial task is y(z) = Σi= (Qx),, where Q e Raxd is matrix with standard normal entries.
Figures 16 and 17 For the experiments on the SVHN dataset, we train four hidden layer neural network with ReLU activations, initialization scale 1.0 in all layers and width 256. For SVHN, subset the dataset to 4000 points. We train for 3000 epochs with learning rate 0.2 for standard training, and 0.3 for SLOO and take NFC measurements every 50 epochs. For SLO, we set Co = 2.5, C₁ = C2 = 0.4, and relaxation parameter e = 0.2. We pre-process the dataset so that each pixel is mean and standard deviation 1. For the experiments on CelebA, we train a two hidden layer network on a balanced subset of 7500 points with Adam with learning rate 0.0001 and no weight docay. We use initialization scale 0.02 in the first layer, and width 128. We train for 500 epochs. We pre-process the dataset by scaling the pixel values to be between 0 and 1.
Code availability We make the code available for the experiments in Figure 2 available at https: //anonymous. .orreouuusooon....mmmm..m wttte/ccee/eeeeeteeeeeeeeeeee NFA-4795. (for MLP+VGG).
31

---

Published in Transactions on Machine Learning Research (11/2024)
K
Additional SLO experiments
We demonstrate the SLO can be applicd adaptively to increase the strength of the UC-NFC in all layers of a deep network on the chain monomial task of rank r = 3. We train a three hidden layer MLP with ReLU activations and an initialization scale of 0.1 by SL.,, and find that all layers finish at the same high UC-NFC (Figure 12). Further, this final UC-NFC value is higher than the highest UC-NFC achieved by any layer with standard training. The generalization loss also lower with SLO on this example, corresponding to better feature learning (through the UC-NFC).
UC-NFC
Train Losses
Test Losses
S
SLD0
Standard
Standard
03.4
Coreaatonn 2000077o
04-
0.1
Layer
Laper2
Layer
Stops
Steps
Figure 12: Training with SLO where the learning speeds are chosen adaptively based on the UC-NFC values of all layers. The dashed lines correspond to training with standard GD.
At every time step we choose C, s for the layer with the smallest UC-NFC correlation value, while setting Cj s- for all other layers, with 20. We again modify SLO by dividing the gradicnts by el V/| for = 0.01. The learning rate is set to 0.05 in SLO and 0.25 for the standard training (gradient descent), and the networks are trained for 500 epochs. We sample n 256 points with d = 32, and use width k = 256.
32

---

Published in Transactions on Machine Learning Research (11/2024)
L
Additional C-NFC/UC-NFC plots across architectures
(MLP on CIFAR-100)
Condo
B. (MLP on MNIST)
10972
(MLP SVHN)
Eoodhe
13079
(MLP on CIFAR-10)
E. (MLP on GTSRB)
5-4
F. (MLP on STL-10)
Liger?
3-
Figure 13: Full trajectories for the C/UC-NFC of a five hidden layer MLP trained on six datasets, averaged over three seeds, with large (1.0) initialization scale. The blue curves are the UC-NFC, while green curves are the C-NFC.
33

---

Published in Transactions on Machine Learning Research (11/2024)
A. (MLP on CIFAR-100 100)
D. (MLP on CIFAR-10)
Layer
Layer
Layer
Layer
Layer
2457
(MLP on MNIST)
E. (MLP GTSRB)
Layer
Layer
Layer
Layer
C. (MLP on SVHN)
F. (MLP on STL-10)
Layer
Layer
Layer
Layer
Layer
Figure 14: Losses and normalized changes in weights across datasets for a five hidden layer MLP. The change in weight is measured as ||W - Wo First column of all subfigures are the losses, while the second columns are the weight changes.
M Additional experiments on real datasets
We replicate Figures and 4 celebrity faces (CelebA) and Street View House Numbers (SVHN). We begin by showing that one can disrupt the NFC correspondence by replacing the PTK feature covariance with a random matrix of the same spectral decay. For this example, measure the Pearson correlation, which subtracts the mean of the image. l.e. p(A, B) = p (A m(A), B - m(B)), where m(A), m m(B) are the average of the elements of and B.
34

---

Published in Transactions on Machine Learning Research (11/2024)
(VGG-11 CIFAR-10)
Layer
Layer 2
Layer
Layer 4
04-
Layer
Layer 6
Layer
Layer 8
23-4
B. (GPT on Shakesepeare)
0.5
0,6
Valure
0.4
3x
Query
600
0003
400
1020
1000
410
1356
Figure 15: Losses and normalized changes in weights across datasets for VGG-11 CIFAR-10 and GPT Shakespeare character generation. The change in weight is measured as |W - Wo Wolll[ First column of (A) and (B) are the losses, while the remaining columns are the weight changes.
W W (NFM)
T KW (AGOP)
WTQW
Figure 16: Various feature learning measures for the CelebA binary subtask of predicting glasses. The diagonals of the NFM (W) (first plot) and AGOP (WT KW) (sccond plot) of a fully-connected network are similar to each other. Replacing K with a symmetric matrix Q with the same spectrum but independent eigenvectors obscures the low rank structure (third plot), and reduces the Pearson correlation of the diagonal from p (diag (F). diag (G)) 0.91 e (diag (F) diag (W (WT)) 0.04.
35

---

Published in Transactions on Machine Learning Research (11/2024)
UC-NFC
C.NFC
C/UC Ratio
Train Lasses
VOUS
Loos
30
200
3100
2100
Steps
Steps
Steps
Steps
Initia ization scale
SLO
GD
(a) Feature learning measurements.
NFM
AGOP
NFM
AGOP
(b) Standard Training
(c) SLO
Figure 17: We demonstrate the SVHN dataset, with 4 hidden layer neural network with large initialization scale, how SLO improve the strength of the UC-NFC, the C-NFC, the ratio of the unnormalized C-NFC to UC-NFC (plot (a)) and the fcature quality (plots (b) and (c)). In plots (b) and (c), we visualize the diagonal of the NFM and AGOP for the first layer of the trained network, where SLO was applied with G = 2.5, C = C2 = 0.4.
36


Reviews:
[{"review_id": "dG8cI7gZ3C", "reviewer": "Reviewer_wFmu", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: The paper is overall well-written and contains many interesting observations about NFA. Understanding feature learning in neural networks is very difficult and so conceptual properties such as the NFA are useful for making progress on this important topic. The origins of NFA have so far been quite mysterious. The experiments and derivations of this work are a positive contribution towards shedding light on the origins of NFA. \n\nWeaknesses: Some of the explanations and claims feel a bit too hand-wavy. For example, the justification for the Speed Limited Optimization was not completely clear to me. I believe these can benefit from more clarifications. The theoretical setting is also fairly limited.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors investigate a concrete structural property of neural networks which is hypothesized to be tightly linked with feature learning. Specifically they seek to the origins of the so called Neural Feature Ansatz (NFA). The NFA is that after training, the Gram matrix of the weights (Neural Feature Matrix or NFA) is highly correlated (as measured by the Neural Feature Correlation or NFC) with the Average Gradient Outer Products (AGOP). Firstly, to more tightly link the NFA and the AGOP the authors observe that the AGOP can be written in terms of the weight matrices and the Preactivation Kernel Matrix (PKM). The authors then use this framing to investigate the origins and implications of the NFA phenomenon, providing several observations: \n\n1. The authors empirically demonstrate that at early times the NFA and AGOP alignment is driven by the alignment of the PKM with the *change* in the weights, which the authors call the centered neural feature correlation (C-NFC).\n\n2. The authors analyze the dynamics of the NFA and AGOP in a simple setting at initialization to provide a heuristic argument for the emergence of the NFA and AGOP correlation. These calculations allow the authors to demonstrate that this emergence is not universal and that a carefully constructed dataset can break it. \n\n3. The authors measure the correlation between the initialization scale and the NFC, showing empirically that higher correlation and better generalization occurs with smaller initialization scales which intuitively is linked with better feature learning.\n\n4. Lastly, the authors use their observations to inspire a new optimization variant called Speed Limited Optimization (SLO). The authors argue and provide some experimental evidence that this optimization algorithm better promotes NFC and hence better feature learning and generalization.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I believe the submission is already quite strong and deserving of acceptance. I would request the following changes to strengthen the work:\n\n1. Provide some introduction and background on free probability for readers who are not already familiar. For example: \"We make the assumption that X and Y are (asymptotically) freely independent of the parameters at initialization\", does this just follow from classical independence?\n\n2. In Eq. (7) why does the NTK limit imply that the Cov term goes to zero?\n\n3. Can the theory be used to understand why the C-NFC is the primary contributor to NFC? Or how the initialization scale affects the correlation?\n\n4. The $\\eta$ in the SLO update should be a $-\\eta$?\n\n4. The motivation for SLO is unclear, as is the statement: \"We expect this rule to increase\nthe strength of the UC-NFC in layers where $C_l$ is large relative to $C_m$ for $m \\neq l$\". Also the appropriate baseline should be to tune all layerwise learning rates, since each layer has a speed. Is there a reason to believe that against this baseline SLO should be better? Is there a way to appropriately set the layerwise speeds?", "broader_impact_concerns": "No concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "HaxQ7gxqGE", "reviewer": "Reviewer_VtKo", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The mechanism of feature learning in neural networks is an important research question. This submission presents a mix of empirical and theoretical investigations of a fairly general feature learning mechanism that goes beyond the idealized problem settings in prior theoretical works, which I believe is a valuable contribution. \n\nI have the following concerns/questions: \n\n1. The analysis of gradient descent dynamics is limited, and the connection between C-NFC and learning needs to be elaborated.\n\n- The current gradient descent analysis is basically restricted to the first gradient step, as it does not account for the parameters depending on the training data. Due to the simplified setting, my impression is that it should be possible to analytically compute the NFC in some idealized student-teacher settings (for the two-layer model at least), using similar Gaussian equivalence calculations similar to (Ba et al. 2022).  Can the authors comment on the technical challenges here?\n\n- The authors wrote: \"*The relation between the NFM and the AGOP is significant, in part, because for a neural network that has learned enough information about the target function, the AGOP of this model with respect to the first-layer inputs will approximate the expected gradient outer product (EGOP) of the target function.*\"  \nCan the authors explain what kind of information is contained in the EGOP, and why it is \"enough\" for learning? The experiment in Figure 1 seems to suggest that this matrix encodes the support of the target function; how does this relate to the low-rank support recovery mechanism in (Damian et al. 2022) (Abbe et al. 2022)? \n\n2. The presentation of theoretical results (especially in Section 4.1) can be improved.\n\n- Can the authors clarify what is the \"*appropriate high-dimensional limit*\" in Section 4.1? Under what scaling of data size and model width do the authors expect the asymptotic freeness to hold? \n\n- The eigenstructure of $\\mathcal{K}$ at initialization is not explicitly characterized, and there are no general principles that decide the extent of alignment between the derivatives in Proposition 3, even in the simple quadratic case.  \n\n3. A few related works are missing.\n\n- How gradient descent induces alignment between the target function and certain feature matrices has been studied in the random matrix theory literature. For example, (Ben Arous et al. 2023) analyzed the Hessian and (Wang et al. 2024) the NFM. It would be a good idea to discuss whether these rigorous analyses apply to the problem setting in Section 4.1.   \n(Ben Arous et al. 2023) *High-dimensional SGD aligns with emerging outlier eigenspaces*.   \n(Wang et al. 2024) *Nonlinear spiked covariance matrices and signal propagation in deep neural networks*.  \n\n- How does the speed-limited optimization relate to the feature learning parameterizations in (Yang and Hu 2022) (Chizat and Netrapalli 2023)? Both works enforced the parameters of the first layer to move faster.  \n(Yang and Hu 2022) *Feature learning in infinite-width neural networks*.   \n(Chizat and Netrapalli 2023) *The Feature Speed Formula: a flexible approach to scale hyper-parameters of deep neural networks*.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This submission studies, both theoretically and empirically, the neural feature ansatz (NFA) proposed in (Radhakrishnan et al. 2024), which states that feature learning results in the Gram matrix of a neural network layer aligning with the average gradient outer product (AGOP) with respect to the input to that layer. The authors identified the alignment between the weight matrix and the preactivation tangent kernel matrix as an indicator of NFA and introduced the centered neural feature correlation (C-NFC) to measure such alignment. The authors also provided some theoretical justification on how NFA may arise in the early stages of gradient descent, and proposed a gradient-based update rule to increase the strength of C-NFC and hence feature learning.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please refer to the weakness section above.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "lDfexZSele", "reviewer": "Reviewer_UL1M", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\nHere is a list of the strengths of the paper, according to me:\n\n- The authors try to tackle an important question : *how a neural network learn features adapted to the data*. To do this, they try to follow the NFA and provide complementing experiments and phenomenology to it. These questions are very important to understand Neural Network training.\n\n- The paper -or at least the story line- is quite clear and easy to follow.\n\n- The authors present many experiments that show the phenomenon they study\n\n*Weaknesses**\n\nHere is a list of the weaknesses of the paper, according to me:\n\n- While the notations and concept are well introduced, and the experiments conclusive, the theoretical contributions trying to explain the phenomenon described are questionable are rather weak compared to the claims the authors make. Here are few important examples of this: \n  - *On Proposition 3*: the mathematical hypothesis are not properly introduced, for example why is it said that $\\bar{W} = 0$, two lines before ? I guess that this corresponds to the fact that all the derivatives in the proposition are taken at $t = 0$ (right after initialization), but it took me quite some time to understand it. Moreover, why the authors claim that *the first non-trivial time derivatives are highly correlated*? There is a $\\mathcal{K}$ vs $\\mathcal{K}^2$ difference between both derivatives and I do not understand how this can be benign. Finally, the alignment is only valid locally at $t = 0$, at second order, this is far from being enough to conclude.\n  - Some mathematical statement are rather rapid, e.g. equations $(6)$ then $(7)$ are provided under certain assumptions, and this is never *clearly stated*. Moreover, from these, conclusion are given as if it were general cases. Another striking example is on the fact that there is a validation of the results in Section 4.2,  under very specific hypothesis (one layer, quadratic activation, $X^\\top X = I$...)", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors claim to provide both an experimental and theoretical study of the Neural Feature Ansatz (NFA), namely that the Gram Matrix of each layer aligns during training to the Average gradient outer product matrix (AGOP). They provide experiments for different architectures/dataset and give theoretical ground for this phenomenon in the case of gradient flow training. Finally, the authors claim to provide an algorithm that boosts the NFA phenomenon in order to do more feature learning.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "*Major changes*\n\nMy main change request is either to forget that the paper is something more than experimental or to largely change the theoretical part: being clear about the assumption, the statement etc...\n\n*Minor changes*\n\n- The discussion on the necessity to deal with the centered and uncentered versions of the feature covariance matrices is rather confusing. I would like the authors to clean this point: notably, what is the link with small initialization on this regard?", "broader_impact_concerns": "I have no concern.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper studies how gradient dynamics may induce feature learning, by analyzing the alignment between the weight gram matrix after training and the AGOP, known as the neural feature ansatz (NFA). This is done through a combination of experiments, theory, and heuristic arguments, which provide some characterization of this phenomenon, particularly near initialization. The reviewers were satisfied with the manuscript and the revisions. The results are interesting and could be useful to the community, therefore I recommend acceptance.
Decision: Accept (poster)