Mixed Dynamics In Linear Networks: Unifying the Lazy and Active Regimes
Authors: Zhenfeng Tu
Santiago Aranguri
Arthur Jacot
Venue: NeurIPS 2024 poster (2024)
Keywords: Linear Networks
Lazy Regime
Active Regime
Training Dynamics
Phase Diagram
Primary area: learning_theory
TL;DR: We derive a simple formula for training dynamics of linear networks that not only unifies the lazy and balanced dynamics, but reveals the existence of mixed dynamics.

Abstract: The training dynamics of linear networks are well studied in two distinct
setups: the lazy regime and balanced/active regime, depending on the
initialization and width of the network. We provide a surprisingly
simple unifying formula for the evolution of the learned matrix that
contains as special cases both lazy and balanced regimes but also
a mixed regime in between the two. In the mixed regime, a part of
the network is lazy while the other is balanced. More precisely the
network is lazy along singular values that are below a certain threshold
and balanced along those that are above the same threshold. At initialization,
all singular values are lazy, allowing for the network to align itself
with the task, so that later in time, when some of the singular value
cross the threshold and become active they will converge rapidly (convergence
in the balanced regime is notoriously difficult in the absence of
alignment). The mixed regime is the `best of both worlds': it converges
from any random initialization (in contrast to balanced dynamics which
require special initialization), and has a low rank bias (absent in
the lazy dynamics). This allows us to prove an almost complete phase
diagram of training behavior as a function of the variance at initialization
and the width, for a MSE training task.

Mixed Dynamics In Linear Networks: Unifying the Lazy and Active Regimes
Zhenfeng Tu Courant Institute New York University New York, NY 10012 zt2255@nyu. edu
Santiago Aranguri Courant Institute New York University New York, NY 10012 aranguri edu
Arthur Jacot Courant Institute
New York University New York, NY 10012 arthur, jacot@nyu........ edu
Abstract
The training dynamics of linear networks are well studied in two distinct setups: the lazy regime and balanced/active regime, depending on the initialization and width the network. We provide a surprisingly simple unifying formula for the evolution of the learned matrix that contains as special cases both lazy and balanced regimes but also a mixed regime in between the two. In the mixed regime, a part of the network is lazy while the other is balanced. More precisely the network is lazy along singular values that are below a certain threshold and balanced along those that are above the same threshold. At initialization, all singular values are lazy, allowing for the network to align itself with the task, so that later in time, when some of the singular value cross the threshold and become active they will converge rapidly (convergence in the balanced regime is notoriously difficult in the absence of alignment). The mixed regime is the 'best of both worlds': converges from any random initialization (in contrast to balanced dynamics which require special initialization). and has a low rank bias (absent in the lazy dynamics). This allows us to prove an almost complete phase diagram of training behavior as a function of the variance initialization and the width, for a MSE training task.
Introduction
Whether in linear networks or nonlinear ones, there has been lot of interest in the distinction between the lazy regime (27] and the active regime [6 43 15 52 13 as the number of neurons grows towards infinity. the lazy regime the training dynamics become linear, so that they can be easily described in terms of the Neural Tangent Kernel (NTK) 27 g 51 36), while the active regime exhibits complex nonlinear dynamics. While our understanding much more limited, it appears to be characterized by the emergence of feature learning 22 52). and form of sparsity [3] [1 2 (the type of sparsity observed depends on the network type [1 19 25 24], but will focus fully-connected linear networks which exhibit rank sparsity in the learned linear map 7 35 28 47) which are both absent the lazy regime.
the active regime remains
Note that even though it is common to talk of the 'the' active regime, we do not know yet whether there only one multiple active regimes. Indeed the term active regime is usually used to describe any regime that differs from the lazy regime and exhibit some form of feature learning. Though we do not have an complete understanding where the lazy regimes ends and the active regime(s) start, we know that the lazy regime requires extreme overparametrization (a large number of neurons in comparison to the number datapoints) 5 20 a 'large' initialization of the weights [15), a small learning rate, and carly stopping when using cross-entropy loss weight decay. Indeed, active regimes have been observed by breaking either of these requirements: taking limits with mild or no overparametrization [10), taking smaller even vanishingly small initializations 35 28), using large
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

learning rates 33 or SGD
42
47
or studying the late training dynamics with the cross cross- entropp
29 26. Though each of these can lead active regimes with significantly different dynamics, they often lead to similar types of feature learning and sparsity.
loss [30 I7 or weight decay 34
39
In this paper, study this transition in the context of linear networks and focus mainly on the effects of the width w and the variance of the weights at initialization o², and give a precise and almost complete phase diagram, showing the transitions between lazy and active regimes. In this setting, we will show that there typically is only "one" active regime, which is the same (up to approximation) as the already well-studied balanced regime
But our result also paint a more subtle picture than the lazy/active dichotomy. We propose a more granular approach, where at certain time some part of the network can be in the lazy regime, while others are the active or balanced regime. More precisely the network is lazy along the singular values of the matrix represented by the network that smaller than o o²w, and in the active regime along the singular values larger than d²w.
1.1 Contributions
We consider the training dynamics of shallow linear networks As = W2W1 and show that for large enough width w (the inner dimension), and a iid N (0, initialization of all weights, the dynamics of Ao(a) as a result of training the parameters for general matrix cost C' with learning rate is approximately given by the self-consistent dynamics
(W,, W2) with GD/GF on the loss [(0) C(Ag)
oeAm(s) = A + VVC(A) Ag Ao + otwPI.
A
(1)
These dynamics contain as special cases both the lazy dynamics
OeAaçe) ~ -2n0 W
when oZw >> Amar (Ag) and the balanced dynamics
O,Ao(t) = - As Aa VC()) - - v/CC)A) ATAa As
when o²w w < Amin (Ag). But it also reveals the whole spectrum of mixed dynamics in between, where some singular values of Ag are below the o²w threshold and some are above it.
This suggests that the lazy/active transition is best understood at a more granular level, where at each time every singular value of Ag either be lazy or active/balanced. The mixed regime is the best
of both worlds: on one hand, since / Ao + o is always positive definite, the network can never get stuck at a saddle/local minimum as can happen in the balanced regime, on the other hand there momentum effect where the dynamics along large singular values is much faster than along the small ones, leading to an incremental learning behavior and a low-rank bias, which is absent in lazy learning. By choosing the threshold d2r adequately, one best take advantage of these two phenomenon.
Finally, we focus on the task of recovering a low-rank d x d matrix A from noisy observations / in the limit the dimension d, width A + E, training the MSE error zz || (A* + E) and variance 2² scale together with scaling laws d' and 0² d= We describe the training dynamics for almost all reasonable scalings Tw, Yoz leading to a phase diagram with two main regimes:
Lazy (1 < Yoa + Tw) where all singular values remain below the threshold 2u throughout training, and where the network fails recover A" due to the absence of low-rank bias.
Active (1 762 + Th) where K Rank singular values pass the threshold and fit A* before the other singular values have time to fit the noise E, leading to the recovery of A*.
There are two other degenerate regimes that we avoid: the underparametrrieed regime when w < d (or The < 1) where the rank is constrained by the network architecture rather than the training dynamics, and the noisy regime 27,² + Tw + > 0 where the variance of the entries of Agco) initialization infinite.

---

MSE arror for diffarent dynamics with Yor - 1.83. Vv. 2.25
Singulaa values for Aso with Vir -1.85, yu 2:25
2:25
2.0
Tra arror, gradien: descent
Tra error, consisenttt dynamics
Test error, gradient descen:
200
1.5
arror, condisen dynamics
Sing. la value
130
Eroor
1.0
o'w threshold
0.5
0.0
50
70
10
Iteraton number
Iteratio number
Figure 1: For both plots, we train either using gradient descent or the coff-consistent dynamics from equation ( with the scaling 702 = - 1.85, Tie 2.25 which lies in the active regime. (Left panel): We plot train and test error for both dynamics. We observe that the train/test error for gradient descent is very close to the train/test error for the self-consistent dynamics. (Right panel): We plot with solid line the singular values of Ag(t) when running the self-consistent dynamics, and use a dashed line for the singular values from running gradient descent. In this experiment, Rank 5. We use different colors for the 5 largest singular values and the same color for the remaining singular values. We can see how the 5 largest singular values 'speed up" they cross the o² w threshold, allowing them to converge earlier than the rest. The minimal test error is achieved in the short period where the large singular values have converged but not the rest.
1.2 Previous Works
Linear networks have been used as a testing ground, stepping stone on the way understand nonlinear networks. Linear networks and their training dynamics are in many ways much simpler than nonlinear ones, but in spite of long research history, our understanding remains limited.
The setting that is best understood is that of diagonal linear networks where the dynamics decouple along the diagonal entries leading to incremental learning behavior and sparsity bias 44
B 45
23 some of this analysis has been extended to include effects of initialization scale While the same decoupling happens in general linear with diagonal initializations and diagonal task, it remains an extremely strong assumption.
[48 and
SGD 42
Some work has been done to prove similar incremental learning dynamics outside the diagonal case 35 28 31 where the incremental aspect can be understood as the parameters going from saddle to saddle. For shallow linear networks, the training dynamics with MSE can be explicitly solved 21 but remain very complex so that one needs to assume some form of alignment to guarantee convergence [4]. For deeper networks there exists explicit formulas in the meaa--field limit where the number of neurons grows to infinity [18), these results can of course be applied to the special case of shallow nets, our paper goes further by giving self-consistent dynamics for the full matrix, revealing the lazy/active transition, and also extends the analysis to finite widths.
A very powerful tool in the analysis of a linear network is its training invariants, and the balancedness condition which greatly simplifies the dynamics 6 7]. Balanced networks exhibit a momentum effect, where the training dynamics along singular value have 'speed' proportional to itself (or 8; to some power). while this momentum effect seems to be key to understand the low-rank bias of linear networks [8). it also means that needs to guarantee that the dynamics never approach zero, which is one the main hurdle towards proving convergence in balanced networks. To solve this issue,
recent work has focused on initialization that slightly imbalanced (49 37 46
38 50). This suggests
that it is key to find the right balance between balancedness and imbalancedness to obtain both fast convergence and low-rank bias.
In a concurrent work 32 a similar transition between lazy and active regimes is observed, and the same mixed dynamics are derived for specific initialization. In contrast, we prove that these dynamics are approximately true with high probability for random Gaussian initializations, which the standard initialization scheme for neural networks.
3

---

1.3 Setup
We will study shallow linear networks (or matrix factorization) where a dout dn matrix Ag is represented as the product of two matrices Ag W, where the weight matrices W and W2 are respectively w x din and dout x w dimensional, for some width w. The parameters the network are the concatenation of the entries of both submatrices 0 = (2).
The parameters are learned in the following manner: they are initialized as i.i.d Gaussian ((, o²). and then optimized with gradient descent to minimize a loss [()) = C(As).. Though most of our analysis works for general convex costs C' Rdoux Xdia - R on matrices, we will in the second part focus the task of recovering low-rank matrix A from noisy observations A + E, by training linear network Ag on the MSE loss
[( d² 11 (A" F
The width allows us to control the over parametrzation, indeed the set of matrices that can be represented by network of width w is the set MM of matrices of rank w or less. The overparametrized regime is when 2 min dows) because all matrices can be represented in this case.
1.4 Lazy Dynamics
The evolution of the weight matrices during gradient descent with learning rate rj is given by Wi 1) (t) - (t)CC(AA)))) W2(t + = W2(t) - (
where we view the gradient VC((A(o) of the cost C as dout din matrix, which for the MSE cost equals VC((e())) = 2d 2d (A* + E)).
But we care more about the evolution of the complete matrix Agce) = W2 induced by the evolution of (t),,, (t),), which can approximated by
Ae(++1) = Ao(e) 77 W2)WW2 (t)VC((A(a) - -VVC W (t) + 0(7²).
(2)
Thus we see that if we can describe the matrices C₁ = WT W, and C2 = WoW? throughout training, then can describe the evolution of Ao(t)-
When very large, we end up in the lazy regime where the parameters move enough up time & to change Ao(a). but not enough to change C1, C. allowing us to make the approximation Ci(() & C.()). Furthermore at initialization these matrices concentrate as to - оo around their expectations E [G]] = o² Id... E [C2] adwlam..- The GD dynamies can then be approximated by the much simpler dynamics:
Ap(c+1) = Ag(z) 2njo (C((g)))
which equivalent to doing GD the cost C directly with learning rate of 2no2w.
One can then easily prove exponential convergence for any convex cost C' following the convergence analysis of traditional linear models. But we can see the absence of feature learning from the fact that the covariance C₁ of the 'feature map' W, is (approximately) constant. More problematic the context of low-rank matrix recovery is the absence of low-rank bias, indeed one can easily solve the dynamics to obtain
Ao(e) = (A" + E) + (1 4d 2d= (Amoo) (A* + E)),
and since (EAo(()) we obtain
E [Aero]
4d 210 (A* F))
'To be more precise the direction in parameter space that change C₁ C2 are approximately orthogonal to those that change As. and GD/GF only along the later direction.

---

The advantage is that there is momentum effect in the form of the prefactor (s(()),, which implies that the dynamics along large singular values are faster than along small ones. As a result, if all singular values are initialized with the same small value, then they will at first grow very slowly until they reach a critical size where the momentum effect will make them converge very fast. The singular values aligned with the top singular values of A' + E will reach this threshold much faster, and they will therefore converge to approximately their final value 8; 8f at a time when the other singular values are still basically zero. If stop training at this time then the linear network will The expected test error E |oAo))) A-1 therefore lower bounded by
[Boc)) - [|(1 4d n w A - + (1 (1 - 4d ))
which never approaches zero.
In linear networks, there is no advantage to being in the lazy regime, as we simply recover simple linear model an additional cost of more parameters and thus more compute. But we will see that a short period of lazy regime at the beginning of training plays crucial role in making sure that the subsequent active regime starts from an 'aligned" state.
1.5 Balanced Dynamics
There has been much more focus on so-called balanced linear networks, which are networks that satisfy the balanced condition Wi WT" W W2- If the network is balanced at initialization, remains so throughout training, because, the difference W,W? W W2 is an invariant of GF (and an approximate invariant of GD with small enough learning rate).
First observe that the balanced condition implies the following shared eigendecomposition W WT = w W2 = USU This implies the following shared SVD decompositions W1 = U UVS in W2 Uout V and Ag = SUT. Furthermore, we have C1 = Un. = Ag. Ag and
C2 SU out / A7 which leads to self-consistent dynamics for Ag(e):
Aoc+11= Ao(e) my Ante) VC(A(()) - v)VC(A()) AT Aoco) + O(7)).
Now these dynamics are quite complex in general, and remains difficult to prove convergence. Indeed one can easily find initializations Ag(o) that will not converge, for example if Aoto) then GD will remain stuck there. A lot of work has been dedicated to finding conditions that guarantee the convergence of the above dynamics 66 14, but these assumptions are often quite strong.
The simplest initialization that guarantees convergence (and the one that will be most relevant to our analysis) is the positively aligned initialization. If at initialization Aι(о) and A + E are 'aligned', i.c shares the same singular vectors Ao(a) SUT and A + E' = Sur S then they will remain aligned throughout training Aa(1) = Unn. S()UUT and the dynamics decouple along each singular value
() x (1)
+ 21y | (1) (5) s,)))) + 0(7))
Since we always have of 2 then for small enough learning rates 1/. see that if (0) E (0, it will grow monotonically and converge to s if (0) > of it will decrease monotonically to sf and if s (0) O it will increase and converge to O. Thus can guarantee convergence if we further assume positive aligmment s; (0) > 0.
have essentially learned only the top K singular values of A + F,, which is good approximation for A leading to a small test error (see 23 for details).
But this analysis relies on the very strong assumption of positive alignment at initialization. If we do not assume positive alignment and assume that the s; random (i.i.d. w.r.t. symmetric distribution), then each s; has probability 1/2 of starting with a negative alignment and getting stuck zero, which means that with high probability training will fail to recover A and will recover only a random subset of the singular values of A". The presence of these attractive saddles shows the complexity of the balanced dynamics.
limitation of this approach is that it requires a quadratic cost and a very specific initialization, and in the case of positive alignment, an initialization that requires knowledge of the (SVD of the)
5

---

true function A", Nevertheless, the positively aligned and balanced dynamics seem to capture some qualitative phenomenon that has been observed empirically outside of this restricted setting. This is the phenomenon of incremental learning, where if the singular values are initialized as very small, they first grow very slowly, but the multiplicative momentum will lead come up one by one in very abrupt manner, and this leads to a low rank bias where the network first only fits the largest singular value, then two largest, and so on. More generally, this can be interpreted as the network performing greedy low-rank algorithm [35].
Our analysis will confirm the fact that positive alignment happens naturally as a result of short period of lazy training, allowing us to prove similar decoupling and incremental learning for general random initialization.
Remark. We can define the time dependent map ((G;t) C2(t)G + GC1 so that the GD dynamics can be rewritten as A0((+1) = Ao( r)((C(AAr))) O(r)). The map θ is none other than the NTK for shallow linear networks, but has also been called the preconditioning matrix in previous work 77. The lazy regime is then characterized by the NTK θ being approximately equal to the time- indeinndeenndent NTK e (G) = 20²wGG whereas the balanced regime characterized by the time- demendennd hhl (G: t) / Amce) AT G /V Ab(e) Au(a), with the distinction that the time
0(1)
dependence is only through Agtt)-
2 Mixed Lazy/Balanced Dynamics
Both lazy and balanced dynamics have the surprising but very useful property that the evolution of the network matrix Ag approximately self-consistent: the evolution of Ag can be expressed in terms of itself. The lazy approximation becomes correct for a sufficiently large initialization, while the balanced one is correct for balanced initialization. However, for most initializations, neither of these approximations are correct.
eooluton of Ag that applies for any initialization scale: We fill this gap by providing self-consistent evolution
O, Ae(rc1))) -"y Aeco) (0)) + totw2 - (VC(A)) A Aocc) + T
This approximation is formalized in the following theorem, denoting Ci(() v Abbie Aoci) ow
and C2(()= Aocco A
Theorem For a linear net Ag = W2W, with width w. initialized with i.i.d. (0, o²) weights and trained with Gradient Flow, we have with high probability that for all time t,
(0:00-(000) (0000CCh0l/o||(/0000 -1200))cc00)1...00) < min (olo2)) (v 1C
Proof. (sketch) The quantity WiW: W W2 is invariant under GF (and approximately so under GD) and approximately equal w P2) for two orthogonal projections P1, P2 (at initialization and for all subsequent times because of the invariance). We therefore have
WF (W, W WW R ~ W (Pr + P2) zu o AwwPCr.
Thus the pairs C1 C2 approximately satisfy the following equations:
0 C AT, A,C C.A, As o w + AT C2A0
0 C2 ArA C2 C2A2A w + 16
The pair Č₁, Ö2 is a solution of the above, and one can show that C1, C2 must approach them and not any of the other solutions.
The takeaway from theorem iis the following.
1. In the lazy regime where ||C; (t)llop + ((C2()))lop S (((Ow)), then [[Ci()) - Ci(() lop < ( << (t)ll)))

---

2. In the active regime where |[Ci(t)lop//oQ f > de >> 1, then [Ci(t) - (t))llpp 1
(()Pw) << (1)
is true that the error does not vanish. However, for our purpose suffices to show that ||ć; C1 C2llop is infinitely smaller than C1 for all times, regardless of the magnitude of [t)(())lop.
We see how both the lazy and balanced dynamies appear as special cases depending on how large the variance at initialization o² is in comparison to the singular values of the matrix Agce):
Lazy: When o²w > Sman (Aoce)): then C₁ as ozwld and C2 Rs obuTa... recovering the lazy dynamics.
Balanced: When pZw & smin (A then Ei 2i V and O2 = / AAcco Ab
recovering the balanced dynamics.
But clearly there can be times when neither conditions are satisfied, when some singular values of Ao(a) are larger than the threshold o²w while others are smaller, in such cases we are in a mixed regime, where the network is lazy along the small singular values of Ao(e) (s; < odw) and active/balanced along the large ones (s; > odw).
At initialization, the singular values are of size Vud This implies that with overparametrization (w >> d), all singular values start in the lazy regime and follow the simple lazy dynamics, which may (or may not) lead to some singular growing and crossing the o²w threshold, at which point they will switch to balanced dynamics (after a short transition period when the singular value around the threshold s; & o²w). Once singular value is far past the threshold s; > o²w, training along this singular value will be much faster than along the lazy singular values (this speed up can be seen in Figure D This allows the newly active singular values to converge while the lazy singular values remain almost constant. Once the active singular values have converged, the slow training of the remaining lazy singular values continues until some of these singular values reaches the threshold, or until GD converges.
This type of behavior is illustrated by the following formula, which describes the derivative in time of the i-th singular value st, of Ar, with singular vectors Ug,e, U4,2
81.-1 S1,1 22 m n.u Ma R -2 / ol uf VC((A())
where the prefactor 21) V s², + ow-22 describes the effective learning rate along the i-th singular value, which depends on the i-th singular value 8;,, itself.
This suggests that it is more natural to distinguish between the lazy and active regime at much more granular level: at every time singular value can either active or lazy (or very close to the transition but this typically only happens for very short time). In contrast, the traditional definition of the lazy regime was defined for a whole network and over the whole training time. avoid confusion, we call this the pure lazy regime, where all singular values remain lazy throughout training. This begs the question of whether a pure balanced regime also exists, but all singular values will always be lazy for at least short time period (assuming w > df), and period plays a crucial role in aligning the network so that the subsequent balanced regime can learn successfully. A pure balanced regime can only be obtained in the underparametrized regime, or by taking a balanced initialization instead of the traditional i.i.d. random initialization.
will see this short lazy
While this challenges the traditional lazy/active dichotomy, it also reinforces it, as it shows that there no fundamentally different third regime, only lazy, active, and some mix of the two. Theorem I thus allows us to revisit previous descriptions of lazy and balanced dynamics and 'glue them together" to extend them to the general case. This simple strategy will allow to almost fully 'fill in the phase diagram', i.c. describe the dynamics, convergence and generalization properties of DLNs for almost all reasonable initialization scales o² and widths w.
Remark. The transition of a singular value s; from lazy to active can be understood as form of alignment happening in the hidden layer: the two vectors Win and WT" un for u, o. the left and right singular vectors of s; are orthogonal in the lazy regime and become perpendicular in the balanced
7

---

Tast error under gradiert descent
Rank of Asld) under gradient descent
2.5
0.25
0.84
64
2.0
0.72
2.0
3.6
0.60
1.5
0.49
1.0
Active threshold
0.36
10-
32
Underparam-eterr and
0.24
24
0:25
0.0
0:35
0.00
0.0
0.8
Log of maximum distaance between gradient descent and f-consistent dynamics
Number iterations until minimal test
32
720
2.5
1.6
2.0
-1
480
disanca
-32
2
1.5
400
-4.8
320
100
10-
-0.4
240
-9.6
0.0
0.0
-3.0
-25
-2.0
-1.5
-10
-0.5
0.0
-300
-2.5
-2.0
1.5
0
-0.5
Ye
Figure 2: As function of Te2, Tw: we run GD and plot different quantities. Our theoretical results only apply the top left region for The 1 and below the red line, although these plots suggest that some results may extend to smaller Sw (Top left panel): We plot the smallest test error + ||Ae(t)) l in the whole run. The active region (below the black line) has a small error while the lazy region does not. (Top right panel): We plot the stable rank of Asçe) (defined as |l / Agçe) llep) at the time of minimal test error. In this experiment, we took Rank = 5. We see that the active region has approximately the correct rank while the lazy region overestimates (Bottom left panel): We plot the number of iterations until minimal test error, illustrating the trade-off
between test error and training time. (Bottom right panel): We compute In
(z | - Agce) |1
where Ap(e) comes from GD and Aς(s) from the self-consistent dynamics. We observe that this distance not only small for the region where our theoretical results apply but also almost everywhere outside this region.
regime. Indeed the normalized scalar product of these two vectors satisfies
2 W2WW
Si
W u, W|a|l
v uf Czusy ut
Ciu;
+o
which is close to zero for lazy singular values 8; < d and close to one for active ones 8; >> 2
2.1 Phase Diagram MSE
To illustrate the power of Theorem we provide phase diagram of the behavior of large shallow networks MSE task, for almost all (reasonable) choices of width and variance r² scalings.
We want recover rank K and d d-dimensional matrix A with s=(A") da, for some a1 > 2 ... 2 independent of the dimension d. We however only observe a noisy version A + for some E such that | < cad One could imagine E' to have iid random Gaussian entries (0, 1) in which case El I v with high probability.

---

As the dimension d grows, the size of the network needs to scale too, as well as the initialization variance, but it is unclear what the optimal way choose w and o². We will therefore consider general scalings w = dh and x² = d'oo2 We will now describe the (Th; To2 J-hase diagram which features regimes: underparametrrzed, infinite-noise, lazy and mixed/active.
We can identify a region of 'reasonable' pairs (o22,7wl by ruling out degenerate behavior. First, the width needs to be larger than the dimension d, since a network of width w can only represent matrices of rank w or less, this means that we need 7 2 1. Another constraint comes from the variance of Ag at initialization: the entries Ag(o),; at initialization have variance o'w. We want this variance to go to zero as d grows which implies that we need 27 The < O.
Now within this reasonable region observe two regimes, the pure lazy regime for 1 < </2 + +Y where the network simply fits A + E thus failing to learn A and the mixed regime for > 762 7 where the dynamics are lazy for short amount of time until K singular values grow large enough to switch to the balanced dynamics and fit the true matrix A".
Theorem 2. For pairs Ywu Yo2 such that "wu > 1 and 20 + +ww < 0, we have two regimes:
L we have that for all time t,
Lazy (1 < Yo2 + Tw): with a learning rate n < 1 T ||Aa(e) - 1- 2a
Active (1 702 + +w): with learning rate 1) & s1(A-) d, and at time
d²
A
2max(1,25)
max(1,22)
d dlogg d + -7 Odd
1 aKK
c(a),..., ak)
28
for A 7/02- -oo 0, we have that
Aw log=
o²w
log
d² A - 1- s 0
+ d 2 + d
+
2
d² )
d2
miine. oo a2-0,lax maxe. pay 10 -2]
for cla)...., ax)
Note that all the terms inside the final O(...) term vanish: Aw + 0 because Yoa + Yau < O. o" o'w² log" d + o2w -> 0 since 1 > 7o2 + Ywo and n7? log" d 0 since we assumed 17 < d.
d2
This shows that the lazy regime only appears for very large widths 7o. > (or at least the lazy regime with finite variance at initialization). Indeed the choice Yor = 2, 7o2 = -1 is at the boundary of the lazy regime with the smallest "w.- This could explain why it is rare to observe the lazy regime in practice.
Our theoretical results applies to the overparametrized regime w > d, but actually we only want to fit A which has a much smaller rank r. and so we might only need w >> r. Figure top left panel, confirms this, since we see a good generalization even for small widths w < d, and in particular when w R Aa But to leverage this underparametrized regime, one would need to know the rank of the true matrix A" in advance, which typically not the case in practice. Nevertheless, the interesting behavior we observe in the (mildly) underparametrized regime warrants further analysis, and the fact that our self-consistent dynamics remain good approximation in this regime (Figure 2 bottom right panel), suggests that the analysis we present here could be extended to this regime too.
Finally, we observe trade-off between generalization emror and training time: on one hand the test eтror has terms that scale negatively with 1 - 7=2 -ww which is the distance to the lazy/active transition, on the other hand, the time it takes to reach the minimal loss point scales positively with the same term. This can be seen from Figure 2 bottom left panel, which plots the number of steps required to reach minimal test error, which increases as one goes further into the active regime.
Remark. In general when trying to fit a matrix B (instead of the special case B A' + E), the transition between lazy and mixed regime is when ozw R3 [|o Thus the exact location of the transition is task- dependent,,,, so that the same variance 6² and width w can lead to NTK or mixed regimes depending on the task. For example, let us assume that A" is full-rank instead of finite rank, then we expect || Ilop ~ vd instead of || ~ d, thus the transition would be at 2 = 7o2 + 7 instead of 7 + Tw. This suggests that linear networks are able to adapt themselves to the task:

---

leveraging active dynamics when the true data is low-rank to get better generalization, or remaining in the lazy dynamics the absence of low-rank structure, to take advantage of the faster convergence. Note also that in the absence of sparsity, the lazy regime can be attained with smaller width (Yw >> 1 is already on the boundary of the lazy instead of Ywo > 2), since the choice "/ 1,702 2
regime.
3 Conclusion
We prove a surprisingly simple self-consistent dynamic for the evolution of the matrix represented by shallow linear network under gradient descent. This description only unifies the already known lazy and balanced dynamics, but reveals the existence of a spectrum of mixed dynamics where some of the singular values are lazy while others balanced.
Thanks to this description we are able to give an almost complete phase diagram of training dynamics function of the scaling of the width and variance initialization w.r.t. the dimension.
natural question that comes out of these results is whether nonlinear network also feature similar mixed regimes, and whether they could be the key to understand the convergence of general DNNs.
References
[1] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-staireass property: a necessary and nearly sufficient condition for sgd learning sparse functions on two-layer neural networks. In Conference on Learning Theory, pages 4782 4782-4887. PMLR, 2022.
[2] Emmanuel Abbe, Enric Boix-Adseria Matthew Stewart Brennan, Guy Bresler, and Dheeraj Mysore Nagaraj. The staircase property: How hicrarchical structure can guide deep learning. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[3] Madhu S. Advani and Andrew M. Saxe. High-dimension dynamics of generalization error in neural networks, 2017.
[4] Bloemendal Alex, László Erdos, Antti Knowles, Homg-TTer Yau, and Jun Yin. Isotropic local laws for sample covariance and generalized wigner matrices. 2014.
[5] Zeyuan Alue Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-parameetriizatoo.. pages 242-252, 2019.
[6] Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks. International Conference on Learning Representations, 2019.
[7] Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization.. Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 244 253. PMLR, 10-15 Jul 2018.
[8] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019.
[9] Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with infinitely wide neural net. Advances in Neural Information Processing Systems, 32, 2019.
[10] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. High-dimensional limit theorems for SGD: Effective dynamics and critical scaling. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[11] Francis Bach. Breaking the curse of dimensionality with convex neural networks. The Journal of Machine Learning Research, 18(1):629-681,2) 2017.
10

---

[12] Rajendra Bhatia. Matrix analysis, volume 169. Springer Science & Business Media, 2013,
[13] Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution wide neural networks. Advances in Neural Information Processing Systems, 35:32240-32256, 2022.
[14] Lukas Braun, Clémentine Dominé, James Fitzgerald, and Andrew Saxe. Exact learning dynamics of deep linear networks with prior knowledge. In S. Koyejo. S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 6615-66222 6629. Curran Associates, Inc., 2022.
[15] Lenaic Chizat and Francis Bach. A note on lazy training in supervised differentiable programming. arXiv preprint arXiv: 18/2.076, 2018.
[16] Lénaic Chizat and Francis Bach. On the Global Convergence of Gradient Descent for Over- parameterized Models using Optimal Transport. In Advances in Neural Information Processing Systems 31, pages 3040-3050. Curran Associates, Inc., 2018
[17] Lénaïc Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Jacob Abernethy and Shivani Agarwal, editors, Proceedings of Thirty Third Conference on Learning Theory. volume 125 of Proceedings of Machine Learning Research, pages 1305-1338. PMLR, 09-12 Jul 2020.
[18] Lénaic Chizat, Maria Colombo, Xavier Fernández-Real, and Alessio Figalli. Infinite-width limit of deep linear neural networks. Communicutions on Pure and Applied Mathematics, 2022.
[19] Zhen Dai, Mina Karzand, and Nathan Srebro. Representation costs of linear neural networks: Analysis and design. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[20] Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes prer-parameterized neural networks. In International Conference Learning Representations, 2019.
[21] Kenji Fukumizu. Effect of batch learning in multilayer neural networks. In International Conference on Neural Information Processing, 1998.
[22] Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart. Disentangling feature and lazy training in deep neural networks. Journal of Statistical Mechanics: Theory and Experiment, 2020(11)) 1): 13301, 2020.
[23] Gauthier Gidel, Francis Bach, and Simon Lacoste-Julien. Implicit regularization of discrete gradient dynamics in linear neural networks. In H. Wallach, H. Larochelle, A. Beygelzimer, d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[24] Arthur Jacot. Bottleneck structure in learned features: Low-dimension regularity tradeoff. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 23609.23629. Curran Associates, Inc., 2023.
[25] Arthur Jacot. Implicit bias of large depth networks: a notion of rank for nonlinear functions. ICLR, 2023.
[26] Arthur Jacot. Implicit bias of large depth networks: notion of rank for nonlinear functions. In The Eleventh International Conference on Learning Representations, 2023.
[27] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural Tangent Kernel: Convergence and Generalization in Neural Networks. In Advances Neural Information Processing Systems 31, pages 8580-8589 8589. Curran Associates, Inc., 2018.
[28) 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, 2022.
11

---

[29] Arthur Jacot, Eugene Golikov, Clément Hongler, and Franck Gabriel. Feature learning in lz-regularized dnns: Attraction/repulsionn and sparsity. In Advances in Neural Information Processing Systems, volume 36, 2022.
[30] Ziwei and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks, CoRR. abs/ 18 810.02032, 2018.
[31] Jikai Jin, Zhiyuan Li, Kaifeng Lyu, Simon Shaolei Du, and Jason D. Lee. Understanding incremental learning of gradient descent: A fine-grained analysis of matrix sensing. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardd, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 15200-15238, PMLR, 23-29 Jul 2023.
[32] 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. 2024.
[33] 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 arXiv:2003.022166 2020.
[34] Aitor Lewkowycz and Guy Gur-Ari. On the training dynamics of deep networks with 1_2 regularization. Advances in Neural Information Processing Systems, 33:4790-4799, 2020.
[35] Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. In International Conference on Learning Representations, 2020.
[36] Chaoyue Liu, Libin Zhu, and Mikhail Belkin. Toward a theory of optimization for over- parameterized systems of non-linear equations: the lessons of deep learning. arXiv preprint wwXiv:2003.00307 2020.
[37] Hancheng Min, Salma Tarmoun, René Vidal, and Enrique Mallada. On the explicit role of initialization on the convergence and implicit bias of overparametrized linear networks. In International Conference on Machine Learning, pages 7760-7768. PMLR, 2021.
(38) Hancheng Min, Rene Vidal, and Enrique Mallada. On the convergence of gradient flow multi-layer linear models. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 24850-24887 PMLR, 23-29 Jul 2023.
[39] Greg Ongie, Rebecca Willett, Daniel Soudry, and Nathan Srebro. A function space view of bounded norm infinite width relu nets: The multivariate case. In International Conference Learning Representations, 2020.
[40] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang. Zach Veto,o, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-perte setbbperformance deep learning library. 2019.
[41] Scott Pesme and Nicolas Flammarion. Saddle-to-sadee dynamics in diagonal linear networks. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 7475-7505. Curran Associates, Inc., 2023.
[42] Scott Pesme, Loucas Fillaud-Vivien, and Nicolas Flammarion. Implicit bias of sgd for diagonal linear networks: provable benefit of stochasticity. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances Neural Information Processing Systems, volume 34, pages 29218-22 19-230. Curran Associates, Inc., 2021.
12

---

[43] Grant Rotskoff and Eric Vanden- Vanden-Eijnden. Parameters interacting particles: long time convergence and asymptotic error scaling of neural networks. In Advances in Neural Information Processing Systems 31, pages 7146-7155. Curran Associates, Inc., 2018.
[44] Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks, 2014.
[45] Andrew M. Saxe, James L.. McClelland, and Surya Ganguli. A mathematical theory of semantic development in deep neural networks. Proceedings of the National Academy of Sciences, 2019.
116(23):11537-11544
[46] Salma Tarmoun, Guilherme Franca, Benjamin D Haeffele, and Rene Vidal. Understanding the dynamics of gradient flow in overparameterrzed linear models. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning. volume 139 of Proceedings of Machine Learning Research, pages 101311 PMLR, 18-24 Jul 2021.
[47] Zihan Wang and Arthur Jacot. Implicit bias of SGD in lz-regularized linear DNNs: One- way jumps from high to low rank. In The Twelfth International Conference on Learning Representations, 2024.
[48] Blake Woodworth, Suriya Gunasekar, Pedro Savarese, Edward Moroshko, Itay Golan, Jason Lee, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models, 2020.
[49] Nuoya Xiong, Lijun Ding, and Simon Shaolei Du How ver-parameterization slows down gradient descent in matrix sensing: The curses of symmetry and initialization. In OPT 2023: Optimization for Machine Learning. 2023.
[50] Ziqing Xu, Hancheng Min, Salma Tarmoun, Enrique Mallada, and Rene Vidal. Linear convergence of gradient descent for finite width wver-parametrized linear networks with general initialization. In Francisco Ruiz, Jennifer Dy, and Jan- Willem van de Meent, editors, Proceedings of The 26th International Conference Artificial Intelligence and Statistics, volume 206 of Proceedings of Machine Learning Research, pages 2262-2284. PMLR, 25-27 Apr 2023.
[51] Greg Yang. Tensor programs ii: Neural tangent kernel for any architecture. arXiv preprint arXiv:2006. 14548, 2020.
[52] Greg Yang and Edward J. Hu. Feature learning in infinite-width neural networks, 2020.
[53] Yi Yu, Tengyao Wang, and Richard Samworth. A useful variant of the davis -ahan theorem for statisticians. Biometrika, 102(2):315-33) -23,,, 2015.
[54] Difan Zou, Philip M Long, and Quanquan Gu. On the global convergence of training deep linear resnets. arXiv preprint 2rXiv:2003.01094, 2020.
13

---

to represent element-wise product of two matrix of the The appendix is structured as follows.
In section A
we introduce the notation, and establish several results about how perturbing a matrix would impact its singular vectors.
In section C
we study the gradient flow dynamics of A, in the active regime and prove that A will be approximately aligned with A throughout the toall--t--Sallee regime.
In section
we prove theorem D for general cost.
In section D
we study the gradient flow dynamics for Agçs) in the lazy regime.
In section we show that Agce) is also approximately aligned with A throughout the
regime, using results in section [ and section [ In subsection [33
summarize the approximate dynamics of Ao(s) throughout training. In section E.4 the final test error.
we bound
In section F we bound the error from gradient descent and prove theorem 2
In section we describe the experimental setup.
A Preliminaries
A.1 Convention and Notation
Constants. d is the dimension of the input and the output layer, K is the rank of matrix A", and w is the dimension of hidden layer. and C will usually denote constants that are independent of d, and depending on the context, the value of and C might be different. If x is scalar that depends on d and y is scalar, then x = O(y) means there exists a constant e independent of d, such that for sufficiently large we have S cy. If A is matrix, then A Oyy) means there exists constant e independent of d such that for d sufficiently large, ||A|||og S cy. Ogy) can be either a matrix or a scalar, and its meaning will be always clear from context.
Matrix. We use || Ilop. ||., norm of a matrix, and the L norm for vectors. For every matrix A, we define max as the L norm, m We use / to denote identity matrix, the dimension of which determined by context. We shall assume that the signal singular values of A" are si..... ... Vi = 1,2.... K, and 5f = a,d where a1 2 02 .... 2 ag are constants independent of By selecting proper basis in the input and output space, assume that A S*. where S* is the diagonal matrix consisting of singular values of A", ordered from largest to smallest. The p, q-th element of matrix denoted Apg. We reserve the notation A(i,j) for the i, j-tt block matrix of A, which we shall define below.
to denote the operator norm of matrix, the Frobenius
Submatrix. Assume that O, and and let Rmm+1 d. For a matrix U, we define the k, j-th sub-block (kk of a matrix U as Un, 1 H1lmm11.0 I with both sides included. Notice that U2 (k, Üj k) In this notation, we can write the singular value decomposition of a matrix A
an ani+1 ...0
Orem OKK
A(i.j)
Σ U(.Kk))(K,A))GGG. Ul, 1)5(m 1, 1,m + 1)V (m 1,5).
Aussggall
We call index & (of sub-block) "signal", if k < m. Index m 1 is called "noise". Let S k))* be block matrix A (" 04+1, m4+1). Then A = diag(S(1, (1,( 1)"...,, S m))* and S(k, k)* Sna is the k-th sub-block of A". Each matrix has only finitely many blocklocks.
Indexing Conventions. Entries of matrices will usually be indexed by p, q, r and sub-blocks of matrices will usually be indexed by i, k, E. Usually k ranges from to m, and j usually ranges from to m + 1.
Element-wise Product. We use same shape.
Important Assumptions. Throughout the paper, shall always assume that
14

---

1. Yo > 1 (i.e., w >> d).
2. 2702 the < 0. (i.e., o'w << 1).
A.2 Matrix Inequalities
In the proof of main theorems, will work extensively with inequalities of matrix norms and inequalities that involves element-wise product. The element-wise product appears naturally in the derivative of singular vectors of matrix.
Lemma A.1. Assume that A, B and R are square matrices. Let Rmaz = maxi,j (Rill Then
tr|A(R B)] triBRR + A],
and
ItrA(( O B)] S Rmaxx (rr(A" A) (r((BB
In particular, Vp q, Rig 2 Rmin then
tr| MIAA(RO A)] 2 Rmintr (A A.A.
Proof. All are simple computations.
Ir[AR O B = Σ
Ap Rap Bap
P-g
BR A Σ
Apq
P.q
The two equations above prove the first claim.
|r|A(R B)] S tr(A A) (T((RO A) BYT B)
< Rmaax / (r(AT AI
This completes the second claim.
tr (RO A)] AiiR
in
2 Rmintr tr(AA'A
This completes the third claim.
Lemma A.2. min B\ & ||B|| S Fmaz Bl|
Proof. This is lemma B.3 of 55
A.3 Perturbation of Singular Values and Singular Vectors
We will often use the following variant of the Davis-Kann Kahan sin theorem.
Theorem A.3 (DK-sin 0 Theorem). Let Σ, E Rpxp be symmetric, with eigenvalues A1 2 and let V = (Vr, Us) and
Ao and '₁ Z ... So. Fix 1 < S S p. let d have orthonormal columns satisfying Σv, Ayo, and Σi, Åzu,, Let 1e. 08 be the singular values of VTV. Let (V, V) be the diagonal matrix with cos (V, V)jz = Jj and sin (V, V\ be defined entry-wise. Then
V= -6,......, 0.)
2min Vdl. Σ' E'lopp S E[lF)
|| sin OVV V)lF
min(Ar-11 dr. do A
Proof. This is theorem in 53
15

---

The implication of the theorem is that if two matrices are sufficiently close, then their singular vectors are also close each other. In the case where and X, of multiplicity the theorem reduces to saying that the sine value of the angle between Vy and ü, is very small.
The term || sin (V, E is complicated to take derivative. In this paper we will use the following characterization of alignment, which easier to take derivatives.
Lemma A.4. Let E. be Assume that 81 ... Snz Sn1+1 has m blocks in total. Assume that || SSllop = d Let = USV and define
d diagonal matrix, and let 81,--- -K.--- be its diagonal entries. Sn2 2 ... 8m SK 5K+1 .... Sa- Then Σ
x 4K-
Σ to(U" (KU, -V (
Resignal
Then
2minn V ((IIX| [1[[loop); X F |(](r) 2
XIlop
V S K
+
m881
||
$1123
SK
SK
ming ( S1z41)
Remark. To better understand x, consider the special case where each signal singular value is of multiplicity 1. Since matrix X "close" to diagonal matrix E, matrix U and matrix should also be close to identity along signal directions: [UEx] s 1, |VEal Rs 1, URa VEs. ss for 1 S k S K. Quantity x captures how much (Uxs], [Vikl, URA de VER deviate from I.
VS is symmetric matrix. Let V(t, (Eme-itl:.... Une) be the singular vectors X corresponding to sm.-tt,...., Sng. There is some freedom choose V. but for simplicity we pick ( k) (enz- Cn,) where e, is the i-th coordinate vector. As a result, ((, ((,,)) Vkk, k). Apply DK-sin theorem to V and V, we see that
Proof. Let X USV Then X
|| sin V(, 2min( /ng nh X X
ops | X E
min n +11111111111
2 +1)
Let op be the singular values of (K, for 1 S ni 7 Here for sub-block of big matrix. Then sin O(V (t, k), ((,,))) is the diagonal matrix, whose diagonal
using the notation
entries given by oz So sin OVV, V) = E, of)) tr (k,, VK, (k)
Now observe that
|x x - 1opp S(|X oo + -111ap)11 Ś'1op
|| X EB11F S(IIXI F +|(1+))|1X Ellop
We conclude that
2
tr| ()) (2min
n4 nz- 1
[X
||1i|| p. || F + (1E1(E)
||x Eillop
min
1
2
1% - s2,+1)
Similar conclusion is true for U (k; k). Next we bound tr(1
[E( k))
Σ
(Eggeo
g2
(nx mi m -X113,
16

---

[(,,k) - Ak. =K) LUMDSSG.J
3
-| -UN.DSSUNVOD-00UAAAS (
j#k
(IU( (A,
2 |I ((R,
Σ (I J#K
11X - EIP
28 1I U(k, (
sîtr(1 - (K, k)² U -
((k, ((k, (k),
||x - E11
For every
11k
T1k
X nk nk-1 matrix M, we have (r(A) S Ep S
VTa Th-1 Σ Xp1 P)
$
Vn 7
M Using this inequality, we conclude that
S(nu mx-1) |/ X
1lop
1r(1 - AV
8nz
2min( ma mk 1 1 (1I |E(lop), |X + ||E1(e)
|1 EDllpp
8yk
min
( 8 1 22 s2a
Summing on k we conclude that
[ X
llpp
ms1
lop + +1E1lop), || + ($(1=)) ) ||x 5:113p
x < K
+
8K
SK
ming (s 22 2
In the case where d + 00 m is a constant, a < 1, 81 = O(sk), SK > csk+1. K is a constant,
||X Ellop o(11_1lop) and (in ) 2 co for some constant we have S
0(1 1X- loe). 12
Lemma A.5. Les X UUSV and define
4K
Σ
(KUU"
kasignal
Let be d x d diagonal matrix whose diagonal entries are given by by,... dy. O.... where
bi =... bn, bm =bn, 2 bnm bg. Then
a --T (m + 1
V + 2(m + 1)²x)
17

---

Proof. First observe that || UTEVIlo llpp - Eii=1 |E UTEV( If + j, then [) 0,
m+1
110 VV,
=1
Σ
E
8=1
<b (m + 1)Vo.
If = then (0, = L
| bn, -lop Slbb,,
- bn, Alop + bi Σ
E+
<m 1) Jar (m 1)xr
Therefore
[ - UT V ((m + 1)3/2 + 2(m + 1)3x).
This bound is not optimal in m, but throughout the paper, is of constant order and it is fine miss a constant factor when estimating error.
Proof of Theorem I
B.1
Weak bound
We prove the following weak bound.
Proposition B.1. For every e 0 and every < T, we have with high probability,
| W1 - ot S (1+
(3)
Analogous results holds for Wz.
Our main tool is the following lemma. Lemma B.2. For every cost C and every time I, we have
W Wt()) (,WF((t W (0) (1)))
(4)
Proof. Let L = ||A" - All² be the loss function. Then
o,W; = 2w VC
a.W, = 2VOWT
We that
a.WWWFF) = 2W3 VCWT + W. W.VC = a.W?
Next, we show that at initialization, W WT and W W2 are approximately orthogonal projections, up to factor. Stating precisely, we have the following lemma.
Lemma B.3. There exists two projections P1 and P₂ R + R such that the following are true.
1. The image of P1 and Pa orthogonal to each other;
2. With high probability.
W WW) - 02 lop = 0 wd log d)
(5)
W (>(T)(()) o² 3w 2/ooVollogg d)
(6)
18

---

Proof of Lemma B.3
At initialization, the rank of wis wo with probability 1. Therefore W1 WT has eigenvalues that are O, and the non-zero eigenvalues equals the eigenvalues of d matrix WTW1.
WT (0) is (scaled) Wishart ensemble, whose limiting distribution is given by the Marchenko- Pastur law. The Marchenko-Pasee law, as stated in [4), proves the following. Let X be an M x N matrix with complex-valued independent entries Xz, such that
[[X 0:
2. E[[X []]
VMN
3. for every p E N, there exists a constant Cp such that
E ([(NNM)+Xx) X | S SCp.
Here, MM satisfies
log M
0 << x<C- <
log N I < 00
for some constant C' independent of M and N. Let Φ = M Then the eigenvalues of N N matrix X* has the same asymptotics
N which may or may not depend on N.
vΦ
[x
(7+ x)
poldz) =
dx + (1 - 5(dx)
2T
x2
where
71 :=:=
22
VΦ
In our situation, W1 is d matrix with independent and identically distributed Gaussian entries whose variance is 0². Let M = w, N = and therefore W W1(W) has the same distribution as o² Vwdxx X. Notice that for this choice of M and N, the asymptotic distribution eigenvalues
of X is py (dæ). Notice that py supported on interval
+ - 2, / + V
from which we conclude that in the limit, all eigenvalues of W (0) is approximately V
By Theorem 2.10 of [4. we have eigenvalue rigidity results for X Let X/ be the k-th largest eigenvalue for X X. Vk 21,2.... w) we have
I - 7kl < d 3+c
(8)
with high probability. Here Ya defined through
p po(dx)
(9)
Let de be the ke-th largest eigenvalue of WT Wi (0). By the relationship between WT ())) and X" X, we know 0g has the same law as 22v wdd Then for every k.
DA rzwl S V (1x6 - +
W
V d
(10)
00 Vwd)
(11)
with high probability. We conclude that the first w eigenvalues of W/ W2 (0) is at most O(o² / wwd= away from 1 and all other eigenvalues are O. Therefore there exists a projection P1 such that
|| W (0) - PP 0(0² 2
(12)
Similarly, there exists a projection P3 such that
|| W2(0) o² p2wP2lopp O(0² Vwd)
(13)
19

---

Notice that P. is not exactly orthogonal to P1. and it remains to find a projection P2 that is orthogonal
Pi and is close W (0). Assume that the column vectors of W1 (0) are UI..... : uu E R= and column vectors of 2 are 01,...,, E I t. For k = 1,2..... ..... d, we define vector o/g as
d = Uy Pi
(14)
We claim that Pau is very small. By law of large numbers, |uull - a a log d with high probability.
1
(0) + Ogy o.v (ulll) W1 W
=
(15)
o²w
w
1 (u,, uk/21 +... + (w 2k uw ll + O(oVdlog d)
(16)
س²w
Notice that (uy, 2k) uy, Vj is a family of independent and identically distributed random vectors. For each of these random vectors, all entries have zero mean. The variance of any one of the entries is given by
E [un, vx) ( ex) O()
(17)
We conclude that for each we have, by CLT.
(u1, (u)(u, +... + (ww: Ub) (uh ce)
1d0, N()
(18)
Vd
In particular,
P (mmax
Kus vx) (x) ee)l
> 100g0 Vdw++ log df SwP(
>
Ku,,
> 100o log
i=1
<2wP 1) > 100
50(d
Therefore with high probability, max, E Ku,, Ux) (ux, eeil I 10003 log d. This implies that with high probability,
[K ub)u1 +...+ + (w vk) /| S Vwo Vdw log
(19)
| I av ovolog d
(20)
Now let W be the matrix with column vector v/,,.. r',...,, onn and let P. be the projection to the column
space of W By construction I[ -W2llap I O l d) so WWTW; - W Willop S
O w/ log d). Since the nonzero eigenvalues of W W2 is at most O / wd) from 1, we know that the nonzero singular values of W? T W/ is also at most O Vwd) from 1. We conclude that ||2 W Wellop I O(02 d).
Proof of Proposition
B.1
We have
W W W (i(() WT W WoW, (W, W W2(O)W,
A A W WT (0) W3 (2))
From lemma B.3 we see that as positive semi-definite matrix, for every constant E > O,
0 wwww (0) < (1+=)o2wl
(21)
with high probability. Therefore
-(1 + =)o2 Wi < (T P2(0)) (1 we)
(22)
By moving terms around and corollary, we have
(WfW))² (1 +1+e)oo (1+))o WT W, + (1 + 2o
(23)
SAT A (1 (1+ ~
(24)
<(WTW, + + W1 (1+ 2
(25)
4
20

---

Theorem V.1.9 of [12 states that the square-root function is operator monotone, which implies that if A 2 B then VĀ 2 v B. Taking square-root, we have
w Td S A A +e) w S WW
o²w
wfW, (1 +
(26)
B.2 Strong Bound
The weak bound does not provide useful information if ||] W1 llop << o²w. For this reason we prove strong bound, which provide useful information if [|W] W1 ll << a2w-+0 for some constant
>
O. Recall that the evolution of weight matrix in gradient descent is given by the following.
d Wi (t) vC
(27)
dt
d W2(() = nVCW
(28)
dt
The goal of this section is prove that
WTW ~ AT otw2l
(29)
For simplicity of notations, shall assume that C1 WTW1, C2 WWW2. c V A + Tf 7 and C2 VAAT + otw?T. It easy of see that C1 and Cg invertible. Our main result for this section is the following proposition.
Proposition B.4, For every cost C we have
(V
[WTW, - V V A + o Ilop S min(O(o?w),
(30)
Proof of Lemma B.4
We start from the equations:
WF [W] W2 - Wi = C. AC CIATA Aw + C²C
W2 [(W W WWWF)2 o' WT = = AA - C2AA" -Aw2C + ACIAT.
Our goal to show that V VAAT 7 with
Ci, C2 are close to the solution Gn VATA+ A
C AT - (
0 = 02 - AA"O, - OZAA" - -'
Apriori. the cubic equation for C₁ and C2 might have multiple solutions. We give an intuitive argument explaining why Ć₁ and Ć₂ are the correct solutions. By selecting proper basis, we assume A = diag(a1..... Q2) is diagonal. Assume that (W, W., W; w7)2 w Also assume that Ci and C2 both commute with A. In this case the equations for C₁ and C₂ reduces to cubic equation for scalars. Let AI,..., Xa be the eigenvalues of C₁. Solving the equations for scalars, we have X; = or + / + Aw? By lemma B.2 and lemma B.3 we have
W.WF( (1) W (2()) W,W (0) w (2(0) o w o(o²w). w).
Since W, W/W?' is positive semi-definite, all its eigenvalues are non-negative, and thus
W.WF (t) 2 (a"w oo
Since the top eigenvalues of WFW W, is the same as the top eigenvalues of W, WT. we conclude
that X; 2 w2w(1 + o(1).. This forces X; a otu
Since C1. C2 are not assumed to be aligned with A, we cannot reduce the matrix cubic equations into scalar cubic equations. The high level idea for proving dC, = C. C. is small is the inverse function theorem.
21

---

1. Step 1: show that LHS of the equations are small.
2. Step 2: reduce the RHS of the equations to linear function of dC1 and dC2. The system of equations is thus reduced
(small small
*
dCi
u dC2)
The matrix is now the "Jacobian" matrix, and uyy U/ are left and right singular vectors of A.
3. Step 3: prove that the "Jacobian" matrix (*
is strictly positive definite, thus proving
that o dC₁ and uf dC2 have small magnitude for all i.
For gradient flow, W 2 W1 W7 is preserved and there exists projections P1 and P2 such that W?W2 WiWT o'w(P2 Pi) + O(o² Vud). Therefore for every unit vector have
WT [w] W₂ w,w - 1' Wil - ||Ciloppo Ad w
Substracting the second pair of equations from the first pair and denoting dC, = C. C, we obtain:
d -
[Cillop 2
C
AT AdC; A + A A'AA
o'
(31)
w
d
|CCllop 2²
C2
C9 AAT dC2AAT o2 + AdC,
(32)
Now since
C - C = CfdC, + +
we substitute the above relation to equation 31 and obtain
||C 1('" V
(09 -54-AAA) (AA)di dc, C, dc, (CP A" dC, + A dC
(33)
C dC1 -Ct A dC22
(34)
For any singular value 8; of A, with left and right singular vectors wi, Us, we multiply equation 34 to the left by of and divide both sides by o²w, to obtain an equation for of dC,. Similarly, we obtain an equation for uf dC2: for uf dC1 and u² dC2:
and similarly for equation
32
||Cil 0 O = u/ dCi
(V
8)² + +1C + σ2
( A'A)) + (1²²,)
8;
u T dC2A
00 wd) = upf dC2,
(v ( (v(=))))
)
+ 102 + F2w
(C2 _ - A47) + (
81
U. dC1AT.
Notice that C? - ATA WT W = PB) In the two equations above, by replacing O A with oww P3) we making
|C wd).
an error of at most |[Cr |loplddCillop 1oop 01 4 From weak bound know that |CCCllop S O(o2w).
Therefore the error we made by making the approximation on the right hand side can be absorbed into left hand side.
To show that |oo dCi and l|u dC2ll are small, suffices to show that the (din + dout) (din + dout) block matrix
+1C + WF 13)
0210 ) A
(x2))) AT
+ 1C2 + W2 P.)WT
22

---

is strictly positive definite. This matrix can be further simplified to
WT
+1 P2
Wi (%)
0
0
W2
(2)) T
+ +1 12 Ps
W
The inner matrix can then be rewritten as RR where R is defined as
(+3m)) +1+ 119 + V ( +1 112
R =
+1 +1 12
#
+1+
WT
0
Let
W2 The "Jacobian" matrix is then QR(QR)T. As described in the strategy,
0
we need to show that the singular values are strictly positive. The smallest nonzero singular value of QR(QR)" is the same as the smallest nonzero singular value of (QR))QM QR. We expand (QR)T QR
as follows.
(QR)QRR QR
81
+1+1111
+1++
PWAW + ((5) P.WiTTPs o2
+
8i
(v(2²))
31
2
+ +1-1
P3Wi
@2w
+
8i
+ +1+11
W2F2 +
8: (
P.W, W2P2
(2w
+
8;
Paw W2P1 +
(V(22))
8i
+1-1 1 P.W.W.P.
02 w
Si
2
+ +1+1
(PI P2)o
+ 2(V
8i
+1
8i W.P P.W
(2
(2
+ 8i
(W:WT + +W2 W2 - o2wpp - + vwd)
22w
8i
+1+1
(2 )
(P+ (P)+Pa)
2
+ 2(v
(
Si
2
8: (PWW +ONTT
+1
(
+ v (W1WT + W W.T W2) O(o3Vwd).
where we used the fact that
WWWF= = P1+ RWW
P2+ A - Ps.
The (din + dout)-th eigenvalue of the above lower bounded by o²w, ww (220) 2 + + 2
oww Fw (5²+) 2 ow.. We conclude that
|CdCilop + ||dC2llop S O /
23

---

Compared to lemma B.3 lemma B.4 gives a tighter bound on ||C1 -illpp when ||Cillp S ow
As suggested by anonymous referee, it is possible to obtain the same approximated dynamics of Aα()) by imposing a non hon-oooogensuus balance condition (in a different setup). Assume that Wi W,W7 WT W. W2 w. Then
+ - A"A= =0;
C 20 wo2 - AAT =0.
Therefore Ci -oww + VA w2 and C2 ozw + VAA" + toww2 Substituting into Gradient Flow equation, we have
dA
= AAT + ACVATA+ouu' dt
The advantage of the setup is that significantly simplifies the proof. The The limitation of the setup that W, W W,'W, otw2 if the initial variance of entries of W1 and W2 are comparable. In this case, W, W W22 W2 will have w positive singular values and w negative singular values, and the absolute value of positive and negative singular values are comparable, In the setup of our problem, the variance of entries of W1 and W2 equal.
C Gradient Flow Dynamics of A, in Active Regime
C.1 Saddle to Saddle Regime
Let A, have the following dynamics:
d A1
A.= (A" Ac)y AT A. otwaI AAA, o'w2I(A" - -A).
The goal of this section is to prove that the singular vectors of A₁ is well-aligned with the singular vectors of A", throughout the Saddle-to- -o-Sadlle regime. In the rest of the section, we will assume the dependence of A, and use A to represent A,. If initialization, A, commutes with A", then throughout the training, A, will always commute with A". this section, we use a delicate stability argument show that if A, almost commute with A at the beginning of the Saddle to Saddle regime, then it will continue to be almost commutative with A throughout the training process.
Definition C.1. Define P1 be the family of d x d matrices that satisfies the following conditions.
88 2 CoR 8K+1 S C² and 9x+1 ak
S < for some d-independent constants C
and C².
Ifaa ak++, then 8k Sk+1 for some d-independent constant
Define Pi to be the family of w x w matrix such that 83 88411 2 S if ax+1 < at.
and šK 2 co2w. Let y > 0 be constant. Define P2(C, be the family of w x w matrices A satisfying the following conditions.
(alignment of signals). Let = USVT. Define
=4K
Σ
kssignal
P2(C, v) is the family of matrix A such that I < Cd
Theorem C.2, Assume that
dA dt
AA
AAT + 1 A)
A(O) P2
Let T Oddlog d). Then Vt [0, we have A, E Pi n (2 min(L,
24

---

Proof. A simple result of the induction lemma
C.3
Jemma
C.4
lemma
C.6 and lemma C.7
We will use the following induction lemma to show that the singular vectors of A, are roughly aligned with A".
Lemma C.3. Assume that Pi and P2 be families of increasing sets, and let P Pt. Assume that we have a family of matrices At, I < T for some fixed number T. T does not depend on the family of matrices. Assume that Ao E P1 LOP. Ps. Let ((t),2) = {A t S t S t2). Assume the following are true.
If p..) E P1 then there exists a constant & > 0 independent of Ag such that Ato,te P/-
2. Let t 12. If Ato,all P4 and Ala,i) Pz. then Ato,tal P2.
3. Let t1 < t2. If Ap. C P2 and Ala E Pi then Asp,tal C P1.
Then A, E Pi and A E P2, 4A TT
Since Ao satisfies Pi, use condition 1 we have A (0.a] P/- Using condition 2, we know that Alu.e C P2. By condition we know that A C Pi and by condition 2, Aso.el C Ps. Iterate the argument.
Proof.
Lemma C.4. Af P1, then there exists odv such that Aga+t)) Pj.
Proof. Let STS.. Using the approximation time [t, + t] we have
d2 dd
ds
=13 (UTAVVVS oz + + AV 22 o+w?
For j 1,2,.... K
de dsj dt
de
< Cdls; +Fw)
for some constant C independent of w and A,. The conclusion follows from Grönwall.
Lemma C.5. Assume that
d² dA AA ATA+ ouul + AAT + A E).
dt
Then for signal k,
d2 di (K(UU =r( [((K)))((
jzk
Here for j + ki
k(j,k) P(kk)) C(iA) R(j,)) @ C(K,))
with C UT + E)V, Ryg spāptep8u for s p.q S w. p # q, and R(k.j)
Roy myjmnyt being the k. j-th block of R.
25

---

Proof. Use SVD dervative.
dU
-U F U (rddA
VS+ SVT
dt d//1)
(35)
dt
_U
(UT(A* (A* EVV
w -
(36)
+U(Fo ( S + AV + U - +)+ (A' + =.) - -SV 52 + o u2
(37)
-U - (FO (**AA*++ + +EVES + V 1 (A' E]VS
(38)
+U F S + (A' + BY (A SV E USe + o w2
SV
(39)
Let D = F U VS + SV dA u] Since F is anti-symmetric and the term in square bracket is symmetrie, know D is anti ayti-symmetrie. Then dú UD. Let $ vS² otwo
C= (A* + EVV. As result, v1 $ $.95 w.
dUm pd dt rarzta
Σ Up
[s
S + (6,5, + Crg]
$²
Apgptegd Let Rpg
if p # q and Ryp = 0. Then
8/ 5/
dU
kk)
Σ
C(K.))
dl
jistk
D(j,)) = & -RAA)))))))))
As a result, for ny P. < 1k+1.
UU (k. (k),
E
Upa
2 dt
p
Σ
Upa Drg +
Σ
Um
Σ
Upr Dra
p.q.e(mu.
D.PE (
Σ
Upg
Σ
Upr
p.gelna,m441)
rElmzmm241)
-t (Ie, k))
Σ
(((
J
Lemma C.6. If A C P1 and Ad, E P2(Cs,,) then
L dr
di < -ddx + O(Vdr),
where Vc = infreje...)) min( (o2 (sy (3, (st (o In particular, Astuita C # - P2(C, min(7: 1)) for some constant C. 1f < 1 then we can take C = Cs-
Proof. Use previous lemma. Let
x. = 4K
R))))
kisignal
26

---

measures the alignment of singular vectors of A with singular vectors of A". From UU = we have
and for every j + k,
- 5()))
In particular, [[ S 50(v)). We first estimate C. For all j = 1,2,...,, ,,, aj # az.
Σ (.'(V.)) EjjK
+O(dz)
(U = EEG + GUGGEEAAA
(U E CV
+o(V))
=0(Va)
D(j,)) -R(k)) 2
-R(k,)) @ (V, R(K,J) (TA
kR(j,)) O
+ +Odd + Vd)
By UUT = / we know that (( = 0. Therefore
U(.D)U(K.))) +UG,,)U(K,)) (0 +0(z) = 0
Similarly we have for V.
We can rewrite D(j, k) as
D(j,)) ( (13.:
+0(dz + Vd)
27

---

d d (0
2t R) [((K))D((
-2 tr CHUNAATUUUJAAAAANAAAAAAAAAAAAAAA
2
+ JUUUAFUHHHJJAAAOAAAAFSSSUUUUUUUD
+ +dddr + VdV=)
2 MA
+ Σ STWAAUU UUU
2 Σ SUL.)UUHAAA70007755A000000000
+
We know that
(S (k, -50444)) - - (K,
(0 - (8x + 8,) + +z)) (sk- - [5) + 25x 8)),
min
*x
2x 8
2ed
for some positive constant e independent of w. We conclude that
de d (r b)] 2 edr + O(dr))+
dt
The same trick applies to V. We similarly have
de @ (k, k) (k. k)] edæ + +((d)))+0(Vz))
It remains to show that tr (U) k)] bounded below.
de d dt [VU dt (k, k)] + (k, dU dt (k,
dV
(k) d + (U((,,) d U(k,)) - RTT d ((KKK)) de
dt
dt 2kr d dt
U (k., - UK,,, ((KK,))] d 1 (U(K
20 + cdæ + 0dzz O(Vdz)
2ade + ( 0(Vdz)
Similarly,
kk, d V(k, 2 cdr +
dt
28

---

This proves that
d² dr
dt S -cdo O(dr2) + O O(vdr).
(40)
Observe that if x 2 d- for some & > 0, then oē < - and therefore æ decrease exponentially. On the other hand, if x Σ d-1-é for some > 0, then da 2 O(vdr) and therefore z might increase. We therefore conclude that A Apr P2CC, min(7, 1)).
Lemma C.7. Assume that A[0.22] C P2(C1,,7) and Alo. C P1. Then Ato.sl C P..
Proof. By assumption we know that t = Odgl
de dS dt =10 (UT EVV + o + / + o (A*+))) - 28V +
d dsskt1 2 (O(Vd) + 0(d -K411 V skt1 + o'w2
dt
1sk+1 (12) (41) S O(o2w(d d ) log << o
It remains to verify the gap between different families. Let T be the first time when ||A|lop 1 and assume that t2 T. Assume that ak >> aj. then there exists some constant such that
d² dsk
dt 2 (5 E) - -sk)sk
d dsj < (17(1 - - [5] + +o? dt
An application of Grönwall inequality on 9y and 8y implics that there exists some constant > 0 that depends only aj. ag such that 7,(7) O(d Next deal with the case when to T. One
sx(7)
important observation is that the in T $1 s1()) addnaid) } T (d). simple result of Grönwall. Moreover, (2 + Odd) z 5 This implies that 81 - 8% > csi for some constant c. Repeat this argument for all remaining signal singular values completes the proof.
C.2 First NTK Regime
At initialization, the singular values of A are of order / by using Mechenko-Pastur law on A which is infinitely smaller than o As result, there very short period when the dynamics is very close to NTK. This section is again a delicate stability argument to show that at the end of the first NTK regime, A, almost commutative with A"
Lemma C.8. Let t be the first when the first singular value of A hits wdd & for 8 2m Then Ad, E P2(C, 1)) for some constant C.
Proof. We approximate the dynamies with NTK dynamics. Assume that
de dB di
= 2(A + E B)o²
with B(0)- = A()). Then
d d( dd(ABB) B)
-2(B A)0² w A A)
dt
+ (V AT A ow ww (V AAT o -2w)(A" + A).
simple application of Grönwall inequality implies that
1
20 nt f 20 (7) ②2w
A (t) _
1 F. -Allop(T)ds
d
||
Ag op
t
<3
s
22
d
29

---

We first check Pg. Clearly, & is of order d' 3. Then we have
$1 0 wA 11 lop + =(O(d=9) Old O(d (d--$ |),
1 o F O(o² vwd)
||A,,
By lemma C.4 we have x S Cd- m for some constant C. Lemma C.9. Let 12 = cd, where e constant to be chosen. Then
T 1 1
A₂ Pi 0 P2C, min(
2
4
Moreover, there exists constants bj ... Das Day+1 +1 for a diagonal matrix whose diagonal entries are given by by.... bx, 0, O, we have ||A El 50low)dd mindddmm(("" ( (2).
bug
Un.. bg, such that
Proof. At time t2, |IB(t2) 2a, co w d(1).. From previous lemma we know that
A 1
|
BlIop op (1) s
0²w si
Let t to we see that [|| (Bld)))lo--l((ttllla) $ 1|A Bloop(t2) S 2 A Atgllerace Therefore as arw long as is sufficiently small, we can guarantee that || Bllop(t) I min --)) [ op (t). This guarantees that the singular values of A grows linearly in [t1, t2l. that sk sk+1 2 сss, if aş > ak+1 for some constant c, that SK+1 < 5Sx and that Ar, E Pt. It remains to check that Ar, satisfies (C ). We are ready to use similar techniques as in lemma C.6 All computation are
similar, and the only difference that [R(kj) = 0
o²w) We have ax
2²w dx + O(dx oLw
) + + dx 2
d² tr( U(K, dt
SK
d² dr
oaw
20 dt
(-edr + dr v vdz)
0x ax(o2 + 0 wd add-3
Therefore A(t C P2 (ann2u1 f)) Let min( 2
f). From lemma C.8 we see that nk, + S P2,P2 S nk. then 18p1 8pp S Od Moreover, the dynamies of sp, 1, 2 are both given by
dsp. de
= 2(O(v)) + 0 + agd - sp.) $2 + otwo
dt
d 8p - 8pal < <0(d - 3pp
dt
1pp Sp2 (tz) I 0(1)ssp, - spx 0² wOd
Let bn, 8,. By lemma A.5 there exists d diagonal matrix E, whose diagonal entries are given
by bag bag bx. such that || Elpp S (
D The Gradient Flow Dynamics of Aaco) in Lazy Regime
In this section, stability argument to show that if ou is infinitely larger than d, then the algorithm cannot converge.
Proposition D.I. With high probability. for all time 6, ||Acco A 2 miinn] A (E
Proof. In gradient flow dynamics, ||An() - A" - ElFF is decreasing with time and therefore for
all time t, || 1 -[ [. In particular, ||Aoc))lop S A[]]A]]| + [)] By
30

---

assumption, |[ Cillop - /11/l// Oldv =),, and the dynamics of Ag(e) can be written
d AA 2(A" + A) + (A* + E 1)0(d / 1 +o + A
dt
Assume that B(0) = Ao(o) and
dB di
de
2(A* E B)ow
Then
-d a Aoct - BI? 20 Al(1) - BlI + + - dl
-
EE A)O )
+r((A - Bd
W (A* + F - ---
d
s - 20 20l2w|lA000 - Bli + || O1I V d
This implies that for all time,
Vo
V W
Aa())) - BlI do(
o2w
The dynamics of B is linear, and we have
B. (A" + E)(1 c - =202wl/dd Bye -20-200wl wo//t/d
|| = |le 12-22021/6 A +(1-c wo/wr/d d } Bye w 11: 0.99|e 20 12099//-2222//8 A' A'+(1-e- w0
2/18
A'+(1-0-20wwt// EE1
Let P be the projection to the image A" Then with high probability,
[B, 11 ||e 20d 2 wdt/d? A" + (1 e wt/// F2 + Boe 20 2 1/
20.99/le 2//6 A (1 20- wr/ wtt dd ElE
=0.99/|e
20 220- wf/d A + (1 - e wt d2
PE 0.99| 20
I[I 21 P)
2 2 min || |
Since ||Arts) - BI? = o(d), we have
||Acc)) A' 2 min (1)1))
E The Gradient Flow Dynamics of Aoce) in Active Regime
In section C we proved stability for gradient flow dynamics for A.. In this section, we prove similar stability statements for Agge), using the approximation results from section B We also summarize the behavior of An(e) in section E.3
E.1 Saddle-to-Saddle Regime
The goal of this section is to show that at the end of the first NTK regime, the singular vectors of Agce) are roughly aligned with A Moreover, the alignment remains to be good throughout the mixed regime. The following generalization of lemma andland lemma C.7 will be useful.
31

---

Lemma E.1. Assume that (at C C1 and AG, E P2(Cs. Let s s
s'k be the first K singular
values of A'. Assume that the dynamics A' the following.
d² 'dA'
- ATA' 0 V A'AT o (A" AAA
O(d (A' + A + (A* + (')
Here, Ogd->&g) is matrix whose operator norm is bounded by O(d->sk)) Then Altr,aa C PaiCs,mini 1, 1,))..
Proof. Let A' USVT. Computing the SVD derivative for A', we have
d² dU dt
-U F
UT AV o5 + 2 + o A AVS)
+ U F S - AT + V AUS + +ut
(F Φ (usrd~A(UT) AVS + SV" 25211)
Assume that a% + ag. Notice that
Fee(O(sLd &V = O(d-~>)
Using the same technique as C.6 we have
d2 d tr(U (k),))))))T UK, di
Zddr O(dx3) + +(Vdxx
+ +rr(U(k, k)² Σ U(k, (F @ (0" A'VS - 25
37k
2odr + Odda + O(d' - vz) O(vdz)
Similar conclusion also holds for tr( (k, VK, k)] and (V V k)]. We conclude that
d2 dx
dt -dd + Oddo Oad V + O(Vdr),
and (1t C P2 min 1,2))).
888 Lemma E.2. Assume that (ertta 72 and Ac, Pi. Let , x be the first K singular values of A'. Assume that the dynamics of A' the following.
d² dA' di (A" A -
+ A'A'
Here, Old ((x)) matrix whose operator norm bounded by O(d (sk).)sk). Then A C Pt.
Proof. Proceeding as in lemma C.7 we have
d² dsp +0( + apd sp)) ow + O(d~))sk)
dt
p is a signal, and
de dsk 2(0 -85 sk + oww2 (dd)))s))
di
Now the conclusion follows from Grönwall.
Theorem E.3. Assume that Aoge) P, 7). Then AppuT P Y') for some
constant Y
32

---

Proof. We use induction lemma C.3 every i ..2,...,2
to prove the theorem. For the first requirement, notice that for
ds;
d²
4|
dt
If Ao(e) E Pi, then Ae((et+r) C Pi for Vud 11
Here e is some small constant that depends only on Pt. This proves the first requirement. For the second requirement, we observe that Agce) satisfies the dynamics described in lemma E.1 at each stage, long |[ć Cillop S Oskdd for some À > 0. suffices to show that there exists some positive independent of w such that ||ć; - Cillop & (dd=As sx) in [o, Т]. Let T, be the first time when ||A|l = o² w. In [0, Ti}. all singular values of the same order. Therefore
|Õ; - Cilop - <O[Ic S = d
Now let T2 be the first time when = ondd =
In (Ti, T2), we have
||ć - Cillop I &O11cill -
~ 0
w
Therefore we can pick X =
Twm T. During [2], we have (0) 2 co²2u for some constant c, and
d2 dok
dt 2 s ( 8K)
dal S (2ar (1) o²w,,, we have T2 Ti 2 Jog for some Therefore w we have SK (T2) 2 w2wd for some constant > 0. Now in T], we have weak bound [|Ć; - Č lop < w()"w) We can therefore pick = v for [72, T). The third requirement verified
Moreover, since d2
in lemma E.2
E.2 First NTK Regime
As in As, we need show that at the end of the NTK regime, Aoce) must be P1 and Ps. Based on what we already have for Ar, this conclusion follows from Grönwall.
Lemma E.4. Assume that Ag = Aocoy- Let Tí be the first time when ||Arllop + ||Aoco)llop reaches pf w. Then
|AT Au(r)) llop O(02
In particular, Aotry) E Pi (C min( 1w-1 1))
Proof. From dynamics of A, we see that TV s for some constant The dynamics of A, is the following.
d2 dA, -(A- - At) ATA o / ALAT + + - -A))
dt
The dynamics of Agce) can be written as follows.
d² dAa(e) di (AA Ao(o)) v Abio Aoc) + + + A0(1) ((1) + +olw21(A' - Aoce))
20 wd) Ar(e)) + (A" (((o-
Observe that
A0(0 Acloop (11Ab1) At)llop + [((Ac)) - AT) A, loop
Soaw|lAore)) Arllop
This implies that we have the following inequality on positive definite matrices.
AF A, o2ullAact) - Aellopd S Al(1)AAcs S AF A, + o² o - Aellop
33

---

Assume that A, = USV then
11 Abi Aoci) + V Af A oA
Aow2 T + N7 A, 02w0l1 Apce)= - Asllopt V / + A7 A, - 02tr|| Ap(e) A Aellop 1lop
=Ily -II AAw2 + - - Acllopt + 52 Votwil - o2wllAa(() - Aellop + +521lop
=max;
(Votue + +
o2wrll + Aellop - Volue + of - orwil Ab(1)
( - A llopp)
<3|( - Aellop
We can control the difference between A, and Ao(t).
de & Acceillop & O(d)llA - Accoll + O(d))((o
Grönwall incquality implies that
|| Aocrr)lop 50(02v
We check that Agery) satisfies the Pt. Since ||Ar; Ascppllop O(oAvwd), Aocryllon 2 20² Let or (Arr) and 0z (AorT))) be the j-th singular value of Ar; and Aiery). If aj # ak j,k K, then oo - o% (AT))) 2 co and therefore
(a,(Ar(T)) (A(ATT)))) 2 co + O(o2 wd).
The noise clearly of order O(o2 This completes the verification of Pl. By lemma C.9 [|Ae(r)) lop S o²"OO(02 wd) (odd:w))mmmm min( Tme
-., By Lemma C.4 we are done.
E.3 Summary of Approximate Dynamics of Aocey at Each Stage
In previous sections we have proved that Aeta) satisfies P1 and P2(C, for some 7 > 0. The Pi and (C) property actually implies that the dynamics of Asco) is such that each group of singular values evolve independently. We state the approximate dynamics for each stage of the dynamics, and show that the alignment will be improved, the alignment is not already good enough.
Initialization. At initialization, Aεco) is a random matrix. The mean of each entry is 0 and the variance of each entry is o² Vu. The gap between singular values is infinitely smaller than the magnitude of singular values, and the singular vectors are not aligned with A".
Initialization to |etcc)lo o²w. Let Ti be the first time when || (Aocc)llop reaches oaw.
[0,Ti Til corresponds the very short NTK regime for the signals, and the dynamics of Aocej is approximately linear. The evolution of signal singular values are roughly linear (i.e., bounded above and below by linear functions), and at Ti, we have Ae(T.) E min (--11 1)) Since the singular values grows roughly linearly, Ti O(d).
[Aocc)) - to ||Aotollop o2w (#) 1/ Let T2 be the first time when |Acc))lop 2²w (w) given by
[Ti, Ty] the early stage of saddle-to-saddle dynamics. The dynamics of 8, is d² ds: dt
= 2(a,((1 + +0()) si) o7 + + (1++
By theorem 1, we have ||ć Cillop S o2w (w) - Let -e o²w (w) we have X = Tw By lemma 4,1, the dynamics of x satisfies
de da
di eda O(dz O(Vid' O(VVdz)
with x(Ti)
2 2
We conclude that
x(Tb) S O(d Tm=11
34

---

|Aoce))lop = ( to sk (x(A(()) = =axd. Let T3. be the first time sk (A(o)))) faxd. time T2, we have ox 2 wdd for some o > 0 that depends only on a. aK. Then we have
||0 - oi < O(skd"))
d² dx
dt S cdl Oddx)) + 0 0(vid' O(Vir).
We conclude that rTTz) O(d O()- The dynamics of s; is given by
(1+ 0(d d² d2ds: 2(a,d(1 O(d 26)) (i) V /3 + 21)
dt
ds;
s:(a,d 8,) s d
dt S 2(a,d Cdl s;) + o²w).
The bounds on da; dt
implies that
dlog e O(d)
1 - 7a² Tw d dlog +ogd+ Odd).
T3 Ti
ak
Final Stage. Lett 2 T3. Since 8K O(d), have
o²w )
|[ Ci lop S 00(sk d
0(skd6o2tw--)
dr v ede + O(dx + Ol O(Vadr22to) O(vdx).
dt
Recall that the constant c in term cdr must satisfy
(*E + (5n + 85) (SE - (5) + 87) + 85)
min jia
8 - 8
of 53
As a result, we can take
c= c=a aa) mink jiou sha, aa a
maxk,s-az a - a3)
Let d' be large constant such that if x > c'd-1 + o2 )(1 then de < c(a ld x. Let T4 be the first time when x I c(d + d2(122 after T3. Then T4 T3 < c(ai.....ak) log(d + (( + O(d). Moreover, for every t > T4. a cannot be larger than c'(d-1 + dh, because de < 0 if by
2
2d
#
c'(d' + d After Tr, the dynamics of each singular value is given
d² ds: = 2(1 + 0 + ( s;) + ((
dt
Let Ts be the first time after T4 when |s; - and| < O d +7 -1)) log d. Then
T3 - T max 1,7222 + Twu -1) d dogg d + Od log d).
2a
We conclude that at time
2max (1, 2(-7o= Nw + 1)) + max (1, 2(-702 Yo + 1)) c(a1,..., ak)
T
d dlog d
OK
2ax
+ Od log d)
we have
a(T") -122
Ao(T-) E Pz(C, max (1,2 702 - [/w)]
and
[s,(T") a,d| O(d + d log
35

---

E.4 Analysis of Testing Error
In section E.3 we proved that Ag(s) is almost aligned with A' throughout the training. With the aligmment in hand, we are ready to give a time to stop training and the testing error at the end of training.
Theorem E.5, Assume that Apt1) follows the gradient flow dynamics. At time
T
Yo?
Tw
21ax(1, 2( 20= - 7 + 1)) + max(11 2(-702
702 - The + 11)
)
dllog Odd log d),
оK
c(01....,ax)
20
the testing error
||AaeT.) - A - =0(o" + O('''0"''0' log" d) + + +
Proof. Let Pkc be the projection to the largest K singular values. Then
Ae A - Al(o) 1 + |I(1 Pk ) 11.
Let st....,8d Sy be the singular values of Au(a). Then II(I - Pk) Px)Aa(e)1E= 11 = Ep2K-1 derivative of $1,..., Sa reads
Sy The
d² ds di =10 (U"(A" + EW S +o + - (A" + EWV - 228 52 + + (A" E.V- - 250((o
+ (U" (A" + E))O(( un) +
If 2 K + 1, then 8p S C' for some constant C", and
d
dsp dt s
(VT + p
d d
Σ
2 5 p
Σ 8p
Σ
UpD V An + (UT Mm O(02w)
dt
P2K+1
p
95K
Σ Sp
Σ
|upl1 O(d))) o(va)) +(d)))
+
O(oFw)
p
1/
$ Σ Upp!")
$/2
Σ
mVx 100 (o
O wd)
QSK
1/2
Σ
(odd-)(?ww) + w())
pp
We conclude that
d d
[I[I - -PK)A < O(o2ud), dt
which implies that
log dl
- Pooool S vdo V wd) + 0
d = 0 V + logd)
- I Aoco 11 S O + +o' low2llog?? 2ggdd. d)
Next we estimate || Aoci) - 1. Notice that A0()) - A = Eis (oco
A (i.) lf (iiti j.. j. then A (1(,5) = 0, and
|P ccc (1.j) <
Σ Resignal,
+ -IUG..
<0(zd²)
36

---

If jí=j== = m + 1, we also have A (m + +. +1 = 0, and
Aeu)) (m 1, m + 1) v Σ 0 ( 1, [V( 50({x)
kosignal
Now assume that i = j are both signals. Then
|| () A (i, ?)
- -1
50(d + - log + O(d
Since there are only finitely many blocks in total, we conclude that at time T",
||Aacr-) A S O + +o'²² log" d + ddd 1,, -)))))))))))))))))))))))))))))) + Ol + 106662+1011)) log d
|| - |I? I 50(a" wd) + O0"aouuuu u log
2 d) + O(d3) + w(d2wd)
(41)
F
Gradient Descent Dynamics and Proof of Theorem 2
In this section we prove that gradient descent dynamies of As(e) is well-approximated by gradient flow dynamics of Ag(()-
F1 Gradient Descent vs Gradient Flow
To study the dynamics of A under gradient flow, we show that if the learning rate is small enough, then the gradient flow dynamics will be close to the gradient descent dynamics.
Lemma F.1. Assume that A is matrix (not necessarily square matrix). F is a function: iddi A + Rddim A, The norm || || satisfies [|AB]] S B| for all and B. In particular, operator norm and Frobenius norm satisfies this property. Assume that sup F S Co and ||V S for some constant. Consider gradient flow dynamics and gradient descent dynamies.
Gradient Flow: dA, = F(AA) dt
Aa((k 1)m) Aa(kn) = =((Aa(kT))
Gradient Descent:
17
Assume that Ay( = Aa)( Then
|Ay (k))(kn) S ((1 + h1 1) no
Proof. Notice that we have
Ay((k + 1)m) Ar(kn) 1-" F(A
["' F(AA(KT)+)) - F(Aa(kn))di
(A, - Aa)(k + 1)m) ((A, - Ako)(
S. ((A + t)) - F(Aa(67))dt== '"" F(A +t) - F(A, (t
+ (FF(()())))- - -F(A
Let G(l) = So ((A + s)) Then
["" F(A t)) -Gn) -C00
G") G
1 2 d
"1 di
(A, +1))
2 OF
(Ay + +)) d
" 0A
dt t - A +1)
OF
"" 0A
(Ay + F()) + E))
37

---

11 J--
- 35100o
(l((F(A)(k)) - = -a(k))))) S (nilllAs(K))
Aa(kn))
We conclude that
[|, 1)n) - Aa((k + 1)m)ll S(1 + (B)))) -
2 200C00
1
Adl((kn) | -all((((( 1)m) + 2 70 5(1 + +))((IAAA) -Cd((())AAd((())) +
s(1 + n/1) 2nc
| Aall(kn) $ ((1 +n, - 1) n
To apply the lemma above we need to prove that |Aore) 13 S (dd)) throughout the training.
Lemma F.2. For both lazy and active regime, we always have || - +0 | throughout the training.
Proof. The gradient descent dynamics of Ao(e) is given by
AB((+1) Ao(1)= 2 + E Ab(e))Ci + C2(A" + E Aa(o))
tr((A* E - Ao(e+1))* (A" (A E Amc+t))) - (A'
E EAoo)) (3Ar))"(AA + E -
21r((A* + AB(1))) (Ap(++1) Ab())) + (|Ag((+1) Ap(e) 112
S - 2ndd tond~2r((( + E. Aa(e)) " + (2) + +F -Aat))) + 17 (d2 21 112p))
From theorem 2, we know that Ci C2 2 2o Therefore 1r((A* F Ap(o)
(C1 A + (A)
w 2 cd2 for some constant c. Therefore for the lazy regime, we always have ||A4 + E A (1++1 S | + E Aeco (1) ((] | <
E - Ae())) )) 2 ||A- + E - 3 0
|A < 7[] 1| + [[,)) which implies that || 10( AA [| |b]))
for all time.
For the active regime, we have
||A- + E - Aoce) (1+ A E - Alco)l (1) S r1)(1).
Since the training has at most 7) [|]])² + O(r7*) S 10 + | (.
steps, we see that Vt, ||A + E AR 11? S 2(11
Proof of main theorem. We first consider the active regime. It suffices to consider the error from considering gradient descent, rather than gradient flow. To apply lemma F. it is more convenient to consider the dynamics for W, and Wz. By lemma G.2 W + [ W Il s ((d) and The gradient flow dynamics of [W1, W is given by ||Wallop + W2llop S O(vd).
d
W 1, W d IW (A' (A* W2W)), W, W2W1)" dt
Let F(W1, W2) = d 2 2 (A* W2W1, W1 WoW) (.. Then sup, || F(W,, W2) (l] S
Old 1). Computing the differential of F, we obtain that dF((W1, W.T = 21 dw: (A" W2W1) + W (-dWWW1 - W2dW1),
dA - W2W1) + W-d-WW? W2 W/ WFdd)))
and therefore |V (dd In the active regime, the total number of training steps is
y -OOd log d) By lemma F1
l flow ddescentt
[F(TT) < Onnd à) b)((1 () Old
-1) O((d log d)
38

---

and the same holds true for W2. We conclude that
[ low Wi Wydesoeniidesseen |
< <1I op +
[ (p)d log
-0 log
||4 gg dessentt 112 1 0(7 log" d).
In the lazy regime, from strong bound we have
WFW, = (1+ NoVoww ATA
Therefore |Willop < Olovw), sup, |F((W1, Wfille S d llop + |W2llop) =
O(d /w). Similarly, || S d =0111A-- W2WillF + 1|22le Wille) d d0
w)
By lemma F.1 at time 100d log
g²w
[|W]
[[W/oww
Wddeseet
100e logd S Ond ow) oo))) ||
100d logd
-1) - =Onnd-to log
We conclude that
W flw w W |F
<(1 W/deeett llop
1nd
=0 o² w l)..).
|| wo Wdeescet Wddseedd ='0)''d2'''llo d).
Recall that in lemma D.I we proved that
| (())
Bell.- = old3).
and at time 100 de log d
o"w [B, A F < d0 which implies that || Agescentt - A" -F S old)).
Before this time, we have |B 1
2 1 || (1))))) After this time, we have
tr((A- - E - deascent 1(++1)
( ((A" + E Adescent) - (A" + E Adescent
( + E
Adescent) ))
A(t)
Ag(t)
= 2tr((A" + E
Afte
deeceet ))² (A
Adeecent Adeacent) + ||4
g((+1) A4(0
A(C)
Adescent 1/?
S 2nd pnnd2tr((A)) "tr((A" E
Ao(t)
descent )T *C1++ 02) + E Ab(t)
Adescent))
+r²tr((A" + E-AAgrarent (100 + E Adescent))
S (2nd 20² - 10 |AA + A0ce)
descent ||2
Therefore ||A- + E
Adescent || is decreasing and therefore ||Adssent A+112 A(1)
minn |1, [}])) for all time.
G Experimental setup
We now describe the experimental setup for the experiments shown in Figures [ and 2 experiments, we used the losses
For all the
1 d
Lusin (0) = d2
A - (A* E)I (0)
| - A 1
where E has i.i.d A 1) entries, A = K K-1/2 ΣK Fasn with u4 ~ ~N Ida) Gaussian vectors in Idd, This means that RankA" = K. The factor K -1/2 ensures that ||A' [| O(d).
39

---

We then either run the self-consistent dynamics (equation T or gradient descent (equation 2 d
Following Theorem 2 we take learning rate rj cwo
for 2 1, and r <lIA-1 ll
d² otherwise,
where is usually but can taken be or 5 faster convergence the cost more unstable training.
For the experiments Figure 1 we took d 500 and K 5. For the experiments in Figure 2 we took d = 200 and K = 5. For making the contour plot, we took grid with 35 points for Toa E 1-3.0, 0.0] and 35 points for Tw E [0, 2.8]. For each of the 35 pair of values for (702 Te), we ran gradient descent (and for the lower right plot the self-consistent dynamics too) until the train error converged. For all the runs, took the same realizations of A' and E.
All the experiments were implemented in PyTorch [40]. Experiments took 12 hours of compute, using GeForce RTX 2080 Ti (1 (IGG memory) and two TITAN V (12GB memory).
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The contribution section accurately describes our contributions, and all theorems/propositions are proven in the main or the appendix.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer:
[Yes]
Justification:
We discuss limitations of results and approach after state them.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
40

---

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairees..
While the authors might fear that complete honesty about limitations might be used by reviewers grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed
penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification:
All assumptions are stated in the Theorem statements.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Thcorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Question:
Does the paper fully disclose all the information needed reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [Yes]
Justification: The experimental setup is described in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
41

---

(c) If the contribution is a new model (e.g.. a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: We use synthetic data, with a description of how to build this synthetic data. The experiments are only there for visualization purposes, we see no particular need to publish it.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. ec/ jublic/guides, Staiai/guuids//oeeeuuniissssss Soode@ubmissionPol111
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https 7/nips. wwwpc//pwppppp...ccee.eee..opeee wttplii/ggiieeeeeeeeeeooee Sttp://oooeoouisiioiiooioooo
sosoooooiisiiiiiciiicicccc for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: Most details are given in the experimental setup section in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
42

---

Answer: [No]
Justification: The numerical experiments are mostl there as a visualization the theoretical results, our main goal therefore clarity, which would be hurt by putting error bars everywhere.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the computer resources (type of compute workers, memory, time of exccution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: In the experimental setup section of the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
We have read the Code of Ethics and see no issue.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there is a special consideration due to laws regulations in their jurisdiction).
43

---

10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The paper is theoretical in nature, it has no direct societal impact that can be meaningfully discussed.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e-g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Not relevant to our paper.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (c.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification: We only use synthetic data.
44

---

Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do release any new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Not relevant this paper.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
45

---

Justification:
Not relevant to this paper.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
46


Reviews:
[{"review_id": "DggQ7PHQf6", "reviewer": "Reviewer_r8Jt", "summary": "The paper studies gradient descent dynamics in two-layer linear networks. It is shown that in the wide-hidden-layer regime with standard random initialization there is a simple self-consistent differential equation describing the network output (Theorem 1). This equation generalizes both lazy training regime and the standard solvable balanced regime. As an application, a phase diagram of the model with respect to different scalings of the hidden layer size and the weight variance is discussed (section 2.1).", "questions": "It seems that a linear network combining lazy and active regime can actually be constructed more directly, within the class of balanced models, by simply considering a non-homogeneous balance condition $W_1W_1^T=W_2^TW_2-b^2I$. This condition is realized, in particular, when the initial $W_1(t=0)=0$ while the initial $W_2$ is isometric up to rescaling: $W_2^TW_2(t=0)=b^2I$. The balance condition is again invariant and, arguing as in the paper, we can obtain the closed-form dynamics for $A=W_2W_1$. Namely,  by multiplying the invariant by the matrices $W_1, W_2$ or their transposes, we get the equations  $C_1^2+b^2C_1-A^TA=0$ and $C_2^2-b^2C_2-AA^T=0$. Finding the roots leads to the self-consistent GF equation $$\\tfrac{dA}{dt}=-\\tfrac{\\eta}{2}[(\\sqrt{4AA^T+b^4}+b^2)\\nabla C+\\nabla C(\\sqrt{4A^TA+b^4}-b^2)]=-\\tfrac{\\eta}{2}[\\sqrt{4AA^T+b^4}\\nabla C+\\nabla C\\sqrt{4A^TA+b^4}].$$ This dynamics also starts in a lazy regime, because initially $W_1=0$ so that $W_2(t)\\approx const$ and learning is linear and occurs only through $W_1$. As deviation of $W_1$ from 0 increases, the dynamics becomes active. \n\nMy impression from the proof of Theorem 1 is that it is essentially a reduction of the wide network model to the algebraically solvable \"balanced-type\" model, but with two channels corresponding to the two approximally orthogonal projectors $P_1, P_2$ appearing in Theorem 1. By restricting to the corresponding subspaces, the matrices $W_1, W_2, A$ are decomposed into two components, say $W_1=Y_1+Z_1, W_2=Y_2+Z_2$ and $A=A_Y+A_Z\\equiv Y_2Y_1+Z_2Z_1$. The initial conditions in one channel are like in the example above, $Y_1(t=0)\\approx 0, Y_2^TY_2(t=0)=\\sigma^2 wI$, while in the other channel they are reversed, $Z_2(t=0)\\approx 0, Z_1Z_1^T(t=0)=\\sigma^2 wI$. However, both channels are described by the same equation with the same initial condition, so yield the same solution $A_Y=A_Z.$ Then the equation for the total $A=A_Y+A_Z$ presented in the paper can be obtained from the single-channel equation above simply by replacing $A$ by $A/2$. I think that the paper would be easier to understand and more convincing if such a two-channel model was explicitly described, and its solution (either this or performed in the paper) explained more carefully.", "limitations": "The model considered in the paper (a two-layer linear network with a wide hidden layer) is very simple and fairly artificial. However, the effects exposed in the paper may be relevant for more complex and realistic models.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is generally nicely and thoughtfully written. \n\nThe main theoretical result described in the paper (the self-consistent evolution equation) is interesting and apparently original. Thanks to the simplicity of this equation, it is likely to be useful in future research. \n\nThis paper uncovers a simple and natural scenario in which lazy training transitions into an active training in an analytically tractable fashion. This scenario can help understand feature learning in more complex models.", "weaknesses": "I find the exposition of the main result (Theorem 1) and surrounding ideas not very clear.\nThe provided sketch of proof is not very convincing, since the derived cubic equations can generally have solutions $C_1,C_2$ other than those indicated. I also couldn't get through the proof in the appendix, particularly Proposition C.4. Some typos there:     \nline 689: \"*Let* $L = ||A^*-A ||^2$ *be the loss function*\"  - not used in the proof   \nline 761: \"*Recall that $\\tfrac{dw}{dt}=$*\" - unfinished sentence   \nline 773: \"*To show that $∥v_i^TdC_1∥$ and $∥u_i^T dC_2∥$*\" - unfinished statement", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the thoughtful review. Regarding the weaknesses you mention:\n\nIt is indeed true that there are many solutions to the set of equations\nwe obtain, and most of the work in the proof is to prove that one\napproaches the `right' solution. This part of the argument is quite\ntechnical and cannot be easily sketched out in the main. We will also\nfix the typos. \n\nRegarding your questions:\n\nThe construction you propose is very interesting, and it can actually\nbe slightly modified to recover the mixed dynamics exactly, if one\ninitializes with weights $W_{1}=(\nI_{d} \\\\;\\\\; 0)^{T}$ and $W_{2}=(0 \\\\\\;\\\\; I_{d})$ then the following three properties are satisfied at initialization\nand for all subsequent times\n$$\nA_{\\theta}C_{1} =C_{2}A_{\\theta}$$ \n$$\nC_{1}^{2} =A_{\\theta}^{T}A_{\\theta}+I_{d}$$\n$$\nC_{2}^{2} =A_{\\theta}A_{\\theta}^{T}+I_{d}$$\nsince it is true at initialization and the derivatives on both sides\nof the first equation match thanks to properties 2 and 3, while the\nderivatives of both sides of 2,3 match thanks to equation 1.\n\nThis initialization/derivation also seems to agree with your `two\nchannels' intuition.\n\nWe added your derivation of mixed regime GF equation at the end of appendix C.", "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": "xxYYK0QPi0", "reviewer": "Reviewer_oPFX", "summary": "This paper consider the GD dynamics for two layer linear network. The authors introduce an approximated dynamics which interpolate between the lazy and the balanced regime. The authors also showed the phase diagram based on the above dynamics for low-rank matrix factorization problem.", "questions": "Please refer to the strengths and weaknesses part. \n\n**Minor points and typos:**\n1. Equation (1) and the equation below Line 204 seems to be different in a $\\eta$ factor.\n2. In Line 284 should be $2 \\gamma_{\\sigma^2} + \\gamma_w < 0$.\n3. Line 939, where is Lemma G.2?", "limitations": "No potential negative societal impact", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. This paper provides an approximated dynamics GD dynamics for two-layer linear network, including different initialization scheme. I believe this is novel in the literature and is an interesting results.", "weaknesses": "1.  Regarding the (pure) lazy training part, it is said in Line 124 \"When $w$ is very large, we end up in the lazy regime where the parameters move enough up to a time $t$ to change $A\\_{θ(t)}$ , but not enough to change $C_1 , C_2$\" . Could you explain more on this or point out some references? In particular why $C_i(t) \\approx C_i(0)$ for any $t$?\n\n2. The statement of Theorem 1 needs to be written more clearly: (1) what is $\\delta$ in the statement? Do you mean that the upper bounds hold for any $0<\\delta<1$? Then how do the upper bounds depend on $\\delta$? (2) what is $C_1$ in the RHS of the bound. From the proof I think you mean $C_1(t)$ not $C_1(0)$?Then how does $C_1(t)$ depend on $t$, in particular, does it blow up in $t$?\n\n3. I also have a few doubts on the utility of Theorem 1: it is not clear to me whether the RHS of the upper bound in Theorem 1 vanishes in $d$ ? In particular, consider the settings in section 2.1 where $\\sigma^2 = d^{\\gamma\\_{\\sigma^2}},w = d^{\\gamma\\_{w}}$, then the first term  $\\mathcal{O}( \\sigma^2 w)$ is not vanishing in $d$ unless $\\gamma\\_{\\sigma^2} +  \\gamma\\_{w}<0$. For the second term $ O( \\sqrt{d/w} ||C_1||\\_{op} )$, it is not clear to me how $ ||C_1(t)||\\_{op}$ depends on $d$ for any $t$. Thus, could you address more on this point? \n\nIn conclusion, so far the utility of the main theorems is not clear to me, I will raise my scores if my questions are addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the thoughtful review. Regarding the weaknesses you raise:\n\n1. Proving this is non-trivial and would require some work to be proven\ndirectly, the intuition is that one needs to take a learning rate\nof order $\\sigma^{-1}$ to get finite size updates to $A_{\\theta}$,\nbut the change to the weight matrices $W_{1},W_{2}$ is then of order\n$1$ which is small relative to their size at initialization (of order\n$\\sigma$). Note however that this fact also follows from our first\ntheorem, if $\\sigma$ is sufficiently large then $C_{1}=\\sqrt{A_{\\theta(t)}^{T}A_{\\theta(t)}+w^{2}\\sigma^{4}I}\\approx\\sigma^{2}wI$\nfor all times $t$. \n\n2. (1) Yes any $\\delta$ can be chosen, actually we will remove $\\delta$ from the Theorem and simply say \"with high probability\"\n(2) It should indeed be $C_{1}(t)$. Note\nthat $C_{1}$ might indeed blow up in time, but we are mainly interested\nin having a small error relative to the size of $C_{1}(t)$ (the learning\nrate $\\eta$ has to be chosen of order $\\left\\Vert C_{1}\\right\\Vert _{op}^{-1}$\nso things only need to be small in comparison to $\\\\left\\\\Vert C\\_{1}\\\\right\\\\Vert\\_{op}$).\nThe fact that in Theorem 2 we are able to control the dynamics up\nuntil convergence using Theorem 1 shows that the approximation of\nTheorem 1 is good enough to be used in a practical setting. We will\nadd a discussion of these aspects after Theorem 1, because we agree\nthat it can be difficult to determine when such an approximation is\n\"good enough\" since everything can vary\nsignificantly in size as $d\\to\\infty$ and throughout training.\n\n3. Again the RHS does not always go to zero in $d$ but it always\nbecomes infinitely smaller than $\\left\\Vert C_{1}(t)\\right\\Vert _{op}$\nwhich is sufficient.", "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": "g9KQqDNOya", "reviewer": "Reviewer_mU2H", "summary": "This paper derives a formula for the training dynamics of two-layer linear networks, encompassing the lazy regime, the active regime, and the mixed regime.", "questions": "- Do the mixed dynamics have a stronger or weaker implicit bias compared to the active regime (such as the low-rank bias or sparsity bias)? Which regime predominates in practical scenarios?\n\n- Can the theory recover other non-trivial training dynamics of linear networks, such as the saddle-to-saddle dynamics [1]?\n\n- How can the main insights from this study be generalized to nonlinear models, such as two-layer ReLU nets? \n\n- Minor: There appears to be a missing '-' in the equation between line 55 and line 56.\n\n[1] Pesme \\& Flammarion. Saddle-to-Saddle Dynamics in Diagonal Linear Networks. (NeurIPS 2023)", "limitations": "The analysis in this study focus on 2-layer linear networks, which is limited.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- In constract to previous works, the authors reveal the existence of mixed dynamics for training two-layer linear networks, which combine the advantages of the lazy regime and the active regime: it converges from any random initialization, and has a low rank bias.\n\n- The authors prove an almost complete phase diagram of training behavior  on the task of recovering a low-rank matrix from noisy observations.", "weaknesses": "The main findings heavily rely on prior insights into lazy training and feature learning. Moreover, the focus on two-layer linear networks, which are simple and lack nonlinearity, limits the broader impact of this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the thoughtful review. Regarding your questions:\n\n- Note that we view the mixed dynamics as part of the active regime,\nwhat we show is that the short lazy dynamic period that always appears\nat the beginning plays an important role, so it is useful to think\nof the active regime as a mix of lazy and active. The balanced dynamics\nthat start after the short lazy period then leads to a low-rank bias.\n\n- We are describing the saddle-to-saddle dynamics actually, and improving\nsignificantly on prior results, since we prove it in the fully-connected,\nnon-diagonal case which is much harder. Furthermore we can determine\nexactly how small the variance $\\sigma^{2}$ needs to be to get the\nsaddle-to-saddle dynamics, in comparison to previous works that required\nan infinitely small initialization (https://arxiv.org/abs/2012.09839\nor https://arxiv.org/abs/2106.15933) and could not really handle the\nsecond saddle making it difficult to determine when one would leave\neach subsequent saddle. We solve all of these issues. We will add a\ndiscussion of the relation to the Saddle-to-Saddle dynamics.\n\n- There are a few high level ideas that could translate to nonlinear\nnets: (1) there could exist mixed regimes in nonlinear networks,\nwhere e.g. some neurons are active while others are still lazy in\nsome sense; (2) the transition from lazy to active could shift depending\non the sparsity of the task, which could explain why we are not yet\nable to fully describe the extent of the lazy regime as a function\nof initialization variance, width, and number of datapoints and how\nit depends on the task at hand.", "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": "AmcxG6NVuY", "reviewer": "Reviewer_pxfL", "summary": "In this paper, the authors study the training dynamics of a two-layer linear network $A = W_1 W_2$ and show that for a wide range of network width $w$ and initialization scale $\\sigma^2$, the dynamics of $A$ can be approximated by the self-consistent equation: \n$\\dot{A} \\approx -\\sqrt{ A A^T + \\sigma^4 w I } \\nabla C(A) - \\nabla C(A) \\sqrt{ A^T A + \\sigma^4 w I }$, \nwhich can be viewed as a mixture of the lazy and balanced dynamics. \nWith this observation, the authors then analyze the behavior of gradient descent with different choices of $w$ and $\\sigma^2$ for the task of recovering a low-rank matrix from noisy observations. In particular, they identify two possible regimes for $w$ and $\\sigma^2$: the pure lazy regime, where the network fails to recover the ground-truth matrix, and a mixed/active regime, where the network first aligns with the ground-truth following the lazy dynamics and then switches to the balanced (and aligned) dynamics and eventually recovers the ground-truth.", "questions": "1. Could you provide some intuition on the proof of Theorem~1? It seems to me to be some clever algebraic manipulations that cannot be easily explained intuitively.\n\n\n2. See item 2 of Weakness.", "limitations": "See Weakness.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "This paper unifies the lazy and balanced regimes of the training dynamics of two-layer linear networks. The equation (cf. (1)) obtained by the authors that approximately characterize the dynamics is surprisingly simple and suggests a way to classify the training behaviors that are more fine-grained than the lazy vs balanced regimes. The proof of Theorem 1 (the theorem on the validity of the approximation (1)) also looks interesting. Instead of controlling the error growth rate like in the proof of many similar results, it relies on the fact that $W_1W_1^T - W_2^T W_2$ is invariant under GF to obtain an equation and the stability of the solutions to that equation. \n\nWith the new equation, the authors also prove a global convergence result for balanced initialization without assuming alignment at initialization. This is a novel result and demonstrates the usefulness of the new characterization, as it leverages the initial short lazy regime to obtain an approximately aligned state.", "weaknesses": "1. The presentation of the paper can still be improved. In particular, some parts of the paper seem to be rushed (the appendix in particular) and contain many typos (e.g. the second equation in line 212, line 761, and line 773). \n\n2. As mentioned by the authors in Section~2.1, in the low-rank matrix recovery setting, the width of the network is assumed to be much larger than the ambient dimension (instead of the rank). What will happen if the network width is much larger than the rank but much smaller than the ambient dimension? It seems that the network will be directly in the active regime and we can no longer rely on the initial lazy stage to align the network. Is this true in theory/practice?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the thoughtful review. Regarding the weaknesses you mention:\n\n1. We will improve the readability of the proofs, thanks to the error/typos\nthat you and the other reviewers have found.\n\n2. We agree that this intermediate width regime is of particular interest,\nand it is probably the regime where one wants to be in practice, and\nwe have thought of it but it seems that it would require different\ntechniques (in particular the invariant approach might not work anymore).\nWhile your intuition that the dynamics would be purely balanced could\nbe correct, there is another possible interpretation: when the task\nhas a low rank structure, then only the dynamics along the low-dimensional\nspan of the true matrix matter, the parameters orthogonal to remain\nessentially constant, at least until the optimal early stopping time,\nthus the dynamics could possibly be approximated by a smaller network\nwith the same hidden layer size but with input and output of size\nequal to the rank of the true matrix $A^{*}$. We hope to be able\nto extend our analysis to this setup to answer which intuition is\ncorrect.\n\nRegarding your question:\n\n1. We will improve the proof and add a intuition section in the Appendix to explain the strategy and how we rule out the other solutions. Depending on how much room we have, we will also explain this in the main. Another to understand Theorem 1\nis to show an equivalence to another dynamic as described in our answer\nto Reviewer r8Jt, we will also add this simple derivation to help with\nthe intuition.", "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:

Meta-review: This paper studies the training dynamics of linear neural networks and identifies a mixed dynamics regime that sits in the middle between the lazy & the active regime studied by prior works. The analysis shows that this regime has theoretical advantages over the previously studied cases. The review reports are generally positive about the paper and I concur with them. The authors are encouraged to incorporate the review comments in preparing the next revision.
Decision: Accept (poster)