Three Mechanisms of Feature Learning in a Linear Network
Authors: Yizhou Xu
Liu Ziyin
Venue: ICLR 2025 Poster (2025)
Keywords: solvable model
feature learning
neural tangent kernel
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Abstract: Understanding the dynamics of neural networks in different width regimes is crucial for improving their training and performance. We present an exact solution for the learning dynamics of a one-hidden-layer linear network, with one-dimensional data, across any finite width, uniquely exhibiting both kernel and feature learning phases. This study marks a technical advancement by enabling the analysis of the training trajectory from any initialization and a detailed phase diagram under varying common hyperparameters such as width, layer-wise learning rates, and scales of output and initialization. We identify three novel prototype mechanisms specific to the feature learning regime: (1) learning by alignment, (2) learning by disalignment, and (3) learning by rescaling, which contrast starkly with the dynamics observed in the kernel regime. Our theoretical findings are substantiated with empirical evidence showing that these mechanisms also manifest in deep nonlinear networks handling real-world tasks, enhancing our understanding of neural network training dynamics and guiding the design of more effective learning strategies.

Published as a conference paper at ICLR 2025
THREE MECHANISMS OF FEATURE LEARNING IN A LINEAR NETWORK
Yizhou Xut, Liu Ziyin
1 Computer and Communication Sciences, École Polytechnique Fédérale de Lausanne
2 Research Laboratory of Electronics, Massachusetts Institute of Technology
3 Physics & Informaties Laboratories, NTT Research
ABSTRACT
Understanding the dynamics of neural networks in different width regimes is cru- cial for improving their training and performance. We present an exact solution for the learning dynamics of a one-hidden-laye linear network, with one-dimensiona, data, across any finite width, uniquely exhibiting both kernel and feature learning phases. This study marks a technical advancement by enabling the analysis of the training trajectory from any initialization and a detailed phase diagram under varying common hyperparameters such as width, layer-wise learning rates, and scales of output and initialization. We identify three novel prototype mechanisms specific to the feature learning regime: (1) learning by alignment, (2) learning by disalignment, and (3) learning by rescaling, which contrast starkly with the dynamics observed in the kernel regime. Our theoretical findings are substan- tiated with empirical evidence showing that these mechanisms also manifest deep nonlinear networks handling real-world tasks, enhancing our understanding of neural network training dynamics and guiding the design of more effective learning strategies.
INTRODUCTION
has been shown that for a neural network under certain types of scaling towards infinite width (or certain parameters), the learning dynamics can be precisely described by the neural tangent kernel (NTK) dynamics (Jacot et al., 2018), or the "kernel regime." We say that a model is in the kernel regime if the NTK of the model remains unchanged throughout training, and the learning dynamics is linear in the model parameters. When the learning dynamics is not linear, we say that the model in the feature learning regime. Since then, a lot works have been devoted to the study of how the kernel evolves during training as it sheds light nonlinear mechanisms of learning (Liu et al., 2020: Huang et al., 2020: Huang & Yau, 2020; Chen et al., 2020; Baratin et al., 2021; Atanasov et al., 2021; Geiger et al., 2021; Bordelon & Pehlevan, 2022; Simon et al., 2023).
Despite this progress, a comprehensive theory that accurately characterizes both the kernel and feature learning dynamics in finite-width models remains clusive. Most existing works focus on infinite-width settings, where the behavior of the network simplifies, but do not extend well to finite configurations. This gap leaves several theoretical and practical questions about the fundamental na- ture of learning in neural networks unanswered. Specifically, there is a lack of analytically solvable models that exhibit both NTK and feature learning dynamics, which are crucial for understand- ing how real-world neural networks learn and adapt. Our main contributions are
We analytically solve the evolution dynamics of the NTK for minimal finite-width model with arbitrary initialization. The model we analyze is a one-hidden-layer linear network, which, de- spite its simplicity, has a non-convex loss landscape and strongly coupled dynamics. Prior to our work, the exact solution for its learning dynamics was unknown.
Based on our exact solutions, we identify three novel mechanisms of learning that are exclusive to the feature learning phase of the network.
Using our exact solutions, we provide phase diagrams that distinguish between the kernel phase and the feature learning phase, for both finite and infinite width models.
AAtthors are listed alphabetically.

---

Published as a conference paper at ICLR 2025
We empirically validate the three mechanisms of feature learning and our phase diagrams in realistic nonlinear networks.
This work is structured as follows. We introduce related literature in Section 2 and Appendix A. Our model and solution are presented Section 3.1, followed by three mechanisms of feature learning in Section 3.2 and 3.3. Section 4 gives the phase diagram that separates feature learning phase and kernel phase. Proofs are provided in Appendix B. Experimental details are given in Appendix C.
2 RELATED WORK
Kernel and feature learning- Under the NTK scaling, it is shown that NTK remains unchanged in the mnfinite-widh limit (Jacot et al., 2018; Lee et al., 2019; Arora et al., 2019b), where the network is asymptotically equivalent to the kernel regression using NTK. Higher order feature learning correc- tions of the NTK have also been studied (Hanin & Nica, 2019; Dyer & Gur-Ari, 2019; Andreassen & Dyer, 2020; Roberts et al., 2022). An alternative to the NTK parameterization is the mean-field (or (PP) parameterization where features evolve at infinite width (Mei et al., 2018; Yang & Hu, 2020; Bordelon & Pehlevan, 2022). Within this literature, the works closest to ours are those computing finite width corrections (Pellegrini & Biroli, 2020; Pham & Nguyen, 2021; Bordelon & Pehlevan, 2024). However, these results are perturbative in nature and applicable when the width is large. Our study has the same goal of understanding the learning dynamics but with a different approach. We analytically solve model that admits analysis both when the model size is finite and infinite.
Linear networks. Linear networks have been extentively used as toy models understand com- plex dynamics of nonlinear networks. For example, they provide significant insights into the loss landscape (Baldi Hornik, 1989; Ziyin et al., 2022), optimization (Saxe et al., 2013; Huh, 2020; Tarmoun et al., 2021; Braun et al., 2022), generalization (Lampinen & Ganguli, 2018; Gunasekar et al., 2018) and learning dynamics (Arora et al., 2018; 2019a; Ziyin et al., 2024) of neural networks. Closely related to are Saxe et al. (2013), Braun et al. (2022), and Atanasov al. (2021), which solve the learning dynamics of linear models under special initializations. Mathematically, previ- ously known results particular solutions the differential equation, whereas solution is general solution. A contemporary paper (Kunin et al., 2024) studies the exact solution in case of one hidden neuron, whereas our result applies arbitrary number of neurons. Another contemporary paper (Beneventano & Woodworth, 2025) analyzes the finite stepsize effect for the same model.
3 AN EXACTLY SOLVABLE MODEL.
3.1 PROBLEM SETTING AND SOLUTION
Let us consider two-layer network with power activations f(x) > (E W where dy is the input size, d is the network width, and W are the weight vector/matrix of the second and the first layer, respectively, and y is a normalization factor. We first show that the learning dynamics of this network can be reduced to 1d dynamics for any 3, and then solve the dynamics exactly for the linear network case where B = 1.
We consider the following network trained on the MSE loss:
d
İu, W) EEa (7 Σ u wyz -
1=1
J-1
where we treated the target y as a function of ã. E denotes the averaging over the training set. The training proceeds with the gradient flow algorithm. We allow the two layers to have different learning rates, 7/ and nw: du: dt
at
dwij = 7m Dwnj
ai
- Ou; dt.
(2)
We restrict to when the data lies on Id subspace, and the following proposition shows that the learning dynamics under (q) is equivalent to that under a simplified loss.
Proposition 1. Let ã = an, where a E is a random variable and n is a fixed unit vector. Let =
E and [(a((a)] Then, the gradient flow of Eq.(1) (1) equals the gradient flow of
[ Σ
de Wijj,,)
2
L(u ui(
(3)
j=1
2

---

Published as a conference paper at ICLR 2025
The following theorem gives a precise characterization of the dynamies of and W for arbitrary initialization and hyperparameter choices.
Theorem 1. Let
(t(t) == 20 Th 2j-1 Way (t)z Bmm (t)) qi(t) = 2p
(4)
1/ Th 2i-1 wiy (t) v pu (t)),
where p :== do
(Analytic Reduction to Id Dynamics) The solution of the gradient flow can be given by the solution of the following ODE
(1.(t) = -27 BTi 2 Bol 700 (37 -1/
1/2 -1,0000---00000000000000-0----00
dt
j=1
(5)
with p,(t)=
for all i,j 1,2, .., d, where de
F((z):= == f x(x+ x(I+cc/z)8-11
(6)
F denotes the inverse function of Fú, (0)) The initial conditions are A,(
(t))
After solving (t),, the weight vector/matrix given by
(t) 1
/ (p) - (t)),
(iy(() wir(() (p:(1) (;(()) Tup
where = F())=F)(AA(t)) ((t)= = p.(1))
2. (Exact Solution of Linear Net) Let := Eリ-1 E) (0) Q= E (0) For = Pf
the solution
a. a.+5(1)a= 1/2
p((1) = (0)
I-((1)
(8)
a.t((a)a 1122
(1(1) = (0)
1-E(()
where
E(1) = 1 1-a,
exp (xp(-41/tc),
(9)
I
to := 1/ (V Th + 4p (=))
(10)
a+ := 2(>2d)p2P TTaaau 7Py +t
(11)
Remark. Theorem does not impose specific assumptions on the width of the network or the ini- tialization of the parameters. Moreover, does not impose more assumptions as long as the data lie a one-dimensional subspace. For the rest of the paper, We focus on the case 3 ==1 1 it admits simpler and more human-ine interpretable solutions (8), which will be the starting point of all following results. Some special cases of 3, including = 2, are also in principle solv- able, but it is difficult write the results in comprehensible form for a large width and we leave to a future work to study these cases in detail. Our focus will be on the optimization dynamics of this network, even though it is also possible to discuss generalization with this solution briefly touch on this in Appendix D.
the data or labels,
Let us begin by analyzing each term and clarifying their meanings. In the theorem, have trans- formed u; and wij into an alternative basis p: and q;, and {(t) is the only time-depende term. Note that & decays exponentially towards zero the time scale fc.
The constants a = Q/P are two asymptotic scale factors. In the limit t + we have that pi( (.(0) (o( (0)\Vaa..
(12)
BBcause a+ 0,t(t) So, is well defined. Also, See Appendix for the case P O.
3

---

Published as a conference paper at ICLR 2025
This directly gives the mapping between the initialization to the converged solution. Unlike strongly convex problem where the solution is independent of the initialization, we see that the converged solution for our model is strongly dependent the initialization and the choice of hyperparameters. Perhaps surprisingly. because O+ in (11) are functions of the learning rates, the converged solution (12) depends directly on the magnitudes of the learning rates. This directly tells us the implicit bias of gradient flow for this problem. Another special feature of the solution is that for any direction orthogonal to x, the model will remain unchanged during training. Let m x, we have that Σ₁ Wiz (t)m, Σ, Way (0)m, Namely, the output of the model in the subspace where there is no data remains constant during training.
In the theorem, what is especially important is the characteristic time scale te. which is roughly the time it takes for learning to happen. Notably, the squared learning speed 1-2 depends on two competing factors:
nu Y p2p"
4p (7" dpQ PQ
contribution from feature learning contribation from kernel learning
The first factor depends on the input-output correlation and learning rate, which we will see is indicative of feature learning. The second term depends only on the input data and on the model initialization. We will see that when this term is dominant, the model is in the kernel regime. In fact, this result already invites a strong interpretation: the learning of the kernel regime is driven by the initialization and the input feature, whereas the learning in the feature learning regime is driven by the target mapping and large learning rates.
Using this theorem, one can compute the evolution of the NTK. Note that when different learning rates are used for different layers, the NTK needs to be defined slightly differently from the conven- tional definition. For the MSE loss, the NTK is the quantity K that enters the following dynamics: df(x) 2K ( x' (f(x)) - y)- This implies that for our problem,
dt
K (x,22'0 = 22== xm (t) (t)) + tull
(13)
which follows from Eq. (2) and (3), where W stands for the matrix with elements way- u(t) and (tt) are obtained via Eqs. (7) to (11). While the overall formula for the NTK dynamics is quite complex, we will provide the conditions concerning when it evolves with an O(1) amount in Section 4. When Teu Tu, our definition agrees with the standard NTK.
3.2 LEARNING BY ALIGNMENT AND DISALIGNMENT
Before we discuss the various phase diagrams implied by Theorem 1, we first focus on an interesting effect predicted by this theorem, which differentiates it from previous results on similar problems. We first set x to be 1d, because Theorem 1 suggests that the dynamics of GD training has only a rank- I effect (i.e. p and q are two effective weight vectors). Numerical results for non-ld are presented at the end of this subsection. A quantity of theoretical and practical interest is ((t) := uz uzw/|l w//|ul/ ||||||||||| which represents the cosine similarity between u and w. Recent works have identified the alignment between the weight and representation after the training starts as a mechanism for feature learning (Everett et al., 2024). Studying the evolution of Ç thus offers a direct clue of how this alignment happens. This quantity is especially interesting to study because it tells us how well-aligned the two layers are during training. Notably, this quantity vanishes as d + oo if and only if the model the kernel regime (Theorem 2), so it serves as a great metric for probing how feature learning happens.
Letting and denoting at))
a_tEa
T-E we have by Theorem 1:
a(1)P-Q/a(1)
<(1) V(a(t)p+Q/a(tt) (( Ep.q.)
(14)
Another pair of important quantities are P and Q, which are essentially the initialization scales of the model. For 1),, Te. a small P implies that (0) (0) for i, which implies that the model close to anti- anti-parallel at the start of training. Likewise, a small Q implies that (0) x (.) for all i; namely, the two layers start training when they are approximately parallel. When both P and Q are small, the model is initialized close to the origin, which is saddle point. We will also see below that in the kernel regime, we always have Q = P.
4

---

Published as a conference paper at ICLR 2025
kernel phase
orthogonal Init. near orthogonal Init.
feature learning phase
Linear
ReLU
Sigmoid
Swish
Leaky ReL.U
alignment
0.5
0.5
0.5
0.5
0.0
0.0
0.0
0.0
1000
1000
0
1000
1000
0
1000
iteration
iteration
iteration
iteration
iteration
disalignment
1.0
1.00
1.00
1.00
0.75
0.75
0.75
0.95
0
2000
2000
2000
2000
2000
iteration Figure 1: The evolution of Ç of two-layer networks with different settings. Specifically, test lincar, ReLU, sigmoid, swish, and leaky ReLU activations for both alignment (upper) and disalignment (lower) cases. For the linear network, we show the theoretical predictions obtained from (14) as lines and experimental results as points. The results for nonlinear networks qualitatively similar.
Iteration
iteration
iteration
iteration
where 4p:g: = un w const does not change during training. In general, the angle evolves by an 0(1)* amount during training". In fact, the angle remains unchanged only in the orthogonal initialization case or the kernel phase, where (x()) 0/1 throughout training (see Appendix C).
For Gaussian initialization, an intriguing fact is that layers tend to align in the feature learning phase, while the alignment remains asymptotically zero for the kernel phase (see Section for more formal definitions of the phases). Assuming x Q and Σ P:q: a 0, which holds for a(t)++++ which monotonously changes from -1
a(t)-11
sufficiently large, (14) leads to ((t) x
ag+ In the feature learning phase, two terms in (10) are of the same order (see Section 4), so we can assume Np 2 2K p² 2² without loss of generality, where K is a certain positive constant. This
(K+1)²-1
further leads to non-zero lower bound of the final alignment ((00) 2 (K+1)=+1' On the other hand, in the kernel phase, a4 x 1, and thus ((t) o(1). Therefore, in the kernel regime, the two layers are essentially orthogonal to each other throughout training. This suggests one mechanism for the failure of the kernel learning phase. For a data point the hidden representation is wxz but predominantly many information in wa ignored after the the layer u. This implics that the model will have disproportionately larger norm than what is actually required to fit the data, which could in turn imply strong overfitting.
When the two layers initialized in parallel way. This setting is often called the "orthogonal initialization" (Saxe et al., 2013). In the orthogonal initialization, u is parallel to w, and so p. Cq: for constant C. In this case, it easy verify that ((t) 1, meaning that u and remain parallel or anti-parallel throughout training.
In general, might be interested in whether the alignment increases decreases. Our solution implies a rather remarkable fact: Ç is always a monotonic function of t. To see this, its derivative is
dc (P+ da [(((IPP+ Q/a(1))2 (3 EP:g:)
(15)
When u and w are parallel, this quantity is zero, in agreement with our discussion about orthogonal initialization. When they not parallel, have that 4PQ (3 Σ p:q:)² 0 by the Cauchy inequality, and thus dc/da O. Because att) monotonically evolves from 1 to a4 0, the evolution of ( is also simple: ((t) monotonically increases if a,, > 1 or, cquivalently, if
Vuulww 2ndpP +
Vhhoy PydpP
Q
(16)
p
and monotonically decreases a. < 1. does not change if aa 1.
'It is because att) changes from 1 to a., during training. angle evolves by O(1). See the proof of Theorem 2 for more details. feature learning phase, a, = =((1, so the
f=x) = Og((z)) means (x) O(g(z)) holds almost surely.
5

---

Published as a conference paper at ICLR 2025
layer 1&2
0.9
layer 2&3
0.8
0.8
layer 3&4
2 0.7
layer 1&2
0.7
layer 283
0.6
layer 3&4
800
1000
lleration
(a) alignment a 4-yer FCN
(b) alignment vs. initialization scale
Figure 2: The alignment angle Ç between different layers of a four-layer FCN with ReLU activation trained on MNIST. (b) shows the final alignment for different initialization scale σ, while (a) shows training curves corresponding to The dashed lines in (h) show the initial alignment. See Appendix C for experiments on a six-layer network.
1500
a 0.1
acc.
0.8
O = 0.1
0= 1
pome
1000
0=5 5
test
0=1
0.6
10
0=
500
10
0.4
100 epoch
150
200
0
100 epoch
150
200
Figure 3: The initialization scale a correlates negatively with the performance of Resnet- 18 on the CIFAR- -10 dataset. Left: test accuracy. Here, a is a constant multiplier we apply to the initialized weights of the model under the Kaiming init. Right: the norm of all weights. While all models achieve 100% training accuracy, models initialized with larger scale converge solutions with higher weight norms, which is sign that the layers are misaligned.
When does condition (16) hold? Let us focus on the case y > O because the theory is symmetric in the sign of y. The first observation is that it holds whenever Q 2 P. which is equivalent to u (0) < O. Namely, if the model is making wrong predictions from the beginning, it will learn by aligning different layers. Moreover, this quantity also depends on the balance of the two learning rates. Notably, when the learning rates for the two different layers are the same, the change in independent of the learning rate. The dependence on the learning rate becomes significant once we use different learning rates on the two layers. For example, when Th >> The (or vice versa), this condition depends monotonically and (essentially) linearly on Yju, and making 7he close to Thu has the effect of making the two layers more aligned.
Why does the alignment effect depend on the ratio Q/P? Because when Q is small, the model layers are initialized to be aligned and are likely to make predictions that are too large, and the learning process necessarily involves decreasing the model output on the data points, which can be achieved in one of the two ways: (1) decrease the scale ||и][//] (2) decrease the aligmment Ç. When condition (16) is not satisfied, GD employs both mechanisms for learning. Lastly, it is also worth noting that for this problem, SGD has been shown converge to a perfectly aligned distribution of solutions (Ziyin et al., 2023). This comparison thus shows a qualitative difference between GD and SGD-using GD, alignment a strong function of the initialization, whereas in SGD, the alignment is quite independent of the initialization. The difference between SGD and GD is of order 1 in this problem, even if the noise is very small.
See Figure 1, where we show that the evolution of ( of two-layer networks with d = 10000. It is trained on a regression task. Similar experimental results are observed for a classification task trained with the cross-entropy loss (Appendix C). We choose ? = 1/vd for the kernel phase and y = 1/d for the feature learning phase. The initial weights are sampled from i.i.d. Gaussian distribution N( 1). Therefore, we have P x Q and the initial ((0) æ O. From (11) we have a. > 1 if y > O, so (t) monotonically increases, but the increase is negligible in the kernel phase. Therefore, in this case, the model learns features by alignment. Meanwhile, the orthogonal initialization remains 1 as predicted. In the near orthogonal case of Figure we set (0) & and the initial model output to be large. As predicted, <(t) monotonically decreases. In this case, the model learns by disalignment. From Figure 1, also see that this phenomenon holds for all non-linear activation functions.

---

Published as a conference paper at ICLR 2025
The layer alignment and disalignment effects can be generalized to higher dimensions, deeper net- works and non- non-linear activations, Here, we define Ç [ WII where U, W are the weight matrices of two consecutive layers and the L2 norm for matrices is used. See Figure 2 for a four-layer fully connected network (FCN) with ReI.U activation and different initialization scales trained on MNIST datasets. Figure 2 (a) shows that the alignment between consecutive layers increases during training, and Figure 2 (b) demonstrates that layers stop being aligned for large initialization. These results are consistent with simpler settings, verifying the generality of the discovered mechanism.
JUWiL
3.3 LEARNING BY RESCALING
Learning can also happen by rescaling the output. The evolution of |/u|| and ||w/|| are given by
||u||2 = dull==d(a._))) P 2 - pig. ||wll d(a, P a+ + Pig.
0+
dnL 4|wll d(p- -.. ) which is positive when ( 0, and negative when 5 < O. and, thus, da, do.
Thus, the rescaling coincides with the alignment, namely, ||u||| and ||د||| become larger when they being aligned (Kc) gets larger), and become smaller when they are being disaligned (() gets smaller). More explicitly, (1) P Q and a+ > 1, or P < Q and ar, < 1, or P P= Q: ||r]|| and |||| monotonically increase. (2) P Q and 1 > a, 2 /QJP, or P << Q and 1 a, S VQ/P: |ul|| and [r|] and monotonically decrease. (3) P Q and a, /Q/P, or P < Q and a, > vQ/P: ||u||] and |||||| first decrease, and then increase. (4) a. 1: everything keeps unchanged. Again, in the kernel phase, the scale change of the model vanishes, In the orthogonal initialization, however, this quantity changes by an O(1) amount. Therefore, the orthogonal initialization essentially learns by rescaling the magnitude of the output.
3.4 HOW DOES FEATURE LEARNING HAPPEN?
The analysis thus suggests three mechanisms for feature learning, all of which are absent in the kernel phase. The first two mechanisms are the alignment and disalignment in the hidden layer, which is driven by the initialization balancing between the two layers. The second mechanism is the rescaling output, which is a simple operation and is unlikely to be related to learning actual features. This argument also agrees with the common technique that even if we normalize the layer output, the performance of the network does not deteriorate (loffe & Szegedy, 2015).
The second question is whether we want alignment or disalignment. The intuitive answer seems to be that alignment should be preferred over disalignment. Because aligned layers require smaller model norm to make the same prediction, whereas a disaligned model requires a very large model norm to make the prediction. Our theory thus offers mechansism of how relatively smaller initial- ization is often more preferable in deep learning - when the model has an overly large initialization, it will learn by disalignment, whereas small initialization prefers alignment. This is in agreement with the common observation that a larger initialization variance correlates strongly with a worse performance (Sutskever al., 2013; Xu et al., 2019; Zhang al., 2020). The parameterization in Yang et al. (2022) ensures that activations are 0(1) at initialization, which could avoid disalignment problems. This is distinct from the NTK/feature-learnin initializations push the model towards the kernel regime, as discussed in Section 4.2 of our paper. In the kernel regime, alignment and disalignment effects are not present. Thus, our explanation is particularly relevant for cases where the initialization is large but not large enough to push the model into the kernel regime. While our explanation complements existing theories, it provides distinct angle on the role of initialization in training dynamics.
explanation, which suggests that larger
numerical result is presented in Figure 3, where a larger initialization leads to worse performance. Note that this example can only be explained through the disalignment effect because (I) the model achieves 100% train accuracy in all settings, yet (2) a larger initialization leads to a larger norm at the end of the training, which also correlates with worse performance. Another piece of evidence is the commonly observed underperformance of kernel models. In the kernel phase, the model norm diverges and the model alignment is always zero, which could be hint of strong overfitting. Therefore, our theory suggests that it would be a great idea for future works to develop algorithms that maximize layer alignment while minimizing the change in the output scale.
7

---

Published as a conference paper at ICLR 2025
4 PHASE DIAGRAMS
Our theory can be applied to study the learning of different scaling limits, where we scale the hyper- parameters with a scaling parameter K co. Here, is an abstract quantity that increases linearly, and all the hyperparameters including the width are power-law function of K. Conventionally, k the model width; however, this excludes the discussion of the lazy training regime in the theory, where the model width is kept fixed and the scaling parameter is the model output scale 7.
We first establish the necessary and suffi- cient condition for learning to happen: the learning time Le needs to be of order @( I).. When it diverges, learning is frozen at ini- tialization. When it vanishes to zero, the discrete-time SGD algorithm will be un- stable, a point that is first pointed out by Yang & Hu (2020). Therefore, we first study the condition for fe to be of order 1. which is equivalent the condition that (assuming x, y are order 1)
Phases of learning in different scaling limits. For brevity, the learning rates of the two layers are set to equal. The first block shows that the models can be frozen or unsta- ble if we do not scale 7 accordingly. The second block shows that one can always choose 1) such that the model training freeze. The third and fourth blocks
Table 1:
is stable and does show that one can always choose a pair of n and y such that the model is either in the feature learning phase or the ker- nel phase. MF refers to the mean-field scaling in (Mei et al., 2018) and lazy refers the scaling (Chizat et al., 2018).
scaling
NTK
MF
Xavier
Kaiming
lazy
(17) 7u + (P=
Cd
0
For Gaussian initializzaiion 4,o ~ ~(
Cy
1/2
-1
0
0
1
Cu
0
0
-1
-1
0
and wio ~ N(..o.)) P and Q random variables with expecta-
Cue
0
0
0
Cip
tion (Nho? + m/oz)/4 and variance
phase
kernel
frozen
learning
unstable
unstable
(Noo + 22 Generally, all hy- perparameters are powers of k; d xc k x²x
Cipy phase
O kernel
1
O
-1
-2
learning lcarning kernel
kerne
Y R o x ke=, of Tw k and T7 x k For simplic- constant.
I
0
0
-1
-1
0
-1
0
phase
learning
learning learning learning learning
ity, we set the input dimension dg to be a
Cy
0
-1
-2
C
-1/2 kernel
-1/2
0
1
Equation (17) implies
phase
kernel
kernel
kernel
kernel
max (2c, + Cy. + Cyu 20,-- + Cd + max Cr_
Cus Cr.
Cw) = 0.
(18)
Whatever choice of the exponents that solves the above equation is a valid learning limit for a neural network. The phase of the network depends on the relative order of the above two terms.
Definition 1. A model is in the kernel phase if (1) Eq. (13) is independent of t as 00 (2) NTK = O1))
When 4. @(1), a model is said to be in the feature learning phase if it is not in the kernel phase.
Theorem 2. When Eq. (18) holds, model is in the kernel phase if and only if lim P/Q = 1, a.s. and
Ca max (Co Cus Cq. + Cw } > Cn. + Cyww
(19)
The necessary condition P a Q for the model being in the kernel phase is interesting and highlights the important role of initialization in deep learning. There are three common cases when this holds:
1. d + co and un and we are independent (standard NTK);
2. is finite and the initial model output is zero: Σ Ege u,wyy = 0 (lazy training)
3. d is finite, ń + oo and C, + Co.. + Cyp + Cyy.;
The first case is the standard way of initialization, from which one can derive the classic analysis of the kernel phase by invoking the law of large numbers. The second case is the assumption used in the lazy training regime (Chizat et al., 2018). (Chizat et al. (2018) assumes C.yy Cr. Cr... -2 and e, Cw O. satisfying the conditions of the exponents (18) and (19).) This case, however, relies on special initialization, and thus our results better illustrate the occurrence of the kernel
'Note that the solution in Theorem is invariant under the transform C. C.. + Cur + Car 0,0, G, - 20, or. Cr, + 20, City Cn 20, corresponding to the abe symmetries in Yang & Hu (2020). "This requires (n W W + -||uu|/ O(1)
8

---

Published as a conference paper at ICLR 2025
mean field
0.55
109
NTK
test acc
IIV Wall-
[Palle
10
Kaiming
0.50
Kalming
10
Xavier
0.45
Xavier
20000
10
10
Figure 4: A two-layer fully connected ReLU net with d neurons trained on the CIFAR-IO dataset for 10000 epochs with batch size 128. The kernel phase is shown in solid lines and the feature learning phase is shown in dashed lines. As the theory predicts, both types of initialization can be turned into either the feature learning or the kernel phase by choosing different combinations of and 1). Left: the best test accuracy during training. Right: relative distance from the initialization.
phase for advanced initialization methods where different weights can be correlated. The third case happens when the learning rate and the initialization are balanced. This suggests that achieve feature learning, one should make sure that the learning rate and the initialization are well balanced:
Cu Ca.
In conclusion, the overall phase is (1) kernel phase, if the first term in (18) is strictly smaller than the second term: 0 2c,, + Cd + max(Cn, + Cus Cr_ + Cwe 2c,, + Cn, + C%. and lim P/Q = 1, (2) feature learning phase if otherwise. A key difference between these two phases is whether the evolution of the NTK is O(1). or equivalently whether the model learns features.
The following corollarie are direct consequences of Eq. (18).
Corollary 1. For any Cap Cu and Cw. choosing Cy. Cn= min(-c,,- -2c, ca max Cw ensures that the model is stable.
Corollary 2. For any Cy and Cw. choosing Cy. = Cn = Onj and cy Ccy with cy. 2 ca + max Cas Cw leads to a feature learming phase.
Corollary 3. Assume limx-sse P/Q- 1. For any and Cws choosing Cy (ca max CCus Go) Cn)) and Cp < Cd + max Cu: Cw leads to kernel phase.
They imply two important messages: for every initialization scheme, (1) one can choose optimal learning rate such that the learning is stable; (2) one can choose an optimal pair of learning rate and output scale such that the model is in the feature learning phase. Point (1) agrees with the analysis in Yang & Hu (2020), whereas point (2) is a new insight we offer. See Table for the classification of different common scalings. We choose scalings according Corollary 2 and 3, to turn each model into the feature learning or the kernel phase. Table is closely related to the Tensor programs framework (Yang & Hu, 2020; Yang al., 2022; Yang & Littwin, 2023). The key difference is that our results apply to winth networks with arbitrary initialization, whereas Tensor programs assume infinite width and Gaussian initialization.
4.1 PHASES DIAGRAM OF INFINITE- WIDTH MODELS
Now, let us focus the case d + oo (ca 1), corresponding the infinite width limit in the NTK and feature learning literature (Jacot et al., 2018; Li et al., 2020; Yang & Hu, 2020).
See Figure 4. We implement two-layer FCN on the CIFAR-I0 dataset with ReLU activation. We run experiments with the scalings of the standard NTK, standard mean-field, Kaiming model, and Xavier model. C. and C, are chosen according to Table 1. Here, we use the superscript + to denote the type of scaling that leads to feature learning phase, and - denotes the kernel phase. For the Kaiming and Xavier model, we choose both oz and et,, and refer them as Kaiming* and Xavier*, respectively. The left figure shows that turning the Kaiming model into the feature learning phase improves the test accuracy by approximately 5%. similar to the gap between the standard NTK model and the mean-field model. Meanwhile, turning the Xavier model into the kernel phase decreases the test accuracy by approximately 10%. This is because the fixed kernel restricts the generalization ability in the kernel phase, and the difference between these models in the kernel phase might be attributed to their different kernels. The right figure asserts that there is power law scaling of the weight evolution ||W-Woll o d= and we can predict that 0 for the feature Wo learning phase, and 0.5, 1, for NTK, Xavier and Kaiming", respectively, which are perfectly consistent with the numerical results (6 0.47, 1.08, 1.93). See Appendix B.4 for details.

---

Published as a conference paper at ICLR 2025
1.0
test train
10
accuracy
0.8
||W ||-- Wolle
||alle
feature learning
NTK
10-4
feature learning
NTK
0.6
10-7.
0.4
5
-5
5
Cw
Cw
Figure 5: A two-layer FCN with different initialization scales trained on the CIFAR-10 10 dataset. We see that finite-width models can also exhibit qualitative differences between the feature learning and the kernel phases when other hyperparameters are scaled toward infinity. Notably, this scaling is different from the lazy training scaling, implying that there are numerous (actually infinitely many) ways for the model enter the kernel phase, even at a finite width.
Thus, in agreement with the theory, choosing different combinations of the output scale and can turn any initialization into the feature learning phase. This insight could be very useful in practice, as the Kaiming init. is predominantly used in deep learning practice and is often observed to have better performance at common widths of the network. Our result thus suggests that possible to keep its advantage even if we scale up the network. Further, our results also imply that any valid learning regimes transfer well when the model gets larger, while the feature learning phase, a larger model generally leads to better performance. This is consistent with Yang et al. (2022).
4.2 PHASE DIAGRAM FOR INITIALIZATIONS
Now, we study the case when d is kept fixed (ca 0), while other variables scale with x co. In this case, the phase diagram is also given by Theorem 2. See Figure 5 for an experiment. We set Cu max(0, Cwe and C., min mww//, 0}. This choice satisfies (18). By Theorem 2, the network is in the kernel phase if and only if c.. > 0.
/2,
One important example for this section is the lazy training regime, where C2 = Cu Ca = 0, and we can choose c, and C, 2 according to Corollary 3, leading to kernel phase in finite width. Another example to consider large initialization, i.e., Cu C c> 0. In this case, we can choose On and e, c according to Corollary 2, leading to feature learning phase. Actually, this choice of the normalization factor 7 cancels out the scaling of the initialization. On the other hand, if we choose 7 to a kernel space. This might be another possible explanation that larger initialization often leads to worse performance empirically. Like before, we implement two-layer fully connected ReLU network on the CIFAR-10 dataset with d 2000. We choose x = 10 for illustration purposes. A clear distinction is observed between the feature learning phase and the kernel phase. (1) In Figure 5(a), the training accuracy can reach 1.0 in the feature learning phase but not the kernel phase, because the NTK in the kernel phase is fixed, and thus the best training accuracy is limited by the fixed kernel. (2) As discussed in the previous section, the test accuracy In Figure 5(a) is about 5% higher in the feature learning phase due its trainable kernel. (3) In Figure 5(b), the weight matrices evolve significantly in the feature learning phase but not the kernel phase.
commonly done, we have choose Cy according Corollary 3, leading
CONCLUSION
Solving minimal models has been a primary approach in natural sciences to understand how con- trollable parameters are causally related to phenomena. In this work, we have solved the learning dynamics of a minimal finite-width model of a two-layer linear network. Through a comprehen- sive analysis of its learning dynamics and phase diagrams, we have uncovered valuable insights into how feature learning happens and the impact of various scalings on the training dynamics of non- linear neural networks. Our theory is obviously limited: the analytical results only hold for inputs lying in one-dimensional subspace. This limitation arises from the inadequacy of conservation laws more general cases (e.g.. Marcotte et al. (2023, Corollary 4.4) suggests that the maximal number of independent conservation laws might be much less than the degrees of freedoms), which potentially implies the impossibility of solving a more general model than ours. Lastly, our results only considers deterministic learning dynamics; feature learning actual models are likely to be also determined by regularization and noise during training (Ziyin et al., 2025a;b), and could be interesting to compare feature learning under noise and without noise.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory, pp. 4782-4887. PMLR, 2022.
Anders Andreassen and Ethan Dyer. Asymptotics of wide convolutional neural networks. arXiv preprint arXiv:2008.08671 2020.
Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization. In International conference on machine learning, pp. 244- 253. PMLR, 2018.
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019a.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov. and Ruosong Wang. On exact computation with an infinitely wide neural net. Advances in neural information processing systems, 32, 2019b.
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference. Journal of Machine Learning Research, 22 (106):1-51, 2021.
Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan. Neural networks as kernel learners: The silent alignment effect. arXiv preprint wrXiv:211..000. 2021.
Jimmy Ba, Murat Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35:37932-37946, 2022.
Pierre Baldi and Kurt Hornik, Neural networks and principal component analysis: Learning from examples without local minima. Neural networks, 2(1):53-58, 1989.
Aristide Baratin, Thomas George, César Laurent, R Devon Hjelm, Guillaume Lajoie, Pascal Vin- cent, and Simon Lacoste-Julien. Implicit regularization via neural feature alignment. Interna- tional Conference on Artificial Intelligence and Statistics, pp. 2269-2277 PMLR, 2021.
Pierfrancesco Beneventano and Blake Woodworth. Gradient descent converges linearly to flatter minima than gradient flow shallow linear networks. arXiv preprini arXiv:2501.09137, 2025.
Alberto Bietti, Joan Bruna, and Loucas Pillaud- Vivien. learning gaussian multi-index models with gradient flow. arXiv preprint arXiv:2310.19799 2023.
Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks. Advances in Neural Information Processing Systems, 35:32240-32256, 2022.
Blake Bordelon and Cengiz Pehlevan. Dynamics of finite width kernel and prediction fluctuations in mean field neural networks. Advances in Neural Information Processing Systems, 36, 2024.
Lukas Braun, Clémentine Dominé, James Fitzgerald, and Andrew Saxe. Exact learning dynamics of deep linear networks with prior knowledge. Advances in Neural Information Processing Systems, 35:6615-6629, 2022.
Shuxiao Chen, Hangfeng He, and Weijie Su. Label-aware neural tangent kernel: Toward better gen- eralization and local elasticity. Advances in Neural Information Processing Systems, 33:15847- 15858, 2020.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. arXiv preprint arXiv: 18/2.09956, 2018.
Hugo Cui, Luca Pesce, Yatin Dandi, Florent Krzakala, Yue M Lu, Lenka Zdeborová, and Bruno Loureiro. Asymptotics of feature learning in two-layer networks after one gradient-step. arXiv preprint arXiv:2402. 04980, 2024.
11

---

Published as a conference paper at ICLR 2025
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. Conference on Learning Theory. pp.
PMLR, 2022.
Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. How two-layer neural networks learn, one (giant) step at a time. arXiv preprint arXiv:2305. 18270, 2023.
Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from feynman diagrams. arXiv preprint arXiv:1909. 110099911344,,0 2019.
Katie Everett, Lechao Xiao, Mitchell Wortsman, Alexander A Alemi, Roman Novak, Peter Liu, Izzeddin Gur, Jascha Sohl-Dickstein, Leslie Pack Kaelbling, Jaehoon Lee, al. Scaling expo- nents across parameterizations and optimizers. arXiv preprint arXiv:2407.05872, 2024.
Mario Geiger, Leonardo Petrini, and Matthieu Wyart. Landscape and training regimes in deep learning. Physics Reports, 924:1-18, 2021.
Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. Advances in neural information processing systems, 31, 2018.
Boris Hanin and Mihai Nica. Finite depth and width corrections to the neural tangent kernel. arXiv preprint arXiv: 190900999,, 2019.
Jiaoyang Huang and Horng-Tzer Yau. Dynamics of deep neural networks and neural tangent hier- archy. In International conference on machine learning, pp. 4542-4551, PMLR, 2020.
Wei Huang, Weitao Du, and Richard Yi Da Xu. On the neural tangent kernel of deep networks with orthogonal initialization. arXiv preprint arXiv:2004.05867, 2020.
Dongsung Huh. Curvatune-corrected learning dynamics deep neural networks. In International Conference on Machine Learning, pp. 4552-4560. PMLR, 2020.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv: 1502.03167, 2015.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kerne:: Convergence and gen- cralization in neural networks. Advances in neural information processing systems, 31, 2018.
Arthur Jacot, François Ged, Berfin Şimşek, Clément Hongler, and Franck Gabriel. Saddle-to-saddle dynamics in deep linear networks: Small initialization training, symmetry, and sparsity. arXiv preprint arXiv:2 02..///..5353 15933, 2021.
Stanislaw Jastrzebski, Maciej Szymezak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Krzysztof Geras. The break-even point on optimization trajectories of deep neural networks. arXiv preprint arXiv:2002.09572 0922000000999,,..... 2020.
Dayal Singh Kalra and Maissam Barkeshli. Phase diagram of early training dynamics in deep neural networks: effect the learning rate, depth, and width. Advances in Neural Information Processing Systems, 36, 2024.
Dayal Singh Kalra, Tianyu He, and Maissam Barkeshli. Universal sharpness dynamics in neural network training: Fixed point analysis, edge of stability, and route to chaos. arXiv preprint arXiv:231 1.02076, 2023.
Daniel Kunin, Allan Raventós, Clémentine Dominé, Feng Chen, David Klindt, Andrew Saxe, and Surya Ganguli. Get rich quick: exact solutions reveal how unbalanced initializations promote rapid feature learning. arXiv preprint arXiv:2406.06//58 2024.
Andrew K Lampinen and Surya Ganguli. An analytic theory of generalization dynamics and transfer learning in deep linear networks. arXiv preprint arXiv:1809. 10374 2018.
Jachoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 32:8572-888 8583, 2019.
12

---

Published as a conference paper at ICLR 2025
Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiw:2003.02218, 2020.
Yuanzhi Li, Tengyu Ma, and Hongyang R Zhang. Learning over parammetrrie two-layer neural networks beyond ntk. In Conference learning theory, pp. 2613-2682. PMLR, 2020.
Chaoyue Liu, Libin Zhu, and Misha Belkin. On the linearity of large non-linear models: when and why the tangent kernel is constant. Advances Neural Information Processing Systems, 33: 15954-15964, 2020.
Philip M Long. Properties of the after kernel. arXiv preprint wrxiv:2105.1058. 2021.
Noel Loo, Ramin Hasani, Alexander Amini, and Daniela Rus. Evolution of neural tangent kernels under benign and adversarial training. Advances in Neural Information Processing Systems, 35: 11642-11657, 2022.
Sibylle Marcotte, Rémi Gribonval, and Gabriel Peyré. Abide by the law and follow the flow: Con- servation laws for gradient flows. 2023.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape two- layer neural networks. Proceedings of the National Academy of Sciences, 115(33):576665-7767), 2018.
Gabin Maxime Nguegnang, Holger Rauhut, and Ulrich Terstiege. Convergence of gradient descent for learning linear neural networks. arXiv preprint arXiv:2108. 020440 2021.
Courtney Paquette, Kiwon Lee, Fabian Pedregosa, and Elliot Paquette. Sgd in the large: Average- case analysis, asymptotics, and stepsize criticality. Conference on Learning Theory, pp. 3548- 3626. PMLR, 2021.
Franco Pellegrini and Giulio Biroli. An analytic theory of shallow networks dynamics for hinge loss classification. Advances in Neural Information Processing Systems, 33:5356-553 5367, 2020.
Huy Tuan Pham and Phan-Minh Nguyen. Limiting fluctuation and trajectorial stability of multilayer neural networks with mean field training. Advances in Neural Information Processing Systems, 34:4843-4855, 2021.
Daniel A Roberts, Sho Yaida, and Boris Hanin. The principles of deep learning theory, volume 46. Cambridge University Press Cambridge, MA, USA, 2022.
Andrew M Saxe, James McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynam- ics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6122 2013.
Inbar Seroussi, Gadi Naveh, and Zohar Ringel. Separation scales and thermodynamic descrip- tion of feature learning in some cnns. Nature Communications, 14(1):908, 2023.
James Simon, Maksis Knutins, Liu Ziyin, Daniel Geisz, Abraham J Fetterman, and Joshua Al- brecht. On the stepwise nature of self- ceff-supervised learning. arXiv preprint arXiv:2303. 15438, 2023.
Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initial- ization and momentum in deep learning. In International conference on machine learning, pp. 1139-1147. PMLR, 2013.
Salma Tarmoun, Guilherme Franca, Benjamin Haeffele, and Rene Vidal. Understanding dynamics of gradient flow in overparameterized linear models. In International Conference Machine Learning, pp. 10153-10161 PMLR, 2021.
Zhichao Wang, Andrew Engel, Anand D Sarwate, Ioana Dumitriu, and Tony Chiang. Spectral evolu- tion and invariance in lincar-width neural networks. Advances in Neural Information Processing Systems, 36, 2024.
13

---

Published as a conference paper at ICLR 2025
Zhi-Qin Chi-Qin John Xu, Yaoyu Zhang, and Yanyang Xiao. Training behavior of deep neural network in frequency domain. In Neural Information Processing: 26th International Conference, ICONIP 2019, Sydney, NSW, Australia, December 12-15, 2019, Proceedings, Part I 26, pp. 264-274. Springer, 2019.
Greg Yang and Edward J Hu. Feature learning in infinite-width neural networks. arXiv preprint arXiv:2011.14522. 2020.
Greg Yang and Etai Littwin. Tensor programs ivb: Adaptive optimization the infinite-wice limit. arXiv preprint arXhv:2308.01814 2023.
Greg Yang. Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ry- der, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203.03466, 2022.
Yaoyu Zhang, Zhi-Qin John Xu, Tao Luo, and Zheng Ma. type of generalization error induced by initialization in deep neural networks. In Mathematical and Scientific Machine Learning, pp. 144-164. PMLR, 2020.
Liu Ziyin, Botao Li, and Xiangming Meng. Exact solutions of a deep linear network. In Advances in Neural Information Processing Systems, 2022.
Liu Ziyin, Hongchao Li, and Masahito Ueda. Law of balance and stationary distribution of stochastic gradient descent. arXiv preprint arXiv:2308.0661 2023.
Liu Ziyin, Mingze Wang, Hongchao Li, and Lei Wu. Parameter symmetry and noise equilibrium of stochastic gradient descent. The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Liu Ziyin, Isaac Chuang, Tomer Galanti, and Tomaso Poggio. Formation of representations in neural networks. International Conference on Learning Representations, 2025a.
Liu Ziyin, Yizhou Xu, Tomaso Poggio, and Isaac Chuang. Parameter symmetry breaking and restoration determines the hierarchical learning in ai systems. arXiv preprint arXiv22502.05300, 2025b.
14

---

Published as a conference paper at ICLR 2025
A ADDITIONAL RELATED WORK
There are also other recent studies focusing on feature learning beyond the kemel regime. Li et al. (2020) demonstrates that two-layer networks outperform kernel methods, while Damian et al. (2022); Abbe et al. (2022) highlight functions learnable by gradient descent in two-layer networks but by kernel methods. Moreover, several studies (Ba et al., 2022; Dandi et al., 2023; Cui al., 2024) show that a few gradient steps help neural networks adapt to dataset features, improving gen- eralization. Our work aims to provide new insights into feature learning through an analytically solvable model.
Another critical regime for feature learning involves large learning rates. Jastrzebski al. (2020); Long (2021); Lewkowycz et al. (2020); Wang et al. (2024) reveal that large learning rates offer bene- fits such as better conditioning of kernel and Hessian matrices and improved generalization. Notably, while the catapult (uv) model (Lewkowycz et al., 2020) shares some similarity with our model, their focus completely different. Lewkowycz et al. (2020) uses the NTK scaling, and shows that neural networks escape the kernel regime with sufficiently large learning rates. On the other hand, we show that how different scalings influence the phases of neural networks with sufficiently small learning rates. Therefore, the mechanisms of feature learning proposed in this paper do not appear Lewkowycz et al. (2020). Moreover, although Kalra et al. (2023); Kalra & Barkeshli (2024) try to extend the catapult (uv) model to finite-wide width settings with different parameterizations, their results remain qualitative, perturbative, or focus solely on fixed points without deriving a complete analytic solution. In contrast, our work provides a fully solvable model with exact dynamics, even for fi- nite widths, under more general initialization and parameterization settings. It is promising future work extend our gradient flow solutions to finite learning rates and adaptive optimizers, which can influence the regime (Yang & Littwin, 2023; Wang et al., 2024).
The concept of "alignment" explored many works. For instance, Lewkowyez et al. (2020) discusses alignment between feedforward activations and backpropagated gradients in the catapult mechanism, feature of the large learning rate regime, absent in our model. Seroussi al. (2023) investigates alignment between backpropagation and weight matrices in neural network Gaussian processes. Some studies (Baratin et al., 2021; Atanasov et al., 2021; Loo et al., 2022; Wang et al., 2024) examine alignment effects between models and data. Our definition of aligmment differs from them. However, we note that our definition of aligmment is inherently equivalent with the "alignment ratio" measured in recent work (Everett al., 2024), which focuses how alignment across layers influences model scaling and supports the importance of analyzing how alignment changes along training.
Finally, while many studies analyze the convergence (Nguegnang et al., 2021; Bietti et al., 2023) or dynamics (Jacot al., 2021; Arous et al., 2021; Paquette al., 2021) of gradient descent without exact solutions, they typically prove limited properties or asymptotic results. We believe an exactly solvable model offers deeper understanding of gradient descent dynamics.
B THEORETICAL CONCERNS
B. PROOF PROPOSITION
Proof. To make the analysis more concrete, we consider the standard loss function
de
i(u, w) = N R=
Σ
Σ 1=
- -2
(20)
j=1
where N is the size of the training set. Data points lie a Id-subspace, meaning that Zyk = agn; for a constant unit vector n. Due to the 1d nature of the data, the training dynamics this loss function is completely identical to training
on the following loss ((,ww) (7 Ei= " (Eja wyyr;) y), where æy. VEM1'
28 ax nj and
15

---

Published as a conference paper at ICLR 2025
Table 2: Varibles used for Theorem 1.
variables
explanation
12 W dedxdo
weight vector/matrix
x R0 .y R. p
In x?
effective data point and signal strength learning rates and normalization factor
do
Ma. ?
(i(t) (/ Ede Wiy (t) pul, (t))
generalized coordinates
9.(1) = 2, (v do way (1) pr (1)
p= Σj (0)² Q := Ejj1 91 (0)²
sufficient statistics of generalized coordinates
Le 1/ 1uol1e 72 p2 Ap (72))
characteristic learning time
Oga 2(x (v 1oy Ypy
characteristic learning scale
{(L) 1-a+ exp 41/16)
characteristic learning
EN of
y=
This is because
V VE of
N
N -2 - 4
da
de uu Σ j=1
N Σ yk
[u, w)
Σ 0
28 (7
Ui
-2
Un
(r)
Wiy
)
k=1
5=1
(21)
L(u, w)
Σ k=1
yk y
Therefore, without loss of generality, the training on the standard loss L is identical the training on L because the difference is only by constant that does not affect gradient descent training, This setting thus equivalent to the case when the dataset contains only a single data point (s, y). As is clearly shown from this example, using the notation in terms of and y is much simpler to understand than using žjk and yk- We believe that this notation is necessary and greatly facilitates the later discussions once the readers accept it.
1 Finally, all these notations can also be written in terms of Eg := NV Ek=1, which is the notation chose for introducing the lemma.
B.2 PROOF OF THEOREM 1
Proof. We summarize the definitions of all variables in Table 2.
To begin with, the gradient flow reads
du: 7u Ou:
ØL
da
22
w
7 Σ us Σ
miyz;) -v
dt
g-1
[-1
(22)
dwij
ØL
de
wirzz) -U) -))
Mw Dwij
2,3 u4
Σ Wiyx,) ) 20
Us Σ
-v)
dt
1=1
j=1
which implies the following two conservation laws
d (1u
Wig 3m 0,
(23)
dt
d
Wii
Wiy
1 dwis
dwiy' 0.
(24)
dt
Iz
dt
"Fssentially, this is because we only need two points to specify line. Also, is trivial to extend to the case when y vector that spans only one-dimensional subspace.
16

---

Published as a conference paper at ICLR 2025
dwis
From Eq. (22), can denote
A. which leads
dg
1
Wij
Σ
wajzzy
dt
Σ/=1 3 j-1
(25)
da
Ede Wijeej
do
A 2
Σ Wajttj 2
2 5-1
1/ =0.
Ege o3
According to the definitions 2p (T Egen wi)(t) (t) 1 (u En Wiy Wiy (t) + /Bmupu,(t) and q:(t) ==
we have
dt (pi(1)a(()))
d
1 d
T/u
N
Wijet,)
4 dt
Σn zjj j-1
(26)
1 d
Thu Σ wij
3m up
=0
4 dt
j-1
Further, substituting (22) into the definition of pt and dt we have
dp: = -YV Banhhapip ( dt
(p: q:)pp)
B-1
Pj Py- qi
(p, + 9:)0 8
(27)
Vuu
BBm
TTu
If we denote q (t), have
dp:
1
dps
(28)
p((p: + c.p.)) 8-1 dt
p,(ps + cg/p,)
dt
which gives
F((p((t)) F.(p:(0))= ())((())))(((()))---))))) -
(29)
where
F1(x)= == f ((x+c///)0TT
de
(30)
Therefore, (27) reduces differential equation of pf pr
2 B3yhhhupPs (Di(t) + (c/D))(t)))) 8-1
dt,(t) dl
VTu
(31)
d Σ
(p)(()++spppp (t( cl/ps(() (p:(t) + cgpps (1)) 1y
VBmp
Thu
Now we denote (t (((p((t). Then we have
dt,(t)
dt -2y Bng 2-
(3 1/2
(32)
where
-0(0)))))
(33)
(32) is an ODE with only one unknown function (t). However, it is in general not possible to solve (32), and this is why we only focus on B = 1.
For the special case 1, (31) reduces to dp((t)
22 (t)
Σ (p)(t)² - -t){)))
PP Vhall -y
(34)
dl
For 8 1, we also have F,(x) log with its inverse F71(x) = and thus (33) reduces to
(t (0) (t)
(35)
(0(
17

---

Published as a conference paper at ICLR 2025
for all i,j = 1,,,, d. Then according (26), also have (;(1) yy( P3 py(1))
Substituting them
into (34), and we obtain a differential equation with only one variable p;
dp: -2p: dt
(r P:)
Qp: p
P.
7PV 7u
(36)
where
d p 7 _ (t) Q-
a
(37)
1=1
This differential equation is analytically solvable by integration
p S
dc
t
(38)
4
(6 p.(0)² 5² -aayy (p2))(pp00))))))))00)))))
Because the denominator as quadratic polynomial has two different roots a the result of the integration is
te (c()))/p,(0))- -a- a t = log 4 (c()))/p;(0)2
const,
(39)
leading to
(1))) 1-a+ P: (ttt)/pp)00) - aa_ a-a
exp (-41/tc),
(40)
which gives (8).
Proposition 2. Under the condition in Theorem if P 0 and Q O, the result becomes
Pr(t)= =0
(41)
(EE(()
- ((0))
(42)
1-t(()
where
<'(1) 1 a' exp (-44thohu
(43)
and
a' YPu
(44)
(r²d)
Specially, if = = 0, we have Pi(() (()) 0, so the gradient flow will be stuck at the trivial saddle point.
Its proof is similar to the proof of Theorem I, because can similarly obtain
dq: -2g
(r²d)p²Q (0)2 q: 1ypyy//T.o MaM
(45)
dt
Its solution gives Proposition 2.
note that the behavior of the solution quite different from & O: when S O, we can obtain solution with zero loss in the end, but when y O, the gradient flow will converge the trivial saddle point p; = q: = 0.
B.3 PROOF OF THEOREM 2
Proof. By definition, {(t) monotonic function. As G+-Sa
1-5 1-g ga-a a. is monotonous to &.
evolves from 1 to a. monotonously. Then according to Equation (13), lim K(x, (t) lim K( (0) if and only if lim a, 1, which holds if and only if lim, P/Q 1 and (19) holds, when we have
PP
lim a. lim R-+00
a.s.
(46)
R-+00
2p²(72d)P
18

---

Published as a conference paper at ICLR 2025
From Equation (18), Equation (19) also implies
2c, + max(Cn. + Cyp + Cu. 0.
(47)
Therefore, we can see that the NTK remains (1)) because
y da, P 0 20, +og++ Cnw tou +ow)) +om))) =(().
(48)
The proof is complete.
n
B.4 SCALING OF WEIGHT EVOLUTION
This section aims to quantitatively analyze the power-law scaling of the right side of Figure 4. (7) and (8) indicate that
Wky (+00) - wi) p:((00) q:(+00) = Vo.pi(0) + (0)/ P(()) (0)
(49)
wiy
(0) + (0)
In the feature learning regime, |a, 11 0(1), and thus the weights evolve in an O(1) amount. In the kernel regime, by using (11), (18) and Theorem 2, we can find that
| a-l o n. (2 t 1/2
(50)
When + cc//2= = Cu 2, we have p=(0)-9(0)) = O(1) and thus P((0)
wio(+ wiy wir( (0)
xc |1 a+l xc x (20g +ro 1/2
(51)
Wiy(0) (0)
When Cw + cu/2 # Cu cw/ we have
oc n 20, +e C /at+ 1///41122/////////// -2 11 11- Wiy (+0)+ (0)
(52)
(i)(1)
[Woll oc 5 In the feature learning regime ò O. In the kernel regime 5 = -2cc + On + Cn.) cw + cu/2 = cu + cw/2 and -(2c, + Cn. + Cna) ) if cw cu/2 C, + c./ Using the values in Table 1, we can verify that 5 = 0.5, 1, for NTK, Xavier" and Kaiming parameterization, respectively.
In conclusion, we have WW-l
19

---

Published as a conference paper at ICLR 2025
10
0.9
layer 1&2
0.9
layer 2&3
layer 1&2
layer 3&4
layer 2
0.8
layer 4&5
0.7
layer 384
0.7
layer 5
layer 4&5
0.6
5&6
0.6
0
100
200
300
400
2
Iteratioo
(a) alignment a 6-layer FCN
(b) alignment vs. initialization scale
Figure 6: The alignment angle ( between different layers of six six-layer FCN trained on MNIST, with the same settings as Figure
ReLU
Sigmoid
Swish
Leaky ReLU
0.5
alignment
0.2
0.2
0.0
0.0
0.0
0.0
near orthogonal Init.
orthogonal init.
feature learning phase
kernel phase
1000
1000
0
2000
2000
iteration
itteraionn
Iteration
iteration
disalignment
1.0
1.00
1.0
1.00
0.98
0.9
0.9
0.98
2000
2000
0
2000
0
2000
iteration
Iteration
Iteration
iteration
Figure 7: The evolution of the alignment angle Ç between u and across two-layer ReL.U, sigmoid, swish, and leaky ReLU networks with d= 10000. The task is classify two Gaussian distributions.
C ADDITIONAL EXPERIMENTAL CONCERNS
C.I ADDITIONAL EXPERIMENTS IN SECTION 3.2
In Figure 1, we choose x = 1. 2, and for others we randomly sample 100 points from N(, 1) data points x, and set y 2x ~N0 1)/10 the target. The learning rates are chosen such that the model converges well within given iterations. For the orthogonal initialization, we initialize the model (O, 104,) and u (0,00.11a)
In Figure 2, to avoid the implicit bias of SGD to make layers aligned (Ziyin et al., 2024), we consider full-batch GD with batch size 2000 and constant learning rate. The learning rates are chosen sepa- rately for each model such that the model converges well in 1000 iterations, with training accuracy above 95%. All models use the standard Kaiming initialization, but we scale each layer by o. The results in Figure 2 also extend to deeper networks, although the training dynamics of deeper FCNs are less stable, as shown in Figure 6.
Moreover, we observe qualitatively the same phenomenon for all kinds of activation functions in the classification task Figure 7, where the task to classify training samples from (0, 1) and (44, 1). Initialization is the same as in Figure 1, but the binary cross-entropy loss is used. From Figure/ 7 we can also see that layers tend to align in the feature learning regime when they are initialized to be disaligned, and vice versa. Note that because of the binary cross-entropy loss, keeps decreasing even after the loss converges. Further, because of the binary cross-entropy loss, deviates from one for non-linear activation functions other than ReLU.
20

---

Published as a conference paper at ICLR 2025
C.2 ADDITIONAL EXPERIMENTS CONCERNING THE ALIGNMENT EFFECT
Figure verifies that the influence of dataset size N and input dimension do on the alignment effect not significant. This consistent with theoretical results, because Theorem characterizes the training dynamics without assumptions on the training set size, distribution, or input dimensions.
Figure includes more detailed ablation experiments, including the influence of large learning rate, data not lying in ID subspace and a three-layer linear network. Together with other figures (e.g. Figure 2), all ablation experiments indicate that our results are not significantly weakened by the following constraints: 1. 1-hidden layer multiple hidden layers, 2. linear vs nonlinear activations, single example (or examples in a I-dimensional subspace) vs examples distributed through space, 4. gradient flow training vs discrete-time gradient descent or stochastic gradient descent.
0.8
0.75
0.6
0.50
kernel phase
kernel phase
feature learning phase
0.4
feature learning phase
0.25
0.2
0.0
250
500
750
1000
250
750
1000
Iteration
iteration
Figure 8: The alignment effect of ReLU network for different dataset sizes and imput dimensions. Left: the dashed line represents datasct size 10 and the solid line represents dataset 100. Right: the dashed line represents input dimension du = 10 and the solid line represents input dimension do = 1. The target chosen be y a² x N(O, 1)/10. 10, where is Gaussian vector. Other settings are the same Figure 1.
large learning rate
not in subquace
three Inyer
Tamell phase
0.75
alighmmnt
0.50
0.25
1000
0
100
iteration
iteration
iteration
1.0
1.000
drsalighiment
TiM Aalooooe Sane Par Sodoouu "He
0.5
0.975
200
1000
iteration
iteration
iteration
Figure 9: More ablation experiments on linear networks. Results about non-linear activation are similar, as in Figures and 2. Left: Learning rate 15 times that of Figure 1. The network jumps out of the kernel regime predicted by Lewkowyez.ed al. (2020), but the alignment and disalignment effects still exist. Middle: is Gaussian vector with do 10, the data the first column of the same as u and the others are zero. Right: three-layer linear network with 7 1/d and d 300. The solid line refers to the alignment between the first two layers, and the dashed line refers to the alignment between the last two layer. For the orthogonal initialization, the first row of the second layer is the same the first layer and the others layer. Other settings are same as Figure This figure, together with Figures and shows that our results are robust to all factors.
in I-dimensional subspace. For the orthogonal initialization.
The last layer the same the first column of the second
C.3 EXPERIMENTS IN SECTION 3.3
In Figure 3, we train Resnet18 network on the CIFAR- 10 dataset with hyperparameters borrowed from https://github. http:://gghuu...oouuuuugg///u// www/kuungllu/pyoooccaoooooooooee wta.gchii/poooorcfar.r...... The only difference is that we scale each layer by and record the test accuracy together with the sum of the norm of all layers.
C.4 EXPERIMENTS IN SECTION 4.1
In Section 4.1, we utilize a two-layer FCN with the ReLU activation and hidden units. The input vectorized and normalized, so the input dimension is do 3072. The cross-entropy loss and the
21

---

Published as a conference paper at ICLR 2025
stochastic gradient descent without moment weight decay used during training. We use batch size 128 and report the best training and test accuracy among all epochs.
We choose 7 v and r7 0.05 for the standard NTK model, y zl and learning rate 0.05d/ 100 for the standard mean-field model, y and f 0.05d/ 100 for the Kaiming model, 100 0.05d/100 for the Kaiming" model, and rj 0.05 (000///)) for the Xavier model. The choice of hyperparameters guarantees that the standard NTK model and the standard mean-field model, the Kaiming" and Kaiming" model, and Xavier* and X model are same for d 100, respectively.
d and
and 1 = 0.05 for the Xavier" model, = 0.01d
C.5 EXPERIMENTS IN SECTION 4.2
The experiment in Section 4.2 is similar to that in 4.1. The only difference is that fix = 2000 and change the initialization scale. More specifically, set 10, of 16,, o max((e.)) and y K- min(0,-o/2) We also fix y 0.005.
D TRAINING AND GENERALIZATION DYNAMICS
This section aims to present the evolution of empirical and population loss.
From the definition pı and qi in Theorem 1, we have
/hoho da u (()))(()); = pi(1)² -a(
(53)
P
Consequently, we have
Σ Σ u. (;(()u;)(1)) 1= F1
70
p
a+ + {(1)a=
Q a+ E(1)a=
(54)
The 1 1-((1) E(t)
{(t)
As E(() is monotonous, o.+g(1)a-
1-(()) monotonously evolves from 1 a... Therefore, the model
7 EE ΣSe1 (t)) monotonously evolves from the initial value ypd
sPpddP- (P- Q) to the fi-
you Thu
nal value
ypd
ripd
(a, P-ap).))))
(a. a.)
(55)
Thelo
where we use the definition of a+ and a.a_ -Q/PP. In conclusion, the empirical loss (see Appendix A.I)
du [(u L7 Σ
Σ u, (ttt (() 11
Σ 1-1
yx y
(56)
evolves from its initial value to its minimal value ΣN v y² monotonously. Notably, the above conclusion does not rely on the choice of all hyperparameters and initialization.
In terms of the population loss, note that when the data subspace, the model output can be written as a linear function f(an) =: as(t), where s a scalar parameter. The population loss is thus a quardratic function of s: E[ "] which takes the minimal at
lies in one-dimensional
Lyyl
In reality, however, s(t) evolves monotonously from its initial value to
ΣN
ΣN a
According to the initial value of and the dataset ag Lk. there are three cases. The population loss might monotonously decrease, monotonously increase, first decrease and then increase. The time
22

---

Published as a conference paper at ICLR 2025
feature learning phase
feature learning phase
15
kernel
kernel phase
10
0.0
0.8
(a) training data point
(b) 10 training data points
Figure 10: The evolution of the training loss (solid lines) and the generalization (dashed lines) for linear network, We choose x 1, 2 (0,11..
scale of the population loss is also te. so there is no grokking for this simple model. Finally, as N increases, Σ 940R
Ejag] B[a
and the population loss converges to its minimum as expected.
2h-1 oi
See Figure 10 for numerical verification. The results include possibilities of the generalization error: monotonously decrease, or first decrease and then increase, as predicted theoretically. The time scales of the training loss and the generalization error are also the same. Moreover, we can see that for different initialization and parameterization, the networks converge to the same MSE. Therefore, our main focus is the training dynamics (e.g., whether NTK evolves), and our main con- tribution lies in analyzing how these dynamics unfold rather than in the final convergence point. The results also indicate that the number samples has no significant impact on the training dynamics but will influence the generalization dynamics.
23


Reviews:
[{"review_id": "7pbFvsyBPy", "reviewer": "Reviewer_Wsc1", "summary": "This paper studies the role of weight initialization on feature learning dynamics in linear networks. The authors provide exact solutions for  two layer linear networks where the input data is one dimensional. In this case, the relevant initial condition is the alignment $\\zeta$ which measures the cosine similarity between the readin and readout vectors. Depending on the scale of the initial weights, the scale of the output multiplier, and the initial alignment between these vectors, the model may or may not exit the kernel regime when fitting the data. The authors show that large weights can lead to worse alignment compared to starting with small weights. They provide a list of possible weight, learning rate and output multiplier scalings that can lead to kernel or feature learning.", "questions": "1. The authors claim larger initialization leads to worse performance. They then provide an experiment with residual skip connections (Figure 3) to try verifying this claim. I agree that with residual skip connections, networks can be initialized with zero weight variance with little consequence to training speed or performance. However, a non-residual feedforward network (such as the linear networks the authors actually analyze) the network is stuck at a saddle point as the weight variance shrinks to zero, leading to slower optimization. This is the reason for the long escape times in Saxe et al 2013. \n2. Are the results/scalings of Table 1 covered by the $a,b,c$ symmetries in the Yang et al [paper](https://arxiv.org/abs/2011.14522) or in this [note](https://mlschool.princeton.edu/sites/g/files/toruqf5946/files/documents/Princeton___Lecture_Notes_0.pdf)? If so, it would be good to comment on. \n3. The authors claim that they characterize the role of width. How does width enter their formulas, is it primarily through the initial alignment $\\zeta$? Or is width to be interpreted as a possible $\\kappa$ scaling quantity?\n4. Could Figure 4 be plotted on a log-log scale? I suspect that $|W-W_0| / |W_0 |$ should scale as a power law in $d$ with exponent that is different for different parameterizations. \n5. In Figure 5, it appears that the best performance is not at *very small* weights, but rather there is some optimal intermediate weight scale. How do the authors reconcile this \n6. Generally throughout the paper some measured metrics are implicitly related to performance/generalization, without this connection being made clear. For example \"In the kernel phase, the model norm diverges and the model alignment is always zero, which could be a hint of strong overfitting.\"", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "This paper studies an important problem of how model initialization impacts learning dynamics in neural networks. They provide exact solutions for the training in a specific setting (linear networks on one dimensional inputs). The solutions indicate a crucial dependence on output scale, learning rates, initialization variance, etc. The authors connect their findings to kernel limits, mean field limits, and $\\mu$P scalings.", "weaknesses": "The assumptions of the model make the scope of the theoretical results fairly limited since it relies on one-dimensional inputs and two layer linear networks (though the authors also have some results for other activations). That said, the authors do make an effort to test their ideas in more realistic settings, which is appreciated. There are also some questions which should be addressed (see below).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ROivsJbDiF", "reviewer": "Reviewer_jPRE", "summary": "The paper studies a linear network setup with single-index inputs and derives exact learning dynamics under gradient flow, and interprets the resulting solutions. In doing so, the paper finds that the time to learn the target function depends on a kernel and feature learning component and that the feature learning component in particular can come from several mechanisms, which are each interpreted and discussed. This results in insights about the role of unbalanced initializations, and output rescaling, the latter being a quantity central to understanding feature learning from past work.", "questions": "- Am I correct in understanding that $p_i, q_i$ are something like \"order parameters\" in the parlance of some statistical physics analyses of network dynamics (DMFT, etc)? ie. you are studying the dynamics of some \"sufficient statistics\" which are informally some rotated/transformed version of our true weights? \n- Why don't either $p_i, q_i$ or initialization variables show up in the feature learning component -- what is the intuition for this? Why is there dependence on the input norm $\\rho$ here?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Overall, I liked this paper. The final rating would be even higher if I was confident in my expertise in this topic (linear networks/optimization trajectories). In the end, I am not such an expert, so my rating reflects the fact that I cannot be sure other related works do not present very similar results. But in my opinion, the strong point of this paper was providing 1) an exact set of solutions, in contrast to the usual perturbative approach from infinite-width dynamics (which is usually extremely complicated and does not yield much insight) and 2) spending several pages interpreting the equations that result -- this is very important: the notion that the feature learning component can come from three distinct \"types of feature learning\" was new and interesting to me.", "weaknesses": "- You never write down the dynamics for the NTK in Equation (13), even though as I understand it you claim the theory implies NTK evolution. Maybe this is immediate from your solution on the weight dynamics, but some interpretation for what the evolving kernel means would be useful. \n- The word \"disalignment\" was confusing to me. As I understand it, it just refers to learning due to asymmetry between either the initialization or learning rate for the two layers, where feature learning doesn't take place under symmetry/orthogonal initialization. I'm not sure what \"disalignment\" means in this context, whereas for instance the meaning of \"alignment\" is intuitively clear. \n- The idea that at some intermediate initialization scale we are not quite in the kernel regime but also not quite aligning layers to the task in some intuitive sense was interesting, and I wish there was more discussion on this topic (ie. Section 3.4). \n- The point made in passing at the end about \"inadequacy of conservation laws\" and the impossibility of stronger analytical results is provocative and insufficiently explained. If you are making such a comment, you must explain it in more detail. \n- Minor spelling/grammar, eg. Theorem 1.1 statement \"Dynamics\" typo\n- I think a set of bullet points of contributions in simple (nontechnical) prose would be useful in the Intro or Section 3. There are many claims made throughout and results derived, and it's sometimes not clear what the overall aim of the paper is. Having one bullet point for each section to illustrate 1) what exactly is the new result and 2) how it all relates to each other would be very useful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wRk06FLdaE", "reviewer": "Reviewer_9ybf", "summary": "The paper studies three mechanisms of feature learning (alignment, disalignment, rescaling) in a one hidden layer (linear) neural network. Under a data assumption, the authors characterize the solution of the network fully in terms of width, initialization scale and learning rate across different regimes. The authors also confirm their findings with empirical evidence.", "questions": "If I’m not mistaken, at line 891, if we take $\\beta=1$, and plug it into equation (30), then $F_i(x) = ln(x) + c$, not $e^x$. I believe this should be the inverse of $F_i$? In which case I do not understand how (33) reduces to (35).\n\nShouldn’t there be $x_{ij}$ with a tilde in Proposition 1? And similarly the loss rewriting in eqn. 3 shouldn’t there be an expectation?\nI don’t understand what the authors mean in line 218 i.e. “the dynamics of GD training only has a rank-1 effect”. \n\nIf I am not misunderstanding, in Line 209, I believe now there is a vector $w$ instead of the previous matrix $W$ because $x$ is taken to be 1d, but this is only mentioned later on in line 217-218, which is confusing. \n\nCould the authors elaborate what they mean in line 239, more specifically why would a larger norm (assuming of the weights?) imply strong overfitting?\n\nWould it be possible as future work to analyze, in the same framework, the solutions for arbitrary nonlinearities? Maybe by using Hermite polynomials?\n\nIs there any intuition for why the alignment increases much faster for Sigmoid/Swish/LeakyReLU in Figure 1?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The main strength of the paper lies in Theorem 1, which completely characterizes the dynamics of the layer weights across time, under the assumption that the data lies on a 1D subspace. The authors obtain solutions for general power activations, but focus on giving closed form solutions for the case of a linear network $\\beta=1$ due to the problem of computing $F_i^{-1}$. By studying the alignment between the layer weights (in the case of $x$ being 1D), the authors characterize the 3 mechanisms of feature learning - which are governed by the angle $\\zeta$. Finally, the authors use the theoretical framework developed to study scaling limits in Section 4.\n\nIn general I think the paper is quite thorough and does a great job covering the mechanisms of feature learning in this toy setting. To my knowledge, the theoretical analysis of these dynamics in the finite width case is novel and interesting, bringing a great addition to the community.", "weaknesses": "In my opinion the main weakness of this paper is the exposition. There are quite a bit of variables involved in the main body which makes it hard to track which quantity is which. Given that the whole dynamics description in the paper revolves around Theorem 1, it might be worthwhile to restate it in simpler terms. \n\nAnother (minor in my opinion, given the analytical intractability of network that is more complicated) weakness is the very simplified setting studied, in particular the restriction over the data lying in a 1D subspace.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AZLPzroYpY", "reviewer": "Reviewer_T1qw", "summary": "This work proposes a solvable shallow (linear) network to gain insight into feature/kernel-regime classification and its parameterization. For a one-dimensional data space, the continuous-time learning dynamics can be reduced to certain variables $(p_i, i=1, ..., \\text{width})$, allowing for explicit solutions in a specific case $(\\beta=1)$. From this solution, particularly from the learning speed $(t_c^{-2})$, the authors identify the parameterization in as general a form as possible, in which the learning dynamics achieve either the feature-learning or kernel regime.", "questions": "### **Difference from uv model**\nHow does the proposed model differ from the well-known catapult (uv) model? Suppose that we rotate the 1-dim data space $n$ to the direction of a unit vector $(1,0,\\dots,0)$ with a rotation matrix $V$. Then, defining the new variable $W(t)= W^*(t) V$, we have:\n$\\frac{dW}{dt} = -\\frac{\\partial L}{\\partial W}$\nequivalent to\n$\\frac{dW^*}{dt} = -\\frac{\\partial L}{\\partial W^*}$\nsince the rotational transformation is orthogonal (that is, the Jacobian $J=\\partial W^*/\\partial W$ satisfies $J^T J=I$). In other words, the proposed model in this work reduces to\n\n$L(u, W^*) = \\left[\\gamma \\left(\\sum_{i} u_i w^{*\\beta}_i  x  - y \\right)\\right]^2$\n\nwith $x:=\\sqrt{E[\\tilde{x}^2]}^\\beta$.\n\nThis suggests that the proposed model is merely a nonlinear parameterization of the uv model, with $v_i = w_i^\\beta$. For $\\beta=1$, which is the focus of the current work, the model is identical to the uv model. Since the authors references Lewkowycz et al. (2020), they  likely know about the solvability of the uv model, yet there is no explanation of this relationship or reference to further developments such as \n\n*Kalra et al., \"Universal Sharpness Dynamics in Neural Network Training: Fixed Point Analysis, Edge of Stability, and Route to Chaos\", arXiv:2311.02076.\n\n*Kalra & Barkeshli, \"Phase diagram of early training dynamics in deep networks: effect of the learning rate, depth, and width\", NeurIPS 2023.\n\nSince  the model is indeed equivalent to the uv model for $\\beta=1$, it is crucial to mention such prior studies and to carefully explain in what aspects the problem setting differs.\n\n### **Difference from muP**\n- It will be better to provide a clearer explanation in Corollaries 1-3 about how these findings differ from the muP parameterization [Yang & Hu (2020)]. The authors state\n>one can choose an optimal pair of learning rate and output scale $\\gamma$ such that the model is in the feature learning phase. Point (1) agrees with the analysis in Yang & Hu (2020), whereas point (2) is a new insight we offer\n\n in lines 471-473. My understanding is that point (2) would be a main difference (in the inifinite width). However, $\\gamma$ also appears in muP as $1/(width)^a$ in the final layer. I recommend that the authors more explicitly state how their \"optimal pair of learning rate and output scale γ\" differs from or extends the muP parameterization.\n\nIf the claim of this work is that the authors have identified a previously unknown feature-learning (or kernel) regime in the infinite-width case, this would imply that certain conditions in prior studies were restrictive, leading to oversight. It seems better to clearly explain which specific conditions from prior work the current work has relaxed or generalized.\n\n- The authors also state\n>A key difference mentioned is that the results apply to finite-width networks with arbitrary initialization. \n\nThis sentence mentioning the finite width appears in Section 4.1, which consider the case of $c_d=1$, that is, the *infinite* width limit. This is confusing. If there's a distinction between their finite-width and infinite-width results, clearly delineate which parts of the paper apply to each case.\n\n### **Typo**\n- Equation 1: $x \\to \\tilde{x}$\n\n###After the rebuttal########\n\nI raised the score because the author(s) revised unclear points and also promised to explain related research in essentially the same model (but different focus) in the final version.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The analysis is technically solid and concrete.\n- This work addresses the challenging yet crucial problem of identifying the feature-learning and kernel regimes and characterizing the dynamics in feature learning. It provides concrete insights from a solvable model perspective.\n- The obtained parameterizations are quite general which cover various width, learning rate, and initialization scales.", "weaknesses": "- The analyzed model seems essentially equivalent to the catapult model (Lewkowycz et al. 2020) or uv model (Kalra et al. 2023). In this sense, the solvability of this model itself seems not so surprising. In addition, It is unclear how the continuous-time analysis of the current work contributes to the original literature of the catapult (uv) model, which claims that a non-trivial dyanamics toword feature learning appears for *discrete* update with a large learning rate.\n\n- Certainly, discrete updates have not been explicitly solved and achieving analytical solvability even in the continuous limit would be novel. However, for alignment, the dynamics is essentially equal to the uv product dynamics, which can already be tracked with any width (Lewkowycz et al. 2020). Therefore, the analytical solvability for the finite width might be a bit straightforward.\n\n- Although this would be a presentation matter, it is not easy to follow how the parameterization derived from the current analysis in Section 4.1 compares to the muP parameterization by Yang & Hu (2020). \n\nPlease see the **Questions** section for more detailed points.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)