Does SGD really happen in tiny subspaces?
Authors: Minhak Song
Kwangjun Ahn
Chulhee Yun
Venue: ICLR 2025 Poster (2025)
Keywords: optimization for deep networks
training dynamics
SGD
Hessian
low-rank subspace
Primary area: optimization
TL;DR: Deep neural networks cannot be trained within the dominant subspace, even though gradients align with this subspace along the training trajectory.

Abstract: Understanding the training dynamics of deep neural networks is challenging due to their high-dimensional nature and intricate loss landscapes. Recent studies have revealed that, along the training trajectory, the gradient approximately aligns with a low-rank top eigenspace of the training loss Hessian, referred to as the dominant subspace. Given this alignment, this paper explores whether neural networks can be trained within the dominant subspace, which, if feasible, could lead to more efficient training methods. Our primary observation is that when the SGD update is projected onto the dominant subspace, the training loss does not decrease further.
This suggests that the observed alignment between the gradient and the dominant subspace is spurious. Surprisingly, projecting out the dominant subspace proves to be just as effective as the original update, despite removing the majority of the original update component. We observe similar behavior across practical setups, including the large learning rate regime (also known as Edge of Stability), Sharpness-Aware Minimization, momentum, and adaptive optimizers. We discuss the main causes and implications of this spurious alignment, shedding light on the dynamics of neural network training.

Published as a conference paper at ICLR 2025
DOES SGD REALLY HAPPEN IN TINY SUBSPACES?
Minhak Song KAIST Math
K Kwangjun Ahn Microsoft Research
Chulhee Yun KAIST AI
ABSTRACT
Understanding the training dynamics of deep neural networks is challenging due their high-dimensional nature and intricate loss landscapes. Recent studies have revealed that, along the training trajectory, the gradient approximately aligns with a low-rank top eigenspace of the training loss Hessian, referred to as the dominant subspace. Given this alignment, this paper explores whether neural networks can be trained within the dominant subspace, which, if feasible, could lead to more efficient training methods. Our primary observation that when the SGD update is projected onto the dominant subspace, the training loss does not decrease further. This suggests that the observed alignment between the gradient and the dominant subspace is spurious. Surprisingly, projecting out the dominant subspace proves to be just as effective as the original update, despite removing the majority of the original update component. We observe similar behavior across practical setups, including the large learning rate regime (also known as Edge of Stability), Sharpness-Awee Aware Minimization, momentum, and adaptive optimizers. We discuss the main causes and implications of this spurious alignment, shedding light on the dynamics of neural network training.
560
0.6
SGD
0.5
SGB
1.0-
Dom-5GD
Dom-SGD
Dom SGD
Bulk: SGD
-0.8
Bul:5GD
-1.0
50D
-1.5
-2.0
1.0
2.5
-2.0
-1.2
6
5000
10000 15000
20000
5000 10000 15000 20000
5000 10000 15000 20000 Training Steps
Training Steps
Training Steps
(a) MLP MNIST-5K
(b) CNN CIFARIO-5K
(c) Transformer SST2-1K 1k
Figure 1: The summary of our main results in Section 3 (training loss in log-scale). For neural network training. Gur-Ari al. (2018) observe that gradients approximately align with the dominant subspace, spanned by the dominant eigenvectors of the training loss Hessian. To see whether such phenomenon lets us train neural networks within the dominant subspace, we implement Dom-SGD, where each SGD update is projected onto the dominant subspace. Surprisingly. training stops after this modification, suggesting that the dominant subspace is not where the learning happens. In contrast, Bulk-SGD, where we project each SGD updates onto the bulk subspace orthogonal to the dominant subspace, is just as effective as the original update, despite removing the majority of original updates. Experimental details are provided in Appendix B.
INTRODUCTION
Understanding the optimization of deep neural networks presents a complex challenge, given their nature and the intricate characteristics of their training loss landscapes. Over the last decade, an abundance of studies has investigated the landscape of training loss L : R - (Li et al., 2018b; He et al., 2019). In this work, we are interested in the following noteworthy phenomena:
Hessian is approximately low-rank. Extensive research (Sagun et al., 2016; 2017; Ghor- bani al., 2019; Papyan, 2019; 2020) has revealed that for k-class classification problems, the loss Hessian v exhibits a "low-rank" structure, characterized by &o dominant eigenval- ues significantly larger than the others. See Figure 2 for details.

---

Published as a conference paper at ICLR 2025
200
150
1000
20
0
0-
5000 10000 15000
20000
0 5000 10000 15000
20000
5000 10000 15000 20000 Training Steps
Training Steps
Training Steps
(a) MLP MNIST-5K
(b) CNN on CIFARIO-5k
(c) Transformer on SST2-1K
Figure 2: Low-rank structure of the Hessian. The plot shows the top eigenvalues of the loss Hessian during SGD training. The blue curves represent the top-k eigenvalues, which are significantly larger than the next top-k cigenvalues, shown in orange. Here, k corresponds the number of classes in the classification task
10 for MNIST-5K, CIFARIO-5K and k 2 for SST2-1k).
1.00
1.00
1.00
0.75
0.75
0.75
0.50
0.50
0.50
0.25
0.25
0.25
0.00
5000 10000 15000 20000
50D0 10000 15000 20000
0
5000 10000 15000 20000 Training Steps
Training Steps
Training Steps
(a) MLP MNIST-5K
(h) CNN on CTFARIO-5k
(c) Transformer on SST2-1k
Figure 3: Alignment of gradients with dominant subspaces. The plot illustrates X (Q.) during SGD training, where is the number of classes for the classification task (see Definition 2). The orange dashed lines represent the exponential moving average (EMA) of X& (v(VL((.)) After few early steps, X (6.) reaches and stays near 1, indicating the alignment between gradients and dominant subspaces.
. Gradients approximately align with the low-rank eigenspace. Gur-Ari et al. (2018) demonstrated that during SGD training, gradients tend to align closely with the low- dimensional subspace spanned by the dominant eigenvectors of the loss Hessian. See Figure 3 for details.
We provide more extensive background on these phenomena in Appendix A. The eigenspace of the top-k cigenvalues of (²L( referred to as the dominant subspace at 0, is the main focus this work. Motivated by Gur-Ari et al. (2018), ask:
Q. Can deep neural networks be trained within the dominant subspace?
(Q1)
This question carries significant practical implications, potentially leading more efficient training methods for neural networks. Furthermore, it offers insights into why deep learning optimization may not suffer from the curse of dimensionality despite operating in high-dimensional spaces.
1.1 SUMMARY MAIN RESULTS
In this paper, we rigorously examine the question (Q1) through systematic experiments. Quite surprisingly, our results reveal that the answer the question is negative, as summarized below.
In Section 3, we demonstrate that the observed alignment is "spurious" in the sense that the aligned component of the gradient is not beneficial for training, even though it constitutes the majority of the gradient. Specifically, projecting each update onto the dominant subspace; we call this Dom-SGD. Unexpectedly, Dom-SGD does further decrease the training loss. Next, we consider Bulk-SGD which projects each update onto the bulk subspace, i.e., the orthogonal complement of the dominant subspace. Despite the fact that the major component of each update is removed, this time the training remains as effective as the original update (see Figure 1).
run critical experiment where we modify SGD by
In Section 4, we identify that the spurious alignment between the gradient and the dominant subspace is caused by the stochastic noise inherent to SGD in the gradient flow (GF)
2

---

Published as a conference paper at ICLR 2025
regime, by showing that alignment disappears when using full-batch GD (see Figure and Figure 5). Additionally, we present a simple quadratic model which also captures the observed phenomena, providing insights into the role of stochastic noise in the spurious alignment (see Figure 6 and Figure 7).
In Section 5, we extend our observations to two other practical settings: (1) GD in the Edge of Stability (EoS) regime (Cohen et al., 2021), and (2) Sharpness-Aware Minimization (SAM) (Foret al., 2021). ). For both of the settings, we again observe that each update approximately aligns with the dominant subspace, yet the aligned component of each update does not contribute to the loss decrement (see Figure and Figure 10). For GD in the EoS regime, the alignment is due to the self-stabilization mechanism (Damian et al., 2023), in contrast to SGD in the GF regime where stochastic noise is the cause of the alignment.
In Section 6, we again observe that for momentum and adaptive optimizers, e.g. Adam, we project the update vector onto the dominant subspace, the training loss fails to decrease. Moreover, we demonstrate that momentum and adaptive methods amplify the bulk subspace component of each update, partially explaining their success in neural network training (see Table and Figure 11).
STARTING POINT: GRADIENT ALIGNS WITH THE DOMINANT SUBSPACE
In this section, set the stage for main results by reviewing the main observation of Gur-Ari et al. (2018). To that end, we first introduce some notations to ease our discussion.
Let [rx] denote the set {1,2..... n). For the Hcssian to be well-defined, let L R2 - R be a twice-differentialee training loss. For E IR", let A1 (0), A2(0), Xp(()) denote the eigenvalues of the loss Hessian L() E R2xp in descending order, and let u (0), (0))0,, ..., (0) denote the corresponding eigenvectors. Given these notations, begin with the most important concept for our discussion, namely, the dominant subspace.
Notations.
Definition (Dominant subspace). The top-k dominant subspace at is denoted by
(x) span (0) : [M]]]
and its orthogonal complement by S (0), referred to the bulk subspace. Unless specified otherwise,
the default choice for is the number of classes for the classification task.
Definition 2 (Dominant subspace projection). The projection matrix onto the dominant subspace Sy(0) is denoted by
k
(x) = Σ (())uu(()) E RPXP 1=1
and the projection matrix onto St by Pr. decompose the vector into v == (0) + Ps (0)z. We say P& (f) is the dominant component of v, and Py. (0)v is the bulk component u.
(B). For a given vector v E TR,,, we can
v. We denote the fraction of u in the dominant subspace by
(v; = (0)x(())ll///1)))
with (u; 0) = of = 0. A vector u E IP is said to (approximately) align with the dominant subspace (0) if (v; is close to 1, and align with St (0) if (u; 0) is close to O.
When clear from context, use shorthand notation such as X == A²( u4 := u. (L) VL S. : Sx (0), S = St (0), Pa == Po) P. = PE (0), and (x(u) == ()(
L())..
Using our notations, the striking observation of Gur-Ari et al. (2018) can be formalized as follows. Phenomenon 1. Gradient approximately aligns with the dominant subspace along SGD tra- jectories. Consider the SGD trajectory {0t) with a constant learning rate. After aefew initial steps, Xx (()))) quickly reaches and remains near 1.
In Figure 3, we confirm the main results of Gur-Ari et al. (2018) for various settings. Notice that X& (v (0,)) reaches 1 after a few early steps, indicating that the gradient (L,),) approximately aligns with the dominant subspace (0,). Given this alignment, it seems that the training can be done within the dominant subspace, which leads to the previously introduced question (Q1). In the next section, conduct set of experiments to investigate (Q1).
3

---

Published as a conference paper at ICLR 2025
3 NEURAL NETWORKS CANNOT BE TRAINED WITHIN DOMINANT SUBSPACES
In this section, we present the first main observation of this paper regarding question (Q1). We start with preliminary analysis using the local quadratic approximation of the neural network landscape.
3.1 WHAT DO WE EXPECT BASED ON QUADRATIC TAYLOR APPROXIMATION?
To analyze the convergence of gradient- based optimization algorithms, a common approach is to use the local quadratic Taylor approximation (see, c.g., (Ghadimi & Lan, 2013)). The "descent lemma" characterizes the one-step progress of the optimizer L(O++1) - ((e,) using this approximation, assuming the training loss I. is smooth. Based the local quadratic Taylor approximation, have:
I(0++1) I(O0) s (VL(0.). 0,+1 0₂) (0.+1 (0++1-0.) v.. 0t)
(1)
ggradieent correlation
Let us denote the first term on the RHS by the gradient correlation term and the second term by the second-order error term. Since the SGD updates are defined
8(+1 f- 8₂ - 79tt
for the stochastic gradient 9r 0,, the gradient correlation term is negative in expectation:
Elgradient correlation] ( -1g)) - | (())l < <.
(2)
For the experiments in Figures la and 1b, we use small learning rates to ensure SGD closely follow the continuous-tinn time gradient flow so that the training loss decreases nearly monotonically, suggesting that the negative gradient correlation dominates the second-order error term in these cases.
Hence, if the quadratic Taylor approximation was accurate enough, based on the above analysis and Phenomenon 1, it is expected that one can decrease the training loss based on updates lying in the dominant subspace, as hypothesized in the question (Q1).
To directly test this hypothesis, we design the following critical experiment.
Our eritical experiment: In the same settings before, whenever Phenomenon occurs, consider the following updates where each update of SGD is projected onto the dominant subspace:
O.+1 + - 1 ()
(Dom-SGD)
By Phenomenon 1, Dom-SGD has approximately same gradient correlation SGD given in (2):
Eggadiient correlation] ([(VV()))
-n -[.)
Therefore, based on the local quadratic Taylor approximation, Dom-SGD should be able to success- fully train neural networks whenever Phenomenon I occurs. Is it really the case?
3.2 THE "SPURIOUS" ALIGNMENT WITH THE DOMINANT SUBSPACE
In the same settings as before, we first train neural networks with SGD until we observe Phe- nomenon 1. Specifically, we track the exponential moving average (EMA) of xx (VLL(()) values (EMA factor set to 0.9), and switch from SGD to Dom-SGD whenever the EMA value exceeds 0.95. Note that we recompute the dominant subspace every step when running Dom-SGD.
For various settings, we plot the training loss of Dom-SGD in Figure 1, comparing it with SGD under the same initialization. We employ constant learning rate and mean squared eтror (MSE) loss for classification (Hui & Belkin, 2021; Cohen et al., 2021). Additional experiments, including those using cross-entropy loss and training standard architectures, are provided Appendix
Consistently, we observe that Dom-SGD fails to further decrease the training loss, unlike standard SGD. Actually, we observe that Dom-SGD even slowly increases the loss long run. This suggests that the dominant component of the stochastic gradient g: is in fact not beneficial for training, despite constituting the majority of gt- We refer to the alignment between the gradient and the dominant subspace in this scenario "spurious," because, based the local quadratic Taylor approximation in Section 3.1, we expect that projecting the update vector onto the dominant subspace should decrease the loss similarly the original update if the update vector is aligned with the dominant subspace.
4

---

Published as a conference paper at ICLR 2025
Remark 1 (This is not the end-of-training phenomenon). Here, note that the switching happens when training accuracy reaches around 53% for CNN on CIFARIO-3K, and 69% Transformer on SST2-Ik. Ik. This indicates that the observed phenomenon is not confined the SGD dynamics near the manifold of local minimizers, which is the basis of many recent theoretical analyses (see, e.g., (Arora al., 2022; Li al., 2022; Lyu al., 2022; Wen et al., 2023; Ahn al., 2024b)).
3.3 BULK SUBSPACE IS WHERE THE LEARNING HAPPENS
To further strengthen our main observation, we conduct another set of experiments, wherein we switch from SGD to the following update scheme:
cc+1 - 0. - 1)
Bulk-SGD)
Essentially, Bulk-SGD discards the majority of the stochastic gradient 9t by removing its dominant component. Consequently, it seems less likely that the remaining bulk component of stochastic gradient would lead to successful training.
Surprisingly, shown in Figure 1, Bulk-SGD is as effective as SGD in decreasing the training loss. This further highlights that it is indeed a small fraction of gradient that aligns with the bulk subspace that contributes to training loss decrease.
One can summarize our observation thus far as follows.
Phenomenon 2. Although the grudient upproximately aligns with the dominant subspuce at each step. the training loss does not decrease within the dominant subspace, suggesting a "spurious" " alignment. Surprisingly, with Bulk-SGD, where each update is projected onto bulk subspaces, the training remains as effective as the original update. This emphasizes the importance of a small component of the update that aligns with the bulk subspace.
Based on our preliminary analysis Section 3.1, Phenomenon 2 appears quite counterintuitive and unexpected. The next section focuses on explaining how this counterintuitive phenomenon occurs, In particular, the seemingly contradictory conclusion from Section 3.1 will be revisited in Section 4.3.
4 WHAT CAUSES THE SPURIOUS ALIGNMENT WITH DOMINANT SUBSPACES?
In this section, we aim explain the spurious alignment discussed Phenomenon 2. To that end, we first distinguish between two differ- ent regimes of SGD dynamics, because the underlying mechanism of the alignment are different.
1.00
0.75
0.50
0.25
HIVE WOS RO1
Definition 3. We say (SJGD is in the GF regime when the sharpness is below the maximum stable sharpness (MSS), where it closely follows the grudient flow. In the GF regime, the sharpness typically increases (progressive sharpening). and the loss decrease is stable. Conversely, (SJGD is in the EoS regime when the sharpness close to MSS. In the EoS regime, the sharpness oscillates around MSS, and the loss decrease is spiky and unstable (Cohen et al., 2021).
0.00
0
10000 20000 30000 40000 Training Stops
Figure 4: X (VL())) when switching from SGD to GD at step 20000 while training MI.P on MNIST-5k.
In this section, focus SGD in the GF regime and present the mechanism of how gradients align with the dominant subspace. The scenario of the EoS regime will be discussed in Section 5.1. For SGD in the GF regime, the spurious alignment is closely tied
0.75
0.50
the landscape of the training loss near SGD trajectories. Impor- tantly, SGD introduces inherent randomness into its trajectories. We investigate how this stochastic noise affects the phenomenon.
0.00
0
10000 20000 30000 40000
Training Steps
4.1 STOCHASTIC NOISE OF SGD THE MAIN CAUSE
Figure 5: X (VL(.) when switching from GD to SGD.
We examine the role of stochastic noise by contrasting the behavior of SGD with (full-batch) GD, isolating the effect of stochastic noise.
MSS is 2/n for GD (Cohen al., 2021), and smaller for SGD (Wu al., 2018).
5

---

Published as a conference paper at ICLR 2025
First, we demonstrate the crucial role of stochastic noise in the alignment by switching from SGD to GD when Phenomenon is observed. Strikingly, as shown in Figure 4, the alignment disappears as as switch to GD. More specifically, X (VL(),) quickly becomes 0 soon as the switch occurs. This sharp transition indicates that the stochastic noise must play a crucial role.
In Section D.I, when training neural networks with GD (from scratch) in the GF regime, observe no alignment between gradients and dominant subspaces, in contrast to SGD. In this case, X& Mn.V ((0)) quickly reaches and remains near O, indicating alignment of gradients with bulk subspaces. Remarkably, despite differences in gradient alignment (GD: X& (VL) L) R 0, SGD: X (VL 1). GD and SGD trajectories closely track each other (see Section D.2). This suggests that the presence of small stochastic noise results in a drastically different behavior in the alignment.
In Figure 5, we switch our optimizer from GD to SGD, complementary the experiment Figure 4. This time we observe that the alignment sharply appears, i.e., X (V [(f))) quickly becomes 1, as soon as the switch occurs. Note that SGD is still in the GF regime as sharpness stably increases, rather than oscillating. This highlights that the gradient alignment during SGD is not due to self-stabilization effect (Damian et al., 2023) the EoS regime. To sum up, one can summarize our findings as follows.
Phenomenon 3 (The spurious alignment is due to stochastic noise). In the GF reigme, the aligmment between the gradient and the dominant subspace quickly disappears when switching the optimizer from SGD GD. Moreover, the alignment quickly reappears when switching GD back to SGD. Hence, the spurious alignment mainly due to the stochastic noise of SGD.
To further demonstrate that noise is the primary driver of gradient alignment for SGD, we conduct experiments using Noisy Gradient Descent (NGD) and SGD with varying batch sizes in Section D.4, We implement NGD by injecting Gaussian noise after each GD update iteration. We observe that either increasing the noise scale or decreasing the batch size lead to the increased gradient aligmment. This observation further demonstrates that noise is the main cause of the spurious alignment for SGD.
Given this observation, one might question how the presence of small stochastic noise leads to a drastically different alignment. We investigate this in the next subsection using simple model.
4.2 UNDERSTANDING THE ROLE OF STOCHASTIC NOISE VIA TOY QUADRATIC MODEL
Towards understanding Phenomena 1 1-3, this section introduces a simple example that recovers all the phenomena.
Given the typical ill-conditioned nature of neural network train- ing, we consider a 2-dimensional ill-conditioned quadratic loss, L(x, y) = (1000x² + y²) where f = (z, y) E R². We define (x, L(x. v) 100ry and (x) L(z, v) 100xy, re- sulting in T(x. = 2 (f. v) +
00.5
SGD
We conduct GD with learning rate 7) as
Initiaiiattion
DGD + gGD - (L
0.0
0.5
1.0
and SGD using random sampling with the same learning rate r as BSGD + V (LLGG)D where k ~ (ni,
Figure 6: GD and SGD
trajectories when training 2-dimensional ill-conditioned toy quadratic model.
In Figure 6, we visualize the optimization trajectories of GD and SGD with initialization OGD = OgGD = (1, 1) and learning rate = 10-4. The Hessian of the quadratic loss remains con- stant during training, with eigenvalues \ 1000 and A2 = 1, and corresponding eigenvectors é (1,0) and e2 = (0, 1). We compute the fraction of gradient in the dominant subspace as Xi (V = (V as shown in Figure 7. [[LL()),,
1.00
0.75
0.50
0.25
SGD
0.00
10000 15000
Notably, this simple quadratic model recovers all the observed phenomena (Phenomena 1-3). In both GD and SGD trajectories,
Trairing Steps
Figure 7: X1 (VI(O))) during GD and SGD for Figure 6.
I, quickly converges to 0 due to the sharper direction along 61
(A) > 22). Subsequently, both trajectories remain close to the y-
axis throughout the remaining of the training. However, in GD, XI (VI(OGP)) quickly approaches

---

Published as a conference paper at ICLR 2025
and remains near (Phenomenon 3), while in SGD, X1 (VL(0))))) stays close to (Phenomenon 1). Notice that if we run Dom-SGD, the updates will done only in the z direction, hence training stops after switching to Dom-SGD (Phenomenon 2). We provide results on NGD in Appendix D.6.
The discrepancy in alignment between GD and SGD arises from the ill-conditioned nature of the loss landscape, where the small stochastic noise of SGD in the direction induces large gradient component along the z direction. For example, if the SGD iterate is at 0 (0.01, 0.5). a slight departure from the y-axis, the gradient alignment X L) is approximately 0.999.
4.3 REVISITING OUR PRELIMINARY ANALYSIS (SECTION 3.1)
At this point, some readers might wonder how we reconcile the results with our preliminary analysis (Section 3.1). Based on our investigations so far, we propose one plausible explanation that the training loss landscape is locally "ill "aaaittionedvvlllee"" causing the spurious alignment:
This landscape has two key features
The landscape is locally valley-shaped, where it is steep along the dominant subspace and flat along the bulk subspace. In particular, the curvature along the dominant subspace is much larger than that along the bulk subspace.
The bottom the valley connected along the bulk subspace. Moreover, there is a direction within the bulk subspace along which the bottom of the valley descends.
To aid readers' understanding, we provide a simple illustra- tion of an iil-coodittooned-vlleey--- landscape in Figure 8, With these features, Phenomena 1-3 can indeed occur:
Dominan
Due stochastic noise, the SGD iterates slightly deviate from the bottom of the valley. Subse- quently, the high curvature along the dominant subspace causes gradients to align with this sub- space, i.e., the iterates exhibit Phenomenon 1.
Bulk Subspace
Figure Illustration demonstrating how spurious alignment can occur with an manditioned-valley--ee
However, Dom-SGI fails to further decrease the training loss, as observed in Phenomenon since it fails to follow the true progress direction along the bulk subspace.
training loss. Dom-SGD iterates (depicted with dots) fail to progress along the bulk subspace where the training loss decreases.
Moreover, without stochastic noise, the iterates quickly approach the bottom the valley, where the alignment disappears, as described in Phenomenon 3.
In Section D.5, we measure the distance of the weights from the step where we switch from SGD to Dom-SGD and Bulk-SGD for the experiments in Figure 1. We observe that weights do not move far from the switching step for Dom-SGD, contrast to SGD and Bulk-SGD. Furthermore, Dom-GI) shows near-zero movement, suggesting it gets stuck the bottom of the valley, consistent with our proposed explanation based on an li-conditioneed-vlleye----
landscape.
Given our results for SGD thus far, one natural question is whether these phenomena are also observed for other practical optimization algorithms.
5 EDGE OF STABILITY AND SHARPNESS-AWAN MINIMIZATION
In this section, we extend our investigations to two other practical settings: (1) (S)GD the Edge of Stability (FoS) regime (Cohen et al., 2021), and (2) Sharpness-Awart Minimization (SAM) (Foret et al., 2021). We show that the same phenomena are observed for the two settings: the update direction aligns with the dominant subspace, but the alignment is again spurious.
5.1 EDGE OF STABILITY
Recent empirical studies (Jastrzębski et al., 2020; Cohen et al., 2021) have observed that when training neural networks using full-batch GD with large learning rates 77. the sharpness A increases until
7

---

Published as a conference paper at ICLR 2025
20
1.00
-1
Dem-GD
CD
0.75
nJlc-DD
2-
10
0.50
0.25
0
0.00
2500 5000 7500
10000
2500
5000 7500 10000
0
2500 5000 7500
10000
Training Steps
Training Steps
Training Steps
(a) Top-20 cigenvalues
(b) X1o LV.(.)
(c) Training loss (log-scalc)
Figure 9: Gradients approximately align with dominant subspaces in the EoS regime. Training MLP on MNIST-5K with GD using a large learning rate yj 0.1. (a) The plot shows the top- 10 cigenvalues in blue and the next top-10 10 cigenvalues in orange. After few steps, GD enters the EoS regime, where the sharpness stabilizes near 2/n, (b) As the sharpness reaches 2/n, X10 ((())) shoots up and remains near 1. (c) We switch the optimizer from GD to Dom-GD and Bulk-GD at step 2500. Dom-GD fails to further decrease the training loss, in contrast GD and Bulk-GD.
reaches the stability threshold, or the maximum stable sharpness (MSS), 2/1,, and saturates around the threshold (see Figure 9a). Cohen al. (2021) call this phenomenon as the Edge of Stability.
In Figure 9b, we observe that gradients closely align with dominant subspaces in the EoS regime. This phenomenon stands in contrast with GD in the GF regime (Phenomenon 3), where Xx V (0.)) remains near 0 (see Section D.1 for details).
Remark 2. We highlight that the mechanisms behind gradient alignment differ between SGD in the GF regime and GD in the EoS regime. For SGD in the GF regime, where sharpness stably increases, alignment is due the stochastic noise and the ill- el-coondttooeed loss landscape. In contrast. for GD in the EoS regime, where sharpness oscillates around MSS, alignment arises from a self-stabilization mechanism (Damian et al., 2023), which illustrates that GD oscillates within the dominant subspace.
Given that the gradient approximately aligns with the dominant subspace, we run experiments analogous Section 3. Specifically, we train neural networks using GD with a large learning rate r) until it reaches the EoS regime. We then switch GD to the following update schemes:
8.41 - 0₁ (Pe)
(Dom-GD)
0141 + 02 - (
(Bulk-GD)
As shown in Figure 9c, observe that Dom-GD fails to further decrease the training loss, unlike GD. Moreover, Bulk-GD is as effective as GD in decreasing the training loss, despite only a small fraction of updates aligning with the bulk subspace.
We provide additional experiments on GD and SGD in the EoS regime in Appendix E. Notably, for SGD in the EoS regime, both stochastic noise and self-stabilization effect affect training dynamics, leading to gradient alignment with dominant subspace, while Dom-SGD still fails to decrease the loss.
5.2 SHARPNESS-AWAN MINIMIZATTION
Sharpness-Aware Minimization (SAM) (Foret al., 2021) is gradient- based optimization method designed to find flat minima. SAM has gained significant attention for its success in practice, especially improving the generalization performance of deep learning models. For concreteness, we focus on the full-batch version of SAM applied to GD as the base optimizer. This leads to the following update equation:
VI.(.))
0 02 VV 0g +
(VL(
where 7) is the learning rate and represents the perturbation radius.
recent study (Long & Bartlett, 2024) highlights that SAM also operates in its own Edge of Stability regime, wherein the sharpness A₁ (6,) saturates near SAM's stability threshold (see Figure 10a). This threshold, denoted as the edge, is defined as:
V (0,)ll2
(VI+ 17 (
8
-1)
(SAM-edge)
2p
8

---

Published as a conference paper at ICLR 2025
SAM-edge
1.00
SAM
15
-1.0
Dom SAM
0.75
Du R-SMM
0.50
-1.5
5
0.25
-2.0
5000 10000 15000 Training Steps
20000
5000 10000 15000 20000
5000 10000 15000 20000 Training Steps
Training Steps
(a) Top-20 cigenvalues
(b) X10 0c)
(c) Training loss (log-scale)
Figure 10: SAM updates approximately lie in dominant subspaces. Training MLP on MNIST-5k with SAM using a learning rate 7/ 0.01 and a perturbation radius 0.1. (a) The plot shows the top-10 10 eigenvalues in blue and the next top-10 eigenvalues in orange. After a few steps, SAM operates in the EoS regime, where the sharpness stabilizes near the SAM-edge. (b) As the sharpness reaches the SAM-edge, X10 (0++1 0,) shoots up and remains near 1. (c) We switch the optimizer from SAM to Dom-SAM and Bulk-SAM at step 5000. Dom-SAM fails to further decrease the training loss, in contrast to SAM and Bulk-SAM.
Note that GD's stability threshold 2/1 remains constant during training, while SAM-edge a decreasing function of the norm of the gradient, so it tends to decrease during training.
As shown in Figure 10b, we observe that update vectors of SAM approximately align with dominant subspaces when sharpness saturates near the SAM-edge, similar GD in the EoS regime. Next, conduct experiments akin to Section 3: we train neural networks using SAM until the alignment occurs. Then, we switch SAM the following schemes:
VL(O.)
O.t1 02 (LL
0, + LVL(0)))
(Dom-SAM)
(I(0.)
0c t 02 1
0, + [L(()))
(Bulk-SAM)
We observe that Dom-SAM fails further decrease the training loss, unlike SAM and Bulk- SAM, depicted in Figure 10c. Our investigation with various practical algorithms suggests that neural networks cannot be trained within the dominant subspace, and the bulk subspace plays an important role in learning.
MOMENTUM AND ADAPTIVE METHODS AMPLIFY UPDATES IN BULK SUBSPACES
In this section, we further extend our investigation to momentum optimizers, e.g., SGD with mo- mentum, and adaptive optimizers, e.g., Adam. In Appendix F.1, we show that momentum and adaptive learning rates lead to less alignment between the update vector and dominant subspace, so Phenomenon 1 no longer holds for this case. However, Phenomenon still holds, i.e., if we project the update vector onto the dominant subspace (running Dom-SGDM or Dom-Adam), the training loss fails to decrease. This observation also demonstrates that bulk space is where the learning happens.
We build our results so far and propose explanations for why momentum and adaptive methods are effective for neural network training. At a high level, we claim that they speed up training by amplifying the bulk component of each update step. Formally, we introduce the following notion.
Definition 4 (Effective learning rate). For given optimization trajectory {0c). we define the dominant effective learning rate (Dom-LR) at step as:
(0+ O L.(
dom
(Dom-LR)
PP. VL(())
and the bulk effective learning rate (Bulk- LR) at step t as:
(0. 0141, Pr (0, (L.)
T bulk
(Bulk-LR)
(L.

---

Published as a conference paper at ICLR 2025
Table 1: Mean effective learning rates over the first 1000 steps (numbers in parentheses show standard deviation). Training Transformer on SST2-1K us- ing GD and Adam with (+m) and without (-m) momentum. GD uses a learning rate of 0.01, and Adam uses learning rate of 0.001. Momentum is set to 3 = 0.9.
-5
GD(-m)
GD(+m)
-10
Adam(-m)
Adam(+
Method
Mean Dom-LR
Mean Bulk-LR
2500
500
750
1000
GD(-m)
0.0100 (0.0000)
0.0100 (0.0000)
Training Steps
GD(+m)
0.0070 (0.0232)
0.0828 (0.0576)
Adam(-m)
0.0325 (0.0054)
0.4672 (0.3555)
Figure 11: Training loss in log-scale for the experiments in Table 1.
Adam(+m)
0.0004 (0.0101)
2.6639 (1.2480)
To understand the above notion better, we first consider the simplest case. For a fixed learning rate
GD has effective learning rates nit doom .tt -njulk 1I for all steps f.
Dom GD has effective learning rates n/ = 17 and ry. = 0 for all steps t.
Bulk-GD has effective learning rates no and rf = for steps t.
Our primary claim in this section is follows: since Dom-GD fails decrease the training loss, while Bulk-GD is as effective as GD in reducing the training loss, we claim that Bulk-LR serves as good indicator for training speed, unlike Dom-LR.
To support this claim, we measure effective learning rates of various optimization methods, including (full-batch) GD with and without momentum, and (full-batch) Adam with and without momentum. Table presents effective learning rates when training Transformer on SST2- Ik, and Figure 11 depicts corresponding training loss plot. Additional experiments other architectures and datasets provided in Appendix F.2.
Across different settings, consistently observe that Bulk-LR positively correlates with the training speed. Moreover, momentum and adaptive methods seem to amplify Bulk-LR. This amplification of the bulk component leads a reduced alignment between the update vector and the dominant subspace, as shown in Appendix F.1. This offers new insights into the effectiveness of momentum and adaptive methods.
CONCLUSION AND DISCUSSION
Motivated by the observation of Gur-Ari et al. (2018) that the gradient aligns with low-dimensional dominant cigenspace of the training loss Hessian, this work investigates the possiblity of training neural networks within the dominant subspace. Our key contributions two-fold:
For every optimizer (e.g., GD, SGD, SGDM, Adam, SAM) we tested, Dom OPT-which projects the update vector onto the dominant subspace-fail fails to decrease the training loss. This indicates that neural networks canmot be trained within the dominant subspace, and bulk subspace plays essential role during training.
We identify distinct mechanisms for gradient alignment across different training regimes. In the GF regime, alignment is primarily caused by stochastic noise from SGD in conjunction with the ill-conditioned loss landscape. In the EoS regime, the alignment arises from the seff-stabilization mechanism, where oscillations within the dominant subspace lead to this behavior.
Discussion. There are remaining questions we leave as a future work. First, providing theory on our empirical findings beyond the toy model would be important direction. One interesting observation we did not discuss is that Dom-SGD decreases sharpness and slowly increases the loss, while Bulk-SGG SGD is less noisy and increases sharpness faster than SGD (see Appendix G). Under- standing such phenomena would provide deeper insights into neural network training. Lastly, this work focuses optimization, and exploring the implications generalization would be important.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was partly supported by the National Research Foundation of Korea (NRF) grants funded by the Korean government (MSIT) (No. RS-2023-002 13522 No. RS-2024-00421203)
REPRODUCIBILITY STATEMENT
We have made significant efforts to ensure the reproducibility of our results. Detailed experimental settings, including byperparameters and training details, are provided Appendix B. Furthermore, to facilitate replication and verification, the source code for the experiments is included in the attached supplementary material. This code contains scripts for reproducing the main results discussed in the paper, along with instructions for running the experiments.
REFERENCES
Atish Agarwala and Yann Dauphin. Sam operates far from home: eigenvalue regularization as dynamical phenomenon. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 152-168, PMLR, 23-29 Jul 2023. URL https: agarwala23a. htmll
https::/proceeediiiiee http:///oocceedeeiee mlr. www..ppees//ee /vieess//22////
Atish Agarwala and Jeffrey Pennington. High dimensional analysis reveals conservative sharpening and a stochastic edge of stability. arXiv preprint arXiv:2404. 19261, 2024.
Kwangjun Ahn and Ashok Cutkosky. Adam with model exponential moving average is effective for nonconvex optimization. arXiv preprint arXiv:2405. 18199, 2024.
Kwangjun Ahn, Jingzhao Zhang, and Suvrit Sra. Understanding the unstable convergence of gradient descent. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference volume 162 of Proceedings of Machine Learning Research, pp. 247-257. PMLR, 17-23 Jul 2022. URL http:::t////mm. Wttttts://poocceeeeeeee httpp//pweeenngn.eeeepe mll. www.mmmt........mm..mm.mmmm www.prraasaeee v1 hwtp///a/22aaaaaaaaaaaaaaaaaa ahn22a html.
Machine Learning,
Kwangjun Ahn, Sebastien Bubeck, Sinho Chewi,
Yin Tat Lee, Felipe Suarez, and Yi Zhang. Learning threshold neurons via edge of stability. In Thirty-sevennth Conference on Neural Information Processing Systems, 2023. URL https: openrevieeeeeeeeeeeeeeeee net.e/ewweeeeeeeeeeeeeeeee.eeeee forum? id= 9cQ6xTonJJ.
/tttppss//ppenevie
Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Linear attention is (maybe) all you need (to understand transformer optimization). In The Twelfth International Conference on Learning Representations, 2024a. URL htttps://oppeeeeeeeee /opnneviee. net/forumt didmmui4411000
Kwangjun Ahn, Ali Jadbabaie, and Suvrit Sra. How to escape sharp minima with random perturba- tions. In International Conference on Machine Learning. PMLR, 2024b.
Kwangjun Ahn, Zhiyu Zhang, Yunbum Kook, and Yan Dai. Understanding Adam optimizer via online learning of updates: Adam is FTRL in disguise. In International Conference on Machine Learning. PMLR, 2024c.
Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware mini- mization. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 639-668. PMLR, 17-23 Jul 2022. html.
URL https::t//// httpht::/ppooceeeiiiiiee// http///oooeedgggee/ee/ee
www.....es/.....m..m.mmmmm.. vl 62/ / www.nddusschhenko222222
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on the edge of stability in deep learning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference Learning, volume 162 of Proceedings of Machine Learning Research, pp. 948-1024 1024. PMLR, 17-23 Jul 2022. URL https: //tp///pr////e ngs. wtt.s...r.rrrs////.. www.ss/sseee. v1 62. 62/arora22 22a html.
Machine

---

Published as a conference paper at ICLR 2025
Gerard Ben Arous, Reza Gheissari, Jiaoyang Huang, and Aukosh Jagannath. High-dimensional SGD aligns with emerging outlier eigenspaces. In The Twelfth International Conference on Learning Representations, 2024. URL
net /ttp://oooomm..ddddaail-i.......
Peter L.. Bartlett, Philip M. Long, and Olivier Bousquet. The dynamics of sharpness-aware minimiza-
tion: Bouncing across ravines and drifting towards wide minima. Journal of Machine Learning Research, 24(316):1-36, 2023. URL httL:ht:/// jmmlr org/ pwwpps/v24/222/// 22//////222233..e/ 043.html.
Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an ornstein-uhlenbeck like process. In Jacob Abernethy and Shiv- ani Agarwal (eds.), Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pp. 483-513, PMLR, 09-12 Jul 2020. URL www.mmrrr..rs....mmmmmmmmmmmm.se
https: httpss///pcceeeinngs. mlr press http/////////aaaa//e html.
Jeremy Cohen, Simran Kaur, Yuanzhi Li, J 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 https: wpwpr/veereeeeeeeeoeoeeee new.wwneeeeeee.e. forumoorummmom j--jh-d-hh-h rTtvkGeM.
Michael Crawshaw, Mingrui Liu, Francesco Orabona, Wei Zhang, and Zhenxun Zhuang. Robustness to unbounded smoothness of generalized signsgd. Advances in neural information processing systems, 35:9955-9966 2022.
Ashok Cutkosky and Harsh Mehta. Momentum improves normalized sgd. In International conference machine learning, pp. 2260-2268. PMLR, 2020.
Yan Dai, Kwangjun Ahn, and Suvrit Sra. The crucial role of normalization in sharpness-aware minimization. In Thirty-sevnnth Conference on Neural Information Processing Systems, 2023.
URL
/twp///peeeeeee.eooo/e//e nett httpr//ewweee///otrrrrree forum? i d dd=z44FneeRi.
Alex Damian, Tengyu Ma, and Jason D. Lee. Label noise SGD provably prefers flat global minimizers. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https: /ttps///ppeeeiieeeee net forum? did= x2 IMPhseWAR.
Alex Damian, Eshaan Nichani, and Jason D. Lee. Self-stabilization: The implicit bias of gra- dient descent the edge of stability. In The Eleventh International Conference Learning Representations, 2023. URL https:
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations, 2021. URL https: opwnreview.ee.ee... net wtt.//ormmmmmi waarrrmmmmmmmmmmmmp.e.mm rM.
Jingwen Fu, Bohan Wang. Huishuai Zhang, Zhizheng Zhang, Wei Chen, and Nanning Zheng. When and why momentum accelerates sgd: An empirical study. arXiv preprint arXiv:2306.c8 09000, 2023.
Martin Gauch, Maximilian Beck, Thomas Adler, Dmytro Kotsur, Stefan Fiel, Hamid Eghbal-zadeh, Johannes Brandstetter, Johannes Kofler, Markus Holzleitner, Werner Zellinger, Daniel Klotz, Sepp Hochreiter, and Sebastian Lehner. Few-shot learning by dimensionality reduction in gradient space. In Sarath Chandar, Razvan Pascanu, and Doina Precup (eds.), Proceedings of The 1st Conference Lifelong Learning Agents, volume 199 of Proceedings of Machine Learning Research, pp. 1043-1064. PMLR, 22-24 Aug 2022. URL https: http:////occceeeeeeeeeee wwt.:ss.......r..... mlt.preess/119/ gauch22a.htir html.
Saeed Ghadimi and Guanghui Lan. Stochastic first-and zeroth-order methods for nonconvex stochastic programming. SIAM journal on optimization, 23(4):2341 12368, 2013.
Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net opti- mization via hessian eigenvalue density. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.). Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 2232-2241. PMLR, 09-15 Jun 2019. URL https: http://preeeee///e ngs ml r.w..r.pp//ssoe... / v / 9b. html.
12

---

Published as a conference paper at ICLR 2025
Frithjof Gressmann, Zach Eaton-Rosen, and Carlo Luschi. Improving neural network training in low dimensional random bases. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 10140-21150. Curran Associates, Inc., 2020.
Guy Gur-Ari, Daniel A Roberts. and Ethan Dyer. Gradient descent happens in tiny subspace. arXiv preprint arXiv: 2018.
Haowei He, Gao Huang, and Yang Yuan. Asymmetric valleys: Beyond sharp and flat local minima. In H. Wallach, H. Larochelle, A. Beygelzimer, d'Alché-Buc, Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Like Hui and Mikhail Belkin. Evaluation of neural architectures trained with square loss vs cross- entropy in classification tasks. In International Conference Learning Represeniations, 2021.
URL https:://///m http:///pppeeeeeeeeeeeeeeeeeeeer net httpi//wwwee//oommeemmm hormm??d--ssNN2ee0EE
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 International Conference on Learning Representations, 2019, URL https: /ttps::///p/e/e/eeeeeeee net/fozee futt/forum?iic id- SaddasSkKaao00 705t
Stanislaw Jastrzębski, Maciej Szymezak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho and Krzysztof Geras*. The break-even point on optimization trajectories of deep neural networks. In International Conference on Learning Representations, 2020. URL https:// openreview n forum foormmiddrlgg 7C4WBB CCAAKW.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantas- tic generalization measures and where to find them. International Conference on Learning Representations, 2020. URL https: wwwpoevieww.w.e.//m
neereeeiieeeeeeeeeeeeeeee..ee.ee net/ forum?id-SJJIIPPS
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. URL https: htttt:://ppeeeeeeeeeeeeeee/eeeee /opnneeeee......
net/ http://ooommiideeeee setprum?dd=HooyRll
Rahul Kidambi, Praneeth Netrapalli, Prateek Jain, and Sham Kakade. On the insufficiency of existing momentum schemes for stochastic optimization. In 2018 Information Theory and Applications Workshop (ITA), pp. 1-9, 2018. doi: 10.1 000/://TA.211880...3
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 141.661412.6900,, 2014.
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. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning. volume 202 of Proceedings of Machine Learning Research, pp. 17684-17744 PMLR, 23-29 Jul 2023. URL hwtps::tt/.............tt...t..t /proced i ngs. ml wtt.g.ml...ppsse. presspresseee. http//22////22 s eer2a .tml.
Alex Krizhevsky. Learning multiple layers of features from tiny images, 2009.
Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt. Noise is the main factor behind the gap between sgd and adam on transformers, but sign descent might be. In The Eleventh International Conference on Learning Representations, 2023. URL https: / //openreiiew.. /tt/ 2id wtt.:ooidaaaoooooeeoooooaea 36550a<<5KK00H5,,,
Frederik Kunstner, Robin Yadav, Alan Milligan, Mark Schmidt, and Alberto Bietti. Heavy-tailed class imbalance and why adam outperforms gradient descent on language models. arXiv preprint 44X1v:2402.19449, 2024.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEFF, 86(11):222 86(11):2278-2) 1):2278-2324, 1998.
13

---

Published as a conference paper at ICLR 2025
Sungyoon Lee and Cheongjae Jang. A new characterization of the edge of stability based sharpness measure aware of batch gradient distribution. In The Eleventh International Confer- on Learning Representations, 2023. URL https: htttps::/apeneeveew..r net forummidumm bH-kCY61.dKg
Chunyuan Li. Heerad Farkhoor, Rosanne Liu, and Jason Yosinski. Measuring the intrinsic dimension of objective landscapes. In International Conference on Learning Representations, 2018a. URL https: wpenreview.eeeeeeeeeeee net/ Eorum?id-ryups WCW.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape neural nets. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018b.
Tao Li, Lei Tan, Qinghua Tao, Yipeng Liu, and Xiaolin Huang. Low dimensional trajectory hypothesis is true: Dnns can be trained in tiny subspaces. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(03):3411-344
mar 2023. ISSN 1939-3539. doi: 10.1 10.109/TPAMM.22223113355
Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss? -a mathematical framework. In International Conference on Learning Representations, 2022. URL
https: ppenreviev net forum? Cte2st5te
Philip M. Long and Peter L. Bartlett. Sharpness-aware minimization and the edge of stability. Journal of Machine Learning Research, 25(179):1-20, 2024. URL http: /mmr org/ haggpppprs/ v25/23-22 1285. html.
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normal- ization layers: Sharpness reduction. Advances in Neural Information Processing Systems, 35: 34689-34708, 2022.
Yurii Nesterov et al. Lectures on convex optimization, volume 137. Springer, 2018.
Vardan Papyan. Measurements of three-level hierarchical structure the outliers in the spectrum of deepnet hessians. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 Proceedings of Machine Learning Research, pp. mlr. mir..ress/ v papyan 9a. html.
PMLR, 09-15 Jun 2019. URL https: httttp::/pproceedingg
Vardan Papyan. Traces of class/cross-class structure pervade deep learning spectra. Journal of Machine Learning Research, 21(25)) 2020. URL http: /tpp:///mmrrrr/////////epeeeee wttpl...or/ppppr///eees///
20 933. html.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, 7achary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc. E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019. URL https: http::ts///cccceeeninnneeeeeeeee neurips cc httppppprrfffit filfee/////i/e /paps/////////////// -apeeaaapppprpaeee
2019 1199ffile http://LLeeeccccaaaaaeeeeeee2222
B.T. Polyak. Some methods of speeding up the convergence of iteration methods. USSR Computational Mathematics and Mathematical Physics, 4(5):1-17, 1964. ISSN 0041-5553. doi: https://cbi.og///oo.oooooo1111s. com/
URL https: htt////www... www.eeeciiiencieett.............
Levent Sagun, Leon Bottou, and Yann LeCun. Eigenvalues of the hessian in deep learning: Singularity and beyond. arXiv preprint arXiv:1611.07476, 2016.
Levent Sagun, Utku Evci. V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the hessian of over-parametrized neural networks. arXiv preprint arXiv: 1706. 1706.04454, 04454,
2017.
14

---

Published as a conference paper at ICLR 2025
Jan Schneider, Pierre Schumacher. Simon Guist, Le Chen, Daniel Haeufle, Bernhard Schölkopf, and Dieter Büchler. Identifying policy gradient subspaces. In The Twelfth International Confer- ence on Learning Representations, 2024. URL https:: https///pppeeeee..eeeeeee net fermmiimmiid i 2Wwxnt22ke.
Dongkuk Si and Chulhee Yun. Practical sharpness-aware minimization cannot converge all the way to optima. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL
net pest/forumiid-ng jiadddiinjJNoHHs JNOLHęM.
Vikrant Singhal and Thomas Steinke. Privately learning subspaces. In A. Beygelzimer, Y. Dauphin, P. Liang. and Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https::////. http/opeevviiieee net/ forum? Bal.dd-BBBBnBBBB
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Y Ng. and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 conference empirical methods in natural language processing, pp. 1631-1642, 2013.
Minhak Song and Chulhee Yun. Trajectory alignment: Understanding the edge of stability phenomenon via bifurcation theory. In A. Oh, T. Naumann, A. Glober- son, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Informa- tion Processing Systems, volume 36, pp. 71632-71682. Curran Associates, Inc., 2023. URL https: hwtp:///nw.........eeeeeeceeecee File/e2a9256bcle 6bd822566d16666 hrw.aa/00.....aa0000000a222...a. otg.:aaaaaaaaa2t1aaaaa---------r Capprr- pdf.
files/ hatps//apppep/////
Runzhe Wang. Sadhika Malladi, Tianhao Wang, Kaifeng Lyu, and Zhiyuan Li. The marginal value of momentum for small learning rate SGD. In The Twelfth International Conference on Learning Representations, 2024. URL https: http///ppeeeeeeeeeeeeeeeeeeeeeee www.eeeweeeeeee.e.eeee n forum fto///odmmmmmiieeeeeeee Pid=3J-JJeeeee
Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How sharpness-aware minimization minimizes sharpness? In The Eleventh International Conference on Learning Representations, 2023. URL https: // /ppenreveeww wortt//fommmmimmm
i iegrrmmidd-sppppWpppp
Jingfeng Wu, Vladimir Braverman, and Jason D. Lee. Implicit bias of gradient descent for logistic regression the edge of stability. In Thinty-sevntth Conference on Neural Information Processing Systems, 2023. URL https: wttpp/peeeeeeeeeeeeeeeeeeeeeeeee
i- 9MWLYNpQ.
Lei Wu, Chao Ma, and Weinan E. How sgd selects the global minima in over- parameterized learning: dynamical stability perspective. In Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018 URL https::/ neurips. cc/ /papprr___ halls//pper// 201///21///// wtt.t////////////ee/eeoeeeeeleee (et.lee//65552266666 /tt.////e/6612226666600000000000
Pawpeeaaaeeeeee.eeeeeeeeeeaaeeee Pff.
Shuo Xie and Zhiyuan Li. Implicit bias of adamw: foo norm constrained optimization. In International Conference on Machine Learning. PMLR, 2024.
Can Yaras, Peng Wang. Laura Balzano, and Qing Qu. Compressible dynamics in deep overparameter- ized low-rank learning & adaptation. In Forty-first International Conference Machine Learning, 2024. URL https::p//... http///pppereeeeeeee net. forum? http://rrummnduuuooooooeaa.. d=uDkXoZMZRv.
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, San- jiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neu- ral Information Processing Systems, volume 33, pp. 15383-15393. Curran Associates, Inc., 2020. /tt.stt/s/// http:///pooeeeennn.n..c/ neurips cc/paper__ www.fffeeesse////ee file /ttle//00SSSTEaaae...00000006668
www.cco0005--PPPP..-.- paper.pps
Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, and Zhi-Quan Luo. Why trans- formers need adam: A hessian perspective. arXiv preprint arXiv:2402.16778 2024.
Yingxue Zhou, Steven Wu, and Arindam Banerjee. Bypassing the ambient dimension: Private (sgd) with gradient subspace identification. In International Conference on Learning Representations, 2021. URL https::////.. opwnrevieeeieeeeee net forum? d kBUJFC
15

---

Published as a conference paper at ICLR 2025
Libin Zhu, Chaoyue Liu, Adityanarayanan Radhakrishnan, and Mikhail Belkin. Catapults in SGD: spikes the training loss and their impact on generalization through feature learning, In Forty-first International Conference Machine Learning, 2024. URL https: http://ppeeeeeeeeeeeeeeeee net/ forum? i JmGJ 01$41.
Xingyu Zhu, Zixuan Wang, Xiang Wang, Mo Zhou, and Rong Ge. Understanding edge-of- stability training dynamics with minimalist example. The Eleventh International Confer- ence on Learning Representations, 2023. URL https:/ / /openrevieww.o.eooooooopp net formmiid- p p7RagBsMAEO.
16

---

Published as a conference paper at ICLR 2025
Appendix
Related work
18
Experimental details
20
B.1 Architectures
20
B.2 Data
20
B.3 Experimental setup
20
Additional experiments for Section
22
C.1 Cross-entropy loss
22
C.2 Standard architectures on CIFARIO-5K
D Additional experiments for Section 4
24
D.1 No alignment with dominant subspaces along GD trajectories
24
D.2 SGD trajectories track gradient flow
25
D.3 Switching between GD and SGD
D.4 Effect of noise scale on spurious alignment
26
D.5 Distance from "switching" step
27
D.6
Toy example
28
Spurious alignment in the EoS regime
29
Additional experiments for Section
31
F.1 SGDM and Adam
31
F.2 Effective learning rates
31
G Sharpness plots for main experiments
34
Test accuracy results for Dom-SGD and Bulk-SGD
35
17

---

Published as a conference paper at ICLR 2025
A RELATED WORK
Gradient descent in tiny subspaces. This work is largely inspired by previous research demonstrating low-rank structures of the training loss Hessian and the gradient in deep neural network training (Sagun et al., 2016; 2017: Gur-Ari et al., 2018). In particular, Jastrzębski al. (2019) also observe that the SGD update direction is highly aligned with the sharpest direction of the loss landscape. Recently, Schneider et al. (2024) observe that policy gradient algorithms in reinforcement learning also seem to operate in low-dimensional subspaces.
Motivated by such prevalent observations, several follow-up works investigate the possibility of training neural networks in low-dimensional subspace. If feasible, it has wide applications, including few-shot learning (Gauch et al., 2022) and differential privacy (Singhal & Steinke, 2021; Zhou et al., 2021). Li et al. 2a) and Gressmann et al. (2020) train neural networks with a small fraction of parameters using random projections, and Li al. (2023) train ResNet8 on CIFARIO with a 15 15-dimensional subspace without sacrificing test accuracy. Note that the results of these work do not contradict our main observation since the low-dimensional subspaces they chose are not the dominant subspace. In particular, Li et al. (2023) construct low-dimensional subspace by sampling parameter trajectories and then using standard PCA to find the low-dimensional subspace that approximately spans the sampled parameter trajectory. This low- owwdiimesiional subspace differs from the dominant subspace we study, as the dominant subspace is defined by the top-k eigenspace of the loss Hessian. Considering our 2D quadratic example (Figure 6) highlights the difference. In this example, our top-1 dominant eigenspace corresponds to the r-axis; in contrast, if we consider the PCA of the SGD iterates, the principle component would align much closer to the y-axis.
From a theoretical perspective, Arous et al. (2024) rigorously prove that the SGD update aligns with the dominant subspace in multi-class logistic regression and XOR classification with two-layer network. More recently, Yaras et al. (2024) theoretically prove that for deep overparameterized low-rank matrix recovery, the learning dynamics of each weight matrix are confined to an invariant low-dimensional subspace. However, they consider scenarios where GD in the GF regime align with the low-rank subspace, which not the case our settings.
Edge of Stability. Most analyses of GD have focused on settings where the learning rate is sufficiently small to ensure that the training loss monotonically decreases in the GF regime. However, recent empirical studies (Jastrzębski et al., 2019; 2020) observe that GD with practically large learning rates decreases the loss monotonically and finds flatter minima. Cohen et al. (2021) call this the Edge of Stability (FoS) phenomenon. Subsequent theoretical works have made progress towards understanding the mechanisms of EoS (Ahn et al., 2022; Damian et al., 2023; Wu et al., 2023). Moreover, several recent works theoretically analyze precise training dynamics under simplified models (Ahn et al., 2023; Kreisler et al., 2023; Song & Yun, 2023; Zhu et al., 2023). The self- stabilization mechanism (Damian et al., 2023) shows that GD in the EoS regime exhibits oscillation along the top eigenvector, while the overall decrease in loss occurs due movement in directions orthogonal to the top eigenvector. Moreover, SGD with large learning rates also operates at the (stochastic) Edge of Stability (Lee & Jang, 2023; Agarwala & Pennington, 2024). Notably, Zhu et al. (2024) observe that catapults in SGD occurs in the low-rank top eigenspace of NTK, which is closely related to the gradient alignment with the dominant subspace.
Sharpness-Aware Minimization. Inspired by prior works (Keskar et al., 2017; Jiang et al., 2020) showing that flat minima often lead to better generalization, Foret et al. (2021) propose an optimization method called Sharpness-Aware Minimization (SAM), designed to find flat minima. SAM has shown great success in practice, and subsequently, several works thcoretically investigate the dynamics of SAM and its convergence properties (Andriushchenko & Flammarion, 2022; Bartlett et al., 2023; Dai et al., 2023; Si & Yun, 2023; Wen et al., 2023). Recently, Agarwala & Dauphin (2023) and Long Bartlett (2024) empirically observe that SAM also goes through unstable dynamics, akin to EoS.
The role of momentum and Adam. Momentum (Polyak, 1964; Nesterov et al., 2018) and adaptive methods (Kingma & Ba, 2014) are workhorses for training deep neural network models. Adaptive methods, such as Adam, have gained renewed interest due to their success in training language models (Zhang et al., 2020). However, the current understanding of their effectiveness for neural network training remains incomplete.
18

---

Published as a conference paper at ICLR 2025
The role of momentum is quite well understood for convex settings, through acceleration mecha- nism (Nesterov et al., 2018; Kidambi et al., 2018). For nonconvex settings, the provable benefits of momentum investigated for variants of SGD, such normalized SGD (Cutkosky & Mchta, 2020) and signSGD (Crawshaw et al., 2022). recent work by Wang et al. (2024) shows that the benefit of momentum is marginal when the learning rate is small and gradient noise is dominant. Moreover, Fu et al. (2023) empirically demonstrate the benefits of momentum for large learning rates from sharpness perspective.
Adam has been observed to be particularly effective in training transformers (Zhang et al., 2024), even for simplified shallow linear transformers trained on linear regression tasks (Ahn al., 2024a). Its superiority over SGD has been attributed to factors such as heavy-tailed class imbalances in language tasks (Kunstner et al., 2024) and block heterogeneity in Hessian (Zhang et al., 2024). recent line of work shows that full-batch Adam is smoothed version of SignGD (Kunstner et al., 2023; Xie & Li, 2024). Additionally, Ahn al. (2024c) and Ahn Cutkosky (2024) study the benefits of Adam from an online learning perspective.
19

---

Published as a conference paper at ICLR 2025
B EXPERIMENTAL DETAILS
In this section, we provide the details of our experiments which are not covered in the main text.
B.1 ARCHITECTURES
The main experiments are conducted on three types of architectures: MIP, CNN, and Transformer. Additional experiments conducted on standard architectures are provided in Section C.2.
MLP: We use 3-layer MLP with width of 200 and tanh activation functions, following
the architecture used Cohen et al. (2021).
CNN: We use 3-layer CNN with width of 32 and ReLU activation functions, also based on the architecture from Cohen et al. (2021).
Transformer: We use a 2-layer Transformer with a hidden dimension of 64 and 8 attention heads, based the architecture used in Damian et (2023).
B.2 DATA
The main experiments conducted on three datasets: MNIST-5K, CIFAR and Sk.2-Ik. The primary task is classification with categorical MSE loss, and additional experiments with cross-entropy loss are provided in Section C.I.
MNIST-5K: We use the first 5000 samples of MNIST dataset (LeCun al., 1998) for multi-class classification. The number classes is 10.
CIFAR10-5k:
We use the first 5000 samples of CIFARIO dataset (Krizhevsky, 2009) for multi-class classification. The number of classes is 10.
We use the first 1000 samples of SST2 dataset (Socher et al., 2013) for binary classification.
SST2-1k:
B.3 EXPERIMENTAL SETUP
Throughout this paper, all experiments conducted using constant learning rate. For experiments using SGD, we use a batch size 50 for all experiments. Below, we provide details on the choice of learning rates for each experiment, which
specified in the main text.
Figure 1, Figure 2, Figure 3, Figure 18, Figure 19, and Figure 20: The training loss, eigenvalues of the loss Hessian, and ( ((.) are computed on the same run of SGD/GD with small learning rates. The learning rates used are:
MLP on MNIST-5K: 0.01,
CNN CIFAR 10-5k: 0.001,
Transformer SST2-1k: 0.001.
Figure 29, Figure 30, Figure 31, Figure 32, Figure 33, and Figure 34: The training loss,
eigenvalues of the loss Hcssian, and Xx ((,) are computed on the same run of (S)GD with large learning rates. The learning rates used are:
ML.P MNIST-5k: 0.1,
CNN CIFAR 10-5k: 0.01,
Transformer on SST2-IK- 0.005.
Figure 4 and Figure 5: We train MLP on MNIST-5k using (S)GD with learning rate of 0.01, under the same initialization.
Figure 12, Figure 13, and Figure 14: The eigenvalues of the loss Hessian, xx (VI(O.)). and the training loss computed on the same run of SGD. The learning rates used are:
-MLP on MNIST-5k: 0.1,
CNN CIFAR 10-5k: 0.001,
Transformer SST2-1k: 0.001.
20

---

Published as a conference paper at ICLR 2025
Figure 15, Figure 16, and Figure 17: The eigenvalues of the loss Hessian, X (V)()))) and the training loss are computed on the same run of SGD. The learning rates used are:
VGG11 on CIFARIO-5K: 0.01,
ResNet8 on CIFARIO-SK: 0.01.
Our experiments were conducted using Pytorch (Paszke et al., 2019), and we referred to the GitHub repository at https: : / ////tttttb/////// comccoomooooooooo Llccullb// wocuslabbedeeooo of stssaabblltt to replicate the ex- perimental setup described in Cohen al. (2021). All experiments were performed on a single server equipped with 4 NVIDIA RTX 3090 GPUs.
21

---

Published as a conference paper at ICLR 2025
C ADDITIONAL EXPERIMENTS FOR SECTION 3
In this section, we provide additional experimental results to support the observations made in Section These experiments demonstrate that critical observation that Dom-SGD fails to further decrease the training loss-also holds when using cross-entropy loss and training with standard architectures.
C.I
LOSS
We use cross-entropy loss instead of MSE loss for classification tasks, and provide the results analgous Figure 1, Figure and Figure
2000
1500
200
1000
100
500
0
1000 2000
3000
4000
5000
0
5000
10000 15000
20000
5.000 10do0 15600 20000 Training Steps
Training Steps
Training Steps
(a) MLP on MNIST-5k
(b) CNN on CIFARIO-5k
(c) Transformer on SST2-1k
Figure 12: Low-rank structure of the Hessian. The plot shows the top eigenvalues of the loss Hessian during SGD training. The blue curves represent the top-k eigenvalues, which are significantly larger than the next top-k eigenvalues, shown orange, where is the number of classes for the classification task.
1.00
1.00
1.00
0.50
0.50
0.50
0.25
0.25
0.25
0.00
15000 Training Steps
Tranning Steps
Training Steps
(a) MLP on MNIST-5k
CNN on CIFARIO-5K
(c) Transformer on SST2-16 1
Figure 13: Alignment of gradients with dominant subspaces. The plot illustrates X V ((,)) during SGD training. The orange dashed lines represent the exponential moving average (EMA) of X ((.)). After few early steps, X (V L(0)) reaches and stays near 1, indicating the alignment between gradients and dominant subspaces.
0.75
0.50
2-
0.25
0.00
SGD
Dem-ca
1000 2000 3000 4000
5000
5000 10000 15000 20000
5000
10000 15000 20000
Training Steps
Training Steps
Training Steps
(a) MLP on MNIST-5K
(b) CNN on CIFAR10-5K
(c) Transformer on SST2-1K lk
Figure 14: Training (log-scale) of SGD and Dom-SGD, Dom-SGD fails further decrease the training loss in contrast to SGD, despite the gradients aligning approximately with the dominant subspace. We switch from SGD to Dom-SGD whenever the EMA value of X (v(VL ((,)) exceeds 0.95.
22

---

Published as a conference paper at ICLR 2025
C.2 STANDARD ARCHITECTURES ON CIFARIO-5K
To ensure the generality of our observations, we conducted experiments on standard architectures, specifically VGG11 and ResNet8, using CIFAR1O-5K dataset, using MSE loss.
60
100
20
10000 20000 30000 40000 Training Steps
0
10000 20000 30000 40000
Training Steps
(a) VGG11 on CIFARIO-SK
(b) ResNet8 on CIFARIO-SK
Figure 15: Low-rank structure of the Hessian. The plot shows the top eigenvalues of the loss Hes- sian during SGD training. The blue curves represent the top-10 eigenvalues, which are significantly larger than the next top-10 eigenvalues, shown in orange.
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
0
10000 20000 30000 40000
0
10000 20000 30000 40000 Training Steps
Training Steps
(a) VGG11 on CIFARIO-5K
(b) ResNetö on CIFARIO-5K
Figure 16: Alignment of gradients with dominant subspaces. The plot illustrates (VL(0)) during SGD training. The orange dashed lines represent the exponential moving average (EMA) of (vo((L((.)) After a few early steps, (0:)) reaches and stays near 1, indicating the alignment between gradients and dominant subspaces.
00.6
Dom-5GD
-1.0-
0.7
-1.5-
0.8
-2.0
Dom-5GD
2500
5000
7500
10000
5000 10000 15000 20000 Training Steps
Training Steps
VGG11 CIFARIO-5K
(b) ResNet8 CIFAR10-5k
Figure 17: Training loss (log-scale) of SGD and Dom-SGD, Dom-SGD fails to further decrease the training loss in contrast to SGD, despite the gradients aligning approximately with the dominant subspace. We switch from SGD to Dom-SGD whenever the EMA value of X (x((V.VVA (VT(()). exceeds 0.95.
23

---

Published as a conference paper at ICLR 2025
D ADDITIONAL EXPERIMENTS FOR SECTION 4
This section provides additional experimental results support the observations made in Section 4.
D.1 No ALIGNMENT WITH DOMINANT SUBSPACES ALONG GD TRAJECTORIES
We run GD in the GF regime under the same settings Figure 3, using the same learning rate and initialization. Figure 18 shows the top Hessian eigenvalues during training. The smooth increase of the Hessian eigenvalues indicate that the training is happening at the GF regime. The corresponding gradient alignment is shown in Figure 19, and we observe that X& (V ((0))) quickly approaches and remains near 0, indicating that gradients do not align with dominant subspaces, unlike in SGD GF regime. Figure 20 shows that Dom-Gi GD fails to further decrease the training loss in this scenario.
100
200
20
50
0.
0-
2500
5000
7500
10000
2500
5000
7500
10000
2500
5000
7500
10000
Training Steps
Training Steps
Training Steps
(a) MLP on MNIST-5K
(b) CNN on CIFARIO-SK
(c) Transformer on SST2-1K
Figure 18: The plot shows the top eigenvalues of the loss Hessian during GD training in the GF regime. The blue curves represent the top-k eigenvalues, which are significantly larger than the next top-k cigenvalues, shown in orange, where k is the number of classes for the classification task.
1.00
1.00
1.00
0.50
0.50
0.50
0.25
0.25
0.25
2500
5000
7500
10000
2500
5000
7500
10000
2500
10000 Training Steps
5000
7500
Training Steps
Training Steps
(a) MLP on MNIST-5K
(b) CNN on CIFARIO-SK
(c) Transformer on SST2-1k
Figure 19: Gradients do not align with dominant subspaces on GD trajectories. The plot illustrates X& (O))) during GD training in the GF regime. After a few early steps, Xx (LL,)) reaches and stays near 0, indicating alignment with bulk subspaces. The same learning rates and initializations as Figure 3 are used.
-0.50
0.75
Dem/GD
-0.8
1.00
-1.5
1.25
-1.0
1.50
2500 Training Steps
10000
2500
10000
10000
Training Steps
Training Steps
(a) MLP on MNIST-5k
(b) CNN on CIFAR10-Sk
(c) Transformer on SST2-1K 1k
Figure 20: Training loss (log-scale) of Dom-GD and GD. Dom-GD fails to further decrease the training loss, unlike GD. We switch from GD to Dom-GD at iteration 2500.
24

---

Published as a conference paper at ICLR 2025
D.2 SGD TRAJECTORIES TRACK GRADIENT FLOW
Figure 3 and Figure 19 suggest that GD and SGD exhibit different alignments with dominant subspaces, even when using the same learning rate and initialization. However, both should track the continuous-time gradient flow trajectory. if the learning rate is sufficiently small. In Figure 21, we confirm that the trajectory of GD [0 and the trajectory of SGD (0SCb) are close to each other. Specifically, observe that ||0 | trajectories are quite similar.
|0GD
0g indicating the |000D
4-
- a'llll
Ie
2
aalll
0
10000 20000
30000
40000
Training Steps
Figure 21: Trajectories of GD and SGD are close to each other.
D.3 SWITCHING BETWEEN GD AND SGD
Here, provide additional plots for the experiments shown in Figure 4 and Figure 5. Figure 22 illustrates x& (VL(O))) training loss, and top eigenvalues of the loss Hessian when switching from SGD GD, corresponding to the experiment in Figure 4. Figure 23 illustrates Xx (V.(O))) training loss, and top eigenvalues of the loss Hessian when switching from GD to SGD, corresponding to the experiment in Figure 5.
1.00
0.75
75
50
0.25
25
0.00
10000 20000 30000 40000
0
10000 20000 30000 40000
10000 20000 30000 40000 Training Steps
Training Steps
Training
(a) X10 L(e.))
(b) Training loss (log-scale)
(c) Top-20 cigenvalues
Figure 22: The left plot shows X (V.().) when training MLP on MNIST-5k. A sharp transition in gradient alignment with the dominant subspace is observed when switching from SGD to GD (same plot as Figure 4). The plots the training loss and top eigenvalues of the loss Hessian change relatively smoothly when switching from SGD to GD, in contrast to X1o (L.)))). In the right plot, the blue curves represent the top-10 cigenvalues, and the orange curves represent the next top-10 eigenvalues.
25

---

Published as a conference paper at ICLR 2025
1.00
75
0.50
0.25
25
10000 20000 30000 40000 Training Steps
10000 20000 30000 40000 Training Steps
10000 20000 30000 40000 Training Soeps
(a) XIo (VL())
(b) Training loss (log-scale)
(c) Top -20 eigenvalues
Figure 23: The left plot shows Xx (V)():) when training MLP MNIST-5K. A sharp transition gradient alignment with dominant subspace is observed when switching from GD to SGD (same plot Figure 5). The plots of the training loss and top eigenvalues of the loss Hessian change relatively smoothly when switching from GD to SGD, in contrast to X10 (L().)). )) In the right plot,
the blue curves represent the top-10 eigenvalues, and the orange curves represent the next top-10 eigenvalues.
D.4 EFFECT OF NOISE SCALE SPURIOUS ALIGNMENT
We train MLP on MNIST-5k using Noisy Gradient Descent (NGD) with varying noise scales, NGD implemented by injecting noise sampled from an isotropic Gaussian distribution after each GD update iteration. We use the same learning rate and initialization Figure 19, We observe that higher noise scales increase the alignment between the gradient and the dominant subspace.
0.75
0.75
0.75
0.50
0.50
0.50
0.00
0.00
0.00
10000 20000 30000 Training Steps
4000p
10000 20000 30000 40000
10000 20000 30000 40000 Training Steps
Training Steps
(a) σ = 0.01
(b) = 0.05
(c) a = 0.1
Figure 24: Effecc of noise scale in NGD on alignment of gradients with dominant subspace. The plot illustrates X1o with varying noise scales NGD is implemented by injecting noise from Gaussian distribution N 12 after each GD update iteration. We observe that higher noise scales increase the alignment between gradient and dominant subspace, supporting our finding that noise causes spurious alignment in SGD.
during Noisy Gradient Descent (NGD) training MLP on MNIST-5K
We train MLP on MNIST-5k using SGD with varying batch sizes by (50, 100, 500, 1000, 5000} under the same initialization and learning rate. We consider training in the GF regime where GD does not exhibit alignment, using the same learning rate as Figure 19. We observe that smaller batch sizes (i.e., higher noise scale) increase the alignment between the gradient and the dominant subspace. Moreover, the loss curves are similar to each other, indicating that the trajectories are closely tracking the continuous-time gradient flow.
26

---

Published as a conference paper at ICLR 2025
1.00
SGD (BS 50)
-1.0
SGD (BS 100)
SGD (BS 500)
0.75
SGD (BS 1000)
-1.
GD (full-batch)
0.50
SGD
SGD (BS 100)
-2.0
0.25
SGD (B5 500)
SGD (BS 1000)
-2.5
0.00
GD (full-batch)
0
5000 10000 15000 20000
0
5000
10000 15000 20000
Training Steps
Training Steps
(a) Training loss (log-scale)
(b) X10 (v(())
Figure 25; Efffet of batch size in SGD alignment of gradients with dominant subspace. The plot illustrates X10( (l)))) during SGD training MLP on MNIST-5k with varying batch sizes by BS E (50, 100, 500, 1000, 5000). We observe that smaller batch sizes increase the alignment between gradient and dominant subspace, supporting our finding that noise causes spurious alignment in SGD.
D.5 DISTANCE FROM "SWITCHING" STEP
We measure the & distance of the weights from the step where we switch from SGD to Dom-SGD and Bulk-SGI) for the experiments in Figure As shown in Figure 26, the distance from the switching step remains relatively small and tends to saturate for Dom-SGD. In contrast, for both SGD and Bulk-SGD, the distance from the switching step increases much faster and in similar manner. This observation is consistent with the explanation based on an ill-conditionee-vvlleeyllli landscape proposed in Section 4.3.
3
SGD
SGD
Dom SGD
Dom Bulk-CD
0.6
Dom
Bulk-5D
aulk-ED
0.4
0.2
5000
10000
15000
20000
10000
15000
20000
14000 16000 18000 20000 Training Steps
Training Sceps
Training Steps
(a) MLP on MNIST-5k
(b) CNN on CIFARI0-5K
(c) Transformer on SST2-16 1k
Figure 26: The plots illustrate the l2 distance of the weights from the step where we switch from SGD to Dom-SGD and Bulk-SGD the experiments Figure 1. We observe that Dom-SGD fails to make further progress in terms of distance, in contrast to SGD and Bulk-SGD.
We also measure the l2 distance of the weights from the step where we switch from GD to Dom-GD for the experiments in Figure 20. As seen in Figure 27, the distance remains near zero for Dom-GD, indicating minimal movement. While Dom-SGD exhibits slight movement due to its stochastic nature, Dom-GD shows near-zen movement, suggesting it gets stuck at the bottom of the valley, consistent with the ill-conditioned--alleelliikkk landscape described in Section 4.3.
27

---

Published as a conference paper at ICLR 2025
2.0
60
6D Dom-SD
0.4
GD
Dom-GD
Dom-GD
1.5
0.2
0.5
0.1
4000
8000 Training Steps
6000
10000
4000
6000
8000
10000
4000
6000
8000
10000
Training Steps
Training Steps
(a) MLP on MNIST-SK
(b) CNN on CIFARIO-SK
(c) Transformer on SST2-1K
Figure 27: The plots illustrate the & distance of the weights from the step where we switch from GD to Dom-GD for the experiments Figure 20, We observe that Dom-GD gets stuck the initial point.
D.6 TOY EXAMPLE
We conduct additional experiments using Noisy GD for toy model we considered Section 4.2. We consider the same 2-dimensional ill-condtiooee quadratic loss, L(x, y) (1000x²² y²), where 8 (x, v) E RR?. We GD with learning rate 77
@DD
and Noisy GD implemented by injecting Gaussian noise after each GD update
NGD
ONGD L(NGGD + 1J62, where E, ~N(00
In Figure 28, visualize the optimization trajectories of GD and Noisy GD with initialization 00 = 0 NGD = (1, 1), learning rate 7) = 10 4, and noise scale o² 10. We also compute the fraction of gradient in the dominant subspace as (VL(O))
(VL(0).ei)] [VL()) In both GD and Noisy GD trajectories, 2; quickly converges 0, and both trajectories remain close the y-axis throughout the remaining of the training. However, in GD, X1 (VL, (VL (VL(B(D) quickly approaches and remains near while in Noisy GD, X1 (VL(O)) stays close
1.0
005
1.00
0.75
Noisy GD
0.50
GD Initialization
0.25
Noisy GD GD
0.0
0.0
0.5
1.0
0.00
0
25000 50000 75000 100000
Training Steps
(a) GD and Noisy GD trajectories when training d-dimensional coldionit toy quadratic model.
(b) Xı (x())(()) during GD and Noisy GD
Figure 28: We train an ill-conditioned quadratic loss L(x, y) 1/ (11000>² y²) using GD and Noisy
GD. Note that spurious alignment is observed for Noisy GD, unlike GD.
28

---

Published as a conference paper at ICLR 2025
E SPURIOUS ALIGNMENT IN THE EOS REGIME
This section provides additional experimental results on GD and SGD in the EoS regime to support our observations made Section 5.1.
Figure 29 shows the evolution of top Hessian eigenvalues during training. The oscillations of the sharpness indicate that the training happening at the EoS regime. The corresponding gradient alignment is shown in Figure 30, and we observe that X (VL((+)) remains near at the EoS regime, indicating that gradients align with the dominant subspace, unlike GD in the GF regime. Figure 20 shows that Dom-GD fails to further decrease the training loss, despite gradients aligning the dominant subspace.
400
100
200
50
100
0
Training Steps
Training Steps
Training Steps
(a) MLP MNIST-5K
(h) CNN CIFARIO-SK
(c) Transformer on SST2-1K
Figure 29: The plot shows the top eigenvalues of the loss Hessian during GD training with large learning rates. The blue curves represent the top-k eigenvalues, which are significantly larger than the next top-k eigenvalues, shown in orange. After a few initial steps, GD enters the EoS regime, where the sharpness stabilizes near 2 2/1).
1.00
1.00
1.00
0.75
0.75
0.75
0.25
0.25
0.25
2500 5000
7500
10000
2500 5000 7500
10000
0
2500 5000 7500 10000
Training Steps
Training Steps
Training Steps
(a) MLP on MNIST-5K
(b) CNN on CIFARIO-5K
(c) Transformer on SST2-1E lk
Figure 30: Gradients approximately align with dominant subspaces in the EoS regime. The plot illustrates X* ( (VL().) during GD training with large learning rates. The orange dashed lines represent the exponential moving average (EMA) of X1o (V.()))).. As the sharpness reaches 2/n, Xio(VL((..) approaches and remains near 1, indicating the gradient alignment with dominant subspace.
CD
GD
CD
Dem-GD
-0.75
Dem-GD
Dem-GD
-2
-1.00
-2
-1.25
-3
11.00
2500
5000
7500
10000
2500
5000
7500
10600
2500
5000
7500
10000
Training Steps
Training Steps
Training Steps
(a) MLP on MNIST-5k
(b) CNN on CIFARIO-SK
(c) Transformer on SST2-1k
Figure 31: Training loss (log-scale) of GD and Dom-GD. Dom-GD fails to further decrease the training loss in contrast to GD, despite the gradients aligning with the dominant subspace. switch from GD to Dom-GD after GD reaches the EoS regime where X (v ((:))) stays near 1.
29

---

Published as a conference paper at ICLR 2025
Similarly, the same set of experiments SGD with large learning rates shown in Figures 32, 33, and 34. In this scenario, both stochastic noise and self-stabilization effect affects the training dynamics. We observe the gradient alignment in this scenario, and Dom- Dom-SGD again fails to further decrease the training loss. Interestingly, Dom-SGD rather increases the training loss.
15
150
300
10
100
200
50
0
2500 Training Steps
10000
10000
2500
7500
Training Steps
Training Steps
(a) MLP MNIST-5K
(b) CNN CIFARIO-SK
(c) Transformer SST2-1K 1k
The plot shows the top eigenvalues of the loss Hessian during SGD training with large learning rates. The blue curves represent the top-k eigenvalues, which significantly larger than the next top-k eigenvalues, shown in orange.
Figure 32:
1.00
1.00
1.00
0.75
0.75
0.75
0.50
0.50
0.50
0.25
0.25
0.25
2500
5000
7500
10000
2500
5000
7500
10000
2500
5000
7500
10000
Training Steps
Training Steps
Training Steps
(a) MLP on MNIST-5K
(b) CNN on CIFARIO-SK
(c) Transformer on SST2-1k
Figure 33: The plot illustrates X (O,))) during SGD training with large learning rates. The orange dashed lines represent the exponential moving average (EMA) of X LVL(0.)). After few initial steps, X V (0()) approaches and remains near 1, indicating the gradient alignment with dominant subspace.
SCB
Dom SGD
-0.75
Dom 9GD
Dom SGD
1.00
2-
-2
-1.25
-3
11.50
2500 5000
7500
10000
2500
5000
7500
10000
2500
5000
7500
10000
Training Steps
Training Steps
Training Steps
(a) MLP MNIST-5k
(b) CNN on CIFARI0-Sk
(c) Transformer on SST2-1K
Figure 34: Training loss (coges of SGD and Dom-SGD. Dom-SGD fails to further decrease the training loss in contrast to SGD, despite the gradients aligning approximately with the dominant subspace. We switch from SGD to Dom-SGI after X& (V,)),) stabilizes near 1.
30

---

Published as a conference paper at ICLR 2025
F ADDITIONAL EXPERIMENTS FOR SECTION 6
This section provides additional experimental results to support the observations made in Section 6.
F.1 SGDM AND ADAM
We train MLP MNIST-5K using SGD with momentum (SGDM) and Adam. The results shown Figure and Figure 36. We observe that SGDM and Adam updates are "partially" aligned with the dominant subspace due to the effects of momentum and adaptive learning rates. Interestingly, the use of momentum and adaptive methods leads to less alignment between the gradient and dominant subspace than SGD, consistent with the observations effective learning rates Section 6. We also implement Dom-SGDM and Dom-Adam, which project each SGDM/Adam update onto the dominant subspace. We observe that Dom-SGDM and Dom- Adam fail to further decrease the training loss, unlike SGDM and Adam, indicating that the update component aligned with the dominant subspace does not contribute to the loss decrease.
Dom-Adam
1.00
40
-1
Dom SGDM
0.75
20
0.50
0.25
0
0.00
2500
5000
7500
10000
2500
5000
7500
10000
2500
5000
7500
10000
Training Steps
Training Steps
Training Steps
(a) Top-20 20 eigenvalues
(b) (10 0,)
(c) Training loss (log-scale)
Figure 35: SGDM updates partially align with the dominant subspace. Training MLP on MNIST: 5k with SGDM using learning rate rj 0.01 and momentum factor 0.9. (a) Top-10 10 cigenvalues (blue curves) are significantly larger than the next top-10 eigenvalues (orange curves). (b) X10(0 02) values remain in the range of (0.5, 0.8), indicating that SGDM updates are partially aligned with the dominant subspace due to the effect of momentum. (c) Switching the optimizer from SGDM com-SGDM at step 2500 shows that Dom-SGDM fails further decrease the training loss, unlike SGDM.
1.00
0.75
10
0.50
0.00
10000
2500
5000
10000
5000 Training Steps
Training Steps
Training Steps
(a) Top-20 cigenvalues
(b) (1 0c)
(c) Training loss (log-scalc)
Figure 36: Adam updates partially align with the dominant subspace. Training MLP on MNIST- 5k with Adam using learning rate 0.001 and momentum factors (31, 32) (0.9, 0.999). (a) Top-10 eigenvalues (blue curves) are significantly larger than the next top-10 10 eigenvalues (orange curves). (b) (10 - 0c) values remain in the range of (0.05, 0.5), indicating that Adam updates are partially aligned with the dominant subspace. Interestingly, the alignment tends to decrease
during training due to the effect of adaptive learning rates. (c) Switching the optimizer from Adam to Dom-Adam at step 2500 shows that Dom- AdmmAdam fails to further decrease the training loss, unlike Adam.
F.2 EFFECTIVE LEARNING RATES
In Figure 37, we plot effective learning rates (smoothed with EMA) as a function over time for experiments in Figure 1.
31

---

Published as a conference paper at ICLR 2025
0.05
4-
0.00
GD(-m)
GD(-m)
0.05
GD/+m)
GD(+m)
Adam(-m)
Adam(-m)
Adaml+m)
Adam(+m)
250
500
750
1000
200
400
600
800
Training Steps
Training Steps
(a) Dom-LR
(b) Bulk-LR
Figure 37: Effective learning rates (smoothed with EMA) for experiments in Figure I when training Transformer on SST2-1k.
We measure the effective learning rates for (full-batch) GD with and without momentum, and (full- batch) Adam with and without momentum on different architectures and datasets. Tables and present the effective learning rates when training an MLP on MNIST-5k and a CNN on CIFARIO-5K. Figures 38 and 40 show the corresponding training loss plots.
Figure 39 and Figure 41 show the corresponding effective learning rate plots.
We consistently observe that (1) higher Bulk-LR positively correlates with increased training speed, and (2) momentum and adaptive learning rates Adam amplify Bulk-LR, resulting in a larger Bulk-LR compared Dom-LR.
Table 2: Mean effective learning rates over the first 1000 steps (numbers in parentheses show standard deviation). Training MLP on MNIST-5k using GD and Adam with (+m) and without (-m) momentum. GD uses learning rate of 0.01, Adam uses a learn- ing rate of 0.001. Momentum is set to 3 0.9.
2
3
GD(m)
GD(+m)
4
Adam(-m)
Adam(+m)
Method
Mean Dom-LR
Mean Bulk-LR
200
400
600
800
1000
GD(-m)
0.0100 (0.0000)
0.0100 (0.0000)
Training Steps
GD(+m)
0.0012 (0.0092)
0.1005 (0.0068)
Adam(-m)
0.3317 (0.0571)
0.4021 (0.0868)
Figure 38: Training loss in log-scale for the experiments in Table 2.
Adam(+m)
0.0356 (0.0194)
0.7301 (0.4717)
1.5
0.4
0.3
1.0
0.2
GDI-m)
GDI-m)
0.1
GDI+m)
0.5
+m)
Adamf-m)
Adamt-mi
0.0
Adamiam ml
0.0
Adami m))
200
400
600
800
1000
0
200
400
600
800
1000
Training Steps
Training Steps
(a) Dom-LR
(b) Bulk-LR
Figure 39: Effective learning rates (smoothed with EMA) for experiments in Figure 2 when training MLP MNIST-5k.
32

---

Published as a conference paper at ICLR 2025
Table 3: Mean effective learning rates over the first 1000 steps (numbers in parentheses show standard deviation). Training CNN on CIFARIO-5K using GD and Adam with (+m) and without (-m) mo- mentum. GD uses a learning rate of 0.001, and Adam uses a learning rate of 10-4, Momentum is set to 3 = 0.9.
-0.5
-1.00
-1.25
GD(-m)
GD(+m)
Adam(-m)
-1.50
Adam(-m)
Method
Mean Dom-LR
Mean Bulk-LR
250
500
750
1000
GD(-m)
0.0010 (0.0000)
0.0010 (0.0000)
Training Steps
GD(+m)
0.0010 (0.0023)
0.0101 (0.0007)
Adam(-m)
0.0191 (0.0020)
0.0289 (0.0053)
Figure 40: Training loss in log-scale for the experiments in Table 3.
Adam(+m)
0.0046 (0.0026)
0.0802 (0.0194)
0.100
0.02
0.075
0.01
GD(-m)
0.050
GDX-m)
GD(+m)
0.025
GD(+m)
Adam(-m)
Adamí-m)
0.00
Adam(+m)
0.000
Adam(+m)
0
250
500
750
1000
250
500
750
1000
Training Steps
Training Steps
(a) Dom-LR
(b) Bulk-LR
Figure 41: Effective learning rates (smoothed with EMA) for experiments in Figure 3 when training CNN on CIFAR10-5k.
33

---

Published as a conference paper at ICLR 2025
G SHARPNESS PLOTS FOR MAIN EXPERIMENTS
In this section, we provide sharpness plots for main experiments when training MLP on MNIST-5K. Figure 42 shows the sharpness plots of Dom-SGD and Bulk-SGD for the experiment in Figure la. Figure 43 shows the sharpness plots of Dom-GD and Bulk-GD for the experiment in Figure Figure 44 shows the sharpness plots of Dom-SAM and Bulk-SAM for the experiment in Figure 10.
Quite surprisingly, we observe that Dom-SGD decreases the sharpness, unlike SGD and Bulk-SGD. Moreover, Bulk-SGD increases sharpness more rapidly than SGD. We believe these phenomena closely interrelated. As Dom-SGD update decreases the sharpness (A Adom < 0), the absence of the dominant component in Bulk-SGD update would lead larger increase of the sharpness compared to the SGD update (AASOD x Adoom + Abbuk < Abbuk).
Based on these observations, we hypothesize that the dominant component in SGD is primarily an effect of stochastic noise, while the bulk component behaves more similarly to continuous time gradient flow. The sharpness reduction during Dom-SGD may be closely related to the phenomenon where SGD noise biases optimization toward flatter regions of the loss landscape, as analyzed in prior works (Blanc al., 2020; Damian al., 2021; Li et al., 2022). However, these prior analyses focus on dynamics near the minimizer manifold (zero loss), whereas the sharpness reduction in Dom-SGD occurs earlier, before reaching the minimizer manifold.
We hypothesize that the noise in Dom-SGI> implicitly reduces sharpness near the manifold of points where the dominant component the gradient vanishes (i.e., xx(V = 0). Further investigation of this intriguing phenomenon is left future work.
the EoS regime, we observe that Bulk-GD increases the sharpness larger than the stability threshold 2/ 2/m. Similarly, Bulk-SAM increases the sharpness larger than SAM-edge, the stability threshold of SAM.
30
60
20
20
10
20
0:
0
5000
10000
15000
20000
5000
10000
15000
20000
0
5000
10000 15000 20000
Training Steps
Training Steps
Training Steps
(a) SGD
(b) SGD to Dom-SGD
(c) SGD to Bulk-SGD
Figure 42:
We train MLP on MNIST-5k using SGD in the GF regime, using learning rate v 0.01 under the same setting as Figure 1a. The plot shows the top eigenvalues of the loss Hessian during SGD and Dom/Bulk-SGh training. The blue curves represent the top-10 eigenvalues, and orange curves represent the next top-10 cigenvalues.
100
20
20
75-
10
50
25:
0
2000 4000 6000
8000 10000
2000 4000 6000 8000 10000
2500
5000 7500
10000
Training Steps
Training Steps
Training Steps
(a) GD
(b) GD to Dom-GD
(c) GD to Bulk-GD
Figure 43: We train MLP on MNIST-5k using GD with large learning rate 7) = 0.1 under the same setting as Figure 9. The plot shows the top eigenvalues of the loss Hessian during GD and Dom/Bulk- GD training. The blue curves represent the top-10 eigenvalues, and orange curves represent the next top-10 eigenvalues.
34

---

Published as a conference paper at ICLR 2025
15
10
0-
2000 4000 8000 8000 10000
2000 4000 6000 8000 Training Steps
10000
2000 4000 6000 8000 10000 Training Steps
Training Steps
(a) SAM
(b) SAM to Dom-SAM
(c) SAM to Bulk-SAM
Figure 44: We train MLP on MNIST-5K using SAM under the same setting Figure 10. The plot shows the top eigenvalues of the loss Hessian during SAM and Dom/Bulk-SAl -SMM training. The blue curves represent the top-10 eigenvalues, and orange curves represent the next top-10 10 eigenvalues.
H TEST ACCURACY RESULTS FOR DOM-SGD AND BULK-SGD
In this section, we provide preliminary results test accuracy of SGD, Dom-SGD, and Bulk-SGD. The experiments are conducted on full MNIST dataset with a learning rate rj = 0.01 and the plots of train loss and test accuracy similar to that of train loss, i.e., Bulk-SGD generalizes as well as SGD, while Dom-SGD fails to generalize well. We believe that more systematic experiments beyond this preliminary results would help to provide deeper insights on the generalization characteristics of Dom/Bulk-SGD, and we leave this question as future work.
provided in Figure 45. The results show that the trend of test accuracy
0.5
SGD
SGl
0.9
0.3
0.2
0.8
SGD
Dom-SGD
0.1
Bulk-5GD
0.7
0
10000 20000 30000 40000 Training Steps
0
10000 20000 30000 40000
Training Steps
(a) Train loss
(b) Test accuracy
Figure 45: We train MLP MNIST using SGD and Dom/Bulk-SGRP with learning rate 7) 0.01. In case of Dom-SGD, training loss does not decrease and test accuracy does not increase. In contrast, for Bulk-SGD, training loss decreases and test accuracy increases similar SGD.
35


Reviews:
[{"review_id": "FdvxfMttWO", "reviewer": "Reviewer_cwmc", "summary": "This paper investigates whether neural networks can be effectively trained by focusing only on the \"dominant subspace\", which is the top eigenspace of the Hessian of the loss, where the gradient seems to align with. The authors projects SGD updates onto this dominant subspace and find that cannot reduce the training loss further. In contrast, removing the dominant subspace from the updates proves equally effective for training, even though it projects out most of the original gradient component. The authors suggest that the alignment between the gradient and the dominant hessian subspace is \"spurious\". This \"spurious\" alignment appears consistently across various practical settings, such as the Edge of Stability, Sharpness-Aware Minimization, and when using momentum or adaptive optimizers. The authors also discuss the possible main causes of the phenomenon, and propose some toy models to understand it.", "questions": "1. Can the authors calculate the full gradient when \"spurious alignment\" happens? Does that also align with the Hessian's top eigenspace? Or it has a very small similarity?\n2. Can the authors experiment with some kind of Gaussian noise for the toy example? And can you empirically show why they believe the stochastic noise has a larger bias in the top-eigenspace?\n3. For Appendix D.1, can the authors (1) report the learning rates (2) train with longer time s.t. EoS happens (3) add experiments with MSE loss?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper systematically investigates the phenomenon of gradient-Hessian alignment in various optimization algorithms, including SGD, GD in the Edge of Stability (EoS) regime, and the Sharpness-Aware Minimization (SAM) algorithm, with a particular emphasis on the analysis of SGD. Previous studies have primarily focused on understanding full-batch algorithms like GD or Adaptive GD without stochasticity. \n\nMoreover, this work takes an initial step toward understanding the 'spurious' alignment where the actual useful component of the gradient is the non-dominant part. They also claim that batch noise is the cause of this alignment by introducing a toy model to provide theoretical intuition. Though it is not very precise to say the mechanism is totally different between GD and SGD (see weakness), I believe the **\"ill-conditioned-valley\" view of loss landscape** captures why dom-GD doesn't decrease the loss but bulk-GD does.", "weaknesses": "First, the so-called 'spurious alignment' is long observed in EoS literature in my opinion. For example, Damian et al. [2] showed that in the EoS regime, (1) the gradient alignment happens (2) the loss decrement in the EoS regime depends on the constrained trajectory **by projecting out the top eigenspace**. It is exactly the finding listed in section 5.1 of this paper that bulk-GD is as effective as GD. I believe the authors should discuss those related works.\n\nThe author may argue that \"The mechanisms behind gradient alignment differ between GD with large learning rates and SGD with small learning rates.\" so their findings are novel. But that is also my primary concern of the paper's conclusion \"Stochastic noise is the cause of spurious alignment\". I agree that \"stochastic noise is part of the cause of spurious alignment\", but I don't think the alignment mechanism is intrinsically different between GD and SGD.\n\nIn section 4. The authors define the \"small LR regime\" and \"large LR regime\" of SGD and claim they are different by setting the threshold of $2/\\lambda_1$ of the learning rate. This threshold is for the GD descent lemma to hold. Then the authors claim in the small LR regime, the cause of the spurious alignment is the gradient noise. They 'verify' their findings by switching SGD to GD, and observe that the alignment disappears. However, [1] found that SGD begins to oscillate within the top eigenspace smaller than $2/\\lambda_1$. An intuitive explanation for this is **injected gradient noise increases the second order term of the Taylor expansion**, making the loss begin to oscillate **when the learning rate is below $2/\\lambda_1$** as the descent lemma predicted. In this case, GD is still stable and descent lemma holds, quickly converging back to the \"bottom of the valley\". That is why the alignment disappears after switching to GD, since after it's in the \"bottom of the valley\" the gradient component in the top eigenspace will be very small. \n\nTherefore, my argument is that stochastic noise indeed reduces the threshold of learning rate for oscillation and makes it easier to have spurious alignment. But the mechanism is the same: the self-stabilization effect induced by the gradient(+noise) within the top Hessian subspace. **To test this argument, the author should also calculate the full gradient when \"spurious alignment\" happens and see if the full gradient aligns with the top eigenspace. Also, the experiments should include various batch sizes to test different levels of gradient noise.** If my conjecture holds, I think this work can be seen as a generalized empirical validation of self-stabilization [2] for SGD, which somehow limits the novelty of this paper. I might also be wrong, so if the conjecture is not correct, I will raise my score.\n\nThe author also did some small learning rate GD in Appendix D.1 to corroborate their argument. I wonder what is the learning rate of the \"small learning rate\"? In the original EoS paper, most of the figures will exhibit EoS due to progressive sharpening when the model is trained for a long time. Also, the authors use the cross-entropy loss, where EoS may not happen due to some margin-maximization effect. What if the authors switch to square loss? I believe when you enter the EoS regime (when trained with enough time),  the gradient alignment phenomenon will also appear.\n\nThe authors' toy examples constructed the two loss functions to introduce the gradient noise. However, I don't think this reflects the real-world gradient noise, since the 100xy term is artificially added to introduce gradient components in $x$-direction unless $y = 0$. But I think injecting some Gaussian noise biasing toward $x$-direction may also work for the toy case. It would be great to have some empirical evidence as to why the authors believe the gradient noise has a larger component in the top eigenspace.\n\n[1] Lee, Sungyoon, and Cheongjae Jang. \"A new characterization of the edge of stability based on a sharpness measure aware of batch gradient distribution.\" The Eleventh International Conference on Learning Representations. 2023.\n\n[2] Damian, Alex, Eshaan Nichani, and Jason D. Lee. \"Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability.\"", "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": "aIjNQf4CYv", "reviewer": "Reviewer_rVK6", "summary": "This paper investigates that neural nets cannot be trained within the dominant subspace (i.e., along sharp directions).", "questions": "- The experiments focus on the optimization comparsion between SGD, Bulk-SGD, and Dom-SGD. \nWhat are the generalization differences among these methods, particularly in the experiments conducted on MNIST and CIFAR, as shown in Figure 1?\n\n- Could the authors provide sufficient theoretical explanation for the main finding?", "limitations": "", "rating": 3, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 1, "strengths": "Many recent papers (Blanc et al., 2022; Li et al., 2022) highlight that the dynamics in bulk subspace (i.e., along flat directions) are crucial for SGD/SAM to move to flat minima, thereby improving generalization.\nIn contrast, this paper emphasizes the significant role of the dynamics in bulk subspace (i.e., along flat directions) in relation to optimization.", "weaknesses": "My primary concern are (i) the paper does not sufficiently explain the main finding, i.e., why only the dynamics in the bulk subspace are crucial for optimization, and (ii) the novelty of many contexts:\n\n- Section 4. This section focus on the alignment between the stochastic gradient and the sharp directions. However,\n\n  - This section fails to adequately explain the main finding: why only the dynamics in the bulk subspace are crucial for optimization, expect for a very toy model.\n\n  - Even regarding the alignment between stochatic gradient and the loss landscape, there are numerous previous works that address this point: $\\mathbb{E}[g_tg_t^T]\\approx 2 L(\\theta)\\nabla^2 L(\\theta_t)$ (Wojtowytsch, 2021; Mori et al., 2022; Wu et al., 2022; Wang and Wu., 2023), implying that stochastic gradient concentrate more along sharp directions of the landscape (i.e., the bump directions). It appears that the authors have overlooked these works.\n\n- Section 5. This section demonstrates that the gradient of GD aligns with the sharp directions in Edge of Stability (EoS), However, \n\n  - The authors do not explain the main finding: why only the dynamics in the bulk subspace are crucial for optimization.\n\n  - Regarding the alignment between the gradient and sharp directions, this has beed well studied in (Arora et al., 2022; Lyu et al; 2022; Damian et al., 2023) that this occurs due to approximate power iterations). However, the authors do not provide new insights.", "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": "nb3bRER918", "reviewer": "Reviewer_LPQC", "summary": "This paper challenges existing assumptions about stochastic gradient descent (SGD) learning within a low-dimensional dominant subspace defined by the top eigenvectors of the Hessian of the loss. The authors find empirically that while gradients appear to align with this dominant subspace during training, projecting updates solely into this subspace (Dom-SGD) halts learning. Instead, projecting updates onto the orthogonal \"bulk\" subspace (Bulk-SGD) is as effective as standard SGD, suggesting that meaningful learning occurs outside the dominant subspace.", "questions": "Why did the authors limit their experiments to a small subset of the datasets and can the authors share any insights for training with the full dataset?\n\nCould the authors provide any test accuracy results for Dom-SGD and Bulk-SGD to offer some insight into their impact on model generalization? When Bulk-SGD performs similarly to SGD in terms of training loss, does this trend hold for test accuracy as well?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper addresses a question of significant interest in the ML community regarding SGD's effectiveness in low-dimensional subspaces. It helps address potential misconceptions arising from the well-cited work of Gur-Ari et al. (2018), which suggests that gradient descent would occur primarily within a tiny subspace.\n\nA convincing quadratic toy model effectively reinforces the authors’ interpretation of the empirical results, lending credibility to their main conclusions.\n\nThe paper is well-structured, with a logical flow that makes the argument easy to follow and the findings readily accessible to readers.\n\nGur-Ari et al. (2018): https://arxiv.org/abs/1812.04754", "weaknesses": "Although the paper includes experiments on three datasets, training is restricted to small subsets (e.g., 5,000 of 50,000 samples of CIFAR10) and primarily uses mean squared error loss instead of cross-entropy, despite focusing on classification tasks. This may limit the generalizability of the findings and should be communicated more clearly.\n\nThe paper exclusively examines the effects on training loss, with no analysis of test accuracy under Dom-SGD and Bulk-SGD. Including at least one plot of test performance would offer valuable insights into the practical implications for readers.", "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": "tBUI9LbVNq", "reviewer": "Reviewer_pDDh", "summary": "This paper is built on recent studies that claim the gradients during neural network training live in a small subspace spanned by the top few eigenvectors of the Hessian. The authors decompose the gradients into the dominant component that lives in the dominant space and the bulk component that is orthogonal to the subspace. They find that dominant component of the gradient does not contribute to decreasing the training loss, while the bulk component, though only accounts for a small portion of the total gradient magnitude, is the main source driving the reduction in the loss. Through different experiments, the authors demonstrate that this observation holds for SGD, GD on the edge of stability, SAM, and adaptive optimizers. The authors also discuss potential explanations from a loss landscape perspective.", "questions": "Instead of computing the subspace based on the Hessian of the full loss, what if the subspace were computed based on the Hessian of the mini-batch loss. Would alignment still hold? Do the findings in the paper still hold?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The paper is well written with a clear structure and compelling experimental results. The design of the experiments is sophisticated, covering both Dom-SGD and Bulk-SGD setups, along with more nuanced experiments like switching between SGD and GD during training. These additional experiments strongly support the paper’s claims, adding robustness to the argument. Additionally, the observations found in this paper is interesting and should contribute to a better understanding of the training dynamics of deep neural networks.", "weaknesses": "Major point:\n\n1.\tGeneralization: The authors acknowledge in the limitation that this paper does not consider the generalization part. However, in modern deep learning, generalization error is usually more important than training loss. And it is widely believed that SGD has an implicit bias that benefits generalization (arxiv: 1611.03530, 2006.08680, 2306.04251, etc.). The absence of results or discussion on generalization weakens the paper. Including an analysis of generalization effects could significantly enhance the paper's impact.\n\n2.\tBulk-SGD and noise: The results suggest that Bulk-SGD is less noisy than SGD (e. g. in Fig. 1) and Bulk-GD is less noisy than GD (Fig. 9). This reduction in noise may explain the speedup observed in Fig. 9 and 10. On the other hand, it is unclear how this will affect generalization. The authors have not adequately addressed this part in the paper, and I suggest including a more thorough discussion around that.\n\n3.\tDom-SGD and loss increase: It appears that with cross-entropy loss or using a standard architecture (VGG-11) (Fig. 14 and 17), the loss increases if training with Dom-SGD. This is a surprising finding and I think the authors should discuss this phenomenon more explicitly in the main text. Additionally, the absence of training curves for Bulk-SGD in these figures raises questions. Is there a reason for not including the curves for Bulk-SGD? I am also wondering whether this increase in the training loss relates to findings in prior works: arxiv: 2107.11774, 2306.04251.\n\n\nMinor point:\n\n1.\tFor results shown in table 1, the authors calculate the effective learning rates over the first 1000 steps. However, the authors also show that the alignment usually happens after some warm-up steps (e. g. Fig. 3). I think it makes more sense to plot the effective learning rates as a function over time (if noise is a concern, then perhaps run some kind of EMA). \n\n2.\tAlthough the toy quadratic model explains the observations well, the source of randomness seems a little contrived. I suggest incorporating the discussion of another model at least in the appendix, where the noise is Gaussian noise directly added on to the gradients, for a more natural comparison.", "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)