Stochastic Collapse: How Gradient Noise Attracts SGD Dynamics Towards Simpler Subnetworks

In this work, we reveal a strong implicit bias of stochastic gradient descent (SGD) that drives overly expressive networks to much simpler subnetworks, thereby dramatically reducing the number of independent parameters, and improving generalization. To reveal this bias, we identify invariant sets, or subsets of parameter space that remain unmodified by SGD. We focus on two classes of invariant sets that correspond to simpler (sparse or low-rank) subnetworks and commonly appear in modern architectures. Our analysis uncovers that SGD exhibits a property of stochastic attractivity towards these simpler invariant sets. We establish a sufficient condition for stochastic attractivity based on a competition between the loss landscape's curvature around the invariant set and the noise introduced by stochastic gradients. Remarkably, we find that an increased level of noise strengthens attractivity, leading to the emergence of attractive invariant sets associated with saddle-points or local maxima of the train loss. We observe empirically the existence of attractive invariant sets in trained deep neural networks, implying that SGD dynamics often collapses to simple subnetworks with either vanishing or redundant neurons. We further demonstrate how this simplifying process of stochastic collapse benefits generalization in a linear teacher-student framework. Finally, through this analysis, we mechanistically explain why early training with large learning rates for extended periods benefits subsequent generalization.

Paper

Similar papers

Peer review

Reviewer DqpF5/10 · confidence 4/52023-07-06

Summary

The paper studies the properties of the stochastic gradient descent using its continuous time SDE approximation. The invariance of the stochastic dynamics in the representation space are investigated and rigorous condition where such invariance sets acts as attractors of SDE are derived. Furthermore, it is empirically validated that neurons collapse towards these invariant sets and the generalization aspects of this phenomenon is studied in two-layer linear networks.

Strengths

a) The paper tackles a very important question of understanding the regularization properties of stochastic algorithms. It formally characteries the two invariant properties of the stochastic algorithms on deep forward networks and established the conditions on the drift and variance terms of the SDE where these sets acts as attractors. In addition, in one-dimensional and two-dimensional settings the conditions are interpreted. b) The above formalism established is novel and is useful to study the empirical behaviour of the deep networks and the role of the learning rate and its influence on generalization.

Weaknesses

a) The paper uses the label noise gradient descent and stochastic gradient descent interchangeably. However, note that there are two different algorithms, hence it should be explicitly stated for which algorithm the results are applicable. Although they share some similar properties, they also differ particularly in the strength of the noise Andriushchenko et al., Ziyin et. al. The paper appears sloppy in this regard. b) Although the paper establishes a more formal treatment the phenomenon of sign invariance and permutation invariance. Andriushchenko et al. provides empirical invariance and theoretical intuitions on same phenomenon.

Questions

a) In the results for evaluation of singular values of linear networks, the set of assumptions used appear strong. Particularly, the assumption that the assumption that noise in labels align with input-output covariance. Can it be verified that the noise of SGD satisfies this assumption ? b) The results in Figure 3 and Theorem 6.1 have been studied when label noise is introduced in the algorithm, how does these changed when using vanilla mini-batch SGD? The conclusion that large step size improve generalization is therefore cannot be generalized to mini-batch SGD.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

The limitations are already addressed in the sections above.

Reviewer nG1u5/10 · confidence 4/52023-07-08

Summary

This paper demonstrate a low-dimensional invariant sets, namely a subset of parameter space, may remain unmodified by SGD. It means that SGD dynamics may lead to simple subnetworks. The theoretical mechanism behind is formally introduced. Moreover, the derived theoretical results revealed that the so-called stochastic collapse may influence generalization in a simplified linear teacher-student setup.

Strengths

- The discovered stochastic collapse is interesting for understanding optimization dynamics of deep learning. Some counterintutive insights are provided to our community. - Formal theoretical analysis with detailed proofs are presented for understanding stochastic collapse. It suggests that stochastic collapse can be precisely analyzed under certain assumptions. - The empirical evidence is direct and informative.

Weaknesses

- I think this paper needs to significantly improve the literature review and discussion. Some closely related studies (e.g. [1,2]) are absent. [1] demonstrated that SGD favor flat minima and SGD dynamics may happen in a low-dimensional parameter space. It will be helpful to explain how the simple subnetworks differs from and relates to the low-dimensional learning subspace in the previous studies. How close are the two low-dimensional subspaces? Will they provides some different insights? - Generalization theoretical analysis is an important part of this paper. While the paper claimed that generalization benefits can be theoretical influenced by stochastic collapse, I did not see really novel or quantitative insights. If the theoretical results are really precise and informative, why not design some novel methods to improve generalization with theoretical supports? For example, if the role of early large learning rates can be well understood in the theory, it is possible to schedule the learning rate in a theory-guided manner that can work well in practice? If not possible, why? Reference: [1] Xie, Z., Sato, I., & Sugiyama, M. (2020, October). A Diffusion Theory For Deep Learning Dynamics: Stochastic Gradient Descent Exponentially Favors Flat Minima. In International Conference on Learning Representations. [2] Gur-Ari, G., Roberts, D. A., & Dyer, E. (2018). Gradient descent happens in a tiny subspace. arXiv preprint arXiv:1812.04754.

Questions

Please see the weaknesses.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

I did see any potential negative societal impact. The main limitation of this work lies its technical contributions, as I mentioned in the weaknesses.

Reviewer kjwn6/10 · confidence 2/52023-07-09

Summary

This paper studies the implicit bias of SGD, and provides a new perspective by characterizing the invariant set in the parameter space along the SGD dynamics, corresponding to specific model architectures. It is shown that SGD noise induces stochastic attractivity towards such invariant sets, which leads to simpler subnetworks. Such simplicity bias is also shown to be beneficial for generalization. Empirical evidence is also provided.

Strengths

This paper provides a novel perspective for characterizing the implicit bias of SGD. The empirical and theoretical analyses in the current paper seem interesting and insightful.

Weaknesses

1. The proof of the stochastic attractivity seems to only apply to very simple networks. For sign invariance, Theorem 5.1 applies to a scalar single neuron model; for permutation invariance, only a two-neuron network is analyzed. 2. Similarly, analysis of the training dynamics only applies to two-layer linear networks. These seem insufficient for justifying what is really happening during actual training of deep neural networks.

Questions

1. The invariant sets for SGD are also invariant sets for GD. Is there evidence that GD doesn't exhibit collapse? 2. Is there any interplay between different types of invariant sets? E.g., which type of invariance does SGD prefer? 3. Is it possible to characterize an invariance type by a corresponding norm of the parameters? If not, why? 4. In general, how to verify the sufficient conditions in Theorem 4.1? Is it possible to verify them empirically for those models used in the experiments in the current paper?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors have adequately addressed the limitations.

Reviewer C15e6/10 · confidence 2/52023-07-12

Summary

The paper explores how stochastic gradient descent (SGD), a common optimization method for deep neural networks, can lead to simpler and more generalizable models. The paper introduces the concept of invariant sets, which are regions of the parameter space that are unaffected by SGD, and shows how SGD can be attracted to these sets under certain conditions. The paper also analyzes how this attraction can result in removing or reducing redundant neurons in the network, and how this can improve the performance on unseen data. The paper provides theoretical and empirical evidence for these phenomena, and explains why training with large learning rates for a long time can help SGD find simpler subnetworks.

Strengths

- The paper is well-motivated, focusing on the study of how gradient noise directs SGD dynamics toward simpler subnetworks. - Several interesting insights are presented, including the relationship between learning rate and stochastic collapse, and the influence of the stochastic collapse on generalization.

Weaknesses

- The paper does not sufficiently clarify the applicability of the attractivity condition to permutation invariant sets in high-dimensional settings. Additionally, it inadequately addresses the extension of the linear teacher-student framework to non-linear models or more general scenarios.

Questions

- Is it possible to (theoretically) analyze stochastic collapse beyond linear models or one-dimensional settings?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

No potential negative societal impact.

Reviewer ZL9R7/10 · confidence 3/52023-07-30

Summary

The authors investigate the implicit bias of the SGD algorithm towards invariant sets -- with a focus on sign and permutation invariance. In particular, they derive a sufficient condition for the stochastic attraction of the SDE describing the first two moments of the SGD process towards these invariant sets, a phenomenon that they call "stochastic collapse". Then, they turn to the study of generalization error dynamics in a one-hidden layer teacher-student linear network, showing that large learning rates increase stochastic collapse, that in turn benefits generalization by implicit low-rank regularization. Finally, they test their theoretical predictions on realistic architectures.

Strengths

The paper presents an interesting perspective on the implicit bias of SGD and its beneficial effect on generalization, by bridging the gap between different approaches -- e.g., the interplay of the Hessian and the effective noise, the tuning of the learning rate, the role of sparsity. The results are solid and well presented.

Weaknesses

- The results are obtained only in continuous time - Assumptions (A1-A4) in Sec. 6 are really restrictive

Questions

- On line 176, the authors comment that stochastic collapse can bias the SGD dynamics towards local maxima of the loss. How can this result in better generalization? A comment in relation to Sec. 6 would be useful.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

Limitations have been adequately addressed.

Reviewer nG1u2023-08-14

Thanks for the response.

Thanks for the response. It helps address some of my concerns. I plan to raise the score at this point.

Reviewer ZL9R2023-08-14

I thank the authors for their careful responses. My concerns have been adequately address and I remain of the opinion that this work is interesting and worth publication. Therefore, I will keep my score of 7.

Reviewer kjwn2023-08-16

Response to the authors

I thank the authors for the detailed response. I don't have further questions, and I'm willing to raise the score.

Authorsrebuttal2023-08-17

Thank you for your consideration in revising the score! Please let us know if any additional questions come up.

Reviewer DqpF2023-08-17

Reply to Author Rebuttal

Thank you for your reply. My score remains unchanged. A few comments on the rebuttal below. Weakness A, Sec 5,6: Yes, I agree that you can set $\zeta$ to mini-batch noise amplitude however then it depends on (t) will not be constant and can be vanishing as the time progresses thus showing that collapse might not be stronger. It would be nice to make this point clear.

Authorsrebuttal2023-08-19

Thank you for your reply and feedback, which has been very helpful in improving our work. The mini-batch noise should not depend on “t” directly, but rather implicitly through a dependency on position. The SDE used in our analysis in section 5 and 6 accounts for an aspect of this positional dependency through the multiplicative structure. That said, you are correct that there can be a further positional dependency unaccounted for in $\zeta$ when considering SGD instead of LNGD. For example, if an interpolating point exists, that is there exists some point $\theta_*$ where $\mathcal{L}(\theta_*) = 0$, then clearly the mini-batch noise amplitude near $\theta_*$ should be much smaller than at a random initialization. One way to account for this would be to insert a dependency on the training loss $\mathcal{L}(\theta)$ into $\zeta$. Unfortunately, this would make the analysis very complex and potentially intractable (we are currently working on if Theorem 5.1 can be extended in this manner). If the training data is not interpolatable, then there is always a positive lower-bound on $\zeta$, and we don’t expect the main lessons from this section would change drastically. As highlighted previously, the experiments for both these sections match well with the theory when run with SGD. This said, we should make it more clear what factors influence $\zeta$ in SGD and why we assume $\zeta$ is positionally independent in order to make the analysis tractable. We hope that the changes we discussed in our rebuttal and the above discussion alleviate your concerns around Weakness A in your review. Have we addressed your other concerns and questions? Is there an aspect of our work that is holding you back from raising your score? Please let us know if you have any more questions. We look forward to hearing from you. And again, thank you for your time and feedback.

Reviewer C15e2023-08-18

Thanks for your helpful responses. The reviewer has no further questions and will keep the positive rating of the paper.

Authorsrebuttal2023-08-19

Thank you for your reply. We are happy to hear you think positively about the paper. Is there an aspect of our work that is holding you back from raising your score? Your comments will help us improve our paper further. Please let us know if you any new questions do come up.

Reviewer C15e2023-08-19

While the paper does contain some restrictive assumptions, it also provides some interesting insights, the reviewer appreciates the value of this work and plans to increase the score from 5 to 6.

Authorsrebuttal2023-08-20

We are happy to hear you appreciate the value of our work and believe that it provides interesting insights. Thank you for revising your score!

Authorsrebuttal2023-08-20

Thank you for considering our response. Please let us know if you have any questions in order to determine the score you plan to raise to before the discussion phase ends on Aug. 21st. Thank you!

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC