Mixed Dynamics In Linear Networks: Unifying the Lazy and Active Regimes

Authors: zhenfeng  tu
santiago  aranguri
arthur  jacot
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 20:05:30

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.

Main paper: 
 introduction: 
 Whether in linear networks or nonlinear ones, there has been a lot of interest in the distinction between the lazy regime [27] and the active regime [16, 43, 15, 52, 13] as the number of neurons grows towards infinity. In the lazy regime the training dynamics become linear, so that they can be easily described in terms of the Neural Tangent Kernel (NTK) [27, 9, 51, 36] , while the active regime exhibits complex nonlinear dynamics. While our understanding of the active regime remains much more limited, it appears to be characterized by the emergence of feature learning [22, 52] , and of a form of sparsity [3, 11, 2, 1] (the type of sparsity observed depends on the network type [11, 19, 25, 24] , but we will focus on fully-connected linear networks which exhibit a rank sparsity in the learned linear map [7, 35, 28, 47] ) which are both absent in the lazy regime .Note that even though it is common to talk of the 'the' active regime, we do not know yet whether there is only one or 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 of 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 of datapoints) [5, 20] , a 'large' initialization of the weights [15] , a small learning rate, and early stopping when using a cross-entropy loss or weight decay . Indeed, active regimes have been observed by breaking either of these requirements: taking limits with mild or no overparametrization [10] , taking smaller or even vanishingly small initializations [35, 28] , using large learning rates [33] or SGD [42, 47] , or studying the late training dynamics with the cross-entropy loss [30, 17] or weight decay [34, 39, 29, 26] . Though each of these can lead to active regimes with significantly different dynamics, they often lead to similar types of feature learning and sparsity.In this paper, we 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 σ 2 , 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 [44, 7, 8] .But our result also paint a more subtle picture than the lazy/active dichotomy. We propose a more granular approach, where at a certain time some part of the network can be in the lazy regime, while others are in the active or balanced regime. More precisely the network is lazy along the singular values of the matrix represented by the network that are smaller than σ 2 w, and in the active regime along the singular values larger than σ 2 w.We consider the training dynamics of shallow linear networks A θ = W 2 W 1 and show that for large enough width w (the inner dimension), and a iid N (0, σ 2 ) initialization of all weights, the dynamics of A θ(t) as a result of training the parameters θ = (W 1 , W 2 ) with GD/GF on the loss L(θ) = C(A θ ) for a general matrix cost C with learning rate η is approximately given by the self-consistent dynamicsEQUATIONThese dynamics contain as special cases both the lazy dynamics∂ t A θ(t) ≈ -2ησ 2 w∇C(A θ )when σ 2 w ≫ λ max (A θ ) and the balanced dynamics∂ t A θ(t) = -η A θ A T θ ∇C(A θ ) -η∇C(A θ ) A T θ A θwhen σ 2 w ≪ λ min (A θ ). But it also reveals the whole spectrum of mixed dynamics in between, where some singular values of A θ are below the σ 2 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 t every singular value of A θ can either be lazy or active/balanced. The mixed regime is the best of both worlds: on one hand, since A θ A T θ + σ 4 w 2 I 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 is a 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 σ 2 w adequately, one can best take advantage of these two phenomenon.Finally, we focus on the task of recovering a low-rank d × d matrix A * from noisy observations A * + E, training on the MSE error 1EQUATIONF in the limit as the dimension d, width w and variance σ 2 scale together with scaling laws w = d γw and σ 2 = d γ σ 2 . We describe the training dynamics for almost all reasonable scalings γ w , γ σ 2 leading to a phase diagram with two main regimes:• Lazy (1 < γ σ 2 + γ w ) where all singular values remain below the threshold σ 2 w throughout training, and where the network fails to recover A * due to the absence of low-rank bias.• Active (1 > γ σ 2 + γ w ) where K = RankA * 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 underparametrized regime when w < d (or γ w ≪ 1) where the rank is constrained by the network architecture rather than the training dynamics, and the noisy regime 2γ σ 2 + γ w + 1 > 0 where the variance of the entries of A θ(0) at initialization is infinite.Figure 1 : For both plots, we train either using gradient descent or the self-consistent dynamics from equation (1) , with the scaling γ σ 2 = -1.85, γ w = 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 a solid line the singular values of A θ(t) when running the self-consistent dynamics, and use a dashed line for the singular values from running gradient descent. In this experiment, RankA ⋆ = 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' as they cross the σ 2 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.Linear networks have been used as a testing ground, a stepping stone on the way to understand nonlinear networks. Linear networks and their training dynamics are in many ways much simpler than nonlinear ones, but in spite of a 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 an incremental learning behavior and a sparsity bias [44, 3, 45, 23, 41] , some of this analysis has been extended to include effects of initialization scale [48] and SGD [42] . While the same decoupling happens in general linear with diagonal initializations and diagonal task, it remains an extremely strong assumption.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 [14] . For deeper networks there exists explicit formulas in the mean-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 a singular value s i have 'speed' proportional to s i itself (or s i 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 one 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 a specific initialization. In contrast, we prove that these dynamics are approximately true with high probability for random Gaussian initializations, which is the standard initialization scheme for neural networks.We will study shallow linear networks (or matrix factorization) where a d out × d in matrix A θ is represented as the product of two matrices A θ = W 2 W 1 , where the weight matrices W 1 and W 2 are respectively w × d in and d out × w dimensional, for some width w. The parameters θ of the network are the concatenation of the entries of both submatrices θ = (W 1 , W 2 ).The parameters θ are learned in the following manner: they are initialized as i.i.d. Gaussian N (0, σ 2 ), and then optimized with gradient descent to minimize a loss L(θ) = C(A θ ). Though most of our analysis works for general convex costs C : R dout×din → R on matrices, we will in the second part focus on the task of recovering a low-rank matrix A * from noisy observations A * + E, by training a linear network A θ on the MSE lossL(θ) = 1 d 2 ∥A θ -(A * + E)∥ 2 F .The width w allows us to control the over parametrzation, indeed the set of matrices that can be represented by a network of width w is the set M ≤w of matrices of rank w or less. The overparametrized regime is when w ≥ min{d in , d out } because all matrices can be represented in this case.The evolution of the weight matrices during gradient descent with learning rate η is given byW 1 (t + 1) = W 1 (t) -ηW T 2 (t)∇C(A θ(t) ) W 2 (t + 1) = W 2 (t) -η∇C(A θ(t) )W T 1 (t)where we view the gradient ∇C(A θ(t) ) of the cost C as a d out × d in matrix, which for the MSE cost equals∇C(A θ(t) ) = 2d -2 (A θ(t) -(A * + E)).But we care more about the evolution of the complete matrix A θ(t) = W 2 (t)W 1 (t) induced by the evolution of W 1 (t), W 2 (t), which can be approximated byEQUATIONThus we see that if we can describe the matricesC 1 = W T 1 W 1 and C 2 = W 2 W T 2throughout training, then we can describe the evolution of A θ(t) .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 21 , allowing us to make the approximation C i (t) ≈ C i (0). Furthermore at initialization these matrices concentrate as w → ∞ around their expectationsE [C 1 ] = σ 2 wI din , E [C 2 ] = σ 2 wI dout .The GD dynamics can then be approximated by the much simpler dynamics:A θ(t+1) = A θ(t) -2ησ 2 w∇C(A θ(t) ),which are equivalent to doing GD on the cost C directly with a learning rate of 2ησ 2 w.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 1 of the 'feature map' W 1 is (approximately) constant. More problematic in the context of low-rank matrix recovery is the absence of low-rank bias, indeed one can easily solve the dynamics to obtainA θ(t) = (A * + E) + (1 -4d -2 ησ 2 w) t (A θ(0) -(A * + E)),and since EA θ(0) = 0 we obtainE A θ(t) = 1 -(1 -4d -2 ησ 2 w) t (A * + E).The expected test error E A θ(t) -A * 2 is therefore lower bounded byEA θ(t) -A * 2 = (1 -4d -2 ησ 2 w) t A * + (1 -(1 -4d -2 ησ 2 w) t )E 2which never approaches zero.In linear networks, there is no advantage to being in the lazy regime, as we simply recover a simple linear model at 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 a crucial role in making sure that the subsequent active regime starts from an 'aligned' state.There has been much more focus on so-called balanced linear networks, which are networks that satisfy the balanced conditionW 1 W T 1 = W T 2 W 2 .If the network is balanced at initialization, it remains so throughout training, because, the differenceW 1 W T 1 -W T 2 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 eigendecompositionW 1 W T 1 = W T 2 W 2 = U SU T . This implies the following shared SVD decompositions W 1 = U √ SU T in , W 2 = U out √ SU T and A θ = U out SU T in . Furthermore, we have C 1 = U in SU T in = A T θ A θ and C 2 = U out SU T out = A θ A T θ , which leads to self-consistent dynamics for A θ(t) : A θ(t+1) = A θ(t) -η A θ(t) A T θ(t) ∇C(A θ(t) ) -η∇C(A θ(t) ) A T θ(t) A θ(t) + O(η 2 ).Now these dynamics are quite complex in general, and it remains difficult to prove convergence. Indeed one can easily find initializations A θ(0) that will not converge, for example if A θ(0) = 0 then GD will remain stuck there. A lot of work has been dedicated to finding conditions that guarantee the convergence of the above dynamics [6, 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 θ(0) and A * + E are 'aligned', i.e. shares the same singular vectors A θ(0) = U out SU T in and A * + E = U out S * U T in , then they will remain aligned throughout training A θ(t) = U out S(t)U T in and the dynamics decouple along each singular values i (t + 1) = s i (t) + 2η |s i (t)| (s * i -s i (t)) + O(η 2). Since we always have s * i ≥ 0, then for small enough learning rates η, we see that if s i (0) ∈ (0, s * i ] it will grow monotonically and converge to s * i ; if s i (0) > s * i it will decrease monotonically to s * i , and if s i (0) ≤ 0 it will increase and converge to 0. Thus one can guarantee convergence if we further assume positive alignment s i (0) > 0.The advantage is that there is a momentum effect in the form of the prefactor |s i (t)|, 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 s i = s * i at a time when the other singular values are still basically zero. If we stop training at this time then the linear network will have essentially learned only the top K singular values of A * + E, which is a 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 a positive alignment and assume that the s i are random (i.i.d. w.r.t. a symmetric distribution), then each s i has probability 1 /2 of starting with a negative alignment and getting stuck at 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.A 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) 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 to come up one by one in a 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 a greedy low-rank algorithm [35] .Our analysis will confirm the fact that positive alignment happens naturally as a result of a short period of lazy training, allowing us to prove similar decoupling and incremental learning for a general random initialization. Remark. We can define the time dependent map Θ(G; t) = C 2 (t)G + GC 1 (t), so that the GD dynamics can be rewritten as A θ(t+1) = A θ(t) -ηΘ(∇C(A θ(t) ), t) + O(η 2 ). The map Θ is none other than the NTK for shallow linear networks, but it has also been called the preconditioning matrix in previous work [7] . The lazy regime is then characterized by the NTK Θ being approximately equal to the time-independent NTK Θ lazy (G) = 2σ 2 wG, whereas the balanced regime is characterized by the time-dependentΘ bal (G; t) = A θ(t) A T θ(t) G + G A T θ(t) A θ(t), with the distinction that the time dependence is only through A θ(t) .2 Mixed Lazy/Balanced Dynamics Both lazy and balanced dynamics have the surprising but very useful property that the evolution of the network matrix A θ is approximately self-consistent: the evolution of A θ can be expressed in terms of itself. The lazy approximation becomes correct for a sufficiently large initialization, while the balanced one is correct for a balanced initialization. However, for most initializations, neither of these approximations are correct.We fill this gap by providing a self-consistent evolution of A θ that applies for any initialization scale:∂ t A θ(t+1) ≈ -η A θ(t) A T θ(t) + σ 4 w 2 I∇C(A t ) -η∇C(A t ) A T θ(t) A θ(t) + σ 4 w 2 I.This approximation is formalized in the following theorem, denoting Ĉ1 (t) = A T θ(t) A θ(t) + σ 4 w 2 I and Ĉ2 (t) = A θ(t) A T θ(t) + σ 4 w 2 I Theorem 1. For a linear net A θ = W 2 W 1 with width w, initialized with i.i.d. N (0, σ 2 ) weights and trained with Gradient Flow, we have with high probability that for all time t,C 1 (t) -Ĉ1 (t) op , C 2 (t) -Ĉ2 (t) op ≤ min O(σ 2 w), O d w ∥C 1 (t)∥ op . Proof. (sketch) The quantity W 1 W T 1 -W T 2 W2 is invariant under GF (and approximately so under GD) and it is approximately equal to σ 2 w(P 1 -P 2 ) for two orthogonal projections P 1 , P 2 (at initialization and for all subsequent times because of the invariance). We therefore haveW T 1 (W 1 W T 1 -W T 2 W 2 ) 2 W 1 ≈ σ 4 w 2 W T 1 (P 1 + P 2 )W 1 ≈ σ 4 w 2 C 1 .Thus the pairs C 1 , C 2 approximately satisfy the following equations:0 ≈ C 3 1 -A T θ A θ C 1 -C 1 A T θ A θ -σ 4 w 2 C 1 + A T θ C 2 A θ 0 ≈ C 3 2 -A θ A T θ C 2 -C 2 A θ A T θ -σ 4 w 2 C 2 + A θ C 1 A T θ .The pair Ĉ1 , Ĉ2 is a solution of the above, and one can show that C 1 , C 2 must approach them and not any of the other solutions.The takeaway from theorem 1 is the following.1. In the lazy regime where ∥C1 (t)∥ op + ∥C 2 (t)∥ op ≤ O(σ 2 w), then ∥C 1 (t) -Ĉ1 (t)∥ op ≤ d/w∥C 1 (t)∥ op << ∥C 1 (t)∥ op .2. In the active regime where ∥C1 (t)∥ op /σ 2 w > d ε >> 1, then ∥C 1 (t) -Ĉ1 (t)∥ op ≤ O(σ 2 w) << d -ε ∥C 1 (t)∥ op .It is true that the error does not vanish. However, for our purpose it suffices to show that ∥ Ĉ1 -C 1 ∥ op is infinitely smaller than C 1 for all times, regardless of the magnitude of ∥C 1 (t)∥ op .We see how both the lazy and balanced dynamics appear as special cases depending on how large the variance at initialization σ 2 is in comparison to the singular values of the matrix A θ(t) :• Lazy: When σ 2 w ≫ s max (A θ(t) ), then Ĉ1 ≈ σ 2 wI dout and Ĉ2 ≈ σ 2 wI din , recovering the lazy dynamics.• Balanced: When σ 2 w ≪ s min (A θ(t) ), then Ĉ1 ≈ A T θ(t) A θ(t) and Ĉ2 ≈ A θ(t) A T θ(t) , recovering the balanced dynamics.But clearly there can be times when neither conditions are satisfied, when some singular values of A θ(t) are larger than the threshold σ 2 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 A θ(t) (s i ≪ σ 2 w) and active/balanced along the large ones (s i ≫ σ 2 w).At initialization, the singular values are of size σ 2 √ wd. 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 σ 2 w threshold, at which point they will switch to balanced dynamics (after a short transition period when the singular value is around the threshold s i ≈ σ 2 w). Once a singular value is far past the threshold s i ≫ σ 2 w, training along this singular value will be much faster than along the lazy singular values (this speed up can be seen in Figure 1 ). 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 s i,t of A t , with singular vectors u i,t , v i,t :s i,t+1 -s i,t ≈ η t u T i,t ∂ t A θ(t) v i,t ≈ -2η t s 2 i,t + σ 4 w 2 u T i,t ∇C(A θ(t) )v i,t ,where the prefactor 2η t s 2 i,t + σ 4 w 2 describes the effective learning rate along the i-th singular value, which depends on the i-th singular value s i,t itself. This suggests that it is more natural to distinguish between the lazy and active regime at a much more granular level: at every time t a singular value can be either active or lazy (or very close to the transition but this typically only happens for a very short time). In contrast, the traditional definition of the lazy regime was defined for a whole network and over the whole training time. To 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 a short time period (assuming w > d), and as we will see this short lazy 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.While this challenges the traditional lazy/active dichotomy, it also reinforces it, as it shows that there is no fundamentally different third regime, only lazy, active, and some mix of the two. Theorem 1 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.e. describe the dynamics, convergence and generalization properties of DLNs for almost all reasonable initialization scales σ 2 and widths w. Remark. The transition of a singular value s i from lazy to active can be understood as a form of alignment happening in the hidden layer: the two vectors W 1 v i and W T 2 u i for u i , v i the left and right singular vectors of s i are orthogonal in the lazy regime and become perpendicular in the balanced where A θ(t) comes from GD and Âθ(t) from the self-consistent dynamics. We observe that this distance is 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 satisfiesu T i W 2 W 1 v i W T 2 u i ∥W 1 v i ∥ = s i u T i C 2 u i u T i C 1 u i ≈ s i s 2 i + σ 4 w 2which is close to zero for lazy singular values s i ≪ σ 2 w and close to one for active ones s i ≫ σ 2 w. 
phase diagram for mse: 
 To illustrate the power of Theorem 1 we provide a phase diagram of the behavior of large shallow networks on a MSE task, for almost all (reasonable) choices of width w and variance σ 2 scalings.We want to recover a rank K and d × d-dimensional matrix A * with s i (A * ) = da i for some 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 is the optimal way to choose w and σ 2 . We will therefore consider general scalings w = d γw and σ 2 = d γ σ 2 . We will now describe the (γ w , γ σ 2 )-phase diagram which features 4 regimes: underparametrized, infinite-noise, lazy and mixed/active.a 1 ≥ a 2 ≥ • • • ≥ a K independentWe can identify a region of 'reasonable' pairs (γ σ 2 , γ w ) by ruling out degenerate behavior. First, the width w 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 γ w ≥ 1. Another constraint comes from the variance of A θ at initialization: the entries A θ(0),ij at initialization have variance σ 4 w. We want this variance to go to zero as d grows which implies that we need 2γ σ 2 + γ w < 0.Now within this reasonable region we observe two regimes, the pure lazy regime for 1 < γ σ 2 + γ w where the network simply fits A * + E thus failing to learn A * and the mixed regime for 1 > γ σ 2 + γ w where the dynamics are lazy for a 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 γ w , γ σ 2 such that γ w > 1 and 2γ σ 2 + γ w < 0, we have two regimes:• Lazy (1 < γ σ 2 + γ w ): with a learning rate η ≪ d 2 σ 2 w we have that for all time t,1 d 2 A θ(t) -A * 2 F ≥ c. • Active (1 > γ σ 2 + γ w ): with a learning rate η ≪ d 2 s1(A * ) ∼ d, and at time t = 1 η ∆ a K + 2max(1, 2∆) c(a 1 , . . . , a K ) + max(1, 2∆) 2a K d log d + η -1 O(d log log d), for ∆ = 1 -γ σ 2 -γ w > 0, we have that 1 d 2 A θ(t) -A * 2 F ≤ O(σ 4 w + σ 4 w 2 log 2 d d 2 + d -1 2 + σ 2 w d + η 2 log 2 d d 2 ), for c(a 1 , . . . , a K ) = min k,j:a k ̸ =a j |a k -aj |a 2 K max k,j:a k ̸ =a j |a 2 k -a 2 j |.Note that all the terms inside the final O(. . . ) term vanish: σ 4 w → 0 because γ σ 2 + γ w < 0,σ 4 w 2 log 2 d d 2+ σ 2 w d → 0 since 1 > γ σ 2 + γ w , and η 2 log 2 d d 2 → 0 since we assumed η ≪ d. This shows that the lazy regime only appears for very large widths γ w > 2 (or at least the lazy regime with finite variance at initialization). Indeed the choice γ w = 2, γ σ 2 = -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 2 , top left panel, confirms this, since we see a good generalization even for small widths w < d, and in particular when w ≈ RankA * . But to leverage this underparametrized regime, one would need to know the rank of the true matrix A * in advance, which is 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 a 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 a trade-off between generalization error and training time: on one hand the test error has terms that scale negatively with 1 -γ σ 2 -γ w , 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 σ 2 w ≈ ∥B∥ op . Thus the exact location of the transition is task-dependent, so that the same variance σ 2 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 ∥A * ∥ op ∼ √ d instead of ∥A * ∥ op ∼ d, thus the transition would be at 1 2 = γ σ 2 + γ w instead of 1 = γ σ 2 + γ w . 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 in 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 a smaller width (γ w > 1 instead of γ w > 2), since the choice γ w = 1, γ σ 2 = -1 2 is already on the boundary of the lazy regime. 
conclusion: 
 We prove a surprisingly simple self-consistent dynamic for the evolution of the matrix represented by a shallow linear network under gradient descent. This description not 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 are balanced.Thanks to this description we are able to give an almost complete phase diagram of training dynamics as a function of the scaling of the width and variance at initialization w.r.t. the dimension.A 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.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 t in the active regime and prove that A t will be approximately aligned with A * throughout the Saddle-to-Saddle regime.• In section B, we prove theorem 1 for general cost.• In section D, we study the gradient flow dynamics for A θ(t) in the lazy regime.• In section E, we show that A θ(t) is also approximately aligned with A * throughout the Saddle-to-Saddle regime, using results in section C and section B. In subsection E.3, we summarize the approximate dynamics of A θ(t) throughout training. In section E.4, we bound the final test error.• In section F, we bound the error from gradient descent and prove theorem 2.• In section G, we describe the experimental setup. 
analysis of testing error: 
 In section E. T * = 1 -γ σ 2 -γ w a K + 2max(1, 2(-γ σ 2 -γ w + 1)) c(a 1 , . . . , a K ) + max(1, 2(-γ σ 2 -γ w + 1)) 2a K d log d+O(d log log d),the testing errorEQUATION2 ) + O(σ 2 wd)Proof. Let P K be the projection to the largest K singular values. Then∥A θ(t) -A * ∥ 2 F ≤ ∥P K A θ(t) -A * ∥ 2 F + ∥(I -P K )A θ(t) ∥ 2 F . Let s 1 , . . . , s d be the singular values of A θ(t) . Then ∥(I -P K )A θ(t) ∥ 2 F = p≥K+1 s 2 p . The derivative of s 1 , . . . , s d reads d 2 dS dt =I ⊙ U T (A * + E)V S 2 + σ 4 w 2 + S 2 + σ 4 w 2 U T (A * + E)V -2S S 2 + σ 4 w 2 + I ⊙ U T (A * + E)V O(σ 2 w 1 ) + O(σ 2 w 1 )U T (A * + E)V -2SO(σ 2 w 1 )If p ≥ K + 1, then s p ≤ C ′ σ 2 w for some constant C ′ , andd 2 ds p dt ≤ (U T A * V ) pp + (U T EV ) pp O(σ 2 w). d 2 d dt p≥K+1 s 2 p ≤ p s p   q≤K U qp V qp A * (qq) + (U T A * V ) pp   O(σ 2 w) ≤ p s p   q≤K |U qp ||V qp |O(d) + O( √ d)   O(σ 2 w) ≤ q≤K p s 2 p 1 2 p |U qp | 2 1 2 max p |V qp |O(σ 2 wd) + p s 2 p 1 2 O(σ 2 wd) ≤ p s 2 p 1 2 O(d 1-γ σ 2 w) + O(σ 2 wd) We conclude that d 2 d dt ∥(I -P K )A θ(t) ∥ F ≤ O(σ 2 wd),which implies that∥(I -P K )A θ(t) ∥ F ≤ √ dO(σ 2 √ wd) + O(σ 2 wd log d d ) = O(σ 2 √ wd + σ 2 w log d) ∥(I -P K )A θ(t) ∥ 2 F ≤ O(σ 4 wd 2 + σ 4 w 2 log 2 d) Next we estimate ∥P K A θ(t) -A * ∥ 2 F . Notice that ∥P K A θ(t) -A * ∥ 2 F = i,j ∥P K A θ(t) (i, j) - A * (i, j)∥ 2 F . If i ̸ = j, then A * (i, j) = 0, and ∥P K A θ(t) (i, j)∥ 2 F ≤ k:signal,k̸ =i ∥U (i, k)∥ 2 F ∥S(k, k)V (j, k) T ∥ 2 F + ∥U (i, j)S(j, j)∥ 2 F + ∥V (j, k) T ∥ 2 F ≤O(xd 2 )If i = j = m + 1, we also have A * (m + 1, m + 1) = 0, and∥P K A θ(t) (m + 1, m + 1)∥ 2 F ≤ k:signal O(d 2 )∥U (m + 1, k)∥ 2 F + ∥V (m + 1, k)∥ 2 F ≤ O(d 2 x)Now assume that i = j are both signals. Then∥P K A θ(t) (i, i) -A * (i, i)∥ 2 F ≤∥U (i, i)(S(i, i) -S * (i, i))V T (i, i)∥ 2 F + O(d 2 )∥U (i, i)V (i, i) T -I∥ 2 F ≤O(d -1 + d 2(γ σ 2 +γw-1) ) log d + O(d 2 √x) Since there are only finitely many blocks in total, we conclude that at time T F Gradient Descent Dynamics and Proof of Theorem 2In this section we prove that gradient descent dynamics of A θ(t) is well-approximated by gradient flow dynamics of A θ(t) . 
 
 Appendix: 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. c and C will usually denote constants that are independent of d, and depending on the context, the value of c and C might be different. If x is a scalar that depends on d and y is a scalar, then x = O(y) means there exists a constant c independent of d, such that for d sufficiently large we have x ≤ cy. If A is a matrix, then A = O(y) means there exists a constant c independent of d such that for d sufficiently large, ∥A∥ op ≤ cy. O(y) can be either a matrix or a scalar, and its meaning will be always clear from context.Matrix. We use ∥ • ∥ op , ∥ • ∥ F , ∥ • ∥ to denote the operator norm of a matrix, the Frobenius norm of a matrix, and the L 2 norm for vectors. For every matrix A, we define maxA as the L ∞ norm, max i,j |A ij |. We use I to denote identity matrix, the dimension of which is determined by context. We shall assume that the signal singular values of A * are s * 1 , . . . , s * K . ∀i = 1, 2, . . . , K, and s * i = a i d where a 1 ≥ a 2 ≥ . . . ≥ a K are constants independent of d. By selecting proper basis in the input and output space, we 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 a matrix A is denoted A pq . We reserve the notation A(i, j) for the i, j-th block matrix of A, which we shall define below.Submatrix. Assume that n 0 = 0, and a 1 = . . . = a n1 > a n1+1 = . . . a n2 > . . . = a nm = a K , and let n m+1 = d. For a matrix U , we define the k, j-th sub-block U (k, j) of a matrix U as U n k-1 +1:n k+1 ,nj-1+1:nj , with both sides included. Notice that U T (k, j) = U (j, k) T . In this notation, we can write the singular value decomposition of a matrix A asA(i, j) = k:signal U (i, k)S(k, k)V (j, k) + U (i, m + 1)S(m + 1, m + 1)V T (m + 1, j).We call an index k (of sub-block) "signal", if k ≤ m. Index m + 1 is called "noise". Let S(k, k) * be block matrix A * (n k : n k+1 , n k : n k+1 ). Then A * = diag(S(1, 1) * , . . . , S(m, m) * ) and S(k, k) * = s n k I is the k-th sub-block of A * . Each matrix has only finitely many sub-blocks.Indexing Conventions. Entries of matrices will usually be indexed by p, q, r and sub-blocks of matrices will usually be indexed by i, j, k, ℓ. Usually k ranges from 1 to m, and j usually ranges from 1 to m + 1.Element-wise Product.We use ⊙ to represent element-wise product of two matrix of the same shape.Important Assumptions. Throughout the paper, we shall always assume that 1. γ w > 1 (i.e., w >> d).2. 2γ σ 2 + γ w < 0. (i.e., σ 4 w << 1).In the proof of main theorems, we 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 a matrix. Lemma A.1. Assume that A, B and R are square matrices. LetR max = max i,j |R ij |. Then tr[A(R ⊙ B)] = tr[BR T ⊙ A],and|tr[A(R ⊙ B)]| ≤ R max tr(A T A) tr(B T B)In particular, if ∀p, q, R pq ≥ R min > 0, thentr[A T (R ⊙ A)] ≥ R min tr(A T A)Proof. All are simple computations.tr[AR ⊙ B] = p,q A pq R qp B qp tr[BR T ⊙ A] = p,q B pq R pq A qpThe two equations above prove the first claim.|tr[A(R ⊙ B)]| ≤ tr(A T A) tr((R ⊙ B) T R ⊙ B) ≤R max tr(A T A) tr(B T B)This completes the second claim.tr[A T (R ⊙ A)] = i,j A ji R ji A ji ≥ R min tr[A T A]This completes the third claim.Lemma A.2. σ min (A)∥B∥ F ≤ ∥AB∥ F ≤ σ max (A)∥B∥ F .Proof. This is lemma B.3 of [54] .We will often use the following variant of the Davis-Kahan sin θ theorem. be the singular values of V T V . Let Θ(V, V ) be the diagonal matrix with cos Θ(V, V ) jj = σ j and sin Θ(V, V ) be defined entry-wise. Then∥ sin Θ(V, V )∥ F ≤ 2min( √ d∥ Σ -Σ∥ op , ∥ Σ -Σ∥ F ) min(λ r-1 -λ r , λ s -λ s+1 ).Proof. This is theorem 2 in [53] .The implication of the theorem is that if two matrices are sufficiently close, then their singular vectors are also close to each other. In the case where r = s and λ r is of multiplicity 1, the theorem reduces to saying that the sine value of the angle between v r and vr is very small.The term ∥ sin Θ(V, V )∥ F is complicated to take derivative. In this paper we will use the following characterization of alignment, which is easier to take derivatives.Lemma A.4. Let Σ be a d × d diagonal matrix, and let s 1 , . . . , s K , . . . , s d be its diagonal entries. Assume that s 1 = . . . = s n1 > s n1+1 = . . . = s n2 ≥ . . . s nm = s K > s K+1 = . . . = s d . Then Σ has m + 1 blocks in total. Assume that ∥X -Σ∥ op = d α . Let X = U SV T , and definex = 4K - k:signal tr(U T (k, k)U (k, k) + V T (k, k)V (k, k) + 2U (k, k)V (k, k) T ).Thenx ≤ K ∥ Σ -X∥ op s K + ms 1 s K 2min( √ K(∥X∥ op + ∥ Σ∥ op ), ∥X∥ F + ∥ Σ∥ F ) min k (s 2 n k -s 2 n k+1 ) 2 ∥X -Σ∥ 2 opRemark. To better understand x, consider the special case where each signal singular value is of multiplicity 1. Since matrix X is "close" to diagonal matrix Σ, matrix U and matrix V should also be close to identity along signal directions:|U kk | ≈ 1, |V kk | ≈ 1, U kk V kk ≈ 1 for 1 ≤ k ≤ K. Quantity x captures how much |U kk |, |V kk |, U kk V kk deviate from 1. Proof. Let X = U SV T . Then X T X = V S 2 V T is a symmetric matrix. Let V (:, k) = (v n k-1 +1 , . . . , v n k ) be the singular vectors of X corresponding to s n k-1 +1 , . . . , s n k .There is some freedom to choose V , but for simplicity we pick V (:, k) = (e n k-1 +1 , . . . , e n k ) where e i is the i-th coordinate vector. As a result, V (:, k) T V (:, k) = V (k, k). Apply DK-sin θ theorem to V and V , we see that∥ sin Θ(V (:, k), V (:, k))∥ F ≤ 2min( √ n k -n k-1 ∥X T X -Σ2 ∥ op , ∥X T X -Σ2 ∥ F ) min(s 2 n k-1 -s 2 n k-1 +1 , s 2 n k -s 2 n k +1 ) Let σ p be the singular values of V (k, k), for 1 ≤ p ≤ n k -n k-1 .Here we are using the notation for sub-block of a big matrix. Then sin Θ(V (:, k), V (:, k)) is the diagonal matrix, whose diagonal entries are given by 1 -σ 2 p . So ∥ sin Θ(V, V )∥ 2 F = p (1 -σ 2 p ) = tr[I -V (k, k) T V (k, k)]. Now observe that ∥X T X -Σ2 ∥ op ≤(∥X∥ op + ∥ Σ∥ op )∥X -Σ∥ op ∥X T X -Σ2 ∥ F ≤(∥X∥ F + ∥ Σ∥ F )∥X -Σ∥ op We conclude that tr[I-V (k, k) T V (k, k)] ≤ 2min( √ n k -n k-1 (∥X∥ op + ∥ Σ∥ op ), ∥X∥ F + ∥ Σ∥ F ) min(s 2 n k-1 -s 2 n k-1 +1 , s 2 n k -s 2 n k +1 ) ∥X -Σ∥ op 2 .Similar conclusion is true for U (k, k). Next we bound tr(I -U (k, k)V (k, k) T ). ∥ Σ(k, k) -X(k, k)∥ 2 F = n k q1,q2=n k-1 +1 ( Σq1q2 -X q1q2 ) 2 ≤ (n k -n k-1 )∥ Σ -X∥ 2 op . ∥ Σ(k, k) -X(k, k)∥ 2 F =∥ Σ(k, k) - j U (k, j)S(j, j)V T (j, k)∥ 2 F =∥ Σ(k, k) - j̸ =k U (k, j)S(j, j)V (k, j) T -U (k, k)S(k, k)V (k, k) T ∥ 2 F ≥∥ Σ(k, k) -U (k, k) Σ(k, k)V (k, k) T ∥ 2 F -∥ j̸ =k U (k, j)S(j, j)V (k, j) T ∥ 2 F -∥U (k, k)(S(k, k) -Σ(k, k))V (k, k) T ∥ 2 F ≥s 2 n k ∥I -U (k, k)V (k, k) T ∥ 2 F - j̸ =k ∥U (k, j)∥ 2 F ∥S(j, j)V (k, j)∥ 2 op -∥X -Σ∥ 2 op ≥s 2 n k ∥I -U (k, k)V (k, k) T ∥ 2 F -s 2 1 tr(I -U (k, k) T U (k, k))tr(I -V (k, k) T V (k, k)) -∥X -Σ∥ 2 op For every n k -n k-1 × n k -n k-1 matrix M , we have tr(A) ≤ p |λ p | ≤ √ n k -n k-1 |λ p | 2 1 2 = √ n k -n k-1 ∥M ∥ F .Using this inequality, we conclude thatEQUATIONSumming on k, we conclude thatx ≤ K ∥ Σ -X∥ op s K + ms 1 s K 2min( √ K(∥X∥ op + ∥ Σ∥ op ), ∥X∥ F + ∥ Σ∥ F ) min k (s 2 n k -s 2 n k+1 ) 2 ∥X -Σ∥ 2 op In the case where d → ∞ , m is a constant, α < 1, s 1 = O(s K ), s K > cs K+1 , K is a constant, ∥X -Σ∥ op = o(∥ Σ∥ op ), and min k (s 2 n k -s 2 n k+1 ) ≥ cs 2 K for some constant c, we have x ≤ O( ∥X-Σ∥op ∥ Σ∥op ).Lemma A.5. Let X = U SV T , and definex = 4K - k:signal tr(U T (k, k)U (k, k) + V T (k, k)V (k, k) + 2U (k, k)V (k, k) T ).Let Σ be a d × d diagonal matrix whose diagonal entries are given by b 1 , . . . , b K , 0, . . . , 0,where b1 = . . . b n1 > b n1+1 = . . . = b n2 ≥ . . . = b nm = b K . Then ∥Σ -U T ΣV ∥ op ≤ ∥Σ∥ op ((m + 1) 3 √ x + 2(m + 1) 3 x) Proof. First observe that ∥Σ -U T ΣV ∥ op ≤ m+1 i,j=1 ∥Σ(i, j) -U T ΣV (i, j)∥ op . If i ̸ = j, then Σ(i, j) = 0, ∥U T ΣV (i, j)∥ op =∥ m+1 ℓ=1 U (ℓ, i) T Σ(ℓ, ℓ)V (ℓ, j)∥ op ≤b 1 m+1 ℓ=1 ∥U (ℓ, i)∥ op ∥V (ℓ, j)∥ op ≤b 1 (m + 1) √ x. If i = j, then Σ(i, i) = b ni I, ∥U T ΣV (i, i) -b ni I∥ op ≤∥b ni U (i, i) T V (i, i) -b ni I∥ op + b 1 ℓ̸ =i ∥U (ℓ, i)∥ op ∥V (ℓ, i)∥ op ≤b 1 (m + 1)x + b 1 (m + 1)x. Therefore ∥Σ -U T ΣV ∥ op ≤ ∥Σ∥ op ((m + 1) 3 √ x + 2(m + 1) 3 x).This bound is not optimal in m, but throughout the paper, m is of constant order and it is fine to miss a constant factor when estimating error.We prove the following weak bound. Proposition B.1. For every ε > 0 and every t < T , we have with high probability,EQUATIONAnalogous results holds for W 2 .Our main tool is the following lemma. Lemma B.2. For every cost C and every time t, we haveEQUATIONProof. Let L = ∥A * -A∥ 2 be the loss function. Then∂ t W 1 = 2W T 2 ∇C ∂ t W 2 = 2∇CW T 1 We see that ∂ t (W 1 W T 1 ) = 2W T 2 ∇CW T 1 + W 1 ∇C T W 2 = ∂ t (W T 2 W 2 )Next, we show that at initialization, W 1 W T 1 and W T 2 W 2 are approximately orthogonal projections, up to a factor. Stating precisely, we have the following lemma. Lemma B.3. There exists two projections P 1 and P 2 : R w → R w such that the following are true.1. The image of P 1 and P 2 are orthogonal to each other; 2. With high probability,EQUATIONEQUATIONProof of Lemma B.3 . : At initialization, the rank of wis w with probability 1. Therefore W 1 W T 1 has w -d eigenvalues that are 0, and the w non-zero eigenvalues equals the eigenvalues of d × d matrixW T 1 W 1 . W T 1 W 1 (0) is a (scaled)Wishart ensemble, whose limiting distribution is given by the Marchenko-Pastur law. The Marchenko-Pastur law, as stated in [4] , proves the following. Let X be an M × N matrix with complex-valued independent entries X iµ such that1. E[X iµ ] = 0; 2. E[|X iµ | 2 ] = 1 √ M N ; 3.for every p ∈ N, there exists a constant C p such thatE (N M ) 1 4 X iµ p ≤ C p .Here, M satisfies0 < C -1 ≤ log M log N ≤ C < ∞for some constant C independent of M and N . Let ϕ = M N , which may or may not depend on N . Then the eigenvalues of N × N matrix X * X has the same asymptotics asρ ϕ (dx) := √ ϕ 2π [(x -γ -)(γ + -x)] + x 2 dx + (1 -ϕ) + δ(dx)whereEQUATIONIn our situation, W 1 is a w × d matrix with independent and identically distributed Gaussian entries whose variance is σ 2 . Let M = w, N = d and therefore W T 1 W 1 (0) has the same distribution as σ 2 √ wdX * X. Notice that for this choice of M and N , the asymptotic distribution eigenvalues of X * X is ρ w d (dx). Notice that ρ w d is supported on intervalw d + d w -2, w d + d w + 2 ,from which we conclude that in the limit, all eigenvalues of W T 1 W 1 (0) is approximately d w .By Theorem 2.10 of [4] , we have eigenvalue rigidity results for X * X. Let λ ′ k be the k-th largest eigenvalue for X * X. ∀k ∈ {1, 2, . . . , w}, we haveEQUATIONwith high probability. Here γ α is defined throughEQUATIONLet λ k be the k-th largest eigenvalue of W T 1 W 1 (0). By the relationship between W T 1 W 1 (0) and X * X, we know λ k has the same law as σ 2 √ wdλ ′ k . Then for every k,EQUATIONwith high probability. We conclude that the first w eigenvalues of W T 1 W 1 (0) is at most O(σ 2 √ wd)away from 1 and all other eigenvalues are 0. Therefore there exists a projection P 1 such thatEQUATIONSimilarly, there exists a projection P2 such thatEQUATIONNotice that P2 is not exactly orthogonal to P 1 , and it remains to find a projection P 2 that is orthogonal to P 1 and is close to W T 2 W 2 (0). Assume that the column vectors of W 1 (0) are u 1 , . . . , u d ∈ R w and column vectors ofEQUATIONWe claim that P 1 v k is very small. By law of large numbers, ∥v k ∥ ≤ σ √ w log d with high probability.EQUATIONNotice that ⟨u j , v k ⟩u j , ∀j 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 byEQUATIONWe conclude that for each ℓ we have, by CLT,EQUATIONIn particular,P{max ℓ w i=1 |⟨u i , v k ⟩⟨u k , e ℓ ⟩| > 100σ 3 √ dw 1 2 log d} ≤wP{ w i=1 |⟨u i , v k ⟩⟨u k , e 1 ⟩| > 100σ 3 √ dw 1 2 log d} ≤2wP{N (0, 1) > 100 log d} ≤O(d -50 )Therefore with high probability, max ℓ wi=1 |⟨u i , v k ⟩⟨u k , e ℓ ⟩| ≤ 100σ 3 √ dw 1 2 log d.This implies that with high probability,EQUATION)EQUATIONNow let W ′T be the matrix with column vector v ′ 1 , . . . , v ′ w and let P 2 be the projection to the column space of W ′ 2T . By construction ∥W 2 -W ′ 2 ∥ op ≤ O(σ √ d log d) , so ∥W T 2 W 2 -W ′ 2 T W ′ 2 ∥ op ≤ O(σ 2 √ wd log d). Since the nonzero eigenvalues of W T 2 W 2 is at most O(σ 2 √ wd) from 1, we know that the nonzero singular values of W ′ 2 T W ′ 2 is also at most O(σ 2 √ wd) from 1. We conclude that ∥P 2 -W T 2 W 2 ∥ op ≤ O(σ 2 √ wd log d).Proof of Proposition B.1. : We have 3 we see that as positive semi-definite matrix, for every constant ε > 0,W T 1 W 1 W T 1 W 1 (t) = W T 1 W T 2 W 2 W 1 + W T 1 (W 1 W T 1 (t) -W T 2 W 2 (t))W 1 = A T A + W T 1 (W 1 W T 1 (0) -W T 2 W 2 (0))W 1 From lemma B.EQUATIONwith high probability. ThereforeEQUATIONBy moving terms around and corollary, we haveEQUATION)EQUATIONTheorem V.1.9 of [12] states that the square-root function is operator monotone, which implies that ifA ≥ B then √ A ≥ √ B.Taking square-root, we haveEQUATIONThe weak bound does not provide useful information if ∥W T 1 W 1 ∥ op << σ 2 w. For this reason we prove strong bound, which provide useful information if ∥W T 1 W 1 ∥ op << σ 2 w 1+□ for some constant □ > 0. Recall that the evolution of weight matrix in gradient descent is given by the following.EQUATION)EQUATION)The goal of this section is to prove thatEQUATION)For simplicity of notations, we shall assume thatC 1 = W T 1 W 1 , C 2 = W 2 W T 2 , Ĉ1 = √ A T A + σ 4 w 2 I and C 2 = √ AA T + σ 4 w 2 I.It is easy of see that Ĉ1 and Ĉ2 are invertible. Our main result for this section is the following proposition. Proposition B.4. For every cost C we haveEQUATIONProof of Lemma B.4 . : We start from the equations:W T 1 (W T 2 W 2 -W 1 W T 1 ) 2 -σ 4 w 2 I W 1 = C 3 1 -A T AC 1 -C 1 A T A -σ 4 w 2 C 1 + A T C 2 A W 2 (W T 2 W 2 -W 1 W T 1 ) 2 -σ 4 w 2 I W T 2 = C 3 2 -AA T C 2 -C 2 AA T -σ 4 w 2 C 2 + AC 1 A T .Our goal is to show that C 1 , C 2 are close to the solution Ĉ1 =√ A T A + σ 4 w 2 I, Ĉ2 = √ AA T + σ 4 w 2 I with 0 = Ĉ3 1 -A T A Ĉ1 -Ĉ1 A T A -σ 4 w 2 Ĉ1 + A T Ĉ2 A 0 = Ĉ3 2 -AA T Ĉ2 -Ĉ2 AA T -σ 4 w 2 Ĉ2 + A Ĉ1 A T .Apriori, the cubic equation for C 1 and C 2 might have multiple solutions. We give an intuitive argument explaining why Ĉ1 and Ĉ2 are the correct solutions. By selecting a proper basis, we assumeA = diag(a 1 , . . . , a d ) is diagonal. Assume that (W T 2 W 2 -W 1 W T 1 ) 2 = σ 2 wI.Also assume that C 1 and C 2 both commute with A. In this case the equations for C 1 and C 2 reduces to cubic equation for scalars. Let λ 1 , . . . , λ d be the eigenvalues of C 1 . Solving the equations for scalars, we have λ i = 0 or ± a 2 i + σ 4 w 2 . By lemma B.2 and lemma B.3, we haveW 1 W T 1 (t) -W T 2 W 2 (t) = W 1 W T 1 (0) -W T 2 W 2 (0) = σ 2 wP 1 -σ 2 wP 2 + o(σ 2 w). Since W 1 W T1 is positive semi-definite, all its eigenvalues are non-negative, and thusW 1 W T 1 (t) ≥ (σ 2 w + o(σ 2 w))P 1 .1 W 1 is the same as the top w eigenvalues of W 1 W T 1 , we conclude that λ i ≥ σ 2 w(1 + o( 1)). This forces λ i = a 2 i + σ 4 w 2 . Since C 1 , C 2 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 i := C i -Ĉi is small is the inverse function theorem.Step 1: show that LHS of the equations are small.Step 2: reduce the RHS of the equations to a linear function of dC 1 and dC 2 . The system of equations is thus reduced to small small = * * * *v T i dC 1 u T i dC 2 .The * matrix is now the "Jacobian" matrix, and u i , v i are left and right singular vectors of A.Step 3: prove that the "Jacobian" matrix * * * * is strictly positive definite, thus proving that v T i dC 1 and u T i dC 2 have small magnitude for all i.For gradient flow, W T 2 W 2 -W 1 W T 1 is preserved and there exists projections P 1 and P 2 such thatW T 2 W 2 -W 1 W T 1 = σ 2 w(P 2 -P 1 ) + O(σ 2 √ wd). Therefore for every unit vector v we have∥v T W T 1 (W T 2 W 2 -W 1 W T 1 ) 2 -σ 4 w 2 I W 1 ∥ ≤ ∥C 1 ∥ op σ 4 d 1 2 w 3 2 .Substracting the second pair of equations from the first pair and denoting dC i = C i -Ĉ, we obtain:EQUATION)EQUATIONNow sinceC 3 1 -Ĉ3 1 = Ĉ2 1 dC 1 + Ĉ1 dC 1 C 1 + dC 1 C 2 1, we substitute the above relation to equation 31 and obtainEQUATIONand similarly for equation 32. For any singular value s i of A, with left and right singular vectors u i , v i , we multiply equation 34 to the left by v T i , and divide both sides by σ 2 w, to obtain an equation for v T i dC 1 . Similarly, we obtain an equation for u T i dC 2 : for v T i dC 1 and u T i dC 2 :∥C 1 ∥ op O(σ 2 √ wd) = v T i dC 1 s i σ 2 w 2 + 1C 1 + 1 σ 2 w (C 2 1 -A T A) + s i σ 2 w u T i dC 2 A ∥C 2 ∥ op O(σ 2 √ wd) = u T i dC 2,i s i σ 2 w 2 + 1C 2 + 1 σ 2 w (C 2 2 -AA T ) + s i σ 2 w v T i dC 1 A T . Notice that C 2 1 -A T A = W T 1 (W 1 W T 1 -W T 2 W 2 )W 1 = σ 2 wW T 1 (P 1 -P 2 )W 1 + ∥C 1 ∥O(σ 2 √ wd).In the two equations above, by replacing C 2 1 -A T A with σ 2 wW T 1 (P 1 -P 2 )W 1 , we are making an error of at most ∥C 1 ∥ op ∥dC 1 ∥ op O( d w ). From weak bound we know that ∥dC 1 ∥ op ≤ O(σ 2 w). Therefore the error we made by making the approximation on the right hand side can be absorbed into left hand side.To show that ∥v T i dC 1 ∥ and ∥u T i dC 2 ∥ are small, it suffices to show that the (d in +d out )×(d in +d out ) block matrix  si σ 2 w 2 + 1C 1 + W T 1 (P 1 -P 2 )W 1 si σ 2 w A si σ 2 w A T si σ 2 w 2 + 1C 2 + W 2 (P 2 -P 1 )W T 2  is strictly positive definite. This matrix can be further simplified toW T 1 0 0 W 2   si σ 2 w 2 + 1I + P 1 -P 2 si σ 2 w I si σ 2 w I si σ 2 w 2 + 1I + P 2 -P 1   W 1 0 0 W T 2The inner matrix can then be rewritten as RR T where R is defined asR =     si σ 2 w 2 + 1 + 1P 1 + si σ 2 w 2 + 1 -1P 2 si σ 2 w 2 + 1 -1P 1 + si σ 2 w 2 + 1 + 1P 2     . Let Q = W T 1 0 0 W 2 .The "Jacobian" matrix is then QR(QR) T . As described in the strategy, we need to show that the singular values are strictly positive. The smallest nonzero singular value of QR(QR) T is the same as the smallest nonzero singular value of (QR) T QR. We expand (QR) T QR as follows.(QR) T QR= s i σ 2 w 2 + 1 + 1 P 1 W 1 W T 1 P 1 + s i σ 2 w P 1 W 1 W T 1 P 2 + s i σ 2 w P 2 W 1 W T 1 P 1 + s i σ 2 w 2 + 1 -1 P 2 W 1 W T 1 P 2 + s i σ 2 w 2 + 1 + 1 P 2 W T 2 W 2 P 2 + s i σ 2 w P 1 W T 2 W 2 P 2 + s i σ 2 w P 2 W T 2 W 2 P 1 + s i σ 2 w 2 + 1 -1 P 1 W T 2 W 2 P 1 = s i σ 2 w 2 + 1 + 1 (P 1 + P 2 )σ 2 w + 2( s i σ 2 w 2 + 1 - s i σ 2 w )(P 1 W T 2 W 2 P 1 + P 2 W 1 W T 1 P 2 ) + s i σ 2 w (W 1 W T 1 + W T 2 W 2 -σ 2 wP 1 -σ 2 wP 2 ) + O(σ 2 √ wd) = s i σ 2 w 2 + 1 + 1 - s i σ 2 w (P 1 + P 2 )σ 2 w + 2( s i σ 2 w 2 + 1 - s i σ 2 w )(P 1 W T 2 W 2 P 1 + P 2 W 1 W T 1 P 2 ) + s i σ 2 w (W 1 W T 1 + W T 2 W 2 ) + O(σ 2 √ wd).where we used the fact thatW 1 W T 1 = P 1 + P 1 W T 2 W 2 P 1 + P 1 W 1 W T 1 P 2 + P 2 W 1 W T 1 P 1 + P 2 W 1 W T 1 P 2 .The (d in + d out )-th eigenvalue of the above is lower bounded by σ 2 w si σ 2 w 2 + 1 + σ 2 wσ 2 w si σ 2 w ≥ σ 2 w. We conclude that∥dC 1 ∥ op + ∥dC 2 ∥ op ≤ O( d w )∥C 1 ∥ op ,Compared to lemma B.3, lemma B.4 gives a tighter bound on ∥C 1 -Ĉ1 ∥ op when ∥C 1 ∥ op ≤ σ 2 w d w .As suggested by an anonymous referee, it is possible to obtain the same approximated dynamics of A θ(t) by imposing a non-homogeneous balance condition (in a different setup). Assume thatW 1 W T 1 -W T 2 W 2 = 2σ 2 wI. Then C 2 1 + 2σ 2 wC 1 -A T A = 0; C 2 2 -2σ 2 wC 2 -AA T = 0. Therefore C 1 = -σ 2 w + √ A T A + σ 4 w 2 and C 2 = σ 2 w + √ AA T + σ 4 w 2 .Substituting into Gradient Flow equation, we havedA dt = -η( AA T + σ 4 w 2 ∇C + ∇C A T A + σ 4 w 2 ).The advantage of the setup is that it significantly simplifies the proof. The The limitation of the setup is thatW 1 W T 1 -W T 2 W 2 ̸ = σ 4 w 2 Iif the initial variance of entries of W 1 and W 2 are comparable. In this case, W 1 W T 1 -W T 2 W 2 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 W 1 and W 2 equal.Let A t have the following dynamics:d 2 d dt A t = (A * -A t ) A T t A t + σ 4 w 2 I + A t A T t + σ 4 w 2 I(A * -A t ).The goal of this section is to prove that the singular vectors of A t is well-aligned with the singular vectors of A * , throughout the Saddle-to-Saddle regime. In the rest of the section, we will assume the dependence of A t on t and use A to represent A t . If at initialization, A t commutes with A * , then throughout the training, A t will always commute with A * . In this section, we use a delicate stability argument to show that if A t 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 P 1 be the family of d × d matrices A that satisfies the following conditions.• s K ≥ Cσ 2 w, s K+1 ≤ C ′ σ 2 w and s K+1 s K ≤ c < 1 2 for some d-independent constants c, C and C ′ . • If a k > a k+1 , then s k -s k+1 ≥ cs k for some d-independent constant c. Define P ′ 1 to be the family of w × w matrix A such that s k -s k+1 ≥ c 2 s k if a k+1 < a k , s K+1 s K ≤ 3 4and s K ≥ cσ 2 w. Let γ > 0 be constant. Define P 2 (C, γ) be the family of w × w matrices A satisfying the following conditions.EQUATIONEQUATIONEQUATIONEQUATIONLet D = F ⊙ U T dA dt V S + SV T dA T dt U . Since F is anti-symmetric and the term in square bracket is symmetric, we know D is anti-symmetric. Then dU dt = U D. Let S = √ S 2 + σ 4 w 2 , C = U T (A * + E)V . As a result, ∀1 ≤ p, q ≤ w, dU pq dt = r:r̸ =q U pr 1 s 2 q -s 2 r [(s q sq + sr s q )C rq + ( sr s r + s r sk )C rq ] Let R pq = sp sp+sp sq s 2 p -s 2 q if p ̸ = q and R pp = 0. Then dU dt (k, k) = j:j̸ =k U (k, j) R(k, j) ⊙ C(j, k) -R(j, k) ⊙ C(k, j) T D(j, k) = R(k, j) ⊙ C(j, k) -R(j, k) ⊙ C(k, j) T .As a result, for n k ≤ p, q < n k+1 ,1 2 d dt tr[U T (k, k)U (k, k)] = p,q∈[n k ,n k+1 ) U pq r U pr D rq = p,q,r∈[n k ,n k+1 ) U pq U pr D rq + p,q∈[n k ,n k+1 ) U pq r / ∈[n k ,n k+1 ) U pr D rq = p,q∈[n k ,n k+1 ) U pq r / ∈[n k ,n k+1 ) U pr D rq =tr[U T (k, k) j̸ =k U (k, j)D(j, k)] Lemma C.6. If A [t1,t2] ⊂ P 1 and A t1 ∈ P 2 (C 4 , γ), then d 2 dx dt ≤ -cdx + O( √ dx),where ∀c = inf t∈[t1,t2] min((s * k -s * j )(s k +sj )(s k +sj ) s 2 k -s 2 j , (s k -sj )(s * k +s * j )(s k +sj ) s 2 k -s 2 j ). In particular, A [t1,t2] ⊂ P 2 (C, min(γ, 1)) for some constant C. If γ < 1 then we can take C = C 4 .Proof. Use previous lemma. Letx = 4K - k:signal tr(U T (k, k)U (k, k) + V T (k, k)V (k, k) + 2U (k, k)V (k, k) T ).x measures the alignment of singular vectors of A with singular vectors of A * . From U U T = I we haveℓ U (k, ℓ)U (k, ℓ) T = I and for every j ̸ = k, U (k, j)U (k, j) T ≤ I -U (k, k)U (k, k) T ≤ O(x)In particular, ∥U (k, j)∥ op ≤ O( √ x). We first estimate C. For all j = 1, 2, . . . , w, a j ̸ = a k ,(U T A * V )(j, k) =U T (j, j)S * (j, j)V (j, k) + U T (j, k)S * (k, k)V (k, k) + ℓ̸ =j,k U T (j, ℓ)S * ℓ V (ℓ, k) =U T (j, j)S * (j, j)V (j, k) + U (k, j) T S * (k, k)V (k, k) + O(dx) (U T EV )(j, k) = ℓ2̸ =k U T (j, j)E(j, ℓ 2 )V (ℓ 2 , k) + ℓ1̸ =j U T (j, ℓ 1 )E(ℓ 1 , k)V (k, k) + U T (j, j)E(j, k)V (k, k) + ℓ1̸ =j,ℓ2̸ =k U T (j, ℓ 1 )E(ℓ 1 , ℓ 2 )V (ℓ 2 , k) =O( √ d √ x) + O( √ d) + O( √ dx) =O( √ d) D(j, k) =R(k, j) ⊙ C(j, k) -R(j, k) ⊙ C(k, j) T =R(k, j) ⊙ (U (j, j) T S * (j, j)V (j, k)) + R(k, j) ⊙ U (k, j) T S * (k, k)V (k, k) -R(j, k) ⊙ (V (k, j) T S * (k, k)U (k, k) + V (j, j) T S * (j, j)U (j, k)) + O(dx + √ d) By U U T = I we know that ℓ U (j, ℓ)U T (ℓ, k) = 0. Therefore U (j, j)U (k, j) T + U (j, k)U (k, k) T + O(x) = 0 U (j, k) = -O(x) -U (j, j)U (k, j) T U (k, k).Similarly we have for V ,V (k, j) = -O(x) -V (k, k)V (j, k) T V (j, j).We can rewrite D(j, k) asD(j, k) =R(k, j) ⊙ (U (j, j) T S * (j, j)V (j, k) + U (k, j) T S * (k, k)V (k, k)) +R(j, k) ⊙ (V (j, j) T V (j, k)V (k, k) T S * (k, k)U (k, k) + V (j, j) T S * (j, j)U (j, j)U (k, j) T U (k, k)) +O(dx + √ d) d 2 d dt tr U (k, k) T U (k, k) =2tr   U (k, k) T j U (k, j)D(j, k)   =2 j tr U (k, k) T U (k, j) R(k, j) ⊙ (U (j, j) T S * (j, j)V (j, k) + U (k, j) T S * (k, k)V (k, k)) + 2 j tr U (k, k) T U (k, j) R(j, k) ⊙ (V (j, j) T V (j, k)V (k, k) T S * (k, k)U (k, k)) + 2 j tr U (k, k) T U (k, j) R(j, k) ⊙ (V (j, j) T S * (j, j)U (j, j)U (k, j) T U (k, k)) + O(dx 3 2 + √ d √ x) ≥ -2 j |S * (j, j)R(k, j) + S * (k, k)R(j, k)| tr(U (k, j) T U (k, j) tr(V (j, k) T V (j, k) + 2 j S * (k, k)tr[U (k, k) T U (k, j)R(k, j) ⊙ (U (k, j) T U (k, k))] + 2 j S * (j, j)tr[U (k, k) T U (k, j)R(j, k) ⊙ (U (k, j) T U (k, k))] + O(dx 3 2 ) + O( √ dx)We know that(S * (k, k)R(k, j) + S * (j, j)R(j, k)) -|S * (j, j)R(k, j) + S * (k, k)R(j, k)| =min( (s * k -s * j )(s k + s j )(s k + sj ) s 2 k -s 2 j , (s k -s j )(s * k + s * j )(s k + sj ) s 2 k -s 2 j ) ≥cdfor some positive constant c independent of w. We conclude thatd 2 d dt tr[U (k, k) T U (k, k)] ≥ cdx + O(dx 3 2 ) + O( √ dx).The same trick applies to V . We similarly haveEQUATION2 ) + cdx + O(dx 3 2 ) + O( √ dx) ≥cdx + O(dx 3 2 ) + O( √ dx)Similarly,tr[U (k, k) T d dt V (k, k)] ≥ cdx + O(dxThis proves thatEQUATIONObserve that if x ≥ d -1+ε for some ε > 0, then dx dt ≤ -c 2 dx and therefore x decrease exponentially. On the other hand, if x ≤ d -1-ε for some ε > 0, then dx dt ≥ O( √ dx) and therefore x might increase. We therefore conclude that A [t1,t2] ⊂ P 2 (C, min(γ, 1)).Lemma C.7. Assume that A [0,t2] ⊂ P 2 (C 4 , γ) and A [0,t1] ⊂ P 1 . Then A [0,t2] ⊂ P 1 .Proof. By assumption we know that t 2 = O(d log d).d 2 dS dt =I ⊙ U T (A * + E)V S 2 + σ 4 w 2 + S 2 + σ 4 w 2 U T (A * + E)V -2S S 2 + σ 4 w 2 d 2 ds K+1 dt = 2 O( √ d) + O(d 1-γ ) -s K+1 s 2 K+1 + σ 4 w 2 |s K+1 (t 2 ) -s K+1 (t 1 )| ≤ O(σ 2 w(d -1 2 + d -γ) log d) << σ 2 w It remains to verify the gap between different families. Let T be the first time when ∥A∥ op = 1 2 s * K and assume that t 2 ≤ T . Assume that a k > a j , then there exists some constant ε such thatd 2 ds k dt ≥ (s * k (1 -ε) -s k )s k d 2 ds j dt ≤ (s * j (1 + ε) -s j )(s j + σ 2 w)An application of Grönwall inequality on s k and s j implies that there exists some constant δ > 0 that depends only on a j , a k such that sj (T )s k (T ) ≤ O(d -δ). Next we deal with the case when t 2 > T . One important observation is that the inf{t > T : s 1 (t) ≥ At initialization, the singular values of A are of order σ 2 √ wd by using Mechenko-Pastur law on A T A, which is infinitely smaller than σ 2 w. As a result, there is a 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 t is almost commutative with A * . Lemma C.8. Let t 1 be the first when the first singular value of A hits σ 2 wd -δ 2 for δ = γw-1 2 . ThenEQUATION)) for some constant C.Proof. We approximate the dynamics with NTK dynamics. Assume thatd 2 dB dt = 2(A * + E -B)σ 2 wwith B(0) = A(0). Thend 2 d(A -B) dt =2(B -A)σ 2 w(A * + E -A) + ( A T A + σ 4 w 2 -σ 2 w) + ( AA T + σ 4 w 2 -σ 2 w)(A * + E -A).A simple application of Grönwall inequality implies that We first check P 2 . Clearly, t 1 is of order d 1-δ 2 . Then we have∥A -B∥ op (t) ≤ 1 d 2 e -2σ∥A t1 - t 1 d 2 σ 2 wA * ∥ op ≤2d -δ ∥A t1 ∥ op + ∥ t 1 d 2 σ 2 wE∥ op + O(σ 2 √ wd) =(O(d -δ ) + O(d -1 2 ) + O(d -δ 2 ))∥A t1 ∥ op By lemma C.4 we have x ≤ Cd -min( 1 2 , δ2 ) for some constant C.Lemma C.9. Let t 2 = cd, where c is a constant to be chosen. ThenEQUATION)).Moreover, there exists constants b ∥B∥ op (t). This guarantees that the singular values of A grows linearly in [t 1 , t 2 ], that s k -s k+1 ≥ cs k if a k > a k+1 for some constant c, that s K+1 < 1 2 s K and that A t2 ∈ P 1 . It remains to check that A t2 satisfies P 2 (C, γw-11 = . . . b n1 > b n1+1 = . . . = b n2 ≥ . . . = b nm = b K ,2 ). We are ready to use similar techniques as in lemma C.6. All computation are similar, and the only difference is that |R(k, j)| = O( σ 2 w s K ). We haveEQUATION)). Let γ = min( γw-1 4 , 1 2 ). From lemma C.8 we see that ifn k1 + 1 ≤ p 1 , p 2 ≤ n k , then |s p1 -s p2 | ≤ O(d -γ )s K .Moreover, the dynamics of s pi , i = 1, 2 are both given byEQUATION).Let b ni = s i . By lemma A.5, there exists a d × d diagonal matrix Σ, whose diagonal entries are given by b1 = . . . = b n1 > . . . = b n2 ≥ . . . = b K , such that ∥A t2 -Σ∥ op ≤ O(σ 2 w)d -γ 2 .In this section, we use a stability argument to show that if σ 2 w is infinitely larger than d, then the algorithm cannot converge. Proposition D.1. With high probability, for all time t, ∥A θ(t) -A * ∥ 2 F ≥ 1 3 min(∥A * ∥ 2 F , ∥E∥ 2 F ).Proof. In gradient flow dynamics, ∥A θ(t) -A * -E∥ F is decreasing with time and therefore for all time t, ∥A θ(t)∥ 2 F ≤ 9(∥A * ∥ 2 F + ∥E∥ 2 F ). In particular, ∥A θ(t) ∥ op ≤ 9(∥A * ∥ F + ∥E∥ F ). By assumption, ∥ Ĉ1 -C 1 ∥ op ≤ O( d w )∥C 1 ∥ op = O(d d w ), and the dynamics of A θ(t) can be written asd 2 dA θ(t) dt = 2(A * + E -A)σ 2 w + (A * + E -A)O(d d w ) + O(d d w )(A * + E -A)Assume that B(0) = A θ(0) andd 2 dB dt = 2(A * + E -B)σ 2 w.Then1 2 d 2 d dt ∥A θ(t) -B∥ 2 F = -2σ 2 w∥A θ(t) -B∥ 2 F + tr((A -B) T (A * + E -A)O(d w d )) + tr((A -B) T O(d w d )(A * + E -A)) ≤ -2σ 2 w∥A θ(t) -B∥ 2 F + ∥A θ(t) -B∥ F O(d 2 d w )This implies that for all time,∥A θ(t) -B∥ F ≤ dO( d σ 2 w d w).The dynamics of B is linear, and we haveB t = (A * + E)(1 -e -2σ 2 wt/d 2 ) + B 0 e -2σ 2 wt/d 2 . ∥B t -A * ∥ 2 F = ∥e -2σ 2 t/d 2 A * +(1-e -2σ 2 wt/d 2 )E+B 0 e -2σ 2 wt/d 2 ∥ 2 F ≥ 0.99∥e -2σ 2 t/d 2 A * +(1-e -2σ 2 wt/d 2 )E∥ 2Let P be the projection to the image of A * . Then with high probability,∥B t -A * ∥ 2 F =∥e -2σ 2 wt/d 2 A * + (1 -e -2σ 2 wt/d 2 )E + B 0 e -2σ 2 wt/d 2 ∥ 2 F ≥0.99∥e -2σ 2 t/d 2 A * + (1 -e -2σ 2 wt/d 2 )E∥ 2 F =0.99∥e -2σ 2 wt/d 2 A * + (1 -e -2σ 2 wt/d 2 )P E∥ 2 F + 0.99∥(1 -e -2σ 2 wt/d 2 )(I -P )E∥ 2 F ≥ 1 2 min(∥A * ∥ 2 F , ∥E∥ 2 F ) Since ∥A θ(t) -B∥ 2 F = o(d 2 ), we have ∥A θ(t) -A * ∥ 2 F ≥ 1 3 min(∥A * ∥ 2 F , ∥E∥ 2 F )E The Gradient Flow Dynamics of A θ(t) in Active RegimeIn section C, we proved stability for gradient flow dynamics for A t . In this section, we prove similar stability statements for A θ(t) , using the approximation results from section B. We also summarize the behavior of A θ(t) in section E.3.The goal of this section is to show that at the end of the first NTK regime, the singular vectors of A θ(t) are roughly aligned with A * . Moreover, the alignment remains to be good throughout the mixed regime. The following generalization of lemma C.6 and lemma C.7 will be useful.Proof. We use induction lemma C.3 to prove the theorem. For the first requirement, notice that for every i = 1, 2, . . . , w,d 2 ds i dt ≤ 4∥A * ∥ 2 op . If A θ(t) ∈ P 1 , then A θ([t,t+τ ]) ⊂ P ′ 1 for τ = c σ 2 √ wd d 2 ∥A * ∥ 2 op .Here c is some small constant that depends only on P 1 . This proves the first requirement. For the second requirement, we observe that A θ(t) satisfies the dynamics described in lemma E.1 at each stage, as long as ∥ Ĉ1 -C 1 ∥ op ≤ O(s K d -λ for some λ > 0. It suffices to show that there exists some positive λ independent of w such that∥ Ĉ1 -C 1 ∥ op ≤ O(d -λ s K ) in [0, T ].Let T 1 be the first time when ∥A∥ op = σ 2 w. In [0, T 1 ], all singular values are of the same order. ThereforeEQUATION).Now let T 2 be the first time when s 1 = σ 2 wd γw -1 4. In [T 1 , T 2 ], we haveEQUATION)Therefore we can pick λ = γw 4 . During [T 1 , T 2 ], we have s K (0) ≥ cσ 2 w for some constant c, andd 2 ds K dt ≥ s K ( 1 2 a K d -s K ) Moreover, since d 2 ds1 dt ≤ (2a 1 -s 1 )(s 1 + σ 2 w), we have T 2 -T 1 ≥ c log w w for some c. Therefore we have s K (T 2 ) ≥ σ 2 wd -ν for some constant ν > 0. Now in [T 2 , T ], we have weak bound ∥ Ĉ1 -Ĉ∥ op ≤ O(σ 2 w). We can therefore pick λ = ν for [T 2 , T ]. The third requirement is verified in lemma E.2.As in A t , we need to show that at the end of the NTK regime, A θ(t) must be in P 1 and P 2 . Based on what we already have for A t , this conclusion follows from Grönwall. Lemma E.4. Assume that A 0 = A θ(0) . Let T ′ 1 be the first time when∥A t ∥ op + ∥A θ(t) ∥ op reaches σ 2 w. Then ∥A T ′ 1 -A θ(T ′ 1 ) ∥ op ≤ O(σ 2 √ wd).In particular, A θ(T ′ 1 ) ∈ P 1 ∩ P 2 (C, min( γw-1 8 , 1 4 )).Proof. From dynamics of A t we see that T ′ 1 ≤ c w for some constant c. The dynamics of A t is the following.d 2 dA t dt =(A * -A t ) A T t A t + σ 2 wI + A t A T t + σ 4 w 2 I(A * -A t ).The dynamics of A θ(t) can be written as follows.d 2 dA θ(t) dt =(A * -A θ(t) ) A T θ(t) A θ(t) + σ 4 w 2 I + A θ(t) A T θ(t) + σ 4 w 2 I(A * -A θ(t) ) + O(σ 2 √ wd)(A * -A θ(t) ) + (A * -A θ(t) )O(σ 2 √ wd)Observe that∥A T θ(t) A θ(t) -A T t A t ∥ op ≤∥A T θ(t) (A θ(t) -A t )∥ op + ∥(A θ(t) -A t ) T A t ∥ op ≤σ 2 w∥A θ(t) -A t ∥ op .This implies that we have the following inequality on positive definite matrices.EQUATION).•∥A θ(t) ∥ op = σ 2 w w d 1 4 to s K (A θ(t) ) = 1 2 a K d. Let T 3 be the first time s K (A θ(t) ) = 1 2 a K d.At time T 2 , we have s K ≥ σ 2 wd δ for some δ > 0 that depends only on a 1 , . . . , a K . Then we haveEQUATION2 ) + O( √ xd 1-δ ) + O( √ dx).We conclude that x(T 3 ) = O(d -2δ ). The dynamics of s i is given byd 2 ds i dt = 2(a i d(1 + O(d -2δ )) -s i ) s 2 i + σ 4 w 2 (1 + O(d -2δ )). s i (a i d -s i ) ≤ d 2 ds i dt ≤ 2(a i d + Cd 1-2δ -s i )(s i + σ 2 w).The bounds on dsi dt implies thatEQUATION2 ) + O( √ xd γ σ 2 +γw ) + O( √ dx).Recall that the constant c in term -cdx must satisfyc ≤ 1 d min k,j:a k ̸ =aj ( (s * k -s * j )(s k + s j )(s k + sj ) s 2 k -s 2 j , (s k -s j )(s * k + s * j )(s k + sj ) s 2 k -s 2 j )As a result, we can takec = c(a 1 , . . . , a K ) = min k,j:a k ̸ =aj |a k -a j |a 2 K max k,j:a k ̸ =aj a 2 k -a 2 j Let c ′ be a large constant such that if x > c ′ (d -1 + d 2(γ σ 2 +γw-1)), then dx dt ≤ -c(a1,...,a K )d 2x. Let T 4 be the first time when x ≤ c ′ (d -1 + d 2(γ σ 2 +γw-1) ) after T 3 . Then T 4 -T 3 ≤ -2d c(a1,...,a K ) log d -1 + d 2(γ σ 2 +γw-1) + O(d). Moreover, for every t > T 4 , x cannot be larger than c′ (d -1 + d 2(γ σ 2 +γw-1) ) because dx dt < 0 if x = c ′ (d -1 + d 2(γ σ 2 +γw-1)). After T 4 , the dynamics of each singular value is given byd 2 ds i dt = 2((1 + O(d -1 + d 2(γ σ 2 +γw-1) ))a i d -s i )s i (1 + O(d 2(γ σ 2 +γw-1) ))Let T 5 be the first time after T 4 when |s i -a i d| ≤ O(d -1 + d 2(γ σ 2 +γw-1) ) log d. Then T 5 -T 4 ≤ - max(-1, 2(γ σ 2 + γ w -1)) 2a K d log d + O(d log log d).We conclude that at timeT * = 1 -γ σ 2 -γ w a K + 2max(1, 2(-γ σ 2 -γ w + 1)) c(a 1 , . . . , a K ) + max(1, 2(-γ σ 2 -γ w + 1)) 2a K d log d + O(d log log d) we have x(T * ) ≤ O(d max(-1,2(γ σ 2 +γw-1)) ), A θ(T * ) ∈ P 2 (C, max{1, 2(1 -γ σ 2 -γ w )}, and |s i (T * ) -a i d| ≤ O(d -1 + d 2(γ σ 2 +γw-1) ) log d E.To study the dynamics of A t 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 a matrix (not necessarily square matrix). To apply the lemma above we need to prove that ∥A θ(t) ∥ Proof. The gradient descent dynamics of A θ(t) is given byA θ(t+1) -A θ(t) = ηd -2 (A * + E -A θ(t) )C 1 + C 2 (A * + E -A θ(t) ). tr((A * + E -A θ(t+1) ) T (A * + E -A θ(t+1) ) -(A * + E -A θ(t) ) T (A * + E -A θ(t) )) = -2tr((A * + E -A θ(t) )) T (A θ(t+1) -A θ(t) )) + ∥A θ(t+1) -A θ(t) ∥ 2 F ≤ -2ηd -2 tr((A * + E -A θ(t) ) T (C 1 + C 2 )(A * + E -A θ(t) )) + η 2 O(d -2 ∥C 1 ∥ 2 op )From theorem 2, we know that C 1 + C 2 ≥ 1 3 σ 2 wI. Therefore tr((A * + E -A θ(t) ) T (C 1 + C 2 )(A * + E -A θ(t) )) ≥ ∥A * + E -A θ(t) ∥ 2 F 1 3 σ 2 w ≥ cd 2 σ 2 w for some constant c. Therefore for the lazy regime, we always have ∥A * +E -A θ(t) ∥ EQUATION2 ). Computing the differential of F , we obtain thatdF ([W 1 , W T 2 ]) = d -2 [ -dW T 2 (A * -W 2 W 1 ) + W T 2 (-dW 2 W 1 -W 2 dW 1 ), -dW 1 (A * -W 2 W 1 ) T + W 1 (-dW T 1 W 2 -W T 1 dW 2 )] and therefore ∥∇F ∥ F ≤ O(d -1 ). In the active regime, the total number of training steps is η We conclude that∥W f low 2 W f low 1 -W descent 2 W descent 1 ∥ F ≤(∥W descent 2 ∥ op + ∥W descent 1 ∥ op )O(ηd -1 σ √ w log d) =O(ηd -1 σ 2 w log d). ∥W f low 2 W f low 1 -W descent 2 W descent 1 ∥ F ≤ O(η 2 d -2 σ 2 w log 2 d).Recall that in lemma D.1 we proved that We now describe the experimental setup for the experiments shown in Figures 1 and 2 . For all the experiments, we used the lossesL train (θ) = 1 d 2 ∥A θ -(A ⋆ + E)∥ 2 F ; L test (θ) = 1 d 2 ∥A θ -A ⋆ ∥ 2 Fwhere E has i.i.d. N (0, 1) entries, A ⋆ = K -1/2 K i=1 u i v T i with u i , v i ∼ N (0, Id d ) Gaussian vectors in R d . This means that RankA ⋆ = K. The factor K -1/2 ensures that ∥A ⋆ ∥ F = Θ(d).To be more precise the direction in parameter space that change C1, C2 are approximately orthogonal to those that change A θ , and GD/GF only moves along the later direction.) + O( √ dx). 


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

Reviews: review: paper_summary: paper_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).
strength_weakness: strength_weakness: I find the exposition of the main result (Theorem 1) and surrounding ideas not very clear.
The 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:     
line 689: "*Let* $L = ||A^*-A ||^2$ *be the loss function*"  - not used in the proof   
line 761: "*Recall that $\tfrac{dw}{dt}=$*" - unfinished sentence   
line 773: "*To show that $∥v_i^TdC_1∥$ and $∥u_i^T dC_2∥$*" - unfinished statement
questions: 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. 

My 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: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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$?

2. 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$?

3. 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? 

In conclusion, so far the utility of the main theorems is not clear to me, I will raise my scores if my questions are addressed.
questions: questions: Please refer to the strengths and weaknesses part. 

**Minor points and typos:**
1. Equation (1) and the equation below Line 204 seems to be different in a $\eta$ factor.
2. In Line 284 should be $2 \gamma_{\sigma^2} + \gamma_w < 0$.
3. Line 939, where is Lemma G.2?
limitations: limitations: No potential negative societal impact
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
questions: 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?

- Can the theory recover other non-trivial training dynamics of linear networks, such as the saddle-to-saddle dynamics [1]?

- How can the main insights from this study be generalized to nonlinear models, such as two-layer ReLU nets? 

- Minor: There appears to be a missing '-' in the equation between line 55 and line 56.

[1] Pesme \& Flammarion. Saddle-to-Saddle Dynamics in Diagonal Linear Networks. (NeurIPS 2023)
limitations: limitations: The analysis in this study focus on 2-layer linear networks, which is limited.
paper_summary: paper_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: 
$\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 }$, 
which can be viewed as a mixture of the lazy and balanced dynamics. 
With 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.
strength_weakness: strength_weakness: 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). 

2. 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?
questions: 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.


2. See item 2 of Weakness.
limitations: limitations: See Weakness.

Comments: comment: Thank you for your detailed response. Now I have a better understanding of the theoretical contribution of this work, and I believe it is an interesting result. Thus I raise my score to 6.
Thanks for the authors' detailed responses to my questions. The authors have addressed my concerns, and I have raised my score.
Thank you for your clarifications. I will maintain my score.