Learning Dynamics of Deep Matrix Factorization Beyond the Edge of Stability

Authors: Avrajit Ghosh, Soo Min Kwon, Rongrong Wang, Saiprasad Ravishankar, Qing Qu
Primary area: optimization
Keywords: edge of stability, deep matrix factorization

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.

Decision: accept (poster)
Reviewer scores: 8, 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1730461956950

thread_type:
official_review

<thread_object>
creator: reviewer_nnre

<content>
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.

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?

- "When the learning rate is large enough to induce catapults, we observe that the training loss decreases rapidly ..." 
The 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"?

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

questions:
The following questions may be related to the weaknesses as the paper may have some unclear points:
- 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.

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

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

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

- "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.
 
- 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".

- L 29-30 (spacing): ... 2020).The -> ... 2020). The
- L415: give a ... -> given a ...
-
</content>
</thread_object>

</thread>
<thread>
creation date:
1730637285286

thread_type:
official_review

<thread_object>
creator: reviewer_n4c2

<content>
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.

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.

---

**References**

[Kreisler et al., 2023] Gradient Descent Monotonically Decreases the Sharpness of Gradient Flow Solutions in Scalar Networks and Beyond, ICML 2023.

[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.

2. **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.

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


Minor comments
- (Section 3, Page 4, Line 169) Typo: unbold subscript $k$ in $W_k$.
- (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].
- (Reference, Page 13) Duplicate entries for Zhu et al. (2024) in the reference (Line 692 and Line 696).

---

**References**

[Cohen et al., 2021] Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability, ICLR 2021.

[Ahn et al., 2022] Understanding the Unstable Convergence of Gradient Descent, ICML 2022.

[Ahn et al., 2023] Learning threshold neurons via the “edge of stability”, NeurIPS 2023.

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.

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

---

**References**

[Jastrzebski et al., 2019] On the Relation Between the Sharpest Directions of DNN Loss and the SGD Step Length, ICLR 2019.

[Jastrzebski et al., 2020] The Break-Even Point on Optimization Trajectories of Deep Neural Networks, ICLR 2020.

[Cohen et al., 2021] Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability, ICLR 2021.

[Ahn et al., 2023] Learning threshold neurons via the “edge of stability”, NeurIPS 2023.

[Song et al., 2023] Trajectory Alignment: Understanding the Edge of Stability Phenomenon via Bifurcation Theory, NeurIPS 2023.

[Karla et al., 2023] Universal Sharpness Dynamics in Neural Network Training: Fixed Point Analysis, Edge of Stability, and Route to Chaos, arXiv 2023.

[Zhu et al., 2024] Quadratic models for understanding neural network dynamics, ICLR 2024.

[Chen et al., 2024] From Stability to Chaos: Analyzing Gradient Descent Dynamics in Quadratic Regression, TMLR 2024.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730697872996

thread_type:
official_review

<thread_object>
creator: reviewer_gbse

<content>
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.

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:

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

*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.

# Experiments
Section 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.

Moreover, 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".

A 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.

# Writing
1. 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.
2. 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. 
3. 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.

[1] Zhang, Y., Saxe, A., & Latham, P. E. (2024). When Are Bias-Free ReLU Networks Like Linear Networks?. arXiv preprint arXiv:2406.12615.

[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.

[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.

questions:
# High-level questions
1. Theorem 1. Why is this called 1-d subspace oscillation?
2. 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?
3. Section 4.2. The original LoRA paper [1] only applies adaptation to the attention weight matrices. Any reason the authors adapt all weights?
4. 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?
5. 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.

# Low-level questions
1. Figure 1. What is the pearson correlation in the last plot? A pointer to the relevant section would help.
2. Line 71. What is "sharpening"? First time this term is mentioned.
3. 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.

[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.


*************
**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.
*************
</content>
</thread_object>

</thread>
<thread>
creation date:
1731120682890

thread_type:
official_review

<thread_object>
creator: reviewer_cf1f

<content>
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.

Using 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.

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?

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? 

2. (15)-(16) on page 19-20 seem to have typos.  $\dot{W}_l(t)$ on the left side should be $\dot{U}_l(t)$?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732031465309

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response to All Reviewers

comment:
Firstly, we would like to thank the reviewers for taking the time to provide insightful comments. We are encouraged that the reviewers find our work “theoretically rigorous and experimentally thorough” (Reviewer CF1F) and “a solid contribution to the theoretical understanding of the Edge of Stability” (Reviewer n4c2). Based on the comments, we have uploaded a revised manuscript, where the revisions can be summarized as follows:
- Toned down our claims on deep nonlinear networks to say that our results shed light on certain phenomena in nonlinear networks rather than directly explaining them. 
- Moved LoRA to the Appendix as suggested by Reviewer n4c2 and edited the manuscript to have a stronger focus on DLNs. 
- Clarificaiton regarding the use of singular vector stationary set. 


In the following, we provide a few more thorough responses regarding shared concerns amongst the reviewers.

**Comment:** Overclaiming results on deep nonlinear networks using our results in deep linear networks.

**Response:** While we believe that our results still shed some light on phenomena that occur in deep nonlinear networks within the EOS regime, we agree that our claims may have been slightly overexaggerated. As the reviewers point out (and stated in our experiments section), DLNs fail to capture catapults that occur in nonlinear networks due to the difference in the loss landscape. To address this point,  we have revised our manuscript in two ways: (i) toned down the contributions such that we say that our results on DLNs contributes to an understanding of phenomena in EOS for nonlinear networks rather than directly explaining them; (ii) provide a fair comparison on the similarities and differences between the two networks in Section 4.2. See [Figure 1](https://imgur.com/3UZNmjP) and [Figure 2](https://imgur.com/7xlMhUA) for landscape visualization for DLNs and complex non-convex function such as Holder table loss at EOS. We chose the Holder table loss function as a proxy for deep non-linear network which is known to have numerous local minimas. Whereas deep linear networks contains only saddle points and global minimas (no spurious local minimas).


**Comment:** The use of the singular vector stationary set is a strong assumption.

**Response:** While this may seem like a strong assumption, we do observe that many of the commonly used initialization schemes ultimately converge to the singular vector stationary set for DLNs. In [Figure 15](https://imgur.com/xON7b4q), we show that the identity, orthogonal, and random initializations all converge to this set empirically. Recall that we do prove that the identity initialization we considered in the original manuscript converged to this set. Since it is exhaustive to prove that all initializations converge to such a set, we used Proposition 1 to cover an extensive base called the SVS set. Theorem-1 for rank-1 subspace oscillation can be easily modified when the initialization is made inside any general point in the SVS set. 


**Comment:** Concerns about LoRA

**Response:** Based on most reviewers' comments, we moved this section to the Appendix and left it as future work for more investigation.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper studies the edge of stability (EoS) phenomenon in deep matrix factorization, also known as deep linear neural networks. The paper analyzes the 2-period oscillation in GD dynamics when the learning rate is larger than the stability threshold, under two assumptions: singular vector stationary set, and strict balanced state.

The reviewers unanimously found the paper to be an interesting submission. To me as well, it is deemed that the paper delivers useful insights into the low-dimensional oscillations in the EoS regime and offers technical tools for analyzing deep matrix factorization. A downside of the paper is that the main results hold under two assumptions: singular vector stationary set, and strict balanced state, as the authors also note in the conclusion.

Although I recommend acceptance of this paper, I am concerned about a slight lack of precision in the authors’ claims regarding their assumptions. It seems to me that the authors assume that some properties that only hold for GF (the continuous time dynamics) also hold for GD with large step sizes. The preservation of balancedness is a good example: For GF, once $W_{j+1}^T(0) W_{j+1}(0) = W_j(0) W_j(0)^T$, the property continues to hold forever. However, for GD the balancedness property slightly breaks every iteration due to the “discretization error”. For example, consider the problem $L(A,B) = 1/2*  \\lVert A B - W \\rVert_F^2$. The GD update for $A$ and $B$ look like $A_+ = A - \eta (A B - W) B^T$ and $B_+ = B - \eta A^T (A B-W)$. Even if $A$ and $B$ satisfy $A^T A = B B^T$, we can realize that $A_+^T A_+ = A^T A - \eta B (A B-W)^T A - \eta A^T (A B - W) B^T + \eta^2 B (A B-W)^T (A B-W) B^T$ and $B_+ B_+^T = B B^T - \eta B (A B-W)^T A - \eta A^T (A B - W) B^T + \eta^2 A^T (A B-W) (A B-W)^T A$, and hence the $O(\eta^2)$ terms do not match in general.

Nevertheless, I believe this issue may not be grounds for rejection because 1) the preservation of balancedness may be true under the singular vector stationary set assumption, 2) the authors may be able to extend Lemma 2 and show that the balance is (approximately) preserved in the balanced initialization case, 3) the empirical observations show balancedness should hold anyway. In any case, I encourage the authors to take a serious look into this matter and make proper adjustments.