Learning Dynamics of Deep Matrix Factorization Beyond the Edge of Stability
Authors: Avrajit Ghosh
Soo Min Kwon
Rongrong Wang
Saiprasad Ravishankar
Qing Qu
Venue: ICLR 2025 Poster (2025)
Keywords: edge of stability
deep matrix factorization
Primary area: optimization
TL;DR: We present a fine-grained analysis of deep matrix factorization beyond the edge of stability regime to explain key phenomena in this field.

Abstract: Deep neural networks trained using gradient descent with a fixed learning rate $\eta$ often operate in the regime of ``edge of stability'' (EOS), where the largest eigenvalue of the Hessian equilibrates about the stability threshold $2/\eta$. In this work, we present a fine-grained analysis of the learning dynamics of (deep) linear networks (DLNs) within the deep matrix factorization loss beyond EOS. For DLNs, loss oscillations beyond EOS follow a period-doubling route to chaos. We theoretically analyze the regime of the 2-period orbit and show that the loss oscillations occur within a small subspace, with the dimension of the subspace precisely characterized by the learning rate. The crux of our analysis lies in showing that the symmetry-induced conservation law for gradient flow, defined as the balancing gap among the singular values across layers, breaks at EOS and decays monotonically to zero. Overall, our results contribute to explaining two key phenomena in deep networks: (i) shallow models and simple tasks do not always exhibit EOS; and (ii) oscillations occur within top features}. We present experiments to support our theory, along with examples demonstrating how these phenomena occur in nonlinear networks and how they differ from those which have benign landscape such as in DLNs.

Published as a conference paper at ICLR 2025
LEARNING DYNAMICS OF DEEP MATRIX FACTORIZA- TION BEYOND THE EDGE OF STABILITY
Avrajit Ghosh"; Soo Min Kwon?*, Rongrong Wang', Saiprasad Ravishankar', Qing Qu? Michigan State University, 2 University of Michigan
ABSTRACT
Deep neural networks trained using gradient descent with fixed learning rate y often operate in the regime of "edge of stability" (EOS), where the largest eigen- value of the Hessian equilibrates about the stability threshold 2/1z. In this work, we present a fine-grained analysis of the learning dynamics of (deep) linear net- works (DLNs) within the deep matrix factorization loss beyond EOS. For DLNs, loss oscillations beyond EOS follow period-doubling route to chaos. We theoret- ically analyze the regime of the 2-period orbit and show that the loss oscillations occur within a small subspace, with the dimension of the subspace precisely char- acterized by the learning rate. The crux of analysis lies in showing that the symmetry-induced conservation law for gradient flow, defined as the balancing gap among the singular values across layers, breaks at EOS and decays mono- tonically to zero. Overall, our results contribute to explaining two key phenom- ena in deep networks: (i) shallow models and simple tasks do not always exhibit EOS (Cohen et al.) (02T); and (ii) oscillations occur within top features (Zhu) (tt al. (023a)). We present experiments to support our theory, along with examples demonstrating how these phenomena occur in nonlinear networks and how they differ from those which have benign landscape such as DLNs.
INTRODUCTION
Understanding generalization in deep neural networks requires understanding of the optimiza- tion process in gradient descent (GD). In the literature, has been empirically observed that the learning rate 1) plays key role in driving generalization (Hayou al. 2024 Lewkowyez et al. 2020). The "descent lemma" from classical optimization theory says that for a 8-smooth loss L(O) parameterized by θ GD iterates satisfy
L(0(1+1) ((O(1) 1/22 1 m18) (
2
and so the learning rate such that n < 2/3, then the loss monotonically decreases. How- ever, many recent works have shown that the training loss decreases even for r > 2/3, albeit non-monotonically. Surprisingly, it has been observed that learning rate beyond the stability thresh- old often provides better generalization over smaller ones that lie within the stability threshold. This observation has led a series works analyzing the behavior of GD within a regime dubbed "the edge of stability" (EOS). By letting + parameterize deep network, formally define EOS as follows:
Phenomenon (Edge of Stability (Cohen et al. 2021)). During training, the sharpness of the loss, defined as S(@)= == [[V continues to grow until it reaches 2/n (progressive sharp- ening), after which stabilizes around 2/T, During this process, the training loss behaves non- monotonically over short timescales but consistently decreases over long timescales.
Using a large learning rate to operate at the EOS is hypothesized to give better generalization per- formance by inducing "catapults" in the training loss (Zhu al., (023a). Intuitively, whenever the sharpness S(@) exceeds the local stability limit 2/1), the GD iterates momentarily diverge (or cat- apults) out of a sharp region and self-stabilizes (Damian et al. 2023) to settle for a flatter region
"Equal contribution; Correspondence ghoshavr@msu.edu. kwonsm @umich.edu

---

Published as a conference paper at ICLR 2025
1:
20p
1-22
24
Singular Oscillatiors
2.-
Elgenvalun Oscillations
200
275
1.8
223
202
0015
3.016
0.07
5.018
0.00
0.000
0.001
0.155
3.016
0.077
5,008
0.19
0.000
0.021
Learning Rate
Learning Rate
Singular Values of Weights
Eigenvalues Hessian
Figure 1: Bifurcation plot of the oscillations in the singular values (left) and the eigenvalues of the Hessian (right) of a 3-layer end-to-end DLN. The bifurcation plots indicate the existence of a period-doubling route chaos in DLNs, which analyze by examining the two-period orbit. Here, 77 > 2/3 corresponds to the EOS regime, where 3 Lo 202-2/L is the sharpness at the minima, L is the depth of the network and 0,, 1 is the first singular value of the target matrix M..
where the sharpness is below 2/1/ This ceff-stabilization mechanism enables GD to auto-regularize and find flatter solution which has shown to correlate with better generalization Keskar et al.) 2017] Izmailov et al.) 2019 Petzka et al. [202 Foret al. 2021 Gatmiry et al.) (023). Of course, the dynamics within EOS differ based on the loss landscape. When the loss landscape is highly non- convex with many local valleys, catapults may occur, whereas sustained oscillations may exist for benign landscapes. When sustained oscillations occur, the sharpness hovers about the local stability limit 2/n rather than settling to a sharpness below 2/n. We refer to this region "beyond the EOS" following existing work (Wang et al. 2023 Zhu et al., 2023b). It is of great interest to understand these behaviors within different architectures further our understanding of EOS.
From a theoretical perspective, there have been many recent efforts to understand EOS. These works generally focus on analyzing "simple" functions, examples including scalar losses (Zhu et al. 2023b Wang et al. 2023, Kreisler et al. 2023), quadratic regression models (Agarwala et al. (023). diag- onal linear networks (Even al. 2024) and two-layer matrix factorization (Chen Bruna 2023).
However, the simplicity of these functions cannot fully capture the behaviors of deep neural net- works within the EOS regime. Specifically, the following observations remain unexplained by ex- isting analyses: (i) mild (or no) sharpening occurs when either networks are shallow "simple" datasets are used for training (Caveat 2 from (Cohen et al. 2021)): and (ii) the oscillations and cat- apults in the training loss occur in the span of the top eigenvectors of the NTK (Zhu al.. (023a).
In this work, we present a fine-grained analysis of the learning dynamics of deep linear networks be partially replicated
(DLNs) beyond the EOS regime, demonstrating that these phenomena and effectively explained using DLNs. Generally, there are two lines of work for DLNs: (i) those that analyze the effects depth and initialization scale, and how they implicitly bias the trajectory of gradient flow towards low-rank solutions when the learning rate is chosen to be stable (Saxe et al.)
2014 Arora et al., 2018, 2019 Pesme Flammarion, 2023 Jacot et al., 2022), and (ii) those that analyze the similarities in behavior between linear and nonlinear networks Wang et al. 2024 Zhang al. 2024 Yaras et al. 2023). Our analysis builds upon these works to show that DLNs exhibit interesting behaviors outside the stability regime and to demonstrate how factors such depth and initialization scale contribute to the EOS regime. Our main results can be summarized as follows:
Walk Towards the Flattest Global Minima Beyond EOS. Similar to the observations made by Chen & Bruna (2023), we adopt the proof techniques of Kreisler et al. (2023) to show that the layers (or weights) of the DLN become increasingly balanced under mild assumptions at EOS. Specifically, we characterize how small the balancing gap at initialization must be for GD to reduce the imbalance the flattest minima in DLNs and our result captures an implicit regularization effect that drives the network toward the flattest minima at learning rates beyond the EOS regime.
iterations. We further show that balanced minima correspond to
2

---

Published as a conference paper at ICLR 2025
Sharpness Scales with the Network Depth. We identify all eigenvalues of the Hessian at the balanced minimum, demonstrating that the sharpness (i.e., the largest eigenvalue) scales with network depth. This rigorously validates the observations made by Cohen et al. (2021) and shows that the learning rate required to enter the EOS regime is depth-dependent, further highlighting its significance in deep networks.
Oscillations in now-Dimensional Subspaces. Once the network goes beyond the EOS regime, we prove that the network undergoes periodic oscillations within r-dimensional subspaces in DINs, where is precisely characterized by the learning rate. For DLNs, a period-doubling route to chaos (Oml 2002) exists in both the singular values of the DLN and the eigenvalues of the Hessian, shown in Figure [
We characterize the case of the two-period orbit, aiming to
contribute to explaining the observations made by Zhu et al. 2023a) Cohen et al. (2021)
2 NOTATION AND PROBLEM SETUP
Notation. We denote vectors with bold lower-case letters (e.g.. x) and matrices with bold upper- case letters (e.g., X). We use I to denote an identity matrix of size N. We use [L]] denote the set {1, 21, L.J. We also use the notation Te, denote the i-th singular value of the matrix We.
We use the notation (r(A) ) to denote the i-th singular value of the matrix A.
Deep Matrix Factorization Loss. The objective in deep matrix factorization is to model a low- rank matrix M. Raxd with rank(M,)) = r via a DLN parameterized by a set of parameters θ = (1, W2, Wt), which can be estimated by solving
argmin f()) = W
W M.
where we adopt the abbreviation W ja W W; to denote the end-to-end DLN and is identity when j < i. We assume that each weight matrix has dimensions We E Rddd observe the effects of overparamettriization. We also assume that the singular values of M, are distinct.
Optimization. Each weight matrix We E ddd
xd is updated using GD with iterations given by
(x(1) We(1 - 1) 7 Vw. (@ 1)), Vl [Z],
(2)
where 7) > 0 is the learning rate and Vw, is the gradient of f(O) with respect to the f-th weight matrix the 1-th GD iterate.
Initialization. In this work, consider both balanced and unbalanced initializations, respec- tively:
We(0) aLd, [eE(L) W1( = O, Wc(0) - = ala. VL -1
(3)
(4)
where a > is a small constant. We assume a is chosen small enough such that a E (0, o..r).
where or, is the r-th singular value of M. Generally, many existing works on both shallow and deep linear networks assume a zero-balanced initialization (i.e., W. (0)
W0) (0)
for i + :). This introduces the invariant W (t)) = w W,( (/WT for all t > 0, ensuring two (degenerate) conditions throughout the training trajectory: (i) the intermediate singular vectors of each of the layers remain aligned and (ii) the singular values stay balanced. For the unbalanced initialization, the zero weight layer can be viewed as the limiting case of initializing the weights with a (very) small constant a' & o, and has been similarly explored by Varre et al. (023) and Xu et al. 2024), albeit for two-layer networks. The zero weight layer relieves the balancing condition of the singular values. Rather than staying balanced, we show that the singular values become increasingly balanced (see Proposition 2 This allows us jointly analyze the singular values of the weights.
Nevertheless, we also show that our analysis is not limited to either initialization but applies to any initialization that converges to a set call the singular vector stationary set (see Proposition I]. To the best of our knowledge, it is common to assume that the singular vectors remain aligned, as
many existing works make the same assumption (Varre et al. 2023 Arora et al. 2019 Saxe al.
3

---

Published as a conference paper at ICLR 2025
Edge of Stability
10a
2/n
Train Loss
Sharpness
Progressive Sharpening
Progressive Sharpening
104
Edge Stability
10a
30
60
70
30
60
70
Iterations
Iterations
Figure 2: Depiction of the two phases learning in the deep matrix factorization problem for a network of depth 3. Left: Plot of the training loss undergoing saddle jumps, followed by periodic oscillations. Right: Plot the corresponding sharpness of the DLN. Upon escaping the first saddle point, the GD iterates enter the edge of the stability regime, where the sharpness hovers just about 2/n-
2014 Gidel et al., 2019, Chou et al. 2024b)] Min Kwon et al. 2024). Hence, throughout the rest of this paper, we refer to the balancing gap the difference singular values across layers and clarify where necessary.
3 DEEP MATRIX FACTORIZATION BEYOND THE EDGE OF STABILITY
When using large learning rate, the learning dynamics can typically be separated into two distinct stages: (i) progressive sharpening and (ii) the edge of stability. Within the progressive sharpen- ing stage, the sharpness lies below 2/1 and tends to continually rise. Our goal is to analyze the FOS stage under the deep matrix factorization formulation. Here, we observe that the training loss fluctuates due to layerwise singular value oscillations, as illustrated in Figure 2
3.1 ASSUMPTIONS AND ANALYTICAL TOOLS
Before we present our main result, we introduce two key analytical tools used in our analyses: the singular vector stationary set and singular value balancedness. First, introduce the singular vector stationary set, which allows us to consider a wider range of initialization schemes. This set defines broad class of weights for which singular vector alignment occurs, simplifying the dynamics of weights to those that only involve singular values.
Proposition (Singular Vector Stationary Set), Consider the deep matrix factorization loss in Equa- tion ) Let M. = U.E,V V and (z(t) = Uz(t) (t) denote the compact SVD for the target matrix and the t-th layer weight matrix at time respectively. For any time Z 0, if Úe(t) Vi(t) 0 for all E [L], then the singular vector stationary (SVS) points for each weight matrix are given by
(Uz, VL) = =(U.Qu)),
SVS(S())
(U.. Vo) =QQ411, Qe),
VE 12. 1),
(U1, =Q2, V.),
where (Qelize can be any orthogonal matrices.
The singular vector stationary set states that for any set of weights where the gradients with respect the singular vectors become zero, the singular vectors become fixed points for subsequent itera- tions. Once the singular vectors become stationary, running GD further isolates the dynamics on the singular values. Hence, throughout analysis, re-write and consider the loss
(t) M1
EL
Σ (a. (o.(EL -
(5)
where Ena1 are the singular values of W f.1- This allows us to decouple the dynamics the singular vectors and singular values, focusing the periodicity that occurs in the singular values within the
4

---

Published as a conference paper at ICLR 2025
Singular Values (5mell (2)
Singular Values (Large ())
Right Singuler our Vectors
Singuler Vectors
Figure 3: Illustrations of the singular vector and value evolution of the end-to-end DLN starting from the unbalanced initialization. The singular vectors of the network remain static across all iterations, as suggested by the singular vector stationary set, regardless of the learning rate. The angle between the true singular vectors and those of the network remains aligned throughout, The first singular values undergo oscillations the large 1) regime, whereas they remain constant in the small 1y regime.
EOS regime. In Propositions and we prove that both the unbalanced and balanced initializations considered here belong to this set respectively, with an illustration in Figure Specifically, we show that the balanced initialization in (3) belongs to the singular vector stationary set for all t 2 0, while the unbalanced initialization in (4) belongs the set for all 2 (far before entering the EOS
regime) with Q = V.. allowing to consider the loss in Equation (5).
Equipped with the loss in (5). notice that the balanced initialization in Equation makes the learn- ing dynamics such that for all t O,
a.( (W) = a. (W)()), Vi [d]. Ve, E [Z]].
where o (We) denotes the i-th singular value of the E-th layer. This allows us to couple the dynamics and analyze the behavior of a single variable: o, (E.1 (t)) = of However, this is certainly not
the case for the unbalanced initialization in Equation 4 Since the singular values of the W, Wz(0) are initialized to zero, there a non-negligible gap of
between the singular values of the L--th layer and the other layers. However, in the following result, we prove that as long as is small, GD will monotonically decrease the gap to balance the singular values across layers at the EOS. This will allow to couple the dynamics in the limiting case for the unbalanced initialization as well.
Proposition 2 (Balancing of Singular Values). Let o, and ot,((t) denote the i-th singular value 2-2. of M. E dxd and W We(t), respectively and define S; = Lo Conssdder GD the i-th index of the simplified loss in with the unbalanced initialization and learning rate g. < S.
2x2 If
1/4
A/L
the initialization scale a satisfies <a<<<
In
2v2 95.
.thn there exists a constant
12.
e E (0, such that for all [T. we have oZ.
1) of, +
o2.(0) - 20,
The proof is available in Appendix C.2.2 This result has been shown to hold similarly for two-layer
matrix factorization (Wang et al. 2022a Ye & Du 2021 Chen & Bruna 2023), and our analysis extends it the deeper case. Precisely, considers the scalar loss for singular value index and states that, as long as a is chosen below a threshold dependent layers will become increasingly balanced. the steady state limit of EOS, this balancing gap will decrease to zero. While this result is presented as tool for the main result, also has interesting implications for the dynamics of GD at EOS.
o..i, the i-th singular value across
Firstly, it is well known that for gradient flow (GF), the balancing gap is conserved throughout its trajectory (see Lemma and Figure 4ap. While GD with small learning rates approximately conserves the gap, GD the EOS (i.c., learning rate close but below the stability limit) breaks this conservation. However, for a learning rate below the stability limit f())ll, the GD iterates may converge unbalanced global minimum (e.g., Figure 4b and the balancing gap will
5

---

Published as a conference paper at ICLR 2025
m
24-
Rraectar
6o Tasjemay
63 Fa
21-
42
41
12
1 15
24 20
24 14
24
(a) GF and Stable GD
(b) GD EOS
(c) GD Beyond EOS
Figure 4: Illustration of the GD trajectories for three different learning rates regimes for minimizing the function f(11, 02) [d2 01 o.)², starting from unbalanced initial point. Gradient flow conserves the balancing gap (t) (t) throughout its trajectory. GD at EOS decreases the gap. but stagnates once the oscillations no longer occur. GD beyond EOS decreases the gap monotonically to zero by oscillating towards and about the balanced minimum.
cease to decrease any further. In contrast, for GD beyond the stability limit (i.e., beyond the EOS),
Proposition 2 states that the balancing gap decreases monotonically to zero. This highlights a key distinction between the learning dynamics of GF, GD, and GD at EOS. To illustrate these differences, we consider toy example of minimizing a two-layer scalar function f(01, = = -0 - -o.²² where 0. = in Figure 4 Once stable GD arrives at an unbalanced global minimum, settles there and the balancing gap do not decrease further. For GD just below the stability limit (Figure Ab), the iterates oscillate, but once they cease oscillating and settle at an unbalanced global minimum, the gap also stagnates. On the other hand, GD beyond EOS (Figure 44) drives the balancing gap strictly to zero, as the GD iterates oscillate toward and around the balanced minimum.
Secondly, for deep matrix factorization, we prove that the balanced minimum (i.e., the minimum where all of the sin- gular values across layers are the same) corresponds to the flattest minimum (see Lemma J. Since GD at EOS mono- tonically decreases the balancing gap. this also implies that GD implicitly walks from a sharper minima to the flattest minima. This also suggests an algorithmic trick: one can initially use a large learning rate to oscillate toward a flatter region and subsequently decrease the learning rate to settle at flat minimum, also highlighted by Chen & Bruna 2023).
10"
10:2
10:4
10
10-4
Gf [Crestant)
Statle CD
10-0
6D S1
6D Beyend 505
169
200
509
Iteratio
Figure 5: Plot of (o] - (1) on a toy example, showing a decaying bal- ancing gap beyond EOS.
Next, note that if the constant in Proposition were to be strictly e < 1, by Lemma 10, the gap would approach zero infinitesimally. Our analysis shows the existence of in two cases: (i) 0; (Eta) < о., and (ii) (Eta) > o..i- While we provably show that e < for the first case, we have that e 1 for the second case. This implies that when the GD iterates are below (o, (E.1)) < o.,i) and approaching the minima, the balancing gap will monotonically decrease, but not guaranteed to decrease to zero when we overshoot above the minima (o, (Eaa) > o.,)). However, note that in the EOS regime, we oscillate below and above the minima as shown in Figure (since for the case of (Eaa) two cases, and hence, the balancing gap will overall decrease to zero as depicted in Figure 5 Since oscillations do not occur or are not sustained in GF and stable GD, the gap does not go to zero in most cases, making this distinct characteristic of GD beyond EOS.
Ja we have e < 1). This indicates that we alternate between the
Finally, we remark that Proposition 2 considers only the loss of single singular value index, whereas Equation 5 is the sum over multiple indices. For Proposition 2 to hold for all indices, we can simply choose a with σ.,1 such that it is the smallest a satisfying the condition for all singu- lar values a...- To this end, in the following sections, we rigorously analyze the behavior of singular value oscillations around the balanced minimum. This can be viewed as the behavior of GD in the steady-state limit, Proposition 2 implies that the singular values become balanced as co.
6

---

Published as a conference paper at ICLR 2025
Rank-1 Oscillation
Rank-2 Oscillation
Rank-3 Oscillation
Magnibuude
Magnitude
399 150 Iteration (x10)
399
150
3K 150
Iteration (x10)
Iteration (x10)
3
2
2
2
81
8g
82
8g
Figure 6: Evolution of the singular values of the end and-to-end 3-layer network for fitting a rank-3 - target matrix with singular values 10, 9.5, and 9. We learning rate of 1y = 2/S, with S; = Lo 2 2/L The oscillations occur as a two-period orbit about the balanced minimum exactly with learning rate ranges specified in Theorem for rank-p oscillations (p 1,)
3.2
MAIN RESULTS
Using analytical tools, we present our main results describing the learning dynamics of DLNs about the balanced solution beyond the EOS. First, we present a result characterizing the set of all eigenvalues Xe of the DLN with respect the flattened Hessian of the training loss.
Lemma (Eigenvalues of Hessian at the Balanced Minimum). The set of all non-zero eigenvalues of the training loss Hessian of the deep matrix factoriiaiio loss f()) defined in Equation ① at the balanced minimum is given by
2- ee LLO "
2-2
U
_
(
L-1
(1.7
U
U
1- be at)
itj
where Txi is the i-th singular value of the target matrix M. Rdxdd a E R. is the initialization scale, L is the depth of the network, and the second element of the set has multiplicity of
The proof deferred to Appendix C.3.2 Let X denote the i-th largest eigenvalue of the Hessian.
By Lemma observe that the sharpness is equal A / Lo the balanced minimum. In Lemma 4 we show that among all the points on the global minima, the sharpness at the balanced minimum is the smallest. Thus, if 1) is set such that 17 > 2/ oscillations in the loss will occur, the step size is large enough to induce oscillations even in the flattest region. Notice that this was alluded to Figure f GD beyond EOS (i.e., when v > 2/x1), there is stable oscillation around the minima, whereas for GD at EOS, the iterates eventually settle down after transient oscillations. Furthermore, notice that all non-zero eigenvalues are a function network depth. For a deeper network, the sharpness will be larger, implying that a smaller learning rate be used to drive the DLN into EOS. This provides a unique perspective on how the learning rate should be chosen networks become deeper and explains the observation made by Cohen al.] (2021), who observed that sharpness scales with the depth of the network. With the eigenvalues, we show in the following result that oscillations occur in a two-period orbit about the balanced minimum within a rank-p subspace, where the rank is dependent on the learning rate.
U.E.V V denote the SVD of Sp ) the target matrix and define Sp = =Lo 2- and K2 max {SptI, 2v .. If we run GD on the deep matrix factorization loss with learning rate 1 = R. where KG K < Sp then the top-p singular values of the end-to-ene DLN oscillate in a 2-period orbit (j (1,2) 2)) around the balanced minimum and admit the following decomposition:
Theorem 1 (Rank-p Periodic Subspace Oscillations). Let M,
W = M u.
Σ 1=p+1
o U. j (1,2)
(6)
oxcillation subspacc
stationary subspace

---

Published as a conference paper at ICLR 2025
10
01
02
Oscillation Range
Balancing Gap
10-11
03
10ma
10-3
10ma
10-4
450
600
1000
1200
1400
0.031
0.032
0.033
0.034
0.035
0.036
Iteration (x10)
Learning Rate
Figure 7: Experimental results on a depth-3 DLN with target singular values 10, 9.5, 9. Left: Plot of the balancing gap decaying monotonically to zero the learning rate is chosen 1j > 2/Sz. Right: Plot of the oscillation range as function of the learning rate. As the learning rate increases, the oscillation ranges also increase.
1/L (0.011) and Pa. (ot!1 (20r,)) g(pt) and
where Pt,
are the two real roots of the polynomial
s(p:) = - 1 [1 (1+ mL(.,, - -pt) 2)-2)) 26-1
1 1 (1 + ( pt) p/--2) 6-1 0+,5.
Theorem I Teplicitly identifies the subspaces that exhibit two-period orbit based the range of the learning rate. It also provides rough characterization of the oscillation amplitude, which is determined by P;, and p1,2-values below and above the balanced minimum, respectively. Since there is no closed-form solution for an arbitrary higher- order polynomial, Pt.1 and Pt,2 are defined as solutions to the polynomial g(p:). Overall, this aims theoretically explain why (i) oscillations occur primarily within the top subspaces of the network, as observed by Zhuu et al. (2023a), and (ii) oscillations are more pronounced in the directions of stronger features, as measured by the magnitudes of their singular values.
The proof is available in Appendix C.3.3
Notice that the range of the learning rate depends on the eigenvalues of the form Sp = Lo 2-2/L rather than all eigenvalues Lemma[]] This is because the eigenvectors associated with the other eigenvalues are orthogonal to the weights of the DLN at the balanced minimum, so oscillations will never occur in those particular eigendirections. They are only orthogonaa in the directions of the eigenvalues of Sp and, hence, oscillations occur only in those specific directions.
We also remark that our result generalizes the recent theoretical findings of [f & Bruna 2023) where they proved the existence of certain class of scalar functions fxx) for which GD does not diverge even when operating beyond the stability threshold. They demonstrated that there exists a range in which the loss oscillates around the local minima with a certain periodicity. These oscil- lations gradually progress into higher periodic orbits (e.g., 2, 4, periods), transition into chaotic behavior, and ultimately result in divergence. In our work, prove that this oscillatory behavior beyond the stability threshold also occurs in DLNs.
4 EXPERIMENTAL RESULTS
4.1 SUBSPACE OSCILLATIONS IN DEEP NETWORKS
Firstly, we provide experimental results corroborating Theorem We let the target matrix be M, E 500 with rank 3, with dominant singular values a 10, 9.5, For the DLN, we consider 3-layer network, with each layer as W, R50x50 and use an initialization scale of = 0.01. In Figure 6 we present the behaviors of the singular values of the end-to-end network under different learning rate regimes. Recall that by Theorem [
the i-th singular value undergoes periodic oscillations when K set to be S, < K < S1+1, where S, = 2-2/L. Figure 6 illus- trates this clearly- we only observe oscillations in the i-th coordinate depending the learning rate.
8

---

Published as a conference paper at ICLR 2025
Interestingly, notice that 02 also begins to oscillate in the rank-1 oscillation region before settling at minimum. This occurs because, while the learning rate is large enough to catapult around an unbalanced minimum, it is not sufficiently large to induce periodic oscillations at balanced minima.
Secondly, in Figure 7 we present an experiment demonstrating the relationship between the os- cillation range and the learning rate by plotting the amplitude of singular value oscillations in the end-to-end network, as well as the balancing gap, to corroborate Proposition 2 in DLNs. Clearly, in Figure 7 (left), we observe that the balancing gap decays monotonically to zero long as the learning rate is chosen such that periodic oscillations occur in the top-3 subspaces. In Figure 7 (right), the oscillations begin occur starting from each region rj = 2/S,, and the oscillation range (or amplitude) increases as the learning rate iccreases This can also be observed in Figure [6| the amplitude of σ1 increases as we move from the rank-1 to the rank-3 oscillation region.
4.2 SIMILARITIES AND DIFFERENCES BETWEEN LINEAR AND NONLINEAR NETS AT EOS
Mild Sharpening. "Mild" sharpening refers to the sharpness not rising to 2/n throughout learning,
and generally occurs in tasks with low complexity as discussed in Caveat 2 of (Cohen et al. 2021
We illustrate mild sharpening in Figure T0, where we plot sharpness in two settings: (i) regression with simple images and (ii) classification with an MLP using a subset of the CIFAR-10 dataset.
For the regression task, we minimize the loss L(O)
[|G(@) Yimagells, where G(@) is UNet parameterized by O, and Yimage denotes one of the images in Figure 1Ob] We observe that when Vimage is a smooth, low-frequency image, the sharpness of the loss generally remains low. However, when Yimage has higher frequency content, the sharpness in- creases and enters the EOS regime (Figure [0a]. Similarly, for the classification task, we train a 2-layer fully connected neu- ral network on N labeled training images from the CIFAR- 10 dataset using MSE loss and plot the sharpness in Figure T0c. The sharpness links to N, the number of data points used for training. For small N values, such as 100 200, the network learns only a limited set of latent features, resulting in mild sharpening, and does not reach the EOS threshold. However, when N exceeds 1000. the sharpness increases and reaches the EOS threshold. Similar observations can also be seen in DL.Ss In Figure g
Sharpness
325
Loi?"
Iteration
Figure 8: DLNs do not enter EOS regime if Loj < 2//11
we show that the sharp- 2- 2/1,
ness reaches T. where ax1 is the singular value of the target matrix. Whenever T the network will not enter the EOS regime. This can be viewed as low-complexity learning, as o-,1 corresponds to the magnitude of the strongest feature of the target matrix. Hence, when σ.,1 is not large enough, the sharpness will not rise to 2/1). While this cannot fully explain mild sharpening, our experiments demonstrate that interpreting sharpness as a measure of complexity, combined with our findings from DLNs, marks an important first step toward fully understanding this phenomenon.
Difference in Oscillation Behaviors. Here, we discuss the differences in os- cillations that arise in DLNs compared to catapults that occur in practical deep nonlinear networks. The main differ- ence lies in the loss landscape-at at con- vergence, the Hessian for DLNs is posi- tive semi-definite, as shown in Lemma meaning there are only directions of pos- itive curvature and flat directions (in the null space of the Hessian). Moreover, the loss landscape of DLNs are known to be benign since they do not contain any spurious local minima, but only saddle points and global minima (Kawaguchi
Figure 9: Loss landscape of the Holder table function and DLNs, respectively (left-right). The Holder table func- tion is non-convex which allows catapulting to other min- ima, whereas DLNs do not have spurious local minima.
(2016)). In this landscape, oscillations occur because the basin walls bounce off, without the di- rection of escape. However, in deep nonlinear networks, it has been frequently observed that the

---

Published as a conference paper at ICLR 2025
Sharpnoss (tigemntoon) [P(]
200
Shargness
100
N- 103
50
203
N- 1002
210
100
1000
1500 2000
(a) Sharpness plots for training im-
(b) Target images (denoted as (c) 2-layer FC network trained
age generator networks using SGD
Yinage) with different frequencies used for training.
with small number N of CIFAR-10 datasct with 1) 10-2
with learning rate 77 2 10
Figure 10:
Illustration of Caveat by Cohen et al. (2021) on how mild sharpening occurs on simple datasets and network. (a) Regression task showing the evolution of the sharpness when UNet (with fixed initialization) is trained to fit a single image shown in (b). (c) Evolution of the minimal progressive sharpening on a classification task of a 2-layer MLP trained on subset of CIFAR-10.
Hessian at the minima has negative eigenvalues (Ghorbani et al. 2019 Sagun et al. 2016). This enables escape direction along the negative curvature, preventing sustained oscillations.
In Figure 9 we demonstrate these two differences by visualizing the loss landscapes and the iterates throughout GD marked in red. The Holder table function Figure 9 (left) exhibits numerous local minima, causing the loss to exhibit sharp "catapult" when a large learning rate is used. In contrast, for DLNs (shown in the right) the loss oscillates in periodic orbit around the global minima since there are no spurious local minima (Ge et al., 2016 Kawaguchi
2016
Lu & Kawaguchi, 2017)
Zhang 2019
Yun al.
2019).
Lastly, Damian et al. (2023) study seff-stabilization, where sharpness decreases below 2/n after ini- tially exceeding 2/1z. Their analysis requires assumptions such as VL(()) uu0) 0 and (S))) lies in the null space of the Hessian, where S(0) and u(O) denotes the sharpness and its corre- sponding eigenvector, respectively. These assumptions do not hold exactly in DI.Ns. Rather, the sharpness oscillates about 2/1 as shown in Figure 2 the condition for stable oscillation holds along each cigenvector of the Hessian. The alignment of (())). u(0) and (() determines the nature of oscillations in deep networks its absence thereof. This alignment usually depends the symmetry of the parameter space and can usually vary across different architectural compo- nents. This work deals with deep linear networks which has rescaling symmetry. however several other symmetries (Kunin et al. (2021) ) can be induced by softmax operator (translation symmetry) or batch-normalizatioo (scaling symmetry), which may further affect these alignments. We leave this study for future work.
5 CONCLUSION AND LIMITATIONS
In this paper, we presented a fne-grained analysis of the learning dynamics of deep matrix factor- ization beyond the EOS, where our analysis revealed two-period orbit within small subspace around the balanced minimum. We showed that as long as oscillations were sustained, the balanc- ing gap. defined as the difference singular values across layers, decreases monotonically to zero. For DLNs, since the flattest minima correspond to the minima where all weights are balanced, this suggests an implicit walk toward flat minima without any explicit regularization, which is distinct characteristic of GD beyond the EOS. Our results also contributed to understanding unexplained phenomena nonlinear networks within EOS, such as mild sharpening or oscillations small subspace.
Since our analysis focuses on the behavior of DLNs around the balanced minima starting from an unbalanced initialization, it technically describes a steady-state limiting behavior of GD once the singular values become balanced. To fully capture the learning dynamics of DLNs, it is great interest to derive the complete dynamics at EOS, where oscillations occur but are not sustained. Furthermore, we focused on cases where the weights belonged to the singular vector stationary set, allowing to isolate the behavior of singular vectors from singular values. is also of great interest to account for how singular vectors align beyond the EOS regime, as this is currently beyond the scope of our paper.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
AG, SR and RR acknowledge support from NSF CCF-2212065. QQ and SMK acknowledges sup- port from NSF CAREER CCF-2143904, NSF CCF-2212066, and NSF IIS 2312842.
We thank Arthur Jacot, Nicholas Flamarrion, Sadhika Malladi, Rene Vidal and Abhishek Panigrahi for their valuable feedback. We also grateful Sungyoon Lee for technical discussions on balancing proof. Special thanks are extended to Molei Tao for his thoughtful contributions, which clarified the nuances between the EOS and sustained oscillations. We appreciate Eshaan Nichani for directing our attention to the work of Kreisler et al. (023), which was instrumental in our analysis of balancing using gradient flow sharpness. Finally, we thank Jeremy Cohen for his support, from early email exchanges in 2023 to in-person discussions at NeurlPS 2023, regarding self seffsstabilization and mild sharpening.

---

Published as a conference paper at ICLR 2025
REFERENCES
Atish Agarwala, Fabian Pedregosa, and Jeffrey Pennington. Second-order regression models exhibit pro- gressive sharpening to the edge of stability. International Conference on Machine Learning. vol- ume 202. pp. 169-195. PMLR, 23-29 Jul 2023. URL htttps://proceed.nngs wwwpinnssmm....p...m mlr www.sppecscccc agarwa 1 a23b. htm.1
Kwangjun Ahn, Sébastien Bubeck, Sinho Chewi, Yin Tat Lee, Felipe Suarez, and Yi Zhang. Learning threshold neurons via edge of stability. Advances in Neural Information Processing Systems, 36, 2024.
Ismail Alkhouri, Xitong Zhang, and Rongrong Wang. Structure Struecure-preserving network compression via low-rank induced training through linear layers composition. Transactions on Machine Learning Research, 2024. ISSN 2835-8856.
forr/f d- KCrVMJoJ9
Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparamelerization. In International Conference Machine Learning, pp. 244-253. PMLR, 2018
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep ma-
trix factorization. In Advances in Neural Information Processing Systems, volume 32, 2019.
https::/proceed neurips. cc..ss.cceeee cc/paper Files/papez/ URL G0013335foo010001111166666662223
Sanjeey Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on the edge of stability in deep learning. In International Conference on Machine Learning, pp. 948-1024. PMIR, 2022.
Lei Chen and Joan Bruna. Beyond the edge of stability via two-step gradient updates. In International Conference Machine Learning, volume 202, pp. 4330-4391, PMLR, 23-29 Jul 2023. URL https:
//proceedings. m
Xuxing Chen, Krishnakumar Balasubramanian, Promit Ghosal, and Bhavya Agrawalla. From stability to chaos: Analyzing gradient descent dynamics in quadratic regression. arXiv preprint urXiv:2310.01687, 2023.
Hung-Hau Hsu Chou, Carsten Gieshoff, Johannes Maly, and Holger Rauhut. Gradient descent for deep matrix factorization: Dynamies and implicit bias towards low rank. Applied and Computational Harmonic Analysis, 68: 101555, 2024a. ISSN 1063-5203. doi: https://doi.or//1100 awt.///..........aaaaaaaaaaaaaaa 2023 101595. URL. (ttpss//
www. www.cceeieeiiccccc rect. com/ scwence/eeeeccccccccoccecccccece caance/artiia aataaatticleepee
Hung-Hsu Chou, Carsten Gieshoff, Johannes Maly, and Holger Rauhut. Gradient descent for deep matrix
factorization: Dynamics and implicit bias towards low rank. Applied and Computational Harmonic Analysis, 08:101595, 2024b.
Jeremy Cohen, Simran Kaur, Yuanzhi Li, Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations, 2021. URL httL/tttts////pceeee//
net/ wotp://oormmddddomm jh- ITtvkGeM
Alex Damian, Eshaan Nichani, and Jason D. Lec. Self-stabilization:: The implicit bias of gradient descent at the edge of stability. In The Eleventh International Conference Learning Representations, 2023. URL httpp:: httpts:///penreview... n forum 1d=nkkA559ggg
Mathieu Even, Scott Pesme, Suriya Gunasekar, and Nicolas Flammarion. (S)GD over diagonal linear networks: Implicit bias, large stepsizes and edge of stability. Advances Neural Information Processing Systems, 36, 2024.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness aware aware minimization for ef- ficiently improving generalization. International Conference on Learning Representations, 2021. URL https://cpenreeveew wetp:/ewwwweeeeeeeeeeeeeeeeeeeee Eorum? 2d= stt.:ttt.t......................
Khashayar Gatmiry, Zhiyuan Li, Tengyu Ma. Sashank Reddi, Stefanie Jegelka, and Ching Yao Chuang. What the inductive bias flatness regularization? A study of deep matrix factorization models. Thirty- seventh Conference on Neural Information Processing Systems, 2023. URL
n feremfrrmm?
Rong Ge, Jason D Lee, and Tengyu Ma. Matrix completion has no spurious local minimum. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 29, 2016. URL htttss://proceeedineeee neuri..pp.....nn............... wwtppppppppppeeeeeeeeeeeeee.eeep tatpapp__eleeepppprrrrrrpe wtl1//iile/fbsceb33. (tt.://eeeddcco15..0666666666660
Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net optimization via hessian cigenvalue density. In International Conference on Machine Learning, pp. 2232-2241. PMI.R, 2019.
12

---

Published as a conference paper at ICLR 2025
Gauthier Gidel, Francis Bach, and Simon Lacoste-Julien. Implicit regularization of discrete gradi- ent dynamics in lincar neural networks. Advances in Neural Information Processing Systems, volume 32, 2019, URL
neur http:ppprr..leeeppppr///0000
pdd
Daniel Gissin, Shai Shalev-Shwartz, and Amit Daniely. The implicit bias of depth: How incremental learning drives generalization. In International Conference on Learning Representations, 2020. URL https:
7/openreview www//eettuuu HdterrmmmmdHH111eeee
Soufiane Hayou, Nikhil Ghosh, and Bin Yu. LoRA+: Efficient low rank adaptation of large models. In Forry- first International Conference on Machine Learning, 2024. URL BBLLattpss/ppeeeevvvv...eeee
id=NEv3YgRROO
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. arXiv preprint arXiv:1803.05400 2019. URL
https://arxiv. www.wwwwwwwww..r.../oo//a// 1803 05407
Arthur Jacot, François Ged, Berfin Simşek, Clément Hongler, and Franck Gabriel. Saddle-to-saddle dy- namics in deep linear networks: Small initialization training, symmetry, and sparsity. arXiv preprint arXiv:2106.1593: 2022.
Stanislaw Jastrzebski. Maciej Szymezak, Stanislav Fort, Devansh Arpit. Jacek Tabor, Kyunghyun Cho and Krzysztof Geras*. The break-even point on optimization trajectories of deep neural networks. Inter- national Conference on Learning Reprexentations, 2020. URL RRLL:ttttps///epeeveee...o Eerorum
id-r Lg87C4KWB
Stanistaw Jastrzebski, Zachary Kenton, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amost Storkey. On the relation between the sharpest directions of DNN loss and the SGD step length. In Interna- tional Conference Learning Representations, 2019. URL htttss://ppeeer////
net/ forum?
Sk-Skgaa
Dayal Singh Kalra, Tianyu He, and Maissam Barkeshli. Universal sharpness dynamics in neural network training: Fixed point analysis, edge of stability, and route chaos. arXiv preprint 0rXiv:2311.0200 2023.
Kenji Kawaguchi. Deep learning without poor local minima. In Advances in Neural Information Processing Systems, volume 29, 2016. URL http:ss//pprccceeigg.ss neuriss cc/ httpp/aperr_ft/ files/ (ttepppppeeepeep paper/
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017. httpphttts:///eeeeee wtw.:peenrr.eeeeeeeeeeeeoorrrrer nst/forum? i 66-11 G-Hloyyll oyRl
Itai Kreisler, Mor Shpigel Nacson, Daniel Soudry, and Yair Carmon. Gradient descent monotonically decreases the sharpness of gradient flow solutions in scalar networks and beyond. International Conference on Machine Learning, pp. 17744. PMLR, 2023.
Daniel Kunin, Javier Sagastuy- Brena, Surya Ganguli, Daniel LK Yamins, and Hidenori Tanaka. Neural me- chanics: Symmetry and broken conservation laws in deep learning dynamics. In International Conference on Learning Representations, 2021. WtRLhttps: /oweeeeviiw/ce net formm? 1d=:@gLLLBBBBup
Aitor Lewkowyzz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari, The large learning rate phase of deep lcarning: the catapult mechanism. arXiv preprint arXiv:2003.022/1 2020.
Haihao Lu and Kenji Kawaguchi. Depth creates no bad local minima. arXiv preprint arXiv: 702.0800008000 2017.
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normalization layers: Sharpness reduction. Advances in Neural Information Processing Syxtems, 35:34689-34708, 2022.
Soo Min Kwon, Zekai Zhang, Dogyoon Song, Laura Balzano, and Qing Qu Efficient low-dimensional com-
pression of overparameterizeed models. Proceedings of The 27th International Conference on Artificial In- telligence and Statistics, volume 238 of Proceedings of Machine Learning Research, pp. 1009-1017, PMLR, 02-(4 May 2024. URL ptttttss:/proceediii www.eedinongigeeo ngs ml.. www.r..peeessv333
v238 wtt.ppeees////mmmmneeeeee ht..a//mmmmmonmm22.aaaattt. html
Leon Mirsky, A trace inequality of John von Neumann. Monatshefte für mathematik, 79(4):303-300 -00-30033,,,,,, 1975.
Edward Ott. Chaos in Dynamical Systems. Cambridge University Press, 2 edition, 2002.
Scott Pesme and Nicolas Flammarion. Saddle-to-saddle dynamics in diagonal linear networks. Advances in Neural Information Processing Systems, 36:7475-7505, 2023.
13

---

Published as a conference paper at ICLR 2025
Henning Petzka, Michael Kamp, Linara Adilova, Cristian Sminchisescu, and Mario Boley. Relative flat-
ness and generalization. In Advances in Neural Information Processing Systems, 2021. URL https:
//openrcview wwweenccview.nne//frummm wtwp:nttt/ooummmddayyooooo
Levent Sagun, Leon Bottou, and Yann LeCun. Eigenvalues of the hcssian in deep learning: Singularity and beyond. arXiv preprint ar 2016.
Andrew M. Saxe, James L.
McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. In 2nd International Conference on Learning Representations, ICLR, 2014. URLpttp:: httttpp:/aaxx...rr///e...
1312. 6120
Minhak Song and Chulhee Yun. Trajectory alignment: Understanding the edge of stability phenomenon via bifurcation theory. In Thirty-sevemth Conference Neural Information Processing Systems, 2023. URL
https: htt.://eeee...tttteooommm.aaaaaa
Aditya Vardhan Varre, Maria-Luiza Vladarean, Loucas PILLAUD-VIVIES and Nicolas Flammarion. On the spectral bias of two-layer linear networks. In Advances in Neural Information Processing Systems, volume 36, pp. 64380- 64414, 2023. URL https: /ttp///ppopppeppppppeppppppo.eap wwwrips cc wwtpppprr_ptepaaaaeeeeeeee
paper/2023 wwtp//il//////o//eooo cwd.caad2fcccccccccc 00222266688119000caaa2aa000 wtt.://0000..00.00.a0.aaaa0-0o00
Peng Wang, Xiao Li, Can Yaras, Zhihui Zhu, Laura Balzano, Wei Hu, and Qing Qu. Understanding deep repre- sentation learning via layerwise feature compression and discrimination. arXiv preprint arXiv:2311 1.029000 2024.
/arxiv.vvv....../v. orgg abs /211 C2960
Yuqing Wang. Minshuo Chen, Tuo Zhao, and Molei Tao. Large learning rate tames homogencity: Convergence and balancing effect. In International Conference on Learning Representations, 2022a. URL https:
/ //opnneeveww net. /utt/forune dtt.rmmmm.dttbbrr7777
Yuqing Wang. Zhenghao Xu, Tuo Zhao, and Molei Tao. Good regularity creates large learning rate implicit biases: edge of stability, balancing, and catapult. In NeurlPS 2023 Workshop on Mathematics of Modern Machine Learning, 2023. URL htt..httss://// /ttp::::////peeeeeee.ee net forum? 2dd=m21513322222
Zixuan Wang. Zhouzi Li, and Jian Li. Analyzing sharpness along gd trajectory: Progressive sharpening and edge of stability. Advances in Neural Information Processing Systems, 35:9983-9994, 2022b.
Jingfeng Wu, Vladimir Braverman, and Jason D Lee. Implicit bias of gradient descent for logistic regression the edge stability. Advances in Neural Information Processing Systems, 36, 2024.
Zhenghao Xu, Yuqing Wang, Tuo Zhao, Rachel Ward, and Molei Tao. Provable acceleration of nes- terov's 's accelerated gradient for rectangular matrix factorization and linear neural networks. arXiv preprint arXiv:2410.09640, 2024.
Can Yaras, Peng Wang, Wei Hu, Zhihui Zhu, Laura Balzano, and Qing Qu. The law of parsimony in gradient descent for learning deep lincar networks. arXiv preprint arXiv:2306.00 055, 154, 2023.
Can Yaras, Peng Wang, Laura Balzano, and Qing Qu. Compressible dynamics in deep overparamcterized low-rank lcarning & adaptation. In Forty-first International Conference Machine Learning, 2024. URL https: httt:::://cpeeeevieee
n wdw.mmdxddddDxdsssse
Tian Ye and Simon S Du. Global convergence of gradient descent for asymmetric low-rank matrix factorization. Advances in Neural Information Processing Systems, 34: 1439, 2021.
Chulhee Yun, Suvrit Sra, and Ali Jadhahaie. Small nonlinearities in activation functions create bad local minima in neural networks. In International Conference Learning Representations, 2019. URL https:
//openreview net hwtp//eet/otttt/o wooormiiddrrre ViRct?
Li Zhang. Depth creates no more spurious local minima. arXiv preprint arXiv:1901.09822 2019.
Yedi Zhang, Andrew M Saxe, and Peter F. Latham. When are bias-free ReIU networks like linear networks? In High-dimenssioaal Learning Dynamics 2024: The Emergence Structure and Reasoning, 2024. URL
lattps://openeeview net. Forum ttte://oooomm---aaVaeeeoCCC
Libin Zhu, Chaoyue Liu, Adityanarayanan Radhakrishnan, and Mikhail Belkin. Quadratic models for under- standing neural network dynamics. arXiv preprint arXiv:2205.1 117877 2022.
Libin Zhu, Chaoyue Liu, Adityanarayanan Radhakrishnan, and Mikhail Belkin. Catapults in SGD: spikes in the training loss and their impact on generalization through feature learning, arXiv preprint arXiv:2306.0488 048815 2023a.
Xingyu Zhu, Zixuan Wang, Xiang Wang, Mo Zhou, and Rong Ge. Understanding edge-of training dy- namics with a minimalist example. In The Eleventh International Conference on Learning Representations, 2023b. htl.:tttp://openreeeee ew. net. forumm (orp
14

---

Published as a conference paper at ICLR 2025
Appendix
CONTENTS
Introduction
2
Notation and Problem Setup
3
Deep Matrix Factorization Beyond the Edge of Stability
Assumptions and Analytical Tools
3.2
Main Results
Experimental Results
Subspace Oscillations in Deep Networks
Similarities and Differences Between Linear and Nonlinear Nets at EOS
5
Conclusion and Limitations
10
Acknowledgments
11
Discussion on Related Work
16
Additional Results
17
B.1
Experimental Details
17
B.2
Initialization Outside Singular Vector Invariant Set
17
B.3
Additional Experiments for Balancing, Singular Vector Invariance, and Theory
19
4
Periodic and Free Oscillations
Deferred Proofs
25
C.I
Proofs for Singular Vector Stationarity
25
C.I.I Proof of Proposition
25
C.1.2
Supporting Results
27
C,2
Proofs for Balancing
29
C.2.1
Supporting Lemmas
29
C22.2
Proof of Proposition
37
C,3
Proots for Periodie Orbits
38
C33.1
Supporting Lemmas
38
Proof of Lemmal
42
333
Proof of Theoremll
44
Auxiliary Results
47
15

---

Published as a conference paper at ICLR 2025
A DISCUSSION ON RELATED WORK
Implicit Bias of Edge of Stability. Edge of stability was first coined by Cohen et al. 2021), where they showed that the Hessian of the training loss plateaus around 2/n when deep models were trained using GD. However, Jastrzebski et al.] (020): Dastrzebski et al. (2019) previously demonstrated that the step size influences the sharpness along the optimization trajectory. Due to the important practical implications of the edge of stability, there has been an explosion of research dedicated to understanding this phenomenon and its implicit regularization properties. Here, we survey a few of these works. Damian et al. (2023) explained edge of stability through mechanism called "self- stabilization"" where they showed that during the momentary divergence of the iterates along the sharpest eigenvector direction of the Hessian, the iterates also move along the negative direction of the gradient of the curvature, which leads to stabilizing the sharpness to 2/1j. Agarwala et al. (2023) proved that second-order regression models (the simplest class of models after the linearized NTK model) demonstrate progressive sharpening of the NTK eigenvalue towards a slightly different value than 2/m.. Arora et al. (2022) mathematically analyzed the edge of stability, where they showed that the GD updates evolve along some deterministic flow on the manifold of the minima. Lyu et al. (022) showed that the normalization layers had an important role in the edge of stability they showed that these layers encouraged GD to reduce the sharpness of the loss surface and enter the EOS regime. Ahhn al. (2024) established the phenomenon in two-layer networks and find phase transitions for step-sizes in which networks fail to learn "threshold" neurons. [Wang et al. (002b)) also analyze a two-layer network, but provide a theoretical proof for the change in sharpness across four different phases. Even et al. (2024) analyzed the edge of stability in diagonal linear networks and found that oscillations occur on the sparse support of the vectors. Lastly, Wu et al.] (2024) analyzed the convergence at the edge of stability for constant step size GD for logistic regression on linearly separable data.
Edge of Stability in Toy Functions. To analyze the edge of stability in slightly simpler settings, many works have constructed scalar functions to analyze the prevalence of this phenomenon. For example,[Chen & Bruna (2023) studied a certain class of scalar functions and identified conditions in which the function enters the edge of stability through two-step convergence analysis. Wang al.. (2023) showed that the edge of stability occurs in specific scalar functions, which satisties certain regularity conditions and developed global convergence theory for a family of non-convex func- tions without globally Lipschitz continuous gradients. Zhu et al. (2023b) analyzed local oscillatory behaviors for 4-layer scalar networks with balanced initialization. [Song Yun (2023); Kalra et al.) (2023) provide analyses of learning dynamics at the EOS in simplified settings such as two-layer networks. Zhu et al.) (2022); Chen et al. (2023) study GD dynamics for quadratic models in large learning rate regimes. Overall, all of these works showed that the necessary condition for the edge of stability to occur is that the second derivative of the loss function is non-zero, even though they assumed simple scalar functions. Our work takes one step further to analyze the prevalence of the edge of stability in DLNs. Although our loss simplifies to loss in terms of the singular values, they precisely characterize the dynamics of the DLNs for the deep matrix factorization problem.
Deep Linear Networks. Over the past decade, many existing works have analyzed the learning dynamics of DLNs as a surrogate for deep nonlinear networks to study the effects of depth and implicit regularization (Saxe et al.) [2114 Arora et al. 2018 2019: Alkhouri et al.) 2024). Generally, these works focus on unveiling the dynamics of phenomenon called "incremental lcarning". where small initialization scales induce a greedy singular value learning approach (Min Kwon al. 2024] Gissin et al.) 2020; Saxe et al.] 2014). analyzing the learning dynamics via gradient flow (Saxe)
jet al., 2014] Chou et al. 2024a; Arora et al.] 2019), or showing that the DLN is biased towards low-rank solution (aara et al. 2024 Arora et al. 2019 Min Kwon et al. 2024), amongst others.
However, these works do not consider the occurence of the edge of stability in such networks. On the other hand, while works such as those by Yaras et al. (024) and Min Kwon et al. (2024) have similar observations in that the weight updates occur within an invariant subspace as shown by Proposition 3 they do not analyze the edge of stability regime.
16

---

Published as a conference paper at ICLR 2025
B ADDITIONAL RESULTS
B.1 EXPERIMENTAL DETAILS
Bifurcation Plot. In this section, we provide additional details regarding the experiments used to
generate the figures in the main text. For Figure | we consider a rank-3 target matrix M. E 125x with ordered singular values 10, 6,3. We use a layer DLN to fit the target matrix. Since 0,,1 10, the network enters the EOS regime at
2
0.0309.
Lo
2-2/L
We show that there exists a two-period orbit after 0. 0.0309/2 = 0.0154, as we do not have scaling of 1/2 in the objective function for the code used to generate the figures.
Contour Plots. In Figure 4 considered the toy example
1 0() - (02 01 o0)²,
which corresponds to a scalar two-layer network. By Lemma 1 the stability limit is computed as ? 0.2, T. = 2 and a, 5. To this end, for GD beyond EOS, we learning rate of = 0.2010, where as we use a learning rate of r = 0999 for GD at EOS. For GF, we plot the conservation flow, and use a learning rate of r 0.1800 for stable GD.
DLN and Holder Table Function Plots. In Figure 9 and I1
we compared the landscape of DLNs
with that of more complicated non-convex function such as the Holder table function. To mimic the DLN, we considered the loss function
= (x' 0.8) 11²-
which corresponds to a 4-layer network. Here the eigenvector of the Hessian at the global minima coincides with the x, y-axis. We calculate the eigenvalues \ and x2 at the minimum (0.8 1) and plot the dynamics of the iterates for step size range A2 > > X
2 and > 2. When X2 >
n > Ai the x-coordinate stays fixed the minima 0.80.25 and the y-coordinate oscillates around its minimum at y 1. This is evident in the landscape figure. Similarly, when > - oscillations occur in both the and y direction. The loss landscape 2 = L(x, 3) does not have spurious local minima, so sustained oscillations take place in the loss basin.
For the non-convex landscape as shown in Figure and T2 we consider the Holder table function:
202 y2
f(x,1) ---(sin() sin(2) cys(y) exp
π
By observation, we initialize near sharp minima and run GD with an increasing learning rate step size as shown in the lefthand side Figure 12 When the learning rate is fixed, we observe that oscillations take place inside the local valley, but when learning rate is increased, it jumps out of the local valley to find a flatter basin. Similar the observations by Cohen et al. (2021), the sharpness of the GD iterates are "regulated" by the threshold 2/7. as it seems to closely follow this value as shown in Figure 12
Overall, these examples aim to highlight the difference in linear and complex loss landscapes. The former consists of only saddles and global minima, and hence (stably) oscillate about the global minimum. However, in more complicated non-convex landscapes, sharpness regularization due to large learning rates enable catapulting to flatter loss basins, where sharpness is smaller than 2/17.
B.2 INIIIALIZATION OUTSIDE SINGULAR VECTOR INVARIANT SET
In this section, we present initialization example that is outside the singular vector stationary set. We consider the following initialization:
W Wz(0) 0,
We(0) oP Vl E [L - 1),
(8)
17

---

Published as a conference paper at ICLR 2025
Oscillation along Y-axis: 2/12 > 1j > 2/A1
Snarpress
2kg
Itoration
baration
Oscillation along both and Y-axis: 1j > 2/
Sharpeone
2/1
Beration
Iteration
Loss Landscape
Sharpness
Oscillatory Components
Figure 11: Demonstration of the EOS dynamics of 2-dimensional depth-4 scalar network as shown in Equation ( X, Y axes are the eigenvectors the Hessian with eigenvalues à and x2 respec- tively. Top: when 7 2/A1. the X component remains fixed, while the Y component oscillates with periodicity of 2. Bottom: for > 2/ X2, the iterates oscillation in both directions.
2/1l
20-
200
400
600
820
3000
Iteration
Figure 12: EOS dynamics at various step learning rates from the Holder table function. Left: plot of the learning rate steps and sharpness, showing that sharpness follows the EOS limit 2/n. Right: Plot showing that the iterates catapult out of a local basin when the learning rate is increased and jumps surface where the sharpness is about 2/17.
where Pt Rdxd is an orthogonal matrix. Note that here for f 1, the singular vectors do not align and lies outside the SVS set we defined in Proposition 3 We consider the deep matrix factorization problem with target matrix M. E Rdxd where d 100, 5, and 0.01. We empirically
18

---

Published as a conference paper at ICLR 2025
obtain that the decomposition after convergence admits the form:
(t)
I P;
(z(t) = U
V
(9)
Li=t11
Wx(I)
[[P.
V
Ex(t)
II P; V 1=f-1
We (2, 1, 1],
(10)
aLe-r
(1)
0
W1(() = PiV*
V
(11)
ald-r
where (.(0) and Wz(0) = aPi, vA E [ 1]. the SVS set the singular vectors align with each other. This demonstrates example where even when the initialization made outside the SVS set, GD aligns the singular vectors such that after certain iterations lies in the SVS set.
The decomposition after convergence lies in
10a
I: M-- I0) (PP....VV) MI-
10a
Subspace Erro
10-1
Subspace Error
10-2
IV (...VVV
[o] (..P.V),, - M+
IV. (O.V I
10
10
10-3
10
10-4
10-5
10-3
10-6
10-4
0
200
400
600
800
1000
1200
1400
200
400
600
800
1000
1200
1400
Iterations
Iterations
Left Singular Vectors
Right Singular Vectors
Figure 13: Empirical verification of the decomposition for initialization with orthogonal matrices (lying outside SVS set) that after some GD iterations, the singular vectors of the intermediate matrices align to lie within SVS set, displaying singular vector invariance.
1.10
10
15
05
70
05
0.01
0.10
0.30
Figure 14: Observing the balancedness between the singular value initialized to 0 and a singular value initialized The scattered points successive GD iterations (going left right). The initial gap between the two values is larger for a larger a, but quickly gets closer over more GD iterations.
B.3 ADDITIONAL EXPERIMENTS FOR BALANCING, SINGULAR VECTOR INVARIANCE, AND THEORY
Our theory relied on two tools and assumptions: balancing of singular values and stationarity of the singular vectors. In this section, we investigate how the dynamics at EOS are affected if these two assumptions do not hold.
Balancing. First, present additional experimental results Proposition 2 and how close the iterates become for different initialization scales. To this end, we consider the same setup from the previous section, where we have target matrix M, Rdd where d 100, 5, and varying
19

---

Published as a conference paper at ICLR 2025
10
Train Loss
Magnitude
110
10th
(x80)
Iteration (x80)
(x80)
250-
22-
33-h
Magnitutt
250
Balancing
Train
234h
sie
315
Iteration (x80)
Iteration (x80)
Iteration (x80)
12-
250
23-1
31-
Balancing Gap
22
Magnitude
150-
Train
2344
Iteration (x80)
Iteration (x80)
Iteration (x80)
Figure 15: Plots of the training loss, singular value magnitude, and the balancing gap over iterations (top to bottom). When the learning rate is stable (1) < 0.031 since the top singular value is 0.,1 10), the balancing gap plateaus, whereas the balancing gap goes strictly to zero when the oscillations occur.
for different learning rates: v 0.030, 0.032, 0.034
initialization a. In Figure 14 we observe that for larger values of a, the balancing quickly occurs, whereas for smaller values a, the balancing is almost immediate. This is also highlight that our bound on in Proposition may be an artifact of our analysis, and can choose larger values of a in practice.
To this end, we also investigate how large can be until Proposition 2 no longer holds. We consider the dynamics of 3-layer DLN to target matrix M. 10x10 of rank-3 with ordered singular values 10, 8, 6. We use learning rate of 7 0.0166, which corresponds to oscillations in the top-2 Figure 16, we show the dynamics of when the initialization scale and a 0.5, where balancing holds theoretically for the former but not for the latter. Clearly, observe that balancing does not hold for a 0.5. However, examining the middle plots reveals that the oscillations in the singular values still have the same amplitude in both cases and for both singular values.
singular values.
0.01
Singular Vector Stationarity. Throughout this paper, we considered two initializations Equa- tions 3 and 4 where balancing holds immediately and where balancing holds for suffi- ciently small initialization scale. In this section, we investigate different initializations with aim to observe (i) if they do not converge to the SVS set and (ii) how they affect the oscillations if they do
20

---

Published as a conference paper at ICLR 2025
13"
250
10-
Balancing Gap
Train Loss
10-11
100
10
10:11
02
10-1
150
Iteration (x80)
Iteration (x80)
Iteration (x80)
12
10.0
10.
Balancing Gap
Train Loxs
Magnitude
1.0
100-
0
30
10-1
-25 Iteration (x80)
250
150
175
123 Iteration (x80)
75 160
25
52
100 122
150
175
Iteration (x80)
Figure 16: Top: EOS dynamics of 3-layer DLN with initialization scale a = 0.01, where balancing theoretically holds. Bottom: EOS dynamics of the DLN with initialization scale 0.5. While the balancing does not hold for a 0.5, the oscillations in the singular values are still prevalent, with the same amplitude.
not belong to the SVS set. To this end, we consider the following:
(z) = 0, We(0) = aId.
vl E [L - 1).
(Original)
W (Z) = 0, Wz(0) = aP&
E [L 1],
(Orthogonal) (Random)
(z = 0, We(0) = aH,
VÉ E [L - 1],
where P: is orthogonal matrix and He is random matrix with Gaussian entries. For all of these initialization schemes, we consider the same setup as in the balancing case, with an initialization scale of a 0.01. To observe if singular vector stationarity holds, we consider the subspace distance as follows:
Subspace Distance
V.r L,
(12)
where Ugr and V are the top-r left and right singular vectors of layer W.. respectively. Since Proposition [ iimplies that the intermediate singular vectors cancel, the initialization converges to the SVS set if the subspace distance goes to zero. In Figure we plot the dynamics for all of the initializations. Generally, we observe that the subspace distance for cases go to zero, validating the use of the SVS set for analysis purposes.
Additional Results. In this section, we provide more experimental results to corroborate our the- ory. Recall that in Lemma [ we proved that the learning rate needed enter the EOS a function of the depth, and that deeper networks can enter EOS using a smaller learning rate. To verify this claim, we provide an additional experiment where the target matrix is M, 5 with the top singular value set to σ,,1 = 0.5. We use an initialization scale of = 0.01. In Figure 18 we can clearly see that shallower networks need a larger learning rate, and vice versa to enter FOS. Here, black refers to stable learning and white refers to regions in which oscillations occur (EOS regime).
B.4 PERIODIC AND FREE OSCILLATIOOS
In this section, we present additional experiments on oscillation and catapults in both deep linear and nonlinear networks supplement the results in the main paper. First, consider a 3-layer MI.P
21

---

Published as a conference paper at ICLR 2025
LVgu; LVzu
137
15-3
Subspace Error
13
Train Loss
22-2
101
12:40
175
1:0 Iteration (x20)
Iteration (x20)
10-1
L
Lv
Subspace Error
Train Loss
22-
13-41
150
175
Iteration (x20)
Iteration (x20)
100
L
L
194
20-1
Train Loss
Subspace
20-1
497
30-1
10m1
50
15
300
125
150
175
50
75
120
125
150
175
Iteration (x20)
Iteration (x20)
Figure 17: EOS dynamics of 3-layer DIN for different initializations where it all converges to the SVS set. The subspace distance defined in Equation (12). Top: Dynamics with the original identity initialization. Middle: Dynamics with orthogonal initialization. Bottom: Dynamics with random initialization.
3
0.691
0.782
0.873
L.055
Learning Rate
Figure 18: Demonstrating that deeper networks requires smaller learning rate enter the EOS regime for DLNs, as implied by Proposition for a target matrix with top singular value 0,,1 = 0.5 and initialization 0.01. Black refers stable learning and white refers to regions in which oscillations in the loss and singular values occur. The EOS limit exactly matches 1) 2/ 2/Lo 2
without bias terms for the weights, with each hidden layer consisting of 1000 units. The network
22

---

Published as a conference paper at ICLR 2025
is trained using MSE loss with learning rate of 17 4, along with random weights scaled by a 0.01 and full-batch gradient descent on a 5K subset of the MNIST dataset, following Cohen et al. (202T). The motivation for omitting bias terms comes from the findings of of et al. (024), where they provably show that a ReI.U network without bias terms behaves similarly to a linear network. With this in mind, aimed to investigate how oscillations manifest in comparison deep linear networks (DLNs). In Figure [9 we plot the training loss, top-5 singular values, and sharpness throughout training. Interestingly, despite the non-convexity of the loss landscape, the oscillations appear to be almost periodic across all three plots. It would be of great interest to theoretically study the behavior of EOS for this network architecture and determine whether our analyses extend to this case as well.
032
Ta
0.0
002
Figure 19: Plot of the training loss, singular values, and sharpness for an MLP network with no bias. Similar to the DIN case, there are oscillations in each of the plots throughout iterations.
Next, we consider the DLN setting to corroborate our result from Theorem I We consider modeling rank-3 target matrix with singular values d+,; (10, 9, 8} with a 3-layer DIN with initialization
2- scale a = 0.1. By computing the sharpness under these settings, notice that 2/ ? es 0.01547 and 2/ se 0.01657. In Figure [1] we use learning rates near these values, and plot the oscillations in the singular values. Here, we can see that the oscillations follow exactly our theory.
Lastly, we provide additional experiments demonstrating stronger oscillation in feature directions measured by the singular values. To this end, we consider 4-layer MLP with ReLU activations with hidden layer size in each unit of 200 for classification on a subsampled 20K set on MNIST and CIFAR-10 10. In Figure 20 we show that the oscillations in the training loss are artifacts of jumps only in the top singular values, which is also what we observe in the DLN setting.
MNIST Dataset with 4-Layer MLP
CIFAR-10 Dataset with 4-yerr MLP
Figure 20: Prevalence of oscillatory behaviors in top subspaces in 4-layer networks with ReLU activations on two different datasets.
23

---

Published as a conference paper at ICLR 2025
130
Tetn Loss
107
Kayer
Blayer
186
26
100
100
350
60
189
200
310
Train Loss (1) = 0.0300)
Layer 1 os (1) 0.0300)
Layer 0, (1) = 0.0300)
Layer 3 a (1) 0.0300)
300
15-
130
Baaye 11
(layer 2)
01ayer
Train Loss 0.031)
Layer 1 (1) 0.0310)
Layer 0. 0.0310)
Layer 3 (1) 0.0310)
flaver
8
69
&
Train Loss (1) 0.0325)
Layer 1 0s (1) 0.0325)
Layer o: (n) 0.0325)
Layer (1) 0.0325)
250
Tain
300
Bayer
0, (Layer
Baaeer
10
Train Loss (1) = 0.0335)
Layer 1 as (1) 0.0335)
Layer di (1 = 0.0335)
Layer 3 a: (n = 0.0335)
Tain
155
160
Baaer
(1
(Layer
65
45
180
380
206
380
Train Loss (1) 0.0350)
Layer 1 0s (1) 0.0350)
Layer as 0.0350)
Layer 3 (1) 0.0350)
350
Tain Loss
300
23
132
Bayer
6. (Layer
(Layer
Train Loss (1) = 0.0360) Layer 1 as (n 0.0360) Layer di (n = 0.0360) Layer 3 a. (n 0.0360)
Figure 21: Depiction of the training loss and the singular values of each weight matrix for fitting a rank-3 matrix with singular values 10, 9.5, 9. The weights enter the EOS regime based on the
learning rate 77 > 2/ K, where K 2-22/4 and 3. For a sufficiently large learning rate (e.g.. 0.04), the singular values start to enter period-4 orbit.
24

---

Published as a conference paper at ICLR 2025
C DEFERRED PROOFS
In this section, we present the deferred proofs from the main manuscript.
C.I PROOFS FOR SINGULAR VECTOR STATIONARITY
C.I.1 PROOF PROPOSITION
Proof. Let us consider the dynamics of Wt(t) in terms of its SVD with respect to time:
We(t) (t (t) (.(t) (t) (tV)).
(13)
By left multiplying by (t) and right multiplying by Vz(t), we have (t) uf (t)()(t)E + (t) (
(14)
where we used the fact that Uz(t) and Ve(t) have orthonormal columns. Now, note that we also have
U (1)),(1)) = I, U, ( + (1),(() 0,
which also holds for V (t). This implies that U, (t) is skew-symmetric matrix, and hence have zero diagonals. Since Σe(() is diagonal, U (L)(,(()E,()) and ET(I)VT (1)) have zero diagonals as well. On the other hand, since Σx(t) is diagonal matrix, we can write
I, (UI(W)(()))))
(t)()))))))))) + (
(15)
where @ stands for the Hadamard product and i, is square matrix holding zeros on its diago- nal and ones elsewhere. Taking transpose of Equation 15 while recalling that U. (I)Ú,(!) and VT (tt) are skew-symmetric, we have
i, ④ (OWT
E,(,),, -
(16)
Then, by right multiplying Equation 15 by E,(t), left-multiply Equation (16) by Ee(t), and by adding the two terms, we get
i, o (OIOW.OV ( (1)).))
(1) (OV,
Since we assume that the singular values of M, are distinct, the top-r diagonal elements of ()() are also distinct (i.e., (1) + 522, for r * r'). This implies that
U (1)),(() H(1)
U (I) (0VO://000000-000O0000000000
where the matrix Htt) E dxd is defined by:
(t) = (t) r+r.
(17)
0.
r=r.
Then, multiplying from the left by Ue(t) yields
Po,a Ü,(() = Ut) (H(I) o U. 200 (18)
with Pu,(t) : UT (t) being the projection onto the subspace spanned by the (orthonor-
mal) columns of Uz()). Denote by PU,)(() the projection onto the orthogonal complement (i...
PU,)()) = I, UZ (t)). Apply PUl)(t) to both sides of Equation 13
Pua Uz(l)= = Puanto Ú,(1) (t) + PU (1VV) (1)
(19)
+ ((()))
(20)
25

---

Published as a conference paper at ICLR 2025
Note that P (t) Uz(t) 0, and multiply from the right by Ve(t) (t) (the latter is well-defined since we have the compact SVD and the top-r elements are non-zero):
(1) Ú.(I) Pu ( ( = (1, U,,) - (O)V) (21)
Then by adding the two equations above, we obtain an expression for Út):
(t(t)= = (o (z(() + Pu. (o Uz(()
= (t(t) ((tt)
@ ( (T() + (WW.e)/(((00000///0///oooo))))))
(I, U.() (O))(V
(22)
We can similarly derive the dynamics for Ve(t) and (t):
Ve(()= = Ve(() (H(I) @ E,((), (1) + V/(W,, T (()(()))))))
(23)
+L, VV.VV((V w. ((),,())) (t).
(24)
(t(t) I, U U (t).(()VV. .
Now, we will left multiply (t) and Ve(t) with UJ (t) and V (t), respectively, obtain
of (1)000)))) = - HH() @ U () (00 v (t) = (t()
+ wttp://woYYY000///////////////0
where we replaced We(t) = by definition. By rearranging and multiplying by (t). we have
Wt(t) is the gradient of f(O) with respect to W
U (1)O,(1));))) - (() = i, @ (TOOV.
(25)
Hence, when Ú.(t) = 0 and (t()) = 0, it must be that the left-hand side is zero and so U. ( is diagonal matrix.
Now, notice that for the given loss function f(O), we have
-te() (w,f(@((t) = WT (t) (Wan(t) - M.) WT
Then, from Equation 25 when the singular vectors are stationary, we have
WT (t) (W M.) (t)
must be a diagonal matrix for all E [L]. The only solution to the above should be (since the intermediate singular vectors need to cancel to satisfy the diagonal condition), is the set
(UL, =(U., Qz), (UG, Ve) (QQ+1, Qe), v [2 12 1],
Svs(f(e))
(U1, VT) =QQ,, V.).
where (Qr are any set of orthogonal matrices. Then, notice that when the singular vectors are stationary, the dynamics become isolated on the singular values:
(1) = I, の ( [UT())((V)())))))
since [U] (1) is diagonal. This completes the proof.
26

---

Published as a conference paper at ICLR 2025
C.1.2 SUPPORTING RESULTS
Proposition Let M, U.E.V V denote the SVD of the target matrix. The initialization in Equation 4 is member of the singular vector stationary set in Proposition 7
where Qz. -...
Q2
Proof. Recall that the initialization is given by Wz(0) and Wz(0) ald E [I. 101111.. 1].
We will show that under this initialization, each weight matrix admits the following decomposition for all t 2 1:
(t) 0 V..
Wx(1) - V. [tt)) O
(1( U.
W IL - 11, (26) V.. ala-r
0
where
(t) = Et - 1) -1. (E - 1) 8 (t 1) - E.rr) ) -56 (t--((11)
E(t = Eit - 1) (I.- 11 Σ (t- 1) (5 (t- - 1). 54-1 (t 1) -1))-1,0) Enr) 1)_r,))) -1-3 (t - 1)) where (t), (t) E xxx is diagonal matrix with Σg (1) 1-l.111 -r,,, and E(1 al,
This will prove that the singular vectors are stationary with Σn = =... = Eg = V.. We proceed
with mathematical induction.
Base Case. For the base case, we will show that the decomposition holds for each weight matrix ( = 1. The gradient of f(O) with respect to We is
Vw.f(e) WLe41
WI-I.
For W, (I()), we have
(1( = Wz(0) 7 Vw, f(((0)
Ww(0) 7) (W.1(0): M.). W..
= nat-iE
= U. (na Σ.) V.
= U. (1) 0 V.T.
0
Then, for each W1(1) in [L
1). we have
W,(1) = wr(0) - Vw,
ala,
where the last equality follows from the fact that W, (0) 0. Finally, we have
= V. (1)
VT.
ala-r]
Inductive Step. By the inductive hypothesis, suppose that the decomposition holds. Then, notice that we can simplify the end-to-end weight matrix to
W La (t) = U, [t)
0
V..
0
0
for which we can simplify the gradients to
Σn(t) (5-((t) Σ,r
0
Σ
T--10 (t)
0 V.
Vw.f(((t))
U,
V
V.
0
0
0
(Ez(t) (t) Σ.,
51-1(t) (t)
= =U.
VT.
0
27

---

Published as a conference paper at ICLR 2025
for the last layer matrix, and similarly, Vw.f(e(tt) V. (1(()- (E2(1) -1-1(1) - (1)- -1) o
V.
EE IL 1],
0
for other layer matrices. Thus, for the next GD iteration, we have
WL(t + 1) = (t( - Ww, (O()))
Σg (t) 7
(t) (t) - -Etr) -..0))))))))))) 11-1())
-
0 v.
= U.
0
= U.
[t
0
0 V..
Similarly, we have
We(t + 1) w.(() n1 Vw,(((1)
-v. E() - n Ez(() (Σr(() -1--(1) -E..) 2..)).54-00 -14-2(1)
0
v!
0
ald-r
E(() (I. 11 Ex(() (Σr(1) 01-1(1) -E.r) (5.))-((((1)
V.
0
V.
ald
= V. E(t B(t+1) 1)
aId-r] v..
0
for all L E [L 1]. This completes the proof.
Proposition 4. Let M. V.E.V. Rdxd denote the SVD of the target matrix. The balanced
initialization in Equation
3 is a member of the singular vector stationary set in Proposition
where U, = Q
V1 V..
Proof. Using mathematical induction, we will show that with the balanced initialization in Equa- tion each weight matrix admits a decomposition of the form
(t( = V. V.E.(
(27)
which implies that the singular vectors stationary for all such that U, Q2 ... =Q²=
V1 = V..
Base Case. Consider the weights iteration t 0. By the initialization scheme, write each weight matrix as
We(0) ala Wi() av.v.. which implies that Wx(0) = V. (OVV with Σ,(0) ala-
Inductive Step. By the inductive hypothesis, assume that the decomposition holds for 2 0. We will show that it holds for all iterations t + 1. Recall that the gradient of f(@) with respect to We is
Vw.l(e) = WL441 (WL. - M.) W
Then, for We(t + 1), we have
We(1 +
We(l) - 71- Vw.S(O(1))
VV. - wTI (t -ton (WL1(t) M.) WT-1a(t)
= V.EV -.. (E/ (t) -.) Σ (t)) VT E.) E x(-(()) V.
V. V.. (Ex(() -n
V.E(VV.
where Σ(1) Ex(() - n. E/-(()) (EE(( -000)-20 EE.) - This completes the proof.
28

---

Published as a conference paper at ICLR 2025
C.2 PROOFS FOR BALANCING
In this section, we present our proof of Proposition 2 along with supporting results. Throughout these results, we the notion of the gradient flow solution (GFS) and the GFS sharpness as presented by Kreisler et al. (2023), which we briefly recap.
Consider minimizing a smooth loss function L IR + R using gradient flow (GF):
(t)) ((w(())
The GFS denoted by (wr(w) is the limit of the gradient flow trajectory when initialized at w. Furthermore, the GFS sharpness denoted by ((w) defined be the sharpness of Ser (w), i... the largest eigenvalue of V (SGr(w)).
C.2.1 SUPPORTING LEMMAS
Lemma (Conservation of Balancedness in GF), Consider the singular value scalar loss
C. ((0e)/2))
1
II
0e- 0+
Under gradient flow, the balancedness between two singular values defined by of(() o (t) for all m, l E [L] is constant for all t 2 0.
Proof. Notice that the result holds specifically for gradient flow and not descent. The dynamics of each scalar factor for gradient flow can be written
i oz(()) (=1
-0. o.) m00(()-0
L - ao.(1))
dz(()
if!
Then, the time derivative of balancing is given
(o)(t) - o (t)) = ox(())e((t) - om
Ot
ax(() o. o.)
L
= =-oe(t)
II V=1
o1)( (1)
[(((())
II ot).
m=1
0.
Hence, the quantity (t) o 2 (t) remains constant for all time 0, hence preserving balanced- ness.
Lemma (Sharpness at Minima), Consider the singular value scalar loss
((1a))111) πI 0+
The sharpness at the global minima given as ||
E1_ Z.
Proof. The gradient is given by
L -
Va =
[1
o1)
10)
1=1
Then,
L V V 11 (t Il oz) +
L
a(t) - o,
((())
II ot)
{
29

---

Published as a conference paper at ICLR 2025
Let x(t) II (t). Then, at the global minima, we have
x²² o² 0:
jj
Thus, the sharpness of the largest eigenvalue is given as
| Σ
Lemma 4 (Balanced Minima is the Flattest). Consider the singular value scalar loss
2
1 II 0. L (for)
i=1
The balanced minimum (i.e., = o. for all i E [L] has the smallest sharpness amongst all global minima with a value of | L122 = Lo 2-2/L
1/4
Proof. By Lemma recall that the sharpness at the global minima given the form
| 2
To show that the balanced minimum is the flattest (i.e., has the smallest sharpness amongst all global minima), we will show that KKT stationarity condition the constrained objective
L
II T0=
min (a))
s.t.
i=1
are only met at the balanced minimum, which gives us the sharpness value /| = 2-22/11
The Lagrangian is given by
I. (i1o--o,)
Σ o
Li L. p)
+ 1
Then, the stationary point conditions of the Langrangian is given by
aL
20
+
= 0,
(28)
0o a
0L
L
II
0,
6,
=0
(29)
0p
From Equation
28
the solution of the stationary point gives
20
20
20
20. o?
II o, =
0,¹ 1
0²00
17
This also indicates that at the stationary point, a; = V2 for all E [L]. which means that the condition only satisfied at the balanced minimum, i.e, o; = = 1/L Furthermore, notice that
V² = 6o² Diag
> 0,
where f(a;) = Ele o immediately. Thus, the balanced minimum has the smallest shaprness (flattest), which plugging into gives a sharpness of ||v2L||2 Lo? 2-2/L
o². indicating that f only has minimum. Notice that Equation (29 holds
30

---

Published as a conference paper at ICLR 2025
Lemma 5. Lets: | 02 σrl R and define the singular value scalar loss as
1 C(s) 2
II 0; 0+ A=1
for some o, 0. Ifo E Rh are initialized such that
(0 =0 0 and or(0) =a, vl E [L - 11.
where 0 < <a< In
2v2 nLo,
and 1j > 0, then the GFS sharpness satisfies (s()) -
2-
1,2.2
21
Tv for some <c<l.
Proof. We will show that the necessary condition for the GFS sharpness to satisfy (s)) < 2yLtc
for some 1 > and 0
1 to hold is that the initialization scale must satisfy
O<a<
2 nLo,
2-
L2.2
21-3
Since the singular values de for all l E [L- 1] are initialized to o, note that they all follow the same dynamics. Then, let us define the following for simplicity in exposition:
y y=a
0 and 01,
and I1 0+ my--1 Then, note that the gradient flow (GF) solution is the intersection between
zy--l= =0_ and y²
where the first condition comes from convergence and the second comes from the conservation flow law of GF from Lemma 2 Then, if we can find a solution at the intersection such that
(#(a), ((a)) = ( =
(30)
x22
3² = =a?
solely in terms of a, we can plug in (i(a). ((a)) into the GFS
of = of ((a)) + L-
1
2/2 n
(a((a), (a)) (s)
(31)
8(a)2
and solve find upper bound in terms of a, where (i) comes from Lemma
3 The strict inequality
ensures that we can find a c in c E 10, such that ti(s < 2/1+c However, the intersection ((a) is a 2L-th order polynomial in j(a) which does not have straightforward closed- form solution solely in terms of a. To this end, we aim to find a more tractable upper bound on (ā(()) ;(a)) by using variational calculus, and use that to find bound instead. Specifically, we will compute the differential dij, upper bound du with tractable function, and then integrate to
obtain our new function 2i for which we use to set e 2/2
Computing the Differentials di and dj. Before computing the differential dub, we need to derive the differentials of ((a) and (ax). We drop the a notation and use á and where applicable. By plugging in è into Equation (30), the solution satisfies
a" 2r of.
Then, by differentiating the relation with respect to a, we obtain the following variational relation:
2Ljz2 2- dÿ 2² I, - 12 dij 2a 21 2 = 0
3 a²(L 1))dý (L
ÿa da. dÿ ( - a² 1))
(32)
NNte that throughout the proof ()(1),((a)) denotes the gradient flow solution function of a. It does refer the GF trajectory.
31

---

Published as a conference paper at ICLR 2025
where we used Lemma 3.10 of Kreisler et al. (2023) to deduce that 9 0 and so g2/-2 0. Then,
notice that we have 1 > V I.a from initialization, and so we di function of a). Then, we also have
> 0, (i.e., j(a) is an increasing
do
lim ((a) = 1/L and lim a(a) = 0 a00
1/L
a-0
as it corresponds to exact balancing. Hence, as increases from 0, j(()) increases from o./4.
Similarly, the intersection at the global minima satisfies the following relation for #:
$ 2+ 1) + -I a²
2
2
((1
dē
2 T.-1
a² 2od x. da =0 0
I,
a
da.
(33)
1.
b-I
I1-1
Note that since à > 0, we have
de
< 0. This implies that as a increases from 0, ((a) decreases
de
from a. 1/L
Computing the Differential dif. Now we are position to derive the differential deb Let us define ((a) ((((a))((a))) we ultimately want the behavior terms of a. Let us simplify Va)
first:
a(a)2 + L-1
V(a) = (((())
o?
(From Equation
;(a)3
=? U(a)² (I. - 1))((a2
(34)
a(a) {())
x2 (a))22²) 1²2 + 11 1 L
(35)
L
Lo?
Then, computing the differential, we have the following:
2
2(L 11--1 1)
dv
o
dē 23
dÿ
(36)
g3
1
200 2
[L (L-1)
2adjo?
da
da
(Substitute dē, dg)
33
La
a
j3
(i2L. a² (1]]
(I-11)
1
2a(l 1))²2 de
(37)
24
L 2
1/1 - a2 (-+))
0
2
12 -/- + (1 - t) 0))
2a L - 1)0
do.
(38)
(i+ a² a20(1=1) (4-1))
L
Then, recall the intersection constraint:
82 22 a²
(0² - 22 (1²² + +2)
a² + +²)
(39)
8¹ ÿt a² (6² + )2))
(40)
By substituting in Equation
44). we can simplify further:
a c
+1-
2a(L 1)0
dv
da.
(f1 + 0² 5 1
L
32

---

Published as a conference paper at ICLR 2025
Now, can plug in Equation
(35 into the numerator:
a + + (1- $22
2a(L -1)02 da
dv
(++ a² (0 o? (1I1))
L
0(00)29²2 Lo7
20r L -1)0²
de
(01 ( 0z
L
V')
20(L 1)0²
da
Lo?( + a2 10 a 22 (4I1))
L
(.)
1
a²
do.
(+
T
Finally, notice that from the conservation flow, we also have
a
L -
--2² ====0 a²
22
L
L
and
(a)
( T T.2
dv
2
dv=
a da
a² da
(+2 + 42
((a)
(+2 + T. )2
T,2
P(a)
P(a)a da.
Upper Bounding the Differential. Note that it difficult to directly solve for a from P(a). as f is also a function of a. Hence, we can upper bound P(a) by function F(a) such that F(a) 2 P(a) for all & > O, and use this to solve for a. We proceed by looking at the derivative of P(a):
4
1 L
dē
2a
P'(a)
2 + 13
L
da
L
-4
(t L
1
2a
2δa
(+2 + 133
L
a² (L-1)ø
L-1
8a
(1 1
1
1--
1
(+ 2
L2
T
a² g2
L
Case 1: T 2. Consider the when T. 2. Then, notice that for all > 0. P'(() 0. Thus, we can choose F(a) as such:
2
1
F= lim P(a)
A/L I
12
0.
which is constant in a that upper bounds P(a).
Case 2: L 2. Now consider the general case. Notice that
a(a) VL(L 2)
a
the only critical point of P(a) (since 0). Furthermore, we have
a(a) < VL(L - 2) = P'(a) < 0,
a
33

---

Published as a conference paper at ICLR 2025
implying that P(a) is decreasing. Then, since 2(a) is also decreasing function in a, this means that there exists an Cterit such that for all a > acnis P(a) is always decreasing. We can find Cerit as such:
x0 ((Ceni))
Ocrit
((Cem) Ocri 1+
VLLL- 2)
L(L 2)
By plugging these into our constraint set, we obtain
L-1
Clerit
1
Ckerit
1+
-0+
VT( 2)
L 2)
L-1
L
(v
(v+ T( -2)
1
=0.V I(L-2
0+ v LLL 22)
Ocnt
(V1+ 2)
4-1
1/b
0
acrit
1/E
1 + L(L-2)
VI(L-2)
Next, also note that for any a < Cerit. P'(a) > 0, and so P(a) is increasing. Hence, P(act) corresponds to the maximum value of P. Therefore, we can choose F' P(aait) as constant function that upper bounds P(a). This leads to
2
1 L
F = P(ac))
(a(a) )2 +
a
L2
2
-1- 12
( 1
L
+
Qill
T.[. 2j
2
1
[ T.2
(1- 1
L(L-2)
2
1
L(L 2) 2
)
(VITI. (vill-2)(
1
(1+
L-1
L/L
L
1
L L-1
L(L - 2)
0+
=.h(L)
(L)
A/T.
0.
Combining Both Cases. To avoid using two separate functions F for different values of L, we can upper bound the function h(L) to encompass both cases. This yields the following upper bound:
(1+ ((L
1
M(L) L3 L.
<
L2.22
2(E-1) b =((
Finally, we left with the new differential
dv (a)
LL(
a da.
o L/L
34

---

Published as a conference paper at ICLR 2025
0.55
0:22
0.29
035
0.32
2
==
L
Figure 22: Plot of P(a) along with its upper bound evaluated at F = P(acatt) for different depths.
The critical point occurs exactly the computed value of Cerit and the function 2 P(a) for all 0.
Finding Upper Bound on a.
Firstly, we integrate the new differential:
dv v(a)
(L) 4/L
f aa do
T
9(L)a
In
Vp
20-
g(L)ao
v Vo exp
20
where v lim (aa Lo Lo?- 2-
Now, we can solve for a:
a-30
2-2 Lo+
((L))a*)
2v2
g(L)a*
2v2
exp
20 1/L
17
exp
201LL 1/L
nLo 2- 7
2011L 1/4
2
a
7
g(L)
Lo
2/2
20 4/L
1/4
a<
2-
T,2
3LE-11
Simplifying further, we obtain
2/2
0+
A/L.
1/4
a
In
Lo.
I
2
which gives us the desired bound. This completes the proof.
Lemma 6. Let (s) = -1 ag denote the end-to-end product of E RL and suppose that each oe O. If the GFS sharpness (s)s) < 2/IFC for some c E (0, 1], then
min (T 7x))
<1+c.
Proof. We consider cases: (i) (S) [0, o.) and (ii) (s) a.. Note that we ignore the case
of (s) = as this occurs with probability zero at EoS.
Case ( E [0, o.)). For this case, notice that we have
min L- 1) (T - T.) 7T)
< x2)))s)
(rT() < a.)
oi_
of of
35

---

Published as a conference paper at ICLR 2025
Then, note that the GFS sharpness is constant for all weights on the GF trajectory, as it is defined to be the sharpness at the limit of the GF trajectory (i.e., the GFS), Hence, we can focus on the weights at the solution, global minima.
Define the GFS as z = Scr (s). By Lemma 2 each coordinate in z E R (and hence s E 12 is balanced across layers under GF, and have that
of am
2 am
ve, E [T].
Hence, it is suffices to show that
min 1)
7f (z)
min 1)
x²
Σ
Sl+c=
Σ
S1+c.
2
o.
i=1
Then, note that (2)) o., since it lies on the global minima, and have
min
(2)
min(2.L-1)
r²00?
Σ
(41)
2
2-
From Lemma 3 the sharpness at the global minima given as
L
o?
(s))
V [())
Σ
(42)
o²
This immediately implies that Therefore, we have
(s) and equivalently, 3B E [0, such that 1
Brb(s).
2..
min L2
}}
of
Σ
S (1-
(43)
Substituting Equations 42 and (43)
into the expression we aim to bound, we obtain
(2 (T( -
(2
1²O²2
1 _ (s) S 11+ 1 S
Σ
of
-7,1
where we used the fact that the maximum of 3(1 B) is when 3 3 1 and (s)) 2/TFc Thus,
(s))s) S 2/I+c then for every weight s R& lying on GF trajectory, we have
min(2,L-1) (T -O. Sl+c.
Case 2 (s) > o.). Consider the case in which rTs) > By assumption, note that we have a; > 0, which implies that each GD update will also remain positive:
1
a - (T((s) - o. )T()))
0.
From this, we get
>> (T((S) - o+)T((s
> 0.
of
This implies that
min
(T((s()) 0. mm) (s)
Σ
< {1+c)),
o7
with c 1.
Putting both cases together, have that
min
(m( 0r 2x))
< (1+c),,
o
for e E (0, 1]. which completes the proof.
36

---

Published as a conference paper at ICLR 2025
C.2.2 PROOF OF PROPOSITION 2
Proof. Consider the i-th index of the simplified loss in (5):
2
1
I oE 0x.
1 2
II
2
and omit the dependency on for ease of exposition. Our goal is to show that the L-th singular value oz initialized zero become increasingly balanced to oe which are initialized to > 0. To that end,
let us define the balancing dynamics between a; and oj as bi. + 1) == (of (o(+t)²²
(o ((2)))
and π (s(t)) I (e(t) for the product of singular values at iteration Then, we can simplity the balancing dynamics as such:
bi (1 + 1)= =
(o)(+++))))) (1,(+))))))(((11)))))
(44)
= (c(()))(((((()) - = (1)(0)-n - (=((((())) o.) a(s(t))
2
a,(t)
oy(t)
(45)
= (ax(t))² - (t)( + (T(S(1)) - o..² ( T
(o.(1))² (o)
(46)
= (lo,(0))- (o,(())) (1 (1-1)(1(6(()-00) - -o.)²
*s(())
(47)
(o:(1)) (a,(1))
(b) (1- (1-1)( o0,² o..) (o.(t)) (o,(t))
*²((())
(48)
Then, in order to show that (t + 1)] < (t)) for some c S 1, we need to prove that
(s(())
1 (=(((((()))))))))))))))))))))))) o.)² (a,(t)) (o)(t)² <c.
for all iterations t. Note that for our case, it is sufficient to show the result for i L and j l for any l # L. WLOG, suppose that the a are sorted such that o 2 02 2 ... 2 σ By assumption, since our initialization scale satisfies
2
A/L
1/4
0<a<
In
nLo 2- 1
L 12.2 2
by Lemma 5 we have that the GFS sharpness (5()) for positive s = lo1
σL) E RR- (i.e., each
element de 0) satisfies (iss) 2v/2 Then, by Lemma 6 we have
min min(22LL11) (T((T()) G. 77)
S1+c,
(49)
(
for some E 1). Then, notice that Equation
49 implies that
(T(s) 0
- xs) <1+c and n7 (T((s) =2.)2=2
1+c
(50)
oF
for all i [T],] E [T. 2) and < j. Notice that the latter inequality comes from the fact that
(T((s) - 23
0i + (T((s) a. (s)
(T((s)) o.) 7
+ (T(((s) a. 2 (s)
of 0
o2
<1+c,
37

---

Published as a conference paper at ICLR 2025
which implies that
1 - - -0.)2m²
< +c (T((s)) - (s)
1 l+
oi
22-2
2
and since a are sorted, holds for all other o. Therefore from Equation i [L 2].
we have for all
(1-1(1+1) < a. (1)41(() and b(t 1)c-2 < c <c.bb-bb bz-2,(1);
as well as
-c- -c- c--b1.1((t) < bL-1, < bc-bL-LLLt).
Then, notice that since we initialized all of the singular values for E [L 1] to be the same, they follow the same dynamics. Since already showed that 1, 1)1 e. (,)) must follow that
[b 1)1 (b, Vi,j e [4].
This completes the proof.
C.3 PROOFS FOR PERIODIC ORBITS
Before presenting our proof for Theorem we first show that the required condition from Chen & also satisfied for DLNs beyond the
Bruna 2023) for stable oscillations occur (see Lemma II
EOS, as shown in Appendix C.3.I
C.3.1 SUPPORTING LEMMAS
Lemma 7 (Stable Subspace Oscillations). Define S, = Loop" and K2 = max (ppt 2 f
2-2
we run GD on the deep matrix factorization loss in 7 with learning rate 1j where Ky K < Sp, then 2-period orbit oscillation occurs in the direction of As,. where Ls denotes the eigenvector associated with the eigenvalue Sp of the Hessian at the balanced minimum.
Proof. Define fa, as the 1-D function the cross section of the loss landscape and the line fol- lowing the direction of A; passing the (balanced) minima, where A; is the i-th eigenvector of the training loss Hessian the balanced minimum. To prove the result,
will invoke Lemma [] which states that two-period orbit oscillation occurs in the direction of A if the minima of fs. satisfies f( 0 and ( /1? A for all of the eigenvector directions, the oscillations only oceur specifically in the directions of As.
for n > X We show that while the condition holds
First, we will derive the eigenvectors of the Hessian of the training loss convergence (i.e., M, Wr.).. To obtain the eigenvectors of the Hessian of parameters (W W....., W2, WT), consider a small perturbation of the parameters:
θ (AW; + wo (WL + AWL We + AW2, W1 AW1).
Given that W f.
M,, consider and evaluate the loss function at this minima:
(( Σ W WW.nAWWWen
(51)
W m- +... AW
(52)
Pcm
38

---

Published as a conference paper at ICLR 2025
By expanding each of the terms and splitting by the orders of AW (perturbation), get that the second-order term is equivalent to
Awwl?
W A...AAWWWW-..
Σ Awwl
W WLH-AWWW-H
Σ
W WLeH&WW 1-1mm1 AWW W Wm-1:1
AW |
2
Σ
L& WLEIAW W- W Wm-
E<m
+ tr
W (WEHHAWWWW-H)
Σ Km<p
W- AW,, Wm-1 W W p- 1:1
The direction of the steepest change in the loss at the minima correspond to the largest eigenvector direction of the Hessian. Since higher order terms such as (E/n | are insignifeant
compared to the second order terms + (E (lAWel?). finding the direction that maximizes the second order leads to finding the eigenvector of the Hessian. Then, the eigenvector correspond- ing to the maximum eigenvalue of 2 f is the solution of
A vev(v (OC(AW 1, AW1)
argmax JAW + IA
S(AW, ,., ,Wi)),
(53)
where
f(). = AWWWL-nt WL3 + A
(54)
While the solution of Equation 53
gives the maximum eigenvector direction of the Hessian, A the other eigenvectors can be found by solving
argmaa
f(AWL, AW1).
(55)
JWW i + JAW F=
A
1A
By expanding f()), we have that
S(WW, AW,) JAW WI. + I|W AWW 1l2 ||
AW
+ tr (W-WW T (W W.3.W W .... W22AW1) +...+
tr (WZ2AW1) T (WLSAW,W,+ + WLSAW,W,)
(56)
We solve Equation
53 by maximizing each of the terms, which can be done in steps:
(i) Each Frobenius term in the expansion is maximized when the left singular vector of AW, aligns with W Li and the right singular vector aligns with We-1.1- This is a result of Von Neumann's trace inequality (Mirsky) 1975). Similarly, each term in the trace is maximized when the singular vector of the perturbations align with the products.
(ii) Due to the alignment, Equation 53 can be written in just the singular values. Let Ast, denote the i-th singular value of the perturbation matrix A AWe Recall that all the singular values of M, distinct (i.e., T..1 ... > r. Hence, it is easy to that Equation (33 is maximized when Asy, = 0 (i.e, all the weight goes to Asy,). Thus, each perturbation matrix must be rank-1.
39

---

Published as a conference paper at ICLR 2025
Now since each perturbation is rank-1, 1, write each perturbation
AW. AscAu,AV E [T.]
(57)
for Ast > and orthonormal vectors Au; E Rd and Av, E TRd with Σ As? 1. Plugging this in each term, we obtain:
W A1s V.O.
Av,
Since alignment maximizes this expression as discussed in first point, we have:
ug Vt v, for all l [2, 1]. then
a = 0,1 Vx,1 and b o,,1 VI,
ab" - - 1-1
V.,
The very same argument can be made for the trace terms. Hence, in order to maximize f(-), we must have
Vr. = Ve and u = Va,1:
V. V.d.
vi E [2, L 1].
To determine uz, and Vı, we can look at one of the trace terms:
-- ir (A) WaL-aa) (Wr.34) W + WL.4,, S L
(
To reach the upper bound, we require uy = u., and v V., Finally, as the for each index, the singular values are balanced, we will have 118 vT. for all [Z]] to satisfy the constraint. Finally, we get that the leading eigenvector is
A vec (
1 VIV
VI V
VI
Notice that we can also verify that f(Ai) = Lo 2- derived in Lemma
which is the leading eigenvalue (or sharpness)
To derive the remaining eigenvectors, we need to find all of the vectors in which AJA; for i * j. where
A: vec W L. A,W,,,
and f(A.) A; where A; is the i-th largest eigenvalue. By repeating the same process above, we find that the eigenvector-eigenvalue pair as follows:
1
u1
1
VIV.
VIVT):
A1 Lo 2- + 2o2
VL
VL
VL
1- +
A vec
U1
V2
VIV2
VIV.
A2
VL
VL
VL
A3 = vec
VL
112
1
V2v
VL
v2v]]
As
0,1
1
1
da -o 2-
4 = vec
VL
VL
V2V
VI
Adrtr vec (1 (VT
1
V Vav,,...,
VI VT -avv
VI
40

---

Published as a conference paper at ICLR 2025
which gives total of dr eigenvectors.
Second, equipped with the cigenvectors, let us consider the 1-D function fa, generated by the cross- section of the loss landscape and each eigenvector A, passing the minima:
fa,(u) L(W& + HAW W2 + 2AW2, W1 mAW1), = n² 1
WWWW L-11 .... W W, + 12
1
Σ Ir (r(W T (W.eH&W.WW W Wm-1a)
W.W.eWWWW W Wm-11
Pcm
Σ tr (WLLAAAWWWW-1)) WW..:WWWWWWWW.. A.
+ nx
k<m<p
Then, the several order derivatives of fa, (H) at μ = 0 can obtained from Taylor expansion as
(0) = W, +.... W WL A.W. + W
(0) = 6
Σ tr
A.WW..A.WWWWW........... W.A. W.A.W W
t=1
W Wi W m--ii |
-6 Σ W.
Σ fcm
6λ; B4
(1 = W +... W.W.
24 Σ tr
Σ WLWTAWWW tim AW W 1p A W, W WP-1:1
t=1
l<mep
12 24X; 6,
where we defined
Σ W WLL+IA W-Weei
(Total (4) terms)
8
WW..LAWWW..... 1-1m AW, W Wn-
(Total
terms)
6, =
Σ L<m<p
AW W
W W
(Total
terms)
and used the fact that tr AB) Then, since & has (5) terms inside the sum, when the Frobenium term is expanded, will have (()++)) number of terms. Under alignment and balancedness, B3 As70, 2-
||A|F B under singular vector alignment.
(G)++)
41

---

Published as a conference paper at ICLR 2025
and 2,6; Asqo, 2- x L.) Thus, we have the expression
(h)
(L)
+1)
2 A,d; = As²0; 2-
2
2
1.
2
2
3
L(L-1) +
= 2- (L)
3
T.
3
L(I -2)
= As30; & 2(5)) L
- x ((L 1)2 +5) > 0,
for any depth I > 2. Finally, the condition of stable oscillation of 1-D function is - 1/ 2
f/4 30 (12) + 24(2)))(6.)) 483 (,o,) 0,
which we have proven to be positive for any depth I. > 2, for all the eigenvector directions cor- responding to the non-zero eigenvalues. Lastly, by Proposition 3 notice that we can write the vectorized weights the form
À = vec (WL, W I......, Wi)
= vec (U, V V.E VT. V.E.V.)
vec (OL,1 u. u. L--1 V. V, 01 V..V..)
Then, A, + 0 only in the eigenvector directions that correspond the eigenvalues of the form S = Lo 2/L Hence, the oscillations can only occur in the direction of As.. where As, are the eigenvectórs corresponding to the eigenvalues S,. This completes the proof.
C.3.2 PROOF OF LEMMA
Proof. By Proposition
3 notice that we can re-write the loss in Equation I as
WL1 - 1F E - E.
where ΣL1 are the singular values of WT.1 L:1. We will first show that the eigenvalues of the Hessian with respect to the weight matrices W. are equivalent to those of the Hessian taken with respect to its singular values Σg. To this end, consider the vectorized form of the loss:
Ι( = Wt. t.1 - M. = vec (W.) taa) - ve.
Then, each block of the Hessian (()(()) e oe Lx with respect to the vectorized parameters is given as
[b) m,t = VeecW...) ... VTo E ddd
By the vectorization trick, each vectorized layer matrix has an SVD of the form vec(We) = vec(U,,EVT)) (V, & U,) (E,). Then, notice that we have
(e(())/(((()) = (V. @ U.) (9()).
which gives us that each block of the Hessian is given by
[V)[())] m.t VworW...) ((e)
= =(V & Um) ViSVE...
T.)
=Hmme
42

---

Published as a conference paper at ICLR 2025
Then, since the Kronecker product of two orthogonal matrices is also orthogonal matrix by we can write the overall Hessian matrix as
Lemma 9
R. R, 22 Ri
2,2 22
H, R,
H
R H. R.
Rr.
for orthogonal matrices (Re)--) Then, by Lemma the eigenvalues are the same as those of H, where H E Rdd Lxd is the Hessian matrix with respect to the vectorized Eg:
H1,1
H1
H L
H2.1
H2
H
H
H
22,
H
Now, we can consider the following vectorized loss:
S( 1 EEn
E.12
vec (En.1 E.)11
(2/1ial @ El.E+1) vve(Ee) - voc(E.)
=A\
Then, the gradient with respect vec (Ee) is given by
V vee = A (Ae vec(Ee) vec(E.))
Then, for m = L, we have
H = V Vi Voo f(e) = AT Ac.
For m + E, we have
Hm, = V Vee((E..) V (ee((E))((() = V vec( A (A,vec(Er) vec(M*))]
V Vvon( (E AT (Aevoc(E)) vcc(MM) +A, V Vooe (Aevoc(Ee) vec(M())
convergence
A Am,
where we have used the product rule along with the fact that Aevec(Er) (E) A Ance (m
Overall, the Hessian at convergence for GD given by
A
AT
A
A A
A
A
A
A2
A. A
H
AJA
A[A₂
AJAL
Now, we can derive an explicit expression for each Am, / by considering the implicit balancing effect we have that convergence,
of GD Proposition] 2
Under balancing and Proposition 3
E// 0
0
VE L - 1}. and En = 1/L
ELA1 =
=.
I
Thus, have
Σ 2(--1 2 Σ.
for m = f
H A
Dm am+t-2 2
for + l.
43

---

Published as a conference paper at ICLR 2025
Now, left with computing the eigenvalues of H L To do this, let block
diagonalize H into H PCP", where P is a permutation matrix and
C
C=
C22
where each (i: entry of Cj E R is the k-th diagonal element of H..- Since C and H are similar matrices, they have the same eigenvalues. Then, since is block diagonal matrix, its eigenvalues (and hence the eigenvalues of H) are the union of each of the eigenvalues of its blocks.
By observing the structure of Hm,e notice that each C₂ is rank-1 matrix. Hence, when considering the top-r diagonal elements of Hm.e corresponding to each Kronecker product to construct C₁, each C₁ can be written as outer product uu" where RR4 is
1 T+.
1 0
1- 1-1 * T x
0
(58)
Then, the non-zero eigenvalue of this rank-1 matrix is simply | which simplifies to
L-1
-e te
t=
Next, we can consider the remaining d - components of each Kronecker product of H Notice that for m
L, we have
2(C
x,1
H
21h-1)
O+r
Id
2L (IL- L
2(L-1)
This amounts to a matrix Cg with a single element 0 212 1) value ".i"
and elsewhere. This gives an eigen-
for all e [r]. with multiplicity d r.
Lastly, we can consider the diagonal components of Hm, that a function of the initialization scale a. For this case, each Cş can be written as an outer product vv where
-
a
o,
1- ||v| which corresponds
0.
1- a²
of
(59)
Similarly, the non-zero eigenvalue is simply
L-1
||
2 = Σ
A
1- -
aa)
t=
This completes the proof.
C.3.3 PROOF THEOREM
Proof. To prove the result, we will consider the GD step on the i-th singular value and show that 2-period orbit condition holds given the learning rate 17 Σ. For of exposition, let denote the i-th singular value of each We as or = Tt,s. Under balancing, consider the two-step GD update on the first singular value:
a,(t + 1) = at()) + nL. (a., of of (1)
(tt)) ai(t 1) пh (o..) of + 1)) -
(By 2-period orbit)
44

---

Published as a conference paper at ICLR 2025
Define == (1 nL (o.,, - (f)) - (t) and by plugging in o,(t + 1) for (t),,, we have
dy(1) - o, nL. (o.,s -0.,1-0000 -1((=)=)) of (1)==-1
= 1 = 2 + nL (o.,) - (t))² // (t)=t-1
= 1 = (1 + nn. (0+.i - (1) (1) + +nI. (0-,i- - (1)= of--(11)2--
= 0= (o... - of (on.s - - of . zt1
Simplifying this expression further, we have
0 Fnss - of (t) 0,,1 - of (t)
(t) + of (t) 21)22)11112 J,, + 02,12
of (t) (1+ 2² 2) 0,1 (1 (1+2"-11
o (1) (1 + )26-1)
(1+2) 21-1) 00,
and by defining P; = (t),, we obtain the polynomial
0,, pi 1 + 22 21
where = (1 + nn(o,.s--ppp - p.--
1
Next, we show the existence of (real) roots within the ranges for Pú,1 and pi,2. We note that these roots only exist within the EOS regime. First, consider Pt.1 E (0, (0,o%/²) We will show that for two values within this range, there is sign change for all L 2. More specifically, we show that there such that
exists Pa E (0,0 1/L 4,
+ 21-1
1 1 + 22-1
1+ 2 0+, 0 and
Pi
1+ + U,, O.
For the positive case, consider p: = (20,,)111, We need to show that
21-1
1 nL 1
1+ 2
2
+
2
1 + 24-1
4
1-1
1+
1 1+
(
2
7
To do this, we will plug the smallest possible value of greater than 2, which gives us
r)
b-
to show that the fraction is still
Lo
21.11
1+ 1 1
2
L(L) 1
(60)
L-I
+ 2
which is an increasing function of L for all L> 2. Since u(2) > Equation (60) must always be greater than 2. For the negative case, we can simply consider P: = 0. Hence, since the polynomial
continuous, by the Intermediate Value Theorem (IVT), there must exist a root within the range
Ps (0,0 (0.0.1)
Next, consider the range Pt.2 E (o!/ 1/L
(20..)
Similarly, we will show sign changes for two
values in Pi.z- For the positive case, consider P (3a..) 1/T. For n, we can plug in the smallest possible value within the range to show that this value of A provides positive quantity. Specifically, we need to show that
2L-1
1
1
2
(o
yo..))
(fon.)) 1-7
1 1+226-1
2
6,
2-7
2
1 1-
4-1
3
1+
1+
2
(0.,5 (0.s) (3o..))
1-7
1-2
2- z ,,t
45

---

Published as a conference paper at ICLR 2025
We can simplify the fraction as follows:
2L-1
マ (o,,1 f0.11) (80.,))' 1-
1+
1 1 1 (2)
(2)
22-1
,,1
L-1
L-1
+ --ビ (0x,11 - 3 00,1) ( (20,,1)
1-2
(9)1-
1
x,2
Then, since we are subtracting by 13)-= i, we can plug in its largest value for L 2 2, which is 3/2. This gives us
+ (-0.5)22-1 2
(-0.5) 4-1 >
3
as for odd values of L, the function increases to starting from = 2, and decreases to 1 for even L. To check negativity, let define
[(p) ((p) = p² (1+ 21
h(p)
1 z 1.-1
We will show that h 1/L)
0:
h' (o111)
f 0x.i
1/L
9 022
1/L
f (0.11) g' 1/L)
1/1)
(o_.
g² (=
f (o.11)
Tndb "g' 111))
2
-o-- 1-t Lo
on,i (2L - 1)
11I2022-
2-2
0, (I) ()11100
4,
+,5
2
1-t (3L 2) 0 Lo
3-
< <0,
as otherwise we need 17 S 32 which is out of the range interest. Since h hp)p) 0, it follows that there exists 8 such that h(p) h(x) for all such that p < < p 8. Lastly, since hp(p) O+,i O for o,, it follows that h(p) o.,i must be negative at 5. Similarly, by IVT, there must exist root within the range P1. E (o!! 1/L (20.,)) This proves that the i-th singular value undergoes two-period orbit with the roots Pt.1 and P.2. Then, notice that if the learning rate is large enough to induce oscillations in the i-th singular value, then is also large enough to have oscillations in all singular values from to the (i - 1)-th singular value (assuming that it is not large enough for divergence). Finally, at the (balanced) minimum, we can express the dynamics as
1/L
WL
Σ phs u.VVV
Σ o u. (E(1,2) WE(L-1 1]. 1
(61)
oscillation subspace
stationary subspace
This completes the proof.
46

---

Published as a conference paper at ICLR 2025
C.4 AUXILIARY RESULTS
Lemma 8. Let (Re E <Rxx be orthogonal matrices and Hig E RR'xxn² be diagonal matrices. Consider the two following block malrices:
H1.1
H1.2 E.1
2
H2.
H.
H
H1,L
H2,
H
H, 1
R,H
H
R H2, H
R
2
R2.1H2.RR.
H
1
R.HL,RR
R,H,,
Then, the matrices H and H are similar, in the sense that they have the same eigenvalues.
Proof. It suffices to show that and have the same characteristic polynomials. Let us define
A
B D
H =
where
A R,,,
B
R,H,
(62)
H22R
R R2,1H
R R2
C
D=
(63)
R R.
R. H2
H,H
Then, have
B D D- det H AII det ([ (
IK
C
det((D - XI) C(A AB)- where the second equality is by the Schur complement. Notice that
det((A AI)
(A (A- (R, R -AI) (R,H,, R7 RR.RR, T)-1
Rr, (H1,1 - A1) R
Then, we also see that,
R
R
C(A - AI-
E
R
R
where
H2 (H1,1 - AI)
H2g
H (H1,1 - AI)- H
-1
H (HI,1
-1 H
H (H1,1 AI) -1 H
Similarly, we can write as
H
H
D V
VT
H
L.L
=
47

---

Published as a conference paper at ICLR 2025
Then, have
det(H - AI) = del(RR (H1,1 AI) RRE) det (V.(E- VT)
1e - AI) det(EE - F),
which is not function of U, V, (R.) By doing the same for H, we can show that both H and H have the same characteristic polynomials, and hence the same eigenvalues. This completes the proof.
Lemma 9. Let A, E IRdxd be two orthogonal matrices. Then, the Kronecker product of A and B is also an orthogonal matrix:
(A + B) = (A 69 B) B) L
Proof. We prove this directly by using properties of Kronecker products:
(A (A@B) A B B
= I & 1 Io.
Similarly, we have
(A @ () 8 T) = AA 2 BB
= I La.
This completes the proof.
Lemma 10. Let {a(t) be a sequence such that a(t) 2 0 for all t. If there exists constant e E (0,1) such that a(I+1) < e. alt) all I, then lim all = 0.
Proof. We prove this by direct reasoning. From the assumption a(t 1) -a((t for some E (0, 1), we can iteratively expand this inequality:
a(t + 1) < a(t), a(t 2) < < + +)) a(t),
and, more generally, by induction:
a(t k) < c a(t), for all k 2 0.
Since e E (0, 1), the sequence {cr 1 converges to as + oc. Hence:
0 < lim all + k) < lim c² a(l) 0.
Therefore, by the squeeze theorem, the sequence {a(t)) converges to 0 as - co.
Lemma 11 (Chen & Bruna (2023) ) Consider any I-D differentiable function f(x) around local
minima F, satisfying (i) f(3 (r) + 0, and (ii) f f" f(4) > 0 at Then, there exists with
sufficiently small |c] and c f f(3
such that: for any point zo between and I there exists
learning rate n such that p2 (m)) = Fo, and
2
2 f"( - c. (z)
["($)
48


Reviews:
[{"review_id": "TToJhv4rq4", "reviewer": "Reviewer_CF1F", "summary": "This paper studies the learning dynamics of deep linear networks trained using gradient descent at a fixed learning rate, especially beyond the edge of stability (EOS) regime.\n\nUsing the singular value stationary set and the strictly balancing assumptions (both can be validated in examples), the paper argues that within the EOS regime, periodic 2-period fixed orbit oscillations occur in a low-dimensional subspace, with the subspace dimension determined by the learning rate. This phenomenon explains the presence of oscillations within top features. Additionally, the paper compares the difference between deep linear networks and diagonal linear networks in the EOS regime, offering insights into how network structure impacts training dynamics. Experiments further reveal that fine-tuning with large learning rates in LoRA for low-rank adaptation of large language models induces \"catapult dynamics,\" which can potentially improve generalization.", "questions": "1. Why is the period of the fixed orbit is always 2? It seems that it comes from the two-step GD update. How about other periodicities? \n\n2. (15)-(16) on page 19-20 seem to have typos.  $\\dot{W}_l(t)$ on the left side should be $\\dot{U}_l(t)$?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This work is both theoretically rigorous and experimentally thorough. The theoretical results are new to me, especially Theorem 1, and the proofs are rigorous and not difficult to follow. Additionally, the experiments are extensive and effectively validate the theoretical insights. Overall, the study contributes valuable advancements in understanding and optimizing the training dynamics of deep networks.", "weaknesses": "Although deep linear network shares some similarities with deep nonlinear networks, the claim that ``Our analysis explains two key phenomena in deep nonlinear networks\" in abstract seems overstated and exaggerated. In the analysis, the initialization (eqn (3)) is very specific, and it is just a member in the singular vector stationary set. It is not clear if all statements work only based on the initialization (eqn (3)) or for any initialization in the singular vector stationary set. For example, will Lemma 1 and Lemma 2 work with general singular vector stationary initial data?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tUzQxdXvim", "reviewer": "Reviewer_GbsE", "summary": "The paper analyzes learning dynamics of deep linear networks a.k.a. deep matrix factorization using Gradient Descent (GD) in the Edge-of-Stability (EOS) regime. The paper proves that, under a \"strict balanced state\" condition of singular values exactly matching across the weight matrices, the top singular values oscillate, leading to oscillations in training loss in the EOS regime. The paper also presents an experiment showing loss oscillations in Low-Rank-Adaptation (LoRA) finetuning on a language modeling task. Results in the paper help explain some observations in existing work.", "questions": "# High-level questions\n1. Theorem 1. Why is this called 1-d subspace oscillation?\n2. Figure 2. There are 2 kinks in the EOS regime in the left plot but only 1 is visible in the right plot. Do the authors know why?\n3. Section 4.2. The original LoRA paper [1] only applies adaptation to the attention weight matrices. Any reason the authors adapt all weights?\n4. Line 450. Pearson correlation between what two quantities? If this is the metric from the STS-B task, then how would this kind of experiment extend to other tasks?\n5. Line 474. Why do smooth, low-frequency images correspond to low task complexity? For the CIFAR dataset, number of samples $N$ is used as a proxy task complexity, why? Both datasets are image datasets. It seems to me that I can pick and choose any proxy/hyperparameter from the experiment setup to argue that sharpness does not rise to EOS regime in low complexity tasks.\n\n# Low-level questions\n1. Figure 1. What is the pearson correlation in the last plot? A pointer to the relevant section would help.\n2. Line 71. What is \"sharpening\"? First time this term is mentioned.\n3. Lines 249-254. Lost here, what are $f_{\\Delta_i}$, $x$, and the line? I think more text is needed to explain how the analysis is done for $r$-dimensional oscillations. A figure would really help the reader.\n\n[1] Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., ... & Chen, W. (2021). Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685.\n\n\n*************\n**Update after author response:** I am increasing my review score from 3 to 6 since the authors addressed my [concerns](https://openreview.net/forum?id=J4Dvxv7WnG&noteId=9ZRUV9lccn). The manuscript has been updated significantly, but not to the extent that I do not understand the central claims. Hence, I'm keeping my original confidence score at 3. Thank you to the authors for a good discussion.\n*************", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The EOS oscillation theorems for deep linear networks are novel, and they directly map to simulations of training such networks using GD (Figs 2,3,6). The paper discusses oscillations in smaller subspaces first (Theorem 1), then extends the analysis to larger subspaces (Theorem 2). The theorems help theoretically ground observations seen in existing work, on loss oscillations in the EOS regime. It is curious that these oscillations can be somewhat explained by singular value oscillations in the top subspaces for deep linear networks.", "weaknesses": "There are 3 main weaknesses:\n\n1. While it is useful to understand learning dynamics of deep linear networks from a theoretical perspective*, I'm not sure the results shed light on dynamics in deep nonlinear networks. In fact, the paper misinterprets the result of [1], which show that only for shallow networks are linear and nonlinear networks (2-layer bias-free ReLU networks) similar---for deep networks there is a strict separation in function approximation. [1] reports similarities between deep linear and nonlinear networks under very strong assumptions on data distribution and structure, which this manuscript does not consider. This submission misses this nuance (in abstract and lines 80-83), and overpromises results for deep nonlinear networks. In fact, the manuscript comments about the differences in loss landscapes in Line 511---nonlinear networks have highly irregular landscapes, and when trained with (minibatch) SGD could lead to catapults across regions not just in the vicinity. This is a major weakness in my view. \n2. The $\\alpha$-scaled initialization scheme is not used in practice, especially for LoRA training where one of the matrix (A) is initialized at 0 and the other (B) with gaussian entries. The paper argues that the results work when initialization leads to convergence to the \"singular vector stationary set\" (line 148). Why is this not a strong and impractical assumption on how initialization is done in practice? Does random initialization satisfy this condition?\n3. As highlighted in Section 5 by the authors, the assumptions on singular values' \"strict balanced state\" and initialization's \"stationary set condition\" are strong. It's not clear if and when they can be met.  \n\n*The authors need not motivate deep linear networks by (wrongly) arguing their similarity to deep nonlinear networks; linear networks have been studied in the ML optimization research for a while now.\n\n# Experiments\nSection 4.2 LoRA initialization does not match practice. In LoRA, one of the matrices is set to 0 and the other to random, so that $AB^T$ is 0 at finetuning initialization. This is different from the alpha-initialization, which is (1) not random and (2) not starting from 0. I don't think this experiment works well with the theoretical setting studied in this paper---the LoRA finetuning experiment seems ad-hoc when the rest of the paper is about deep linear networks.\n\nMoreover, the \"catapult\" observation in LoRA experiments is on very shaky ground, since stochasticity affects the training loss and makes it oscillate. The paper mentions this in Line 431, so I'm not convinced that the oscillations are, in fact, \"catapults\".\n\nA natural experiment to try: In Line 188, step size is suggested to be chosen based on the depth $L$ of the network (as $L$ is larger, smaller values of step size will lead to oscillcations). This is curious to verify.\n\n# Writing\n1. In the introduction, 3 observations are described from [2, 3], but it is not clear to me how or which ones the paper addresses. It is also not clear to me how sharpening and catapults from the two papers are related to each other.\n2. Definition 2 is introduced abruptly without need or context. There is discussion in Lines 191-195, which can be better placed before/after Def. 2. \n3. The paper is missing a related work section. It is included as Appendix A.1, but I'd suggest cutting down on the main sections to include related work within 10 pages.\n\n[1] Zhang, Y., Saxe, A., & Latham, P. E. (2024). When Are Bias-Free ReLU Networks Like Linear Networks?. arXiv preprint arXiv:2406.12615.\n\n[2] Cohen, J. M., Kaur, S., Li, Y., Kolter, J. Z., & Talwalkar, A. (2021). Gradient descent on neural networks typically occurs at the edge of stability. arXiv preprint arXiv:2103.00065.\n\n[3] Zhu, L., Liu, C., Radhakrishnan, A., & Belkin, M. (2023). Catapults in SGD: spikes in the training loss and their impact on generalization through feature learning. arXiv preprint arXiv:2306.04815.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CZuDvuAmli", "reviewer": "Reviewer_n4c2", "summary": "This paper rigorously analyzes gradient descent dynamics of deep linear networks for deep matrix factorization at the Edge of Stability. Assuming initialization within the singular vector stationary set and strictly balanced weights, the authors characterize a 2-period oscillation within top subspaces. The main insight is that, under these assumptions, the loss can be rewritten as a function of singular values, decoupling the effect of singular vectors. The authors also provide empirical support for the strict balancing assumption, showing that singular values become increasingly balanced during gradient descent dynamics, with numerical experiments aligning well with the theoretical findings.", "questions": "1. Could the authors please include the following references in the related works section? Prior to [Cohen et al., 2021], [Jastrzebski et al., 2019] and [Jastrzebski et al., 2020] demonstrated that step size influences sharpness along optimization trajectories. Additionally, [Ahn et al., 2023], [Song et al., 2023], and [Karla et al., 2023] provide rigorous analyses of learning dynamics at the Edge of Stability in simplified settings, such as two-layer linear networks. [Zhu et al., 2024] and [Chen et al., 2024] study gradient descent dynamics for quadratic models in large learning rate regimes where catapults occur.\n\n2. When initialization is outside the singular vector invariant set, how do loss and sharpness behave at the Edge of Stability? Specifically, does the weight converge to a period-2 orbit, or does the loss oscillate while decreasing non-monotonically? Similarly, if oscillations begin before strict balance is achieved, how are the loss trajectory and learning dynamics affected? It would be insightful if the authors could provide additional experimental results exploring these scenarios.\n\n---\n\n**References**\n\n[Jastrzebski et al., 2019] On the Relation Between the Sharpest Directions of DNN Loss and the SGD Step Length, ICLR 2019.\n\n[Jastrzebski et al., 2020] The Break-Even Point on Optimization Trajectories of Deep Neural Networks, ICLR 2020.\n\n[Cohen et al., 2021] Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability, ICLR 2021.\n\n[Ahn et al., 2023] Learning threshold neurons via the “edge of stability”, NeurIPS 2023.\n\n[Song et al., 2023] Trajectory Alignment: Understanding the Edge of Stability Phenomenon via Bifurcation Theory, NeurIPS 2023.\n\n[Karla et al., 2023] Universal Sharpness Dynamics in Neural Network Training: Fixed Point Analysis, Edge of Stability, and Route to Chaos, arXiv 2023.\n\n[Zhu et al., 2024] Quadratic models for understanding neural network dynamics, ICLR 2024.\n\n[Chen et al., 2024] From Stability to Chaos: Analyzing Gradient Descent Dynamics in Quadratic Regression, TMLR 2024.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper makes a solid contribution to the theoretical understanding of the Edge of Stability by analyzing gradient descent dynamics in deep matrix factorization. It extends previous studies on deep scalar linear networks [Kreisler et al., 2023], offering a rigorous proof that oscillations in the EoS regime are confined to top subspaces, as empirically observed in prior work [Zhu et al., 2024]. Additionally, the exact characterization of Hessian eigenvalues at convergence under strict balancing provides a valuable technical contribution to deep matrix factorization problem.\n\n---\n\n**References**\n\n[Kreisler et al., 2023] Gradient Descent Monotonically Decreases the Sharpness of Gradient Flow Solutions in Scalar Networks and Beyond, ICML 2023.\n\n[Zhu et al., 2024] Catapults in SGD: spikes in the training loss and their impact on generalization through feature learning, ICML 2024.", "weaknesses": "1. **My main concern is that Lemma 2 does not ensure convergence to strict balancing.** While Lemma 2 shows that the balancing gap decreases with each step, it does not guarantee that this gap converges to zero. The authors suggest that Lemma 2 justifies the strictly balanced assumption, but this needs further clarification. Although Figure 4 empirically shows the gap converging to zero, Lemma 2 alone does not rigorously establish this phenomenon.\n\n2. **The model fails to capture (non-monotonic) decreasing training loss in the Edge of Stability.** A key characteristic of the Edge of Stability is that training loss oscillates but decreases over time, reflecting unstable convergence (e.g., [Cohen et al., 2021], [Ahn et al., 2022]). In this study’s setup, weights oscillate in a 2-period orbit, causing the loss to oscillate indefinitely without decreasing. This limits the model’s applicability to realistic settings.\n\n3. **The claims on LoRA dynamics in Section 4.2 lack sufficient evidence.** The authors claim that using large learning rates in LoRA dynamics leads to oscillations and catapults, enhancing generalization by increasing Pearson correlation. However, Figures 7 and 8 do not clearly show catapult behavior, as oscillations occur even at lower learning rates (e.g., $\\eta = 10^{-5}, 10^{-6}$), making it difficult to conclude that catapults are unique to $\\eta=10^{-4}$. Additionally, empirical evidence linking catapults to better generalization is insufficient. The claims here seem overstated and less relavant to the main theory. I recommend moving this section to the appendix and avoid making strong claims.\n\n\nMinor comments\n- (Section 3, Page 4, Line 169) Typo: unbold subscript $k$ in $W_k$.\n- (Appendix A, Page 14, Line 723) Wrong citation: \"Ahn et al. (2022) established the phenomenon in two-layer networks...\" is citing the paper [Ahn et al., 2022], but it should be citing another paper [Ahn et al., 2023].\n- (Reference, Page 13) Duplicate entries for Zhu et al. (2024) in the reference (Line 692 and Line 696).\n\n---\n\n**References**\n\n[Cohen et al., 2021] Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability, ICLR 2021.\n\n[Ahn et al., 2022] Understanding the Unstable Convergence of Gradient Descent, ICML 2022.\n\n[Ahn et al., 2023] Learning threshold neurons via the “edge of stability”, NeurIPS 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PvpvUpZhGa", "reviewer": "Reviewer_nnRE", "summary": "This paper analyzes the catapult behaviors of DLNs (Deep Linear Networks) and explains many unexplained observations regarding EOS (Edge of Stability) phenomenon: (i) periodic subspace oscillations (especially in the top subspace that corresponds with more prominent features), (ii) mild (or no) sharpening for simple dataset with low complexity (when the eigenvalue for the strongest feature is too small), (iii) difference between DLNs and DiagLNs (i.e., the architecture matters), and (iii) catapult in LoRA.", "questions": "The following questions may be related to the weaknesses as the paper may have some unclear points:\n- In Thm 1, what is $\\rho_i(t)$ exactly? It only says that $\\rho_i(t)\\in \\\\{\\rho_1,\\rho_2\\\\}$. Does $i$ depend on $t$? As it says that the matrix oscillates, it is likely that $\\rho_i(t)$ changes with $t$. Is it $\\rho_i(t)=\\rho_1$ if $t$ is odd (even) and $\\rho_i(t)=\\rho_2$ if $t$ is even (odd)? If the readers have to guess the meaning, then it is not well-written.\n\n- What do you mean by that \"To show oscillations in two or more subspace, we can easily extend Theorem 1\"? Why do we need to set $\\lambda_2\\geq K>\\lambda_3$? I don't think this is a trivial result from Thm 1. Can you elaborate more on the paragraph L241-248? My understanding is that $\\rho_i(t)$ in the first term of $W_\\ell(t)$ in Thm 1 corresponds to the oscillation, is it right? I didn't fully understand the $r$-subspace oscillation part.\n\n- For the second term $\\sum_j \\sigma_{\\ast,i}u_{\\ast,j}v^\\top_{\\ast,j}$, isn't it $\\sigma_{\\ast,j}$ with $j$, not $\\sigma_{\\ast,i}$ with $i$?\n\n- Is there any reason that we can view the adaptations as individual low-rank matrix factorization problems? Can you elaborate more on this? I hope some \"math\" may help the reader to understand this.\n\n- \"Notably, for ranks $r=4$ and $r=8$, there are catapults ... do not occur for $r=1$ or $r=2$\". What catapults are you talking about? In Fig 7, it is hard to see the catapult (or loss oscillation) and compare it with other ranks ($r=1,2$). I don't fully understand Section 4.2.\n \n- Can you somehow compute or estimate the $\\sigma_{\\ast,1}$ for Fig 10 (a) and for each image in Fig 10 (b)? It would be better to have a quantitative understanding of the \"low-complexity learning\".\n\n- L 29-30 (spacing): ... 2020).The -> ... 2020). The\n- L415: give a ... -> given a ...\n-", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper explains many unexplained observations regarding EOS phenomenon.", "weaknesses": "- Lemma 2 does **not** tell us that the singular values become balanced as $t\\rightarrow \\infty$ (L191-193). The strictly decreasing balancing gap and the difference being lower bounded by zero does **not** lead to the convergence to zero (L307-309). Think about the sequence $0.5+\\frac{1}{t}$. Can you provide a proof to show the convergence to zero?\n\n- \"When the learning rate is large enough to induce catapults, we observe that the training loss decreases rapidly ...\" \nThe statement does not seem right. Large lr may lead to fast decreasing of the training loss because of its large stepsize not because of the catapults. Same for the statement \"when the learning rate is small, convergence takes much longer, as the model seems to bounce around within the same local basin before reaching a low training loss\". Can you provide an experiment to show that \"the catapult implies faster optimization\"?\n\n- It is hard to say \"in DLNs, self-stabilization does not occur\". Actually, also in the original EOS paper, the sharpness oscillates **above** $2/\\eta$ (they say \"it hovers above $2/\\eta$\"). It is true that DLNs do not satisfy the condition of the proposition in the self-stabilization paper, but self-stabilization is a broader concept and the threshold may not be necessarily $2/\\eta$. The statement should be written carefully.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)