Does SGD really happen in tiny subspaces?

Thank you for your efforts and your insightful comments! We appreciate that you found our observations in this paper interesting and should contribute to a better understanding of the training dynamics of deep neural networks.

Let us address your concerns one by one.

Generalization

We definitely agree that considering generalization is an important problem. While this question is out of the scope of our paper, we provide preliminary results on the test accuracy of SGD, Dom-SGD, and Bulk-SGD in Appendix H. The additional experiments are conducted on the full MNIST dataset and the plots of train loss and test accuracy for each method are provided. The results show that the trend of test accuracy is 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 these preliminary results would be necessary to provide better insights on model generalization, and we leave this question as a future work.

Bulk-SGD and noise / Dom-SGD and loss increase

Thank you for your insightful comments! We appreciate your suggestion and have added a discussion on this phenomenon in the Conclusion section, along with additional experimental results in Appendix G. Our experiments consistently show that Bulk-SGD tends to be less noisy than SGD, while Dom-SGD tends to increase the loss over time. Specifically, Dom-SGD exhibits a slow but steady increase in loss in the long run across all experiments, including Figure 1. Additionally, we observe that Dom-SGD decreases sharpness, and we have included sharpness plots for both Dom-SGD and Bulk-SGD.

Interestingly, Bulk-SGD increases sharpness more rapidly than SGD. We believe these phenomena are closely interrelated. As Dom-SGD update decreases the sharpness ($\Delta \lambda_{\text{dom}}<0$), the absence of the dominant component in Bulk-SGD update would lead to larger increase of the sharpness compared to the SGD update ($\Delta \lambda_{\text{SGD}} \approx \Delta \lambda_{\text{dom}} + \Delta \lambda_{\text{bulk}} < \Delta \lambda_{\text{bulk}}$). We currently do not have a fully satisfactory explanation of these phenomena, and we leave this question as a future work.

In the additional plots in the Appendix, we focused on verifying our main claim that training cannot occur in the dominant subspace, which is why we did not include the training loss curves for Bulk-SGD. However, we observed the same behavior for Bulk-SGD in those experiments.

Furthermore, we are uncertain whether the prior works demonstrating that SGD can converge to local maxima (arXiv:2107.11774) and the stochastic collapse of SGD (arXiv:2306.04251) are directly related to the loss increase observed in Dom-SGD. Specifically, we do not believe that the loss increase in Dom-SGD is caused by convergence to a local maximum. Could you elaborate further on the potential connection between our findings and the stochastic collapse?

Plot the effective learning rates as a function over time

Thank you for the suggestion. In Appendix F.2, we added the plots of effective learning rates as a function over time for the experiments training MLP on MNIST-5k, CNN on CIFAR10-5k, and Transformer on SST2-1k. Notably, we observe that Adam (with momentum) tends to gradually increase Bulk-LR during training Transformer on SST2-1k (see Figure 37). This observation further supports the idea that Adam effectively adapts to the loss landscape, enabling efficient optimization of Transformers.

Toy quadratic model with Gaussian noise

Thank you for the suggestion. We added the experiments on the toy quadratic model with Noisy GD (i.e., Gaussian noise) in Appendix D.6. We implement Noisy GD by injecting a noise sampled from an isotropic Gaussian distribution after each GD update. The results are similar to the original toy model. Both GD and Noisy GD trajectories closely follow the y-axis. However, the gradient aligns with the dominant subspace (x-axis) for Noisy GD, unlike GD. Note that we also provide experiment results with Noisy GD under deep learning setups in Appendix D.4, Figure 24.

Subspace computed based on the Hessian of the mini-batch loss

We compute the subspaces with the full loss Hessian because Gur-Ari et al. (2018) observed that full gradients align with the top eigenspace of the full loss Hessian during SGD training. Moreover, they also observed that the top eigenspace of full loss Hessian (i.e., dominant subspace) is mostly preserved over the long period of training. In contrast, the top eigenspace of mini-batch loss Hessian is unstable and quickly changes over time during training. Therefore, the alignment between gradient and top eigenspace of mini-batch loss Hessian is also unstable compared to that of full loss Hessian.

I thank the authors for their further responses. I intend to update the score to 7.

Thank you for your support and for your valuable feedback! We appreciate that you found our work addresses a question of significant interest in the ML community regarding SGD's effectiveness in low-dimensional subspaces.

We would like to address your questions as below.

Why did the authors limit their experiments to a small subset of the datasets?

Let us explain why we decided to conduct experiments on small subsets. Conducting extensive systematic experiments on full datasets is computationally challenging given that Dom-SGD and Bulk-SGD require computing top-10 eigenvalues and eigenvectors of the full loss Hessian at every single iteration. Existing works such as Cohen et al. (2021) did not compute the sharpness of full loss Hessian for the CIFAR10 dataset (see https://github.com/locuslab/edge-of-stability). Instead, they computed the proxy of sharpness by randomly selecting a batch of 5k samples and computing the top eigenvalue of mini-batch loss Hessian. However, such an approximation method could lead to an unstable approximation of eigenvectors, making it inappropriate to implement Dom-SGD and Bulk-SGD. For this reason, we chose to perform experiments on smaller subsets of datasets.

To demonstrate that the trends observed on the subsampled datasets are consistent with those on the full dataset, we added a new set of experiments using the full MNIST dataset in Appendix H. The results confirm that our main observations hold true in this setting as well.

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

Thank you for the suggestion. We provide preliminary results on the test accuracy of SGD, Dom-SGD, and Bulk-SGD in Appendix H. As noted above, the additional experiments are conducted on full MNIST dataset and the plots of train loss and test accuracy for each method are provided. The results show that the trend of test accuracy is 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 these preliminary results would be necessary to provide better insights on model generalization, and we leave this question as a future work.

Primarily uses MSE Loss instead of cross-entropy, despite focusing on classification tasks.

The choice of MSE loss was intended to clearly distinguish the settings between the small LR (GF) regime and the large LR (EoS) regime. When cross-entropy loss is used with a large LR, progressive sharpening initially occurs and enters the EoS regime, but training does not remain in the EoS regime till the end since sharpness decreases at the end of training (Cohen et al., 2021). In contrast, when MSE loss is used with a large LR, sharpness does not decrease after entering the EoS regime. Therefore, in order to conduct systematic experiments investigating the distinct mechanism of spurious alignment depending on LR, we decided to primarily use MSE loss.

Q1. The experiments focus on the optimization comparsion between SGD, Bulk-SGD, and Dom-SGD. What are the generalization differences among these methods, particularly in the experiments conducted on MNIST and CIFAR, as shown in Figure 1?

Although our paper primarily focuses on optimization, we provide preliminary results addressing the reviewer’s question about generalization differences among SGD, Dom-SGD, and Bulk-SGD in Appendix H. These experiments were conducted on the full MNIST dataset, and we include plots of training loss and test accuracy for each method. The results indicate that test accuracy trends closely mirror those of training loss: Bulk-SGD generalizes as well as SGD, while Dom-SGD struggles to generalize effectively.

Q2. Could the authors provide sufficient theoretical explanation for the main finding?

Our paper is an empirical work that uncovers an interesting and counterintuitive phenomenon in NN training dynamics and proposes (through a toy example) a novel hypothesis that can explain it. We believe that providing a theory on our empirical findings beyond the toy model is definitely an important future research direction, but is currently out of the scope of this paper.

Scientific progress is driven by discovery, the formulation of hypotheses, and their subsequent validation. In our work, we have presented an intriguing discovery and hypothesis, which we believe constitutes a meaningful contribution. While we acknowledge that a full theory is not provided, we feel that rejecting the paper solely on this basis is overly harsh. We kindly ask you to reconsider your evaluation of our manuscript.

We appreciate your efforts and valuable feedback. Before addressing each of the reviewers' comments, we would like to highlight the novel contributions of our paper. As Reviewer LPQC noted, our work tackles a potential misconception arising from the influential work of Gur-Ari et al. (2018), which suggests that gradient descent would occur primarily within the dominant subspace. Contrary to this expectation, our empirical findings reveal that projecting each gradient update onto the dominant subspace does not, in fact, decrease the loss—a surprising observation, especially given that gradient vectors appear aligned with the dominant subspace during SGD training. Furthermore, we deepen our analysis of this phenomenon through systematic experiments on real-world datasets, as well as a toy model, providing further insights into the underlying mechanism.

Let us address your concerns one by one.

(Section 4) This section fails to adequately explain the main finding: why only the dynamics in the bulk subspace are crucial for optimization, except for a very toy model.

The primary focus of our paper is to identify and analyze a novel phenomenon in deep learning through systematic experiments. The toy model serves as a tool to provide intuition for the counterintuitive phenomenon we uncovered. Specifically, our ill-conditioned quadratic toy model, though simple, captures the essence of our empirical findings. This is evidenced by its ability to replicate all the key phenomena observed in the deep learning setup (see Section 4.2 for details).

Building on the intuitions from the toy model, we hypothesize that the "ill-conditioned valley" loss landscape of neural networks explains why the bulk subspace dynamics are crucial for optimization (see Section 4.3 for details). In such a landscape, the SGD trajectory remains near the bottom of the valley. The valley is elongated along the bulk subspace, and true optimization progress occurs within this subspace, as the bottom of the valley descends in that direction.

(Section 4) Even regarding the alignment between stochastic gradient and the loss landscape, there are numerous previous works that address this point: $\mathbb{E}[g_tg_t^\top]\approx 2L(\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.

First, we did not compute the alignment between the "stochastic" gradient and the loss landscape. Instead, we measured the alignment between the "full-batch" gradient and the loss landscape along the SGD trajectory. The previous works you mentioned (e.g. Wu et al., 2022) demonstrate that the stochastic noise has a larger bias in sharp directions, but this does not imply that full-batch gradient aligns with top eigenspace because full-batch gradients are "noiseless."

In fact, the "spurious alignment" also occurs even when the stochastic noise is isotropic, i.e., the same bias along every direction. In Appendix D.4, we run Noisy GD in neural network experiments by injecting noise sampled from an isotropic Gaussian distribution after each GD update iteration and observe the alignment between full-batch gradient and top eigenspace.

(Section 5) The authors do not explain the main finding: why only the dynamics in the bulk subspace are crucial for optimization. The alignment between the gradient and sharp directions has been 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.

We agree that the alignment between the gradient and dominant subspace in the EoS regime can be understood through approximate power iteration (also known as self-stabilization). Indeed, we cited Damian et al. (2023) and explicitly highlighted its relevance in Remark 2 of Section 5, as well as in the Related Work section. However, we emphasize that the primary contribution of Section 5 is not the discovery of this alignment but rather our observations on Dom-GD and Bulk-GD in the EoS regime. It is important to note that while Damian et al.'s analysis explains the behavior of GD, it does not extend to Bulk-GD. Specifically, their theory cannot predict or explain the training dynamics when the self-stabilization effect, arising from oscillations along the dominant subspace, is absent.

Furthermore, we hypothesize that the "ill-conditioned valley" loss landscape of neural networks explains why dynamics in the bulk subspace are critical for optimization, even in the EoS regime. Specifically, in the EoS regime, the gradient descent (GD) updates jump across the valley, oscillating within the dominant subspace, while the "true" optimization progress occurs within the bulk subspace, where the bottom of the valley descends. This provides an intuitive explanation for why bulk subspace dynamics play a crucial role in optimization.

Dear Reviewer rVK6,

Thank you for your initial review of our paper. We wanted to check if you have had the chance to read our rebuttal, as it has been a week since we uploaded our response. The discussion period is coming to an end, and we would greatly appreciate your thoughts on whether our responses resolved the concerns you raised.

If there are any remaining concerns, we would be happy to clarify further. Alternatively, if you find that your concerns have been sufficiently addressed, we kindly ask you to consider adjusting your score accordingly.

Thank you for your time and contributions to the review process.

Sincerely,

Authors

I thank the authors for their clarifications; however, my primary concerns remain unresolved.

• Main concern: the paper does not sufficiently explain the main finding, i.e., why only the dynamics in the bulk subspace are crucial for optimization? While I strongly agree that toy models can be valuable for early-stage research, it is essential that they preserve the key characteristics of the problem. For example, neural network training is well-known to occur at the edge of stability (EoS). However, the toy model used in this work—a quadratic function—does not exhibit EoS, as it only allows for two states: convergence or divergence.

• Regarding Section 4: I find this section particularly confusing. As clarified, this work measures the alignment of the full-batch gradient with the loss landscape. Why, then, does it emphasize the critical role of stochastic gradients (e.g., the title of Section 4.1)? I would appreciate clarification on this apparent contradiction.

• Regarding Section 5: While I acknowledge the clarification on the distinctions from related works, the current explanation fails to elucidate the underlying mechanism of Bulk-GD.

• Additionally, a recent paper should be discussed.

Given these points, I will maintain my rating.

[1] Cohen et al. Understanding Optimization in Deep Learning with Central Flows. https://arxiv.org/pdf/2410.24206

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

First, we would like to emphasize that our observation does not necessarily imply that stochastic noise has a larger bias in the top eigenspace, given that we measure the alignment of full-batch gradients, but not mini-batch gradients along SGD trajectories. Indeed, the "spurious alignment" also arises when we run Noisy GD for neural network training, by injecting noise sampled from an isotropic Gaussian distribution after each GD update iteration. The results are presented in Appendix D.4. Notice that even when the noise is isotropic (i.e., same bias to every direction), we can still observe the alignment between full-batch gradient and top eigenspace, which happens due to the ill-conditioned landscape.

We also added the experiments on the toy quadratic model with Noisy GD (i.e., Gaussian noise) in Appendix D.6. We implement Noisy GD by injecting a noise sampled from an isotropic Gaussian distribution after each GD update. The results are similar to the original toy model. Both GD and Noisy GD trajectories closely follow the y-axis. However, the gradient aligns with the dominant subspace (x-axis) for Noisy GD, unlike GD. Note that we also provide experiment results with Noisy GD under deep learning setups in Appendix D.4, Figure 24.

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

First, the learning rates of ALL experiments are reported in Appendix B Experimental Details. For experiments in Appendix D.1, the learning rates are reported in Appendix B.3 and here are the values:

• MLP on MNIST-5k: 0.01
• CNN on CIFAR10-5k: 0.001
• Transformer on SST2-1k: 0.001

Note that the learning rates (and also the initialization) for experiments in Appendix D.1 (Figures 18,19,20) and experiments in the main text (Figures 1,2,3) are the same, but the only difference is the choice of optimizer (GD vs SGD).

Second, every experiment in the paper uses MSE loss, except Appendix C.1 (we use cross-entropy loss here).

Finally, if we train longer until EoS happens, then the gradient alignment phenomenon will also appear after entering the EoS regime. This phenomenon can be naturally expected from our experiments at the EoS regime in Section 5. Specifically, if we run GD with a constant LR $\eta$, the sharpness initially increases (progressive sharpening) until it reaches $2/\eta$. This phase of training is in the GF regime, where alignment does not occur. As training continues, GD enters the EoS regime, where sharpness oscillates around $2/\eta$. In the EoS regime, the alignment between gradient and dominant subspace occurs.

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. (...) Therefore, my argument is that stochastic noise indeed reduces the threshold of learning rate for oscillation and makes it easier to have spurious alignment. However, the mechanism is the same: the self-stabilization effect induced by the gradient(+noise) within the top Hessian subspace.

We definitely agree that stochastic noise reduces the "stable" learning rate (i.e., stable LR of SGD is smaller than stable LR of GD), as found by previous works. However, in the GF regime of SGD, we considered using learning rates smaller than the stable LR of SGD, and hence the mechanism behind the gradient alignment is different from the self-stabilization effect. Let us provide concrete evidence on why this is the case.

1. Sharpness smoothly increases along the SGD trajectory: In Figure 2, we plot top eigenvalues of Hessian during SGD training. Notice that top eigenvalues smoothly increase (i.e., progressive sharpening occurs). If the SGD training were in the EoS regime, then sharpness should oscillate, as observed in e.g., Lee and Jang (2023). This indicates that experiments we did for SGD in the GF regime use LRs smaller than the stable LR of SGD.
2. GD and SGD trajectories are close to each other: In Appendix D.2, we observe that trajectories of SGD (exp in Figures 1,2,3) closely track that of GD (exp in Figures 18,19,20), under the same initialization and LR. If LR were larger than the stable LR of SGD, then SGD trajectory would become unstable and move away from GD trajectory.

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.

We did calculate the full gradient when we computed the gradient alignment in the SGD experiments. In fact, we did not measure the alignment between the mini-batch gradient and top eigenspace. Instead, we measured $\chi_k(\nabla L(\theta))$, the alignment between full-batch gradient and top eigenspace along the SGD trajectory (e.g., see Figures 3, 4, 5, 7).

Also, the experiments should include various batch sizes to test different levels of gradient noise.

Thanks to your suggestion, we conducted additional experiments with various batch sizes and the results are provided in Appendix D.4, Figure 25. We train MLP on MNIST-5k with the same initialization and LR and vary batch sizes by {50, 100, 500, 1000, 5000(full-batch)}. (Here, the alignments are still computed with the full gradients.) We observe that the loss curves are similar to each other. Moreover, a smaller batch size increases the gradient alignment. This observation is consistent with our experiments on Noisy GD in Appendix D.4, Figure 24. These observations further support our finding that noise causes spurious alignment in SGD, and this is different from the self-stabilization effect in the EoS regime.

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

As we mentioned above, we did calculate the full-batch gradients and observed that they align with the Hessian's top eigenspace. For example, in Figure 3, we measure $\chi_k(\nabla L(\theta))$, the alignment between the full-batch gradient and dominant subspace, and observe that the value stays near 1, showing that the full-batch gradient aligns with the dominant subspace along the SGD trajectory.

Thank you for your efforts and valuable feedback. Let us address your concerns and questions below. We look forward to discuss further if anything remains unclear.

First, let us highlight one major update in the revised manuscript.

Definition of small LR regime and large LR regime

In the revised manuscript, we replace the terms "small learning rate regime" and "large learning rate regime" with "GF regime" and "EoS regime", which are defined in Definition 3. In the GF regime, (S)GD closely follows the continuous-time gradient flow. This happens when the sharpness is below the maximum stable sharpness (MSS), where MSS is $2/\eta$ for GD and lower than $2/\eta$ for SGD. We can observe the increase of the sharpness (i.e., progressive sharpening) and a stable decrease of the loss in the GF regime. In the EoS regime, (S)GD oscillates away from the gradient flow, where we can observe oscillation of the sharpness and unstable decrease (spikes) of the loss.

Notably, in Figure 6, GD training initially occurs in the GF regime, but after the sharpness reaches $2/\eta$, training shifts to the EoS regime. This shows that a single run of GD with a constant LR can occur in both of the regimes. Note that, the gradient does not align with the dominant subspace in the GF regime, but after reaching the EoS regime, the spurious alignment occurs. Hence, we thought the original terms small/large LR regime could be misleading given that both GF and EoS regimes are observed with the same LR.

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

You are right, the alignment in the "EoS regime" along the GD trajectory arises from the self-stabilization mechanism by Damian et al. (2023). We did cite Damian et al. and explicitly mentioned its relevance in the paper, in Remark 2 of Section 5 (and also in the Related Work section). We emphasize that the primary contribution of Section 5 is not the discovery of this alignment but rather our observations on Dom-GD and Bulk-GD in the EoS regime.

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

This is not true, our finding that Bulk-GD is as effective as GD in the EoS regime is distinct from the finding in Damian et al (2023). Note that while Damian et al.'s analysis explains the behavior of GD, it does not extend to Bulk-GD. Their theory cannot predict or explain the training dynamics when the self-stabilization effect, arising from oscillations along the dominant subspace, is absent.

Specifically, the self-stabilization mechanism by Damian et al. considers constrained trajectory on the "stable set", the set of points where gradient and top eigenvector is orthogonal and sharpness is below $2/\eta$. They prove that GD trajectory in the EoS regime closely tracks the projected gradient descent (PGD) trajectory on the "stable set" while oscillating along the top eigenvector. This PGD trajectory is different from Bulk-GD we consider. Bulk-GD projects each gradient update on the bulk subspace, while PGD projects each gradient update on the stable set. As a result, sharpness grows larger than $2/\eta$ for Bulk-GD, while in contrast, sharpness remains at (or below) $2/\eta$ along PGD trajectory by the definition of a stable set. In Appendix G, we provide a sharpness plot of Bulk-GD in Figure 43. Note that Bulk-GD increases the sharpness larger than $2/\eta$.