Stochastic Collapse: How Gradient Noise Attracts SGD Dynamics Towards Simpler Subnetworks
Authors: Feng Chen
Daniel Kunin
Atsushi Yamamura
Surya Ganguli
Venue: NeurIPS 2023 poster (2023)
Keywords: Implicit Bias
SGD Dynamics
Implicit regularization
Learning rate schedule
Stochastic Gradient Descent
Invariant set
Attractive saddle points
Stochastic collapse
Permutation invariance
Simplicity bias
Teacher-student
TL;DR: This study reveals how stochasticity from mini-batch gradients biases overparameterized neural networks towards "invariant sets" corresponding to simpler subnetworks with improved generalization.

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

Stochastic Collapse: How Gradient Noise Attracts SGD Dynamics Towards Simpler Subnetworks
Feng Chen"
Daniel Kunin"
Atsushi Yamamura (1))
Surya Ganguli
Stanford University
{fenge, kunin, atsushi3, sgaaguuli)@stanforrd
Abstract
In this work, we reveal a strong implicit bias of stochastic gradient descent (SGD) that drives overly expressive networks much simpler subnetworks, thereby dramatically reducing the number of independent parameters, and improving gen- eralization. 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 simpler (sparse or low-rank) subnetworks and commonly appear in modern architectures. Our analysis uncovers that SGD exhibits property of stochastic attractivity towards these simpler invariant sets. We establish suf- ficient condition for stochastic attractivity based on 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 strength- ens attractivity, leading to the emergence of attractive invariant sets associated with suddle-points 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 stochastic collapse benefits generalization in linear teacher-studed framework. Finally, through this analysis, we mechanistically explain why early training with large learning rates for extended periods benefits subsequent generalization.
1 Introduction
The remarkable performance of modern deep learning systems relies on a complex interplay between training dataset, network's architecture, and an optimization strategy. Contrary traditional statistical learning theory, these highly expressive models exhibit impressive generalization capabili- ties natural tasks, even without explicit regularization, despite having the capacity to memorize random data []] This phenomenon is often attributed to implicit biases introduced in the training process that drive the learning dynamics towards models with low-complexity, thereby improving generalization.
widely believed that a central source of this implicit bias the randomness introduced by stochastic gradient descent (SGD) 2 In this work discuss SGD's role in attracting the dynamics, throughout training, towards subsets of parameter space that correspond to simpler subnetworks. We reveal that the architecture of modern deep neural networks as well the nonlinear activation function plays a crucial role in forming these subsets, thus providing a novel perspective on the source of SGD's implicit bias. Our contributions are as follows:
We introduce invariant sets as subsets of parameter space that, once entered, trap SGD. We characterize two such sets that correspond to simpler subnetworks and appear extensively in modern architectures: one for vanishing neurons and the other for identical neurons (Sec. 3)
* Equal contribution. Ordered alphahetically.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

2. We formulate a sufficient condition for stochastic attractiviry process attracting SGD dynamics towards invariant sets that reveals a competition between the loss landscape's curvature around an invariant set and the noise introduced by stochastic gradients (Sec. 4
3. We apply the attractivity condition to determine when neurons with origin-passing activation functions collapse their parameters to zero, effectively removing the neuron. We empirically show that the parameters of permutable neurons within the same hidden-layer can collapse towards invariant sets corresponding to identical neurons (Sec.
4. We demonstrate how this process of stochastic collapse influences generalization in a linear reacher-student framework. We apply these findings to shed light on the empirically observed
importance of maintaining a large learning rate for an extended period during the early stages of training (Sec. 6
2 Related Work
Implicit biases of SGD. Several recent works have explored the properties of SGD and its implicit biases. Barrett and Dherin [8] showed how SGD can be interpreted as adding implicit gradient norm regularization and Geiping et al. [4] showed that training full-batch gradient descent while explicitly adding this regularization can achieve the same performance as SGD. Kunin et al. [5) showed how the anisotropic structure of SGD noise effectively modifies the loss and HaoChen et al. [6] showed that parameter-dependee noise introduces an implicit bias towards local minima with smaller variance, while spherical Gaussian noise does not. Blanc et al. (7) and Damian et al. [8] studied the dynamics of gradient descent with label noise near a manifold of minima and proved that SGD implicitly
regularizes the trace of the Hessian. Li et al. [9] developed framework to study this bias of SGD by considering the projection of the trajectory to the manifold. Kleinberg et al. [10] showed how
SGD is effectively operating local minima. Zhu et al. [1] further demonstrated how the anisotropic structure of SGD noise helps SGD escape efficiently. Xie al. [12] demonstrated that the covariance matrix of SGD approximates the Hessian around local minima, and thus the escape rate from a minima is linked to the Hessian eigenvalues flatness. While numerous studies have sought to identify the origin of SGD's implicit bias, most have predominantly focused on how SGD introduces an implicit regularization term (35), minimizes a measure of curvature among equivalent minima (7-), or escapes from sharp local minima 10 [3.. In our work, we provide a novel perspective and discuss how stochastic gradients introduce a strong attraction towards regions of parameter space associated with simpler subnetworks, even when this attraction is detrimental to the full-batch train loss.
smoother version of the original loss allowing it to escape sharp
Simplicity biases. Several works have explored implicit biases that encourage some notion of sparsity or low-rankness in neural network learning dynamics. Kunin et al. [14] showed how gradient flow training with an exponential loss can lead sparse solutions via an implicit max-margin process. Woodworth et al. [5] demonstrated how an implicit IA penalty occurs for diagonal linear networks trained with gradient flow in the limit of small initialization. Nacson al. [16) and Pesme et al. [17] showed how large step sizes and stochastic gradients further bias these networks towards the sparse regime. Andriushchenko et al. [18] observed that longer training with large learning rates keeps SGD high in the loss landscape where an implicit bias towards sparsity is stronger and theoretically analyzed diagonal linear networks, showing that an associated SDE has implicit bias towards sparser representations. Vivien et al. [19] studied the role of label noise in inducing an implicit Lasso regularization for quadratically parameterized models. Kunin al. 20 and Ziyin et al. [2]] showed how weight decay induce bias towards rank minimization for linear networks. Galanti et al. [22 and Wang and Jacot 23 extended this observation to deep settings trained with SGD. Jacot 24
theoretically, and Huh et al. 25 empirically, showed an implicit bias towards learning low-rank functions with increasing depth. Gur-Ari et al. (26) showed that gradients of SGD converge to a small subspace spanned by few top eigenvectors of the Hessian. Ziyin et al. (27 demonstrated how data augmentation can promote dimensional collapse of representations in self-supervised lcarning. While many works have explored the emergence of sparsity or low-rankness during training, much of the analysis is confined to a particular architecture or consists of general empirical observations lacking clear underlying mechanism. In our work, we focus on fundamental characteristic of neural networks trained with SGD that leads to novel empirical observations.
Our analysis is closely related to recent work studying failure modes of SGD follow up work studying the stability of SGD near a fixed point 22).. In App. we further discuss these works and other relevant works.
28 and concurrent
2

---

3 Invariant Sets of SGD Generated by Reflection Symmetries Throughout this work, we consider a feed-forward network
fa w (ui(m=1).... ...o (wox))x== + b1) 6(m b(m-1)) + B(m)
(I)
parameterized by 0 (w:1),6 ((1) ...... w mm 6(m-1), w(m), b(m)) E R. where and b denote weights and biases respectively, the superscript indexes the m layers, and a-.) is an activation function. The network is trained by stochastic gradient descent (SGD) over a training dataset {(21, 91),..., (In.Ya }n of size n, where 2, E R' and Yi E R The network parameters are randomly initialized at ((0) and iteratively updated according to
@(++1) = -(())
Σ (EB())
Vol (8()) :2.11.
(2)
B
where vi > O is the learning rate, B C [n] is randon index set (mini-batch) of size / at step f, and 8(0; z;, M4) is a loss function. We let [()) A Lie(m) 10;;1;;0 ri, vi) denote the train loss averaged over the entire training dataset. Despite the simplicity of SGD's optimization procedure, understanding how it can navigate through complex, high-dimensionall non-convex loss landscapes to find generalizing solutions is still mystery. Here, we take a dynamical systems perspective and identify subspaces of the parameters that are preserved under the SGD update equation (Eq. (2)
Definition 3.1 (Invariant Set of SGD), A Borel-measuran
set A C Rd is an invariant set of SGD if
given any initialization g(0) E A, all future iterates of SGD g(t) fort 2 are contained within A almost surely, for any batch size 3 E N, learning rate 1) > O, and mini-batcher (B()) C [m] E N).
From this definition, it is immediately clear that all of parameter space = Rd and any interpolating point A = } such that C(0.) = 0 are invariant sets of dimension d and 0 respectively. However, the highly over- sarr-paramtteizeed and layer-wise structure of neural network architectures lead to many additional, non-trivial invariant sets. We will focus on two fundamental invariant sets that appear ubiquitously in neural networks and correspond to simpler (sparse or low- low-rank) subnetworks.
Proposition 3.1 (Sign Invariant Sets). Consider hidden neuron p within layer l of a feed-forward neural network. Let (wop)p
in.p b and Wond denote the parameters directly incoming and outgoing
from the neuron respectively If the non- non-linearity is origin-passing (a(()) = 0), then the axial subspace A {0 E |
11+1) 0, b1) 0, Wout.p = 0} is an invariant set.
The subspace corresponding to a sign invariant set represents the parameter space of a sparse subnetwork obtained by removing a hidden neuron. Most modern neural network architectures employ origin-passing activation functions, which includes linear, hyperbolic tangent, Rectified Linear Unit (ReL.) S0), Leaky ReL.U (31, Exponential Linear Unit (ELU) [3], Swish [33. Sigmoid Linear Unit (SiLU), and Gaussian Error Linear Unit (GELU) (34). Networks with any of these functions will exhibit invariant sets of this nature for each hidden neuron. A fully-connected network of depth y and width k possesses (m - 1)k distinct sign invariant sets.
Proposition 3.2 (Permutation Invariant Sets). Consider two hidden neurons p. q within the same layer of a forward neural network. Let (w bl (w!)) b(") denote the parameters directly incoming to the neurons, and wontpp Wow.g the parameters directly outgoing from the neurons. The (1) (I+1) ((+1)) an invariant set.
affine subspace A {0 E IR ( Wpn.gg bp ba W Wout.g
(1) (1)
The subspace corresponding to a permutation invariant set represents the parameter space of a low- rank subnetwork obtained by constraining two neurons be identical. All feed-forward networks possess this permutation invariance among the neurons in each hidden-layer. For fully-connected network of depth yn and width l there are (m - ()) distinct permutation invariant sets.
Invariant sets generated by symmetry. The presence of sign and permutation invariant sets within neural networks is a direct result of reflection symmetries inherent in their architectural design. In this
"This notation encompasses fully-connected and convolutional networks, excluding more complex architec- tures such as transformers for simplicity, although much of our theory would directly apply.
FFrr simplicity assume sampling with replacement such that there is no dependency between batches. 'The incoming and outgoing weights are related by: (we(i) ) (w--1))
3

---

context, symmettry' is defined as any transformation of the parameter that preserves the network's s input-output mapping. regardless of the input. As stated in the following theorem, any approximate linear symmetry defined by symmetric or orthogonal matrix generates an invariant set:
Theorem 3.1 (Symmetry Induced Invariant Sets). Let Q Rdxd be a symmetric (Q = QT) or orthogonal matrix (QQT= = Ia). The affine subspace A = {0 E R= = O is an invariant set if the loss function E(0; x.s ye), for any i E [n], is approximately Q- 2-symmetric around A, i.e., for any
e > 0, there exists 6 > 0 such that [e{Q0; 24, yi) - f(0; zt. vs)l ed(0 A) for any 0 E Rd satisfying d(0, A) < 8 where d(0, A) is the Euclidian distance between 0 and A
Remarkably, the sign and permutation invariant sets defined earlier can be understood through this symmetry perspective. In both cases, the symmetry defined by a symmetric and orthogonal matrix. This class of matrices describes the generalized notion of reflection in d-dimensional space and the invariant sets are the axes of reflection, which can be at most (d - -)dimennsiona.. In the case of permutation invariant sets, any two hidden neurons in the same layer of a network can be permuted without changing the input-output function of the network, and thus the affine subspace A where these two neurons are identical is invariant set. On the other hand, in the case of sign invariant sets, the symmetric transformation Q corresponds to a simultaneous sign flip of a neuron's 's input/output weights. This transformation does not change the input-output function of the network, as long as the activation function is odd. However, even for origin-passing activation functions that are not exactly odd, the loss function is approximately Q-symmetric around A, since the activation function is approximately odd around its origin. Therefore, Theorem [3.T iimplies that the axial subspace A, where input and output weights are exactly zero, an invariant set.
Additional invariant sets. In App. B we discuss other classes of invariant sets, associated with softmax nonlinearities, low dimensional structure in data, and provide further discussion the connection between symmetry and invariant sets. We also discuss how our definition of invariant sets relates to other works 38) where symmetry and invariance used to explore geometric properties of the loss landscape and the presence of critical points.
Gradient Noise Attracts SGD Dynamics towards Invariant Sets
To study the dynamics of SGD we will approximate its trajectory with a stochastic differential equation (SDE), a common analysis technique applied in many works 5 43] SGD can be interpreted as full-batch gradient descent plus a per-step gradient noise term introduced by the random minibatch. Let (f) represent the per-sample gradient noise at O. It is easy check that (0) 0 for all 0. This motivates studying the SDE, which we refer to as stochastic gradient flow (SGF),
d0,- = - (L + TE
T V
8g = =()
Here B, denotes a standard d didimensional Wiener process and (())) = / (O), The drift term of this SDE is the negative full-batch gradient V L(,.) while the diffusion is determined by a spatially-dependen diffusion matrix ((0,) = Z (0.) (),)) Although this SDE cannot be interpreted as the continuous time limit of SGD, it matches the first and second moments of the SGD process. See Appendix ( for further discussions on the relationship between SGD and SGF. Additionally, SGF preserves all affine invariant sets of SGD:
23
Proposition 4.1 (Informal). All affine invariant sets of SGD are also affine invariant sets of SGF.
Like deterministic dynamical processes, the stochastic processes of SGF can be attracted toward the invariant sets. To describe this attraction, we adopt the concept from stochastic control theory [44). Definition 4.1 (Stochastic Attractivity). An invariant
A Rd of stochastic process {0, Rd t 2 0} is stochastically attractive if for any > O and e. >> O, there exists 8 > 0 such that for any 0 E I2d with the Euclidian distance d(0, A) < 5,
R
P sup d(O (0,, A) 2 c 0 = 0 Sp
(4)
120
'The Euclidian distance between 0 and A is defined as d(0, A) = minge |0 -0-al all.
"ss discussed in Li et al. [39] Eq. F is an order 1 weak approximation of Eq. 2
7 "eee App. C for the definition of invariant sets for continuous processes and the formal statement.
'This property is called stochastic stability in 44
4

---

Intuitively, a set is stochastically attractive if there always exists close initial condition that guaran- tees all future iterates remain close to the set with high probability. Understanding when the dynamies are attracted by particular invariant set becomes fundamental for understanding the implicit bias of SGD. To gain insight into this process lets consider a canonical example one-dimension.
Geometric Brownian motion. Geometric Brownian Motion (GBM) is a linear stochastic process given by the SDE, d0,- = p0,dt + (0,d where μ E IR represents the drift rate, < R the volatility rate. Denote 0g E IR as the initialization. The set A = {0} is an invariant set of GBM, which, depending on the relationship between μ and Ç, can be stochastically attractive. As the trajectory of GBM approaches the invariant set, its movement diminishes. To account for this effect, we take the logarithmic transformation 2
(6+), where we assume without loss of generality that fn > 0. The stochastic process 2n obeys the transformed SDE, dz,- (n Ç² /2) dt + [dB,, which has an additional drift term -C2/2 /2 from Itô's lemma that attracts the process towards the invariant set. The total drift is now determined by competition between the original drift rate 1 and the additional attractive force driven by the diffusion. This competition controls whether the invariant set A {0) of GBM is stochastically attractive (i.e. when μ < 5² / the invariant set is attractive).
Stochastic attraction in one-dimension. Stochastic attractivity is local property and thus might expect SGF to act like GBM around an invariant set, enabling us to derive broader conclusions about stochastic attraction. If we consider SGF with a smooth loss and diffusion in one-dimension, where without loss of generality we assume A = {0} is an invariant set, we can Taylor expand the gradient and the diffusion around the set, yielding geometric Brownian motion near the set:
dig- xc (0) dD),,,,
(5)
Here we use
= (0) = 0, which is a consequence of {0} being an invariant set. From Eq. 5 we can formulate a necessary/sufficiet condition for stochastic attractivity in one-dimension:
Theorem 4.1 (Necessary/Sufficient Condition for Stochastic Attraction in One- One-Dimension). Let
A = {0} E E be an invariant set of {0, R 2 0} obeying Eq. 3 Suppose (') R + R. ((0)
R + R are C2 functions such that (0) > O. Define rate of attractivity a = C²(()) 1 DD). A
is stochastically attractive if > 0, while it is not stochastically attractive if a < 0. 9
One of the most surprising implications of Theorem 4.1 is that SGF can potentially converge to a saddle-poin moitle-piinteee or local maxima of loss landscape, an observation previously made by Ziyin et al. (28). To see this, recall that D(#)= 2 by definition, and (0)) = 0 as {0} is an invariant set. As a result, (0 0 and (0) 2 0 by the continuity assumption of D. The collapsing condition, therefore crucially depends on the curvature of the loss function L at 6 = 0. When D²(() is strictly positive, the collapsing condition can still be satisfied with negative curvature provided that L" (0) > (0)) Given that D is proportional to n/ 3, the learning rate to batch size ratio determines the maximum attainable steepness for an invariant set to be attractive.
An illustrative example. Consider SGF in a deulle-well potential L(0) = 16 nم²² with multiplicative diffusion 10 such that the dynamics are d0, - (0 - 1ute )dt + <0,dB, where μr, C > 0. The minima of the potential are located 0 + -VT,, while 0 0 is local maximum and forms an invariant set. This example is special as the dynamics have an analytical steady-state distribution Pas (0), given any initialization. Thus, determining the stochastic attractivity of {0} can be achieved by examining the steady-state distribution. As in GBM, we assume, without loss of generality, that 6g > and consider the logarithm process 2, log(6.). In this new coordinate, the noise term is constant, which allows us to determine the steady-state distribution. Transforming back to the original coordinate, we find that Pss (0) is given by a Gibbs distribution Pas xx 6 (x) for modified potentia [I] (0) 8 0 (2n/c² 2) log(f) with constant = 2/5² and partition (0) d0. When μ S <²/2, the partition function diverges and pas (0) collapses
function Z - xa
to Dirac delta distribution Theorem 4.1
0. This transition agrees with the collapsing condition from
Whee a = O. higher-order derivatives of the loss and diffuusion determine the attractivity of A.
Heer build a substantial literature, originating from the seminal work by Kramers (45). that investigates the escape rate between minima a particle in double-well potential with constant diffusion (46).
'TThe interpretation of modified potential determined by the gradient noise that drives the learning dynamics was explored in 2 and [5]. Stochastic collapse is when the partition function for this modified loss diverges.
5

---

1.00
Theaiogd
-0
+/7
0.75
0
0.50
30000
0.25
20006
42000
0.50
164 0.00 0.50 1.00
324
Time
Figure 1: Stochastic collapse in quartic loss. The left three plots show sample trajectories with the same, random initializations driven by the SDE d0, = - (0² n0² )dt + Ç0,dB, for different values of Ç. The analytic stationary solution is plotted on the side plot. Leftmost: for small ( the steady-state distribution is two bumps around the minima 0 = -VØ. Middle left: with increasing < the distribution spreads out and is biased towards collapsing condition 22 the steady-state distribution collapses to dirac delta distribution at = 0. Rightmost: the empirical probability of the sample trajectories at different steps being within c of the origin as a function of increasing gradient noise. As we see empirically. there is sudden phase transition which aligns with the collapsing condition from Theorem
O. Middle right: when C surpasses the
Stochastic attraction in high-dimensions. To extend the collapsing condition derived in Theorem 4.1 nigh-dimensional cases, natural idea would be consider all one-dimensional slices of parameter space orthogonal to the invariant set. However, one challenge is that some of these slices might satisfy the collapsing condition while others do not. This result in complex dynamics near the boundary of the invariant set making it difficult to derive a necessary and sufficient condition for attractivity in high-dimensions. Nonetheless, we can derive sufficient condition:
Theorem 4.2 (A Sufficient Condition for Stochastic Attraction in High- high-Dimensioos). Let A Rd be a d da-dimensional affine subset, and a stochastic process {0g E Rd :t 2 0} obey Eq. 3 in Ag. open c-neighborhood of A with some >> O. Suppose L A. RR is C Co-function whose first and second-order derivatives are L-Lipschitz continuous. D A. - Raxd is the diffusion matrix such that the second-order derivatives of its elements are L-Lipschitz continuous. Furthermore, we assume that all the elements of vD Ac + Raxd are T--Lipschizz continous. Let Dı PDDP1 where P1: : R.d > Rd projects to the normal space of A. If there exists 5 > O such that V2 > 6 and
v2 4- - + (1 51--5) 0. 2
for any unit normal vector á Rd perpendicular to and 0 A, then A is stochastically attractive.
To prove Theorem 4.2 (see App. Ep. we consider a family of distance measures between O and the invariant set A. Eq. 6 eensures that one of the distance measures super-martingale in neighborhood
of the invariant set, allowing us to apply version of Doob's 's maximal inequality (Lemma in [44).
5 Attractivity of Invariant Sets in Deep Neural Networks
In this section we explore theoreti- cally and empirically the stochastic attractivity of invariant sets in deep neural networks.
Loss C(W), W2)
1.00
Output weight
1
0.75
Staps
400
0.50
850
1600
0.25
3200
6400
Sign invariant sets. Sign invari- ant sets correspond to the parame- ter space of a sparse subnetwork ob- tained by removing a hidden neu- ron. Here, we demonstrate the stochastic collapse to a sign invari- ant set in a minimal model of neu- ral networks scalar single neuron
-2 -1 Input weight W1
0.0
0.5
1.0
15
2.0
noise level n5²
o'()/2
Figure 2: Stochastic collapse of single neuron. The SDE dynamics of single-neuron model are simulated with var- ious noise levels Left: A scatter map of trained weights (wn, w2). The colors of the dots represent the noise level, corresponding to the vertical lines in the right panel. For each noise level, we plot trained models with different noise re- alizations. The background heatmap shows the loss landscape. Right: The colored lines represent the empirical probability that / + w < c = 10 (from 10 samples) after a given number of update steps. We observe a sudden transition that aligns with the collapsing condition from Theorem 5.1
model: [2; W1, W2) = W20 W22 (w) with w1. 32,
R. Suppose this model is trained via label-noise
gradient descent and using the L²_
loss [(w), W2) = 1 Eie
f W1, 2² on non-empty data

---

Layer with little stochastic collapsing My: conv5 1256)
Layers with strong stochastic collapsing
Weuti convél (256)
We: Conv32 (512)
Weut conv12 (512)
We: Fc14 (532)
Maxi 1c15 (512)
CIEAR-100 ResNet-18 CIFAR-18 vCcc-66
2.5
2.0
10
lyerr conv2 1128 layer2. convi (126) syerl conel (64) layer 1 come (64) leyerd 3 [cov]]
w leyer com22 (512)
2.0
15
0.5
00
Figure 3: Evidence of stochastic collapse towards permutation invariant sets in deep neural networks. Each pair of plots shows the normalized distance matrix between incoming parameters (Win) and outgoing parameters (Wom) of neurons within the same hidden layer. We show two pairs of plots with little stochastic collapsing and four pairs of plots with strong stochastic collapsing: three pairs for three different layers in VGG-16 trained for 10 steps on CIFAR-10 (top row), and three pairs for three different layers ResNet-18 trained for 10 steps on CIFAR- 100 (bottom row). For each pair of plots, we sort the neurons by hierarchical clustering according to normalized distances between incoming parameters only (left plot in each pair), and then show the distance matrix between outgoing parameters using the same sorting of neurons (right plot in each pair). The similarity between each plot in a pair indicates that clusters of neurons with similar incoming parameters also have similar outgoing parameters. See App. for experimental details.
set {(11, y1), (In, bn)} such that Exen x? + Exploiting Theorem 4.1 we can determine the attractivity of the sign invariant set A = {00 0)} for this network.
Theorem 5.1 (Informal). Let {(WLe, w2,1) e RR 2 0} be a process obeying label-noise gradient flow with loss of a single scalar neuron in a neighborhood U of A {(0, 0)}, with learning rate 7) 0 and noise amplitude Ç <> O. Suppose the activation function R - IR is smooth satisfying
(0) 0, o o'(0)= >
The invariant
stochastically attractive if o' (0)mc.
Eveto #7
Eielm) s.96 Eiejn =?
The term stochastic attractivity is determined by balancing the signal and noise idea we will later in Sec. as the origin of generalization benefits from stochastic collapse. We leave the poof in App.
can be thought of the signal of the dataset. Thus, Theorem 5.1 states that
Permutation invariant sets. Permutation invariant sets correspond to subnetworks with fewer unique hidden neurons. An important question to what extent are permutation invariant sets attractive as this would indicate an implicit bias towards low-rank model. We provide intuitive insights into the attractivity of permutation invariant sets through a toy example of a two-neuron neural network in App. G More importantly, we empirically test whether permutation invariant sets be attractive in deep neural networks via experiments using VGG-16 [7 and ResNet-18 [48) trained on CIFAR-10 and CIFAR- 100 respectively 49 Since the ReLU activation has additional symmetry resulting in a potentially larger invariant set [35], we replace ReLU with GELU activation in all our experiments below. To detect stochastic collapse to permutation invariant set, compute the normalized pairwise distance between parameters for neurons from the same layer, defined by 21) [1) -
Strikingly, hierarchical clustering based this distance reveals dist(w:") 10 I| 13+
multiple large clusters of many neurons with near identical incoming and outgoing parameters after training (Fig. 3). Such clustering implies SGD implicitly drives weight matrices towards low-rank structures in many layers via attraction towards permutation invariant sets.
To explore the influence of the learning rate and batch size on the attraction towards sign and permutation invariant sets, we trained VGG-16 and ResNer-18 on CIFAR-10 and CIFAR-100, respectively, while varying these hyperparameters. We defined vanishing neurons as those with incoming and outgoing weights under 10% of the maximum norm for that layer, and identified the number of independent con-vanishing neurons by clustering the concatenated vector of incoming

---

Indopindent Neurons
CIAAR100 10 VGG-16
CIFR-1000 Restiet-1 18
1.0
Fraction of
0.8
Learning Rate
0.8
0.6
0.2
0.2
0.0
0:10
convi
conv2
conv3
cond
comed
convi
come7
comul
condd
conv10
convi1
12
R:13
1:14
R:15
convi
a.cond
comve
Leonve
a.conll
a.coma
Laad
Leonz
Geanv
Lennv
L.comng
a.ond
Loonz
Layer
Layer?
Lays
Layer
Larger learning rates intensiff stochastic collapse. This figure illustrates how the fraction of independent neurons per layer in VGG-16 trained on CIFAR-I0 (left column) and ResNet- 18 trained on CIFAR- 100 (right column) varies with changes in learning rates, The networks are evaluated at training steps of 10 A reduced percentage of independent neurons indicates stronger stochastic collapse. See App. for further details.
Figure 4:
and outgoing parameters based on normalized distance. Two neurons were considered identical if their distance in parameter space was 10% of their norms, stringent criterion given the high- dimensionality of the weight vectors.
We found that increased learning rates typically intensify the stochastic attraction to invariant sets by reducing the number of independent neurons (Fig.4). A large reduction in the fraction of independent neurons was observed in VGG-16 between conv? and conv8, where the number of channels increase from 256 to 512, indicating this excess model capacity is counteracted by stochastic collapse. Also, we note that Proposition 3.2 and 3.1 do not apply to neurons with residual connections. More strict definitions are required, as described Proposition the source of syyertt-llyeer oscillations in the fraction of independent neurons for ResNet18 in Fig. 4 Surprisingly, we noticed that unlike learning rate effects, changes in batch size produce complex shifts in the number of independent neurons. See Fig. 10 for further discussion.
[..2 [4.2 H.T] We believe, this stricter constraint
How Stochastic Collapse Can Improve Generalization
To understand how stochastic collapse benefit generalization, will theoretically analyze learning dynamics in a teacher-student model for linear networks, a well established framework for studying generalization (50 This analysis not only provably connects the phenomenon of stochastic collapse to generalization benefits, but also suggests stochastic collapse based mechanism for why successful learning rate schedules improve generalization. We then provide direct evidence for this theoretically suggested mechanism in deep neural networks.
Insights from two-layer linear neural network. We build the analysis of full-batch gradient learning dynamics of training error 53 and test error (51) in the limit of infinitesimal learning rates for two-layer linear neural networks in student-teacher setting. In new analysis here we incorporate the important new ingredient of stochastic gradients, which dramatically alters the learning dynamics. The details of our analysis are presented fully in App.
teacher neural network with low-rank weight matrix W generates a training dataset by providing noise corrupted outputs to a set of random Gaussian inputs. The input-output data covariance matrix Eye then drives the learning dynamics of a two-layer student with composite weight matrix W- = W, We analyze the learning dynamics of the student under set of four assumptions precisely stated in App. [ These assumptions are similar to the ones made in [51 53 and are roughly stated as: (A1) whitened inputs to the teacher, (A2) structured gradient noise related to input-output covariance; (A3) spectral initialization with student singular vectors aligned to input- output covariance singular vectors; (A4) balanced initialization with equal power in the two weight layers of the student. Under these assumptions, the learning dynamics of the student weight matrix W = W2 can be decomposed into independent learning dynamics for each singular value 8; of W. An associated singular value 3, of Eye drives the dynamics of & via the SDE,
(7) where Ç >> 0 is the amplitude of the gradient noise. In the small noise limit of Ç - 0, this SDE aligns with the nonlinear ODE corresponding to the student network's 's dynamics under gradient flow, where exact solutions exist [51 53 Important insights gained from this previous analysis are that larger data singular modes are learned earlier [5X] and larger data singular modes, when learned, help the student generalize, while smaller singular modes, when learned, impede generalization by
då; = 28, ((5, + 22 8.) dt + 2 2.,dB,
8

---

X10
X10
Theary (-0
x10
2.0
1.2
(-5
1.2
Train Loss
5.0
0.9
Steps entrr anp
4.8
Test
1.0
0.6
0.3
0.3
4.6
0.5
0.0
0.0
0
K1d 0.0
0.4
0.8
1.2
0.0
0.4
0.8
1.2
Steps x10'
Steps
Figure 5: Demonstrating generalization benefits of stochastic collapse in teacher-student setting. We show the train loss (leftmost) and test loss (middle left) during the training of the student with different gradient noises. Dashed lines in the leftmost and middle left panels indicate the step where learning rate is dropped. Training with larger gradient noises ( generalizes better. Middle right: we show the noisy teacher signals against the learned student signals before dropping the learning rate. Larger noises have a stronger implicit bias towards zero. Rightmost: same as the middle right panel but we show the learned signals (for the brightest green in the middle right panel) at different training steps after the learning rate drop. All the results are averaged over 256 replicates.
overfitting to noise in the training data [5T] Both these insights will apply in our analysis of SGF. However, we will also uncover third behavior unique to SGF: stochastic collapse effectively governs the minimum learnable singular mode, serving as a natural mechanism to prevent overfitting. Our analysis in App. D reveals that the student singular values 8; evolve according the process:
Theorem 6.1 (Stochastic Student Dynamics). Under assumptions Al A4, the dynamics of (c)
given by Eq. 7 are governed by the stochastic process,
(1) = e28: ( 2222
(2
2
(28: ( ( dr (;(0))
(8)
In terms of the per layer balanced student weight WA = V the above process can be thought of as originating from the quartic loss landscape (,(w,) = } (w? 52) equivalent to the example presented in Sec. 4 Thus, there exists an invariant set for each student singular value corresponding to the affine space 8; wy = O where both drift and diffusion terms vanish. Stochastic collapse to this invariant set would thus induce low-rank regularization on the student in which some number of singular modes in the data are never learned despite having nonzero data singular values $,. To confirm stochastic collapse to this invariant set, we derive:
Corollary 6.1. In the limit (0) + O, for any > O, then (t)))(((() + c(28, nc (222)+227227
Corollary 6.1 has important implications. First, if < " the distribution will exhibit stochastic collapse, converging delta distribution the origin, consistent with Theorem 4.1 (Fig. S middle right panel). This in turn implies that early training with large learning rates promotes stochastic collapse of more student singular modes. Indeed, even after the training loss plateaus, further extended training with a large learning rate 7) will continue to drive stochastic collapse of small singular modes. After subsequent learning rate drop, some of these modes will cease to obey the collapsing condition and as consequence, the associated invariant set will repuise rather than attract the dynamics. Nonetheless, the repulsive escape will take longer the closer it starts from the invariant set (Fig. 5 rightmost panel). This results in dynamics in which generalization error is reduced by the student not learning many data singular modes with large y). because 8; < " and then learning only the subset with large ă that have time to escape the vicinity of the invariant set after a learning rate drop. The overall effect is that the student selectively learns modes with higher signal-to- -onoie ratios, resulting in improved generalization after the learning rate drop as shown Fig. () middle left panel. In summary, the linear teacher-studed schedule, noisy student can effectively avoid overfitting through stochastic collapse and achieve superior generalization compared to a noiseless student.
model highlights that with an appropriate learning rate
Stochastic collapse explains why early large learning rates helps generalization. The analyses in the linear teacher student network provide valuable insights into how stochastic collapsing can enhance generalization. Importantly. it sheds light on the mystery of why training at large learning rates for long period of time (even after training and test loss have plateaued) helps generalization after subsequent learning rate drop. The key prediction is that large learning rate induces stronger

---

Freeltin of Inckppr Nerons traction b-dkpendkmm dendddmettt Noon
CIFAR-10 VGG-16
93-
Taining Lood
1.0
Stepa
Teat Accurily
Test Accurcyy
92
0.8
10200
0.7
20303
40303
0..4
60
8
50
Steps after Leamnngg
CIEAR- 100 RosNet-18
1.8
65
76
1.0
Training
1.4
Staps
Looge
Accurary
Test Accurecy
0.8-
1.0-
20000
40000
8
0.6
0.2
10000
20000
30000
40000
5000
10000
Steps
Staps after Learning Rate Drop
Figure 6: Large learning rates aid generalization via stochastic collapse to simpler subnetworks. Left: The training loss (grey) and test accuracy (blue) during an initial training phase using high learning rate (Ir=0.1). Vertical dashed lines highlight the training steps at which we will drop the learning rate. Middle: Test accuracy preceding and following a learning rate drop to 1r-0.01. Colors indicate when the learning rate drop shown in the left plot occured. Despite plateauing of the training loss and test accuracy, a later learning rate drop yields higher final test accuracy. Right: The fraction of independent neurons per layer evaluated different learning rate drop times (indicated by the color) during the initial high learning rate training phase. Prolonged training with a large learning rate induces implicit regularization by reducing the number of independent neurons. All curves represent the average across eight replicates. See App. ffor experimental details.
stochastic collapse, thereby regularizing the model complexity. Furthermore, remaining in phase of larger learning rates for a prolonged period drives SGD closer to the invariant set. Consequently, when the learning rate is eventually dropped, overfitting in these specific directions is mitigated.
To test this predicted mechanism, we trained VGG-16 and 1e on CIFAR-10 and CIFAR-100, respectively. The training loss and test accuracy had already plateaued at 10 steps (Fig. 6 left column). We dropped the learning rate after different lengths of the initial high learning rate training phase and confirmed that training with large learning rates for longer periods helps subsequent generalization (Fig. 6 middle column). We then tested our prediction that stochastic collapse is occurring during the early high learning rate training, thereby enhancing implicit bias towards simpler subnetworks. In particular, we computed the fraction of independent neurons at each step where we drop the learning rate. We found, as predicted, that models with a longer initial training phase collapse towards invariant sets with fewer independent neurons (Fig. right column).
Conclusion
In this study, we demonstrated how stochasticity in SGD biases overly expressive neural networks towards invariant sets that correspond to simpler subnetworks, resulting in improved generalization. We established sufficient condition for stochastic attractivity in terms of the Hessian of the loss landscape and the noise introduced by stochastic gradients. We combined theory and empirics to study the invariant sets corresponding to vanishing neurons with origin-passing activation functions and identical neurons generated by permutation symmetry. Furthermore, we elucidated how this process of stochastic collapse can be beneficial for generalization using a linear teacher-student framework and cxplain the importance of large learning rates carly in the training process.
Limitations and future directions. Activation function with additional continuous symmetries, such as RELU, might have larger class of invariant sets than those studied in our work, which remains to be explored. As all the invariant sets in our study are affine, exploring the possibility of curved invariant sets and how curvature affects the analysis is an interesting direction for future work. Furthermore, the interplay between symmetries in data and invariant sets warrants investigation. Identifying the necessary conditions for stochastic attractivity general affine invariant sets is an important goal for future work. Extending our analytic results from the continuous SGF to the discrete SGD updates an interesting theoretical direction. Lastly, designing new optimization algorithms based on our insights into stochastic collapse is major goal for our future work.
10

---

Acknowledgments and Disclosure of Funding
We thank Nan Cheng, Shaul Druckmann, Mason Kamb, Itamar Daniel Landau, Chao Ma, Nina Miolane, Mansheej Paul, Allan Raventós, Ben Sorscher, and Daniel Wennberg for helpful discussions. D.K. thanks the Open Philanthropy AI Fel for support. S.G. thanks the James S. McDonnell and Simons Foundations, NTT Research, and NSF CAREER Award for support.
Author Contributions
D.K. and F.C. initiated the project. F.C. formulated an initial analysis on stochastic collapse and is primarily responsible for the experiments associated with deep neural networks. D.K. is primarily responsible for the linear teacher-student analysis and the original observation that SGD can collapse saddle-point introduced by symmetry. A. primarily contributed to the theoretical formulations and proofs, which include the concept of invariant sets and conditions for stochastic attractivity. S.G. advised throughout the work and provided funding for computation. All of the authors have worked on writing the manuscript.
References
[1] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communn ACM, 64 feb 2021. ISSN 0001-0782. doi: 10.1145/3446776 URL httptttt///.o...o
10. 1145//4467/77
[2] Fengxiang He, Tongliang Liu, and Dacheng Tao. Control batch size and learning rate to gener- alize well: Theoretical and empirical evidence. In H. Wallach, H. Larochelle, A. Beygelzimer, d'Alché-Bus Buch E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Sys- tems, volume 32. Curran Associates, Inc., 2019. URL httpsttttt////poceeeiiigg neurips. cc/ pdf
files/
[3] David Barrett and Benoit Dherin. Implicit gradient regularization. In International Con- ference Learning Representations, 2021. URL http::/://opemreeie.eeeee///eeee 3q5IqUrkcF]
[4] Jonas Geiping, Micah Goldblum, Phil Pope. Michael Moeller, and Tom Goldstein. Stochastic training is not necessary for generalization. In International Conference on Learning Represen- lations, 2022. URL htt./htttp::/oppreevieee http:/aaerrrrreettefoooruu/////
met:foormmidd2an00 -idiZEie--BBeeeee IUB5k
[5] Daniel Kunin, Javier Sagastuy-Brent Lauren Gillespie, Eshed Margalit, Hidenori Tanaka, Surya Ganguli, and Daniel L.K Yamins. Limiting dynamics of sgd: Modified loss, phase space oscillations, and anomalous diffusion. arXiv preprint arXiv:2107.09133 2021.
[6] Jeff Z. HaoChen, Colin Wei, Jason Lee, and Tengyu Ma. Shape matters: Understanding the implicit bias of the noise covariance. In Mikhail Belkin and Samory Kpotufe, editors, Proceedings of Thirly Fourth Conference on Learning Theory, volume 134 of Proceedings of Machine Learning Research, pages htmll //proceedings. mlr 3rress/v134/aacchen22
PMLR, 15-19 Aug 2021. URL https
[7] Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an cmstein-uhlenbee like process. In Jacob Abernethy and Shivani Agarwal, editors, Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pages 483-513. PMLR, 09-12 Jul 2020. URL https://prrceedding mlr prwss/ www11////lan20000 httml20a..hmml html
[8] Alex Damian, Tengyu Ma, and Jason D Lee. Label noise sgd provably prefers flat global minimizers. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 27449-27461. Curran Associates, Inc., 2021. URI https: http::::/ neurips. cc/paper baaper_fiees files/

---

[9] Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss? -a mathematical framework. In International Conference on Learning Representations, 2022.
URI.https: httpss://ppenreeieews http://ppmaaeeeeeeeeeeeooonm/maa
[10] Bobby Kleinberg, Yuanzhi Li, and Yang Yuan. An alternative view: When does SGD escape local minima? In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 2698-2707. PMIR, 10-15 Jul 2018. WtR.httts://pprceeddiii htt.eeeeeennn.mm...ee.e.e/s//s// press/v80 kleinberg1@a html
mlr
[11] Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 7654-7663. PMLR, 09-15 Jun 2019. URL https: pttttp:://proceeddngge mlr press/v97/ zhu19e html
[12] Zeke Xie, Issei Sato, and Masashi Sugiyama. diffusion theory for deep learning dynamies: Stochastic gradient descent exponentially favors flat minima. arXiv preprint MxXi:20022034493 2020.
[13] 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. arXiv preprint arXiv: /60900486, 2016.
[14] Daniel Kunin, Atsushi Yamamura, Chao Ma, and Surya Ganguli. The asymmetric maxi- mum margin bias of quasi-homogeneous neural networks. In The Eleventh International Conference on Learning Representations, 2023. URL BRLhttpps: kd=1M4xp7kGI5V
net/forum?
[15] Blake Woodworth, Suriya Gunasekar, Jason D. Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparamettrized models. In Jacob Abernethy and Shivani Agarwal, editors, Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pages 3635 3673. PMLR, 09-12 Jul 2020. URL https::/ htt.:i/.........ss/eeeee
mlr prwss/ v/// https//v1////ooo// hot5///ooooott
html
[16] Mor Shpigel Nacson, Kavya Ravichandran, Nathan Srebro, and Daniel Soudry. Implicit bias of the step size in linear diagonal neural networks. In Kamalika Chaudhuri, Ste- fanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceed- ings of the 39th International Conference Machine Learning, volume 162 of Proceed- ings of Machine Learning Research, pages 16270-16295 PMLR, 17-23 Jul 2022. URL htttps:/pprccediiii
mlr
[17] Scott Pesme, Loucas Pillaud- Vivien, and Nicolas Flammarion. Implicit bias of sgd for di- agonal linear networks: a provable benefit of stochasticity. In M. Ranzato, A. Beygelz- imer, Y. Dauphin. P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neu- ral Information Processing Systems, volume 34, pages 29218-29230. Curran Associates, Inc., 2021. URL /ttp wtw.:/occeediinnnnnniiepeeeeeeee cc/ http/pppr__ffile www.fflllssppppee paalle//ppper/22///1i1. f 146119880saaaaaaffeese0000111111 Paper. .pff
[18]
Maksym Andriushchenko, Aditya Varre, Loucas Pillaud-Vivies and Nicolas Flammarion. Sgd with large step sizes learns sparse features. arXiv preprint arXiv:22 10.05337, 2022.
[19] Loucas Pillaud Vivien, Julien Reygner, and Nicolas Flammarion. Label noise (stochastic) gradient descent implicitly solves the lasso for quadratic parametrisation. In Po-Ling Loh and Maxim Raginsky, editors, Proceedings of Thirty Fifth Conference Learning Theory, volume 178 of Proceedings of Machine Learning Research, pages 2127-2159. PMLR, 02-05 Jul 2022. URL https://prooceeding. mlr.
[20] Daniel Kunin, Jonathan Bloom, Aleksandrina Goeva, and Cotton Seed. Loss landscapes of regularized linear autoencoders. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors,
12

---

Proceedings of the 36th International Conference on Machine Learning, volume 97 of Pro- ceedings of Machine Learning Research, pages 3560-3560 3569. PMLR, 09-15 Jun 2019. URL httpp:://roocsedings mlr.. wtt..press///////unnn//////
html
[21] Liu Ziyin, Botao Li, and Xiangming Meng. Exact solutions of deep linear network. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URLhtttst httpp:://pperreveee u X6bp8r18dV
net/ferummie
[22] Tomer Galanti, Zachary S Siegel, Apama Gupte, and Tomaso Poggio. Sgd and weight decay provably induce low-rank bias in neural networks. arXiv preprint arxiv.org/abs/2206.00799// 2022.
[23] Zihan Wang and Arthur Jacot. Implicit bias of sgd lg-reuularized linear dnns: One-way jumps from high to low rank. arXiv preprint arXiv:2305. 16038, 2023.
[24] Arthur Jacot. Implicit bias of large depth networks: a notion of rank for nonlinear functions. arXiv preprint arXiv:2209. 15055, 2022.
[25] Minyoung Huh, Hossein Mobahi, Richard Zhang, Brian Cheung, Pulkit Agrawal, and Phillip Isola. The low-rank simplicity bias in deep networks. arXiv preprint MxXiv:2103.1022: 2021.
[26] Guy Gur-Ari, Daniel A. Roberts, and Ethan Dyer. Gradient descent happens in a tiny subspace, 2018. URL TRRLhttps:// http::tttt///iiiiee http://axxvvvvve org/ www./////////////ee//o.o//eeeee/ abaaaaaaaaaaa 11822. 181..0472411.11111
[27] Liu Ziyin, Ekdeep Singh Lubana, Masahito Ueda, and Hidenori Tanaka. What shapes the loss landscape of self supervised learning? In The Eleventh International Conference on Learning Representations, 2023. URL httpss:/oppereviieee www.:neeeoofuneee http://pnneeeeeeeeeeeeerrrrrrrdd 3o33S
[28] Liu Ziyin, Botao Li, James Simon, and Masahito Ueda. Sgd with constant large learning rate can converge to local maxima. arXiv preprint arXiv:2 2021.
[29] Liu Ziyin, Botao Li, Tomer Galanti, and Masahito Ueda. The probabilistic stability of stochastic gradient descent. arXiv preprint arXiv:2303.13093, 2023.
[30] Vinod Nair and Geoffrey E. Hinton. Rectified linear units improve restricted boltzmann machines. In Proceedings of the 27th International Conference on International Conference on Machine Learning, ICML'10, page 807-814, Madison, WI, USA, 2010. Omnipress. ISBN 9781 105558907777
[31] Andrew Maas, Awni Hannun, Andrew Y Ng, et al. Rectifier nonlinearities improve neural network acoustic models. In Proc. ieml, volume 30, page 3. Atlanta, Georgia, USA, 2013.
[32] Djork-Amé Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (clus). arXiv preprint arXiv:1511.07289. 2015.
[33] Quoe Le Prajit Ramachandran, Barret Zoph. Searching for activation functions, 2018. URL http://wwww......ooo///d///////a
[34] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint
arXiv:1606.088/ 08011100000811.... 2016.
[35] Daniel Kunin, Javier Sagastuy-Brent Brena, Surya Ganguli, Daniel LK Yamins, and Hidenori Tanaka. Neural mechanics: Symmetry and broken conservation laws in deep learning dynamics. In International Conference on Learning Representutions, 2021. URL fett/forum?id=qqLLALB
[36] Johanni Brea, Berfin Simsek, Bernd Illing, and Wulfram Gerstner. Weight-space symmetry in deep networks gives rise permutation saddles, connected by equal-loss valleys across the loss landscape. arXiv preprint arXiv:1907.02911 2019.
13

---

[37] Berfin Simsek, François Ged, Arthur Jacot, Francesco Spadaro, Clement Hongler, Wulfram Gerstner, and Johanni Brea. Geometry of the loss landscape in overparameterized neural networks: Symmetries and invariances. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages proceedings... mlr. wttps:/v39//immseeeeeeeeeee html
PMLR, 18-24 Jul 2021. URL https://
[38] Haikun Wei, Jun Zhang. Florent Cousseau, Tomoko Ozeki, and Shun-ichi Amari. Dynamics of learning near singularities in layered networks. Neural Comput., 20(3):813-84) 843, mar 2008. ISSN 0899-7667. doi: 10.1 162/neco.2007. 10-000000..1--1444 URL https::tp///...... :/tp::://ddod. www.oo. 1162/ neco n 2000 10-06-41-
[39] Qianxiao Li, Cheng Tai, and Weinan E. Stochastic modified equations and adaptive stochastic gradient algorithms. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 2101-2110. PMLR, 06-11 Aug 2017. URL htttpss://pooceeeiiiniee mlr press/v70/11175 html
[40] Stephan Mandt, Matthew Hoffman, and David Blei. A variational analysis of stochastic gradient algorithms. In Maria Florina Balcan and Kilian Q. Weinberger, editors, Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages 354-363, New York, New York, USA, 20-22 Jun 2016. PMLR. URL attps://prroeeedings mlr press/ ptress/v48/nse v48/mandt1e httmaldt16.hmm111111
[41] Stanislaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in sgd. arXiv preprint arXiv:1711.0422 2017.
[42] Pratik Chaudhari and Stefano Soatto. Stochastic gradient descent performs variational inference. converges to limit cycles for deep networks. In 2018 Information Theory and Applications Workshop (ITA), pages 1-10, 2018. doi: 11..110...........3....
[43] Alnur Ali, Edgar Dobriban, and Ryan Tibshirani. The implicit regularization of stochastic gradient flow for least squares. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 233-244, PMLR, 13-18 Jul 2020. URL https:ts////reoeediii
mlr. pwwss/v119/al1200./ html
[44] Harold J Kushner. Stochastic stability and control. Academic Press, 1967.
[45] Hendrik Anthony Kramers. Brownian motion in a field of force and the diffusion model of chemical reactions. Physica, 7(4):284-304, 1940.
[46] Nils Berglund. Kramers' law: Validity, derivations and generalisations. Markov Processes And Related Fields, 19(3):459-490, October 2013. URL https::/hal..cciieeeee http://aaal..ceeeecc/ala-000o99/ 26 pages.
[47] Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv: 1556, 1556, 2014.
[48] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for im- age recognition. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition (CVPR), June 2016.
[49] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. University of Toronto, 2009.
[50] Madhu S. Advani, Andrew M. Saxe, and Haim Sompolinsky. High-dimensional dynamics of generalization emror in neural networks. Neural Networks, 132:428 446, 2020. ISSN 0893-6080. doi: https://ddoi.or///0.000......i http://00...........e.e.0....... http://ccinee//rrrceee///s/////0
URL. https: //www sci sciencedirect.
14

---

[51] Andrew K. Lampinen and Surya Ganguli. An analytic theory of generalization dynamics and transfer learning deep linear networks. In International Conference on Learning Representa- lions, 2019. URL
[52] Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. In H. Wallach, H. Larochelle, A. Beygelzimer. F. d'Alché Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, neurips. httpscccppppppppeeeeeeee httpc///aee...ieeeeppp////////// pdf
Inc., 2019. URL https: ccb77855502201111555552222222522
[53] Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv: 2013.
[54] Pierre Baldi and Kurt Hornik. Neural networks and principal component analysis: Learning from examples without local minima. Neural Networks, 2(1):53 58, 1989. ISSN 0893-6080. doi: http://l..ggg///.0/0///0/9666000 com http://cceeee/tttccceeee/////o00
URL https http:////wwwwwweeeeeeeeeeeeeeeee scieecedrrrcr
[55] Kenji Fukumizu and Shun-ichi Amari. Local minima and plateaus in hierarchical structures of multilayer perceptrons. Neural networks, 13(3):31-337,,, 2000.
[56] Shun-ichi Amari, Hyeyoung Park, and Tomoko Ozeki. Singularities affect dynamics of learning in neuromanifolds. Neural computation, 18(5):1007-106; 2006.
[57] Christopher DiMattina and Kechen Zhang. How modify neural network gradually without changing its input-output functionalityy Neural computation, 22(1):1-47, 2010.
[58] Bernt Oksendal. Stochastic differential equations: an introduction with applications. Springer Science Business Media, 2013.
[59] Florent Benaych-Georger and Raj Rao Nadakuditi. The singular values and vectors of low rank perturbations of large rectangular random matrices. Journal of Multivariate Analysis, 111:120- 135, 2012. ISSN 0047-259X. doi: http::/ddbbiooo/11.0011111.aama https://www sciencedirect. httpa///ww....ooooooooooooooaaaa
URL
[60] Olivier Roy and Martin Vetterli. The effective rank: A measure of effective dimensionality. In 2007 15th European Signal Processing Conference, pages 606-610, 2007.
15

---

A Extended Related Work
SGD's noise covariance shape matters. Our work is closely related to HaoChen et al. 6 work that showed how parameter-dependent noise introduces an implicit bias towards local minima with smaller variance, while spherical Gaussian noise does not. particular, their work conjectured that the parameter dependent noise introduced by SGD has an implicit bias effect" parameters where the noise covariance is smaller". They studied this effect in two-layer model introduced by 15 Our work generalizes this effect to broad range of models by introduc- ing invariant sets and demonstrates that this implicit bias towards these regions regulated by a competition between the curvature of the full-batch loss and the magnitude of the noise.
SGD can converge to a local maximum. Our work is closely related to a recent work studying failure modes of SGD 28 and a concurrent follow up work studying the stability of SGD near fixed point 29 Ziyin al. (2 studied the behavior of SGD in one-dimensional quadratic and quartic losses and showed that SGD can converge to local maxima, escape from saddle-pon -oodnts slowly, and prefer sharp minima due to the gradient noise. strength of their analysis that they directly derive their results from discrete SGD updates, while we use a continuous formulation. However, their analysis is limited to one-dimensional scenarios and is used to illustrate behaviors of SGD they implicitly deemed unfavorable. In our work, we describe the mechanism of stochastic collapse generally and connect to a simplicity bias that can benefit generalization. Concurrent to our work, notion of stability for SGD near fixed point called probabilistic stability, which overlaps with our work in discussing stochastic collapse for single neurons but not for permutation invariance. They spotlight SGD learning phases using this stability concept, while we explore how stochastic gradients incline networks toward simpler subnetwork invariants.
Ziyin et al. [29] proposes
SGD learns sparse features. Our work is closely related to Andriushchenko et al. [18] recent work studying a sparsity bias of SGD with large learning rates. The novelty of their work can be summarized as: 1. They revealed that large-step SGD dynamics have effective slow dynamics with multiplicative noise during loss stabilization. 2. They theoretically analyzed diagonal linear networks, showing that the SDE has implicit bias toward sparser representations. 3. They conjectured that deep nonlinear networks also show this phenomenon, which is supported by the empirical results. 4. They argued that SGD first learns sparse features and then fits the data after the step size is decreased. With this foundation laid, we would like to emphasize two fundamental aspects that set our work apart:
First, our analysis goes far beyond the basic understanding of diagonal linear networks, offering a broader perspective by introducing the invariant sets with a more general theorem (Theorem [.2)). In their work. the authors conjectured an implicit bias towards sparser features beyond their diagonal linear model, leaving this exploration open as an "exciting avenue for future work". We believe that our paper makes substantial progress along this avenue by understanding the implicit bias towards sparse representations as stochastic collapse to the invariant sets. Furthermore, our theory is applicable to general deep neural nets. We also provide theoretical framework that predicts the quantitative conditions for this attractiveness in a general setting. While applying this theorem to arbitrary neural nets is challenging future work, our work utilized this theorem to quantitatively analyze the collapsing threshold of simple models. Second, our results shed light on the collapsing phenomena of weight vectors, a perspective that contrasts with their emphasis neuron similarity based on activation patterns. Given that weight vector similarities often lead to similarity in activation patterns, our findings suggest that neural networks adhere to stronger condition.
16

---

B Further Discussion on Invariant Sets of SGD
In our work, start by describing subsets of parameter space that are unmodified by SGD, which we term invariant sets. Several works have explored geometric properties of the loss landscape through the lens of symmetry and invariance. (35 demonstrated how the continuous symmetries the loss function create level sets of the loss and impose geometric constraints on the gradients. [36] and [37] showed how permutation symmetries in the hidden neurons can lead to many potentially isolated minima with saddle points between them. (38] studied the learning dynamics of gradient flow near singular regions of parameter space, which for a multilayer perceptron, include permutation and sign invariance. A series of works [54 have leveraged the hierarchical structures of neural networks along with properties of their activation functions to identify critical points of the loss landscape.
B.1 Proof of Proposition 3.1 and Proposition 3.2
Proof. A is affine and 0 E A, thus it suffices to show that Øt A and is independent of the training 80 data. Denote the activation at layer as h()) (w!l)"((e (1-1)) + b(1). The gradients associated
with the outgoing weights are,
de
08 (1+1)
Ow (1+1) out,
(h(+1
(1+1)
de o
(1) (u(+t)) h + by )
amm ha-1)) 6!")
(9)
ah(+1)
0 E A,
at Swisep
0. The gradients associated with the incoming weights are,
0e
de
0h++1) Oh Orcoor
(1)
Ohu++j²
an
(1)
(
OE
(d((witt)))
((+1) Woul.
h3 + b(tt)) o w
o'
(1) T (1-1) + 6.00 o (a-1))
@h(++1j'
w
(10)
If 0 E A,
88
0. The gradients associated with the biases are,
Swin.
08
Oe
ah(+
ah
06!1
Oh(+1)²
ah!
0b
(o'
( out h + ((++1)
(1+1)
C w ( ) o' w h0-1) (0-1)+b())
(11)
0hu+1)
If 0 E A,
88
at = 0. Therefore,
A.
86
06
Proof.
We can directly see from Eq. 9 Therefore,
that if E A,
at
ag and to
Dwan.
(1+1)
Dw
[c+1]
Du
(2)
ar 06
0t 36
A.
06
B.2 Proof of Theorem 3.1
Proof. It suffices show that Vl(0; y) for all 0 E A and training data. the following, we
use the notation of E(O) as shorthand for E(f; 2),8. First, we show Vlg(O) VE(0) for 0 E A. where (0) = I(Q0). The derivative of lg along an arbitrary unit vector n is given by
(0) lim E(Q(O + hn)) E(QO) lim R0
l(o hn) E()) 1(Q(0 hn)) E(O + hn)
h
is tangent vector of A, the second term zero. If n is a normal vector of A, then by the approximate Q-symmetry assumption, the second term converges to zero. These two facts imply that
17

---

Vio(0) = V.(0). Then, for any E A. V Vio(0) Q QTVE(Q0). Since Q is symmetric (Q QT) or orthogonal (QQT Ia), we obtain
QV x, ve(0;x,)))
which implies V4(0;2, v) A. Hence, regardless of the sequence of mini-batches, starting from ((0) A, the SGD trajectory remains within A, demonstrating A is invariant set.
Theorem 3.1 provides an alternative proof of Proposition 3.1 and Proposition 3.2 simply identify
symmetric or orthogonal matrix Q that generates these sets. For the sign invariant sets, Q is the identity matrix with negative diagonals restricted to the index of the parameters associated with the invariant set. For the permutation invariant set, Q is the identity matrix with the diagonal blocks for the parameters associated with the two permutatable neurons swapped to being off diagonal.
B.3 Additional Invariant Sets
Beyond the classes of invariant sets already presented, there can exist other invariant sets depending on the structure of the architecture. For example, for a network using softmax layer, the columns of the weight matrix and the vector of biases immediately preceding this layer can be shifted by constant without changing the outputs. As a result, the dynamics for these parameters are constrained to an affine space with constant sum determined by their initialization. This subspace is an invariant set. Similarly, if the training data has non-trivial kernel space, as in the case of overparameterized linear regression, then the rows of the first-layer's weight matrix can be translated by any vector in the kernel space without changing the output the network. Again, this constrains the learning dynamics to affine space, which is invariant set. In general, any translational symmetry should generate invariant sets.
18

---

C Invariant sets of SGF
C.1 Definition of invariant sets for continuous processes
Here, we provide the formal definition of invariant sets for continuous process and a formal statement of Proposition[4.1
Definition C.1 (Invariant sets of continuous processes). For a given stochastic process (O, E Rd t t2 0}, Borel-measurable set A C I& is defined as invariant for every initial point Oo in A, the probability that the process stays in for all 0 is I, e.,
P(O, E A for any t 2 0 EA=1
(12)
C.2 Proof of Proposition 4.1
Proposition 4.1 (Formal). Consider SGD process (0(t) E Rd :t NJ given by Eq. where the gradients V Ve(;,a1 Vo) Rd >> R in Eg. are Lipschitz continuous and bounded for any E [n]. If :NENN NJ. then A is also an invariant set of an SGF A C Rd is affine invariant set of SGD (000 :t process (0. 2 0} given by Eq. 3
proof. By assumption, Ve(0: zz, y:) is L-Lipschitz in f and Ve(0; Nı, 36)l S L with some constant 0. Then we can see that VL is also Lipschitz continuous and bounded as follows:
1
[VV)()))) - v C(0)))
Ve(01;:1; v.) Ve(02: v:))
S 1 02l
EE(m]
|VC(())
- 11
Σ |1(00 xz, w:)1 s
Let E(0, Fz. v4) E00, Fz, U:) [(B, v,, v1). Vē(,, Fz. M4 is clearly 21-Lipschitz continuous and is bounded by 2L for any E [n]. Then,
mm (01) Dtm(02)) (02)
Σ
wi) 2.) v:))
iE(m)
(2((()(((0;;2;,,)) + 0n E(02 10110 - -221
(E]m]
<
ATT.110: - 0211
This implies that || [(()) - D(f2)llac < 4L - 02ll?, where || 1|oo denotes the operator norm. On the other hand, the Powers-Stormer inequality gives,
Il ((O)) V((2)) vāviID(0)) D(O2)]l dVID(())-
where || || and || l| denote the Frobenius norm and Schatten 1-norm, respectively. Thus, v (())
is Lipschitz continuous with respect to the Frobenius norm. This establishes the existence of a unique
strong solution the SDE in Eq. It now suffices show that the SDE has a solution {0, 2 0} satisfying Eq. 12 Let P Rd
dd be the projection operator to the affine subset A, and
construct a projected SDE:
dØ, -VV(((d /
By construction, this projected SDE also has a unique solution o, which clearly satisfies the con-
dition Eq. [2] Moreover, given that A is an invariant set of the original SGD process, we have PV Fi, vi) Vé(0; xx, yı) for all B and almost surely. Hence, we conclude that A is also an invariant set of the solution to Eq 3
A. This means that {0, : 2 0} satisfies Eq 12
19

---

D SDE formulations of learning dynamics with gradient noise
D.1 The relationship between SGD and SGF
Recall the SGD update rule from Eq. 2
(((+1) g(1) Σ (EB()
Vol (010 2.))
(13)
We factorize the right-hand side into the negative full-batch gradient and the deviation from the full-batch gradient,
( = 0 - ((00)) -V
Vol (0(0) ;00) -
(14)
KEB
gradient descent
gradient noäse
Let Eι (0) represent the gradient noise at 0. Assuming we sample with replacement, the gradient noise be disentangled into of i.i.d. per-sample gradient noises: {s(0) ン Eies &
where & VE(0; z; y4) (0)).. It is easy to check that EE[; = for all 0. Assuming that the gradient noises are gaussian random variables, Eq.
14 turns into
Ag V( 2D(
(15)
where AW ~ (N(0, and the diffusion matrix D = Var(Es) = 23 Var(š.). Eq. 15 is Suler-Maruyany discretization of the following continuous-time stochastic differential equation, which we denote as stochastic gradient flow (SGF).
d0: = - dt + d2D,)))
0o = (
(16)
The diffusion matrix can be expressed as:
((0,)
17
Σ KE[N]
Vol (0:; Me) Vol (0: "i V.(0.)VZ((0))
(17)
2.3
N
The diffusion matrix can be decomposed into a product of parameter-independen magnitude term D m and parameter-depender shape term (0) such that D(#) D,, D,) where
Dm 233
(18)
(+) Σ SEIN]
Vol (0l y1) Vol (0, v. VC(0.)V2(())
(19)
The optimization hyperparameters learning rate and batch size affect the magnitude term D.,, while the network architecture, the training data, and the loss function affect the shape term (0). For generic network architectures, determining the exact form of the shape term is quite complicated. For the purpose of analytical assessments, we sometimes pivot considering label-noise gradient descent (LNGD) as an alternative to stochastic gradient descent (SGD).
D.2 A simpler SDE formulation via Label-Noise Gradient Descent
In the following, will restrict our discussion to the loss function with the mean square error. Our aim is to achieve analytical results by simplifying the diffusion matrix. In its most basic form, using the deterministic gradient descent, the diffusion matrix is simply 0. Expanding when we deliberately introduce noise into the gradients? One common method to achieve this is by adding noise to the dataset labels. If we introduce i.i.d. mean-zero, variance ç² Gaussian label noise per training label, denoted as & it results an update rule we refer to as label-noise gradient descent (LNGD),
this, what happens
g(++1) =(( - L)((()))))
Sow (x:)
(20)
NE
20

---

This is discretization of the following continuous time stochastic differential equation, which we denote as naise gradient flow (LNGF),
d0,- -1) (2DLN(0.JdBB where the diffusion matrix of LNGF can be expressed as,
0g (
(21)
S1 1 N
DLN
(22)
KE[N]
which can also be decomposed into magnitude term and a shape term,
Dm S
(23)
2
D. = N
Σ NE[N]
(24)
Compared with SGF, the magnitude term here depends on the learning rate and the variance of the label noise. Compared with SGF, the shape term here much simpler, which allows us to get analytical results in Sec. To compare with the theory, we also conducted experiments with LNGD. We study label-noise in theory because it can replicate the implicit regularization effects of the mini-batch noise. Essentially, label noise serves as a foundational model that paves the way for more comprehensive grasp of the more complex dynamics inherent to mini-batch noise.
21

---

E Conditions for stochastic attractivity
E.1 Proof of Theorem 4.2
Theorem 4.2 holds as a corollary of the following slightly more general statement where the drift term of the SDE is not necessarily a gradient of a function.
Theorem E.1. Ler A C B be d--dimensional affine subset, and a stochastic process {0, E R t 12 0} obey the following SDE Ac. open c-neighhorhood of A with some e. >>
d0, - p(O)dd + (2D(0.),Wr,
(25)
where WW:tz 0f is the d-dimensional standard Wiener process. Here μ A. + R is a L- Lipschitz C²-function with its first-order derivatives being I.-Lipschitz as well. D : A. - Rmxm the diffusion matrix such that the second-order derivatives of its elements are I.-Lipschitz continuous. We further assume that all the elements of /D(0) are L-Lipschitz continuous. Let D1 = DPD where P1 Th + TRe projects to the normal space of A. Suppose A is an invariant set of Or. is stochastically attractive if there exists >> such that for any unit normal vector E TR. perpendicular to A and any 0 A.
VANDA >
(26)
Tr ( +11-5)nT
VAn p V
8) Dí > 0.
(27)
Proof. We exploit a local version of Doob's 's maximal inequality (Lemma. 44). )) By the Lipschitz
property of μa and vD, the SDE Eq. 25 has a unique strong solution up to the exit time T. = inf(t 2 0: 0g & Ac}. The stopped process {0, Senn. i 2 0} is a continuous strong Markov process. It suffices to show that A is stochastically attractive for θ.. We first show that AI) < 0 for any and À >> O, where Á is the infinitesimal operator of the stochastic
E AAAA with some 0 < process O, and A, denotes the open neighborhood of A within distance >> 0. We consider the Euclidian distance ((et) between the process 0g and the affine invariant set A.
The SDE for (0,) is given by,
a (0.) = (-21~~2000)0 PLU+ A (0,) Tr + A(A 2)r* 2)6* (O.DL a (0.)- / 22222__0 dW,.
This means that for any E Ac,
( = ->~~22 / + X1^ (x) (r) Dr D1 + A(A - -2)1-~1(xr) (x (x x). (28) By the fact that A is an invariant set, (1)
holds for any I A. Therefore, we obtain Tr Dr) = 0. Since Dı is positive semi definieee (1(r)= = 0 for any E A. Furthermore, since Áf(x) 0 with (0.,) PiO, 0, for any E A, PI (1)) = 0 for any E A.
Going back Eq 28 we obtain
F(x)
AL = 221 2
1 + 1/1/2)
2G(x))
where
( Pix
||P1a ) ((x) Pl
Piæ
f(o):=(=<TP)) (TDD))))) G(x)
Notice that G(x)= > for any E A with small enough e > 0 because Eq. 26 holds and the such that for
second derivatives of D(x) are L-Lipschitz. We want show that there exists
any & E A, [0(=) 1 + X/2 < 0 with some À > O. Parameterize x by xxE) = æi + εú, where Tij E A and â is the unit normal vector perpendicular to A. From I. I''Héiital's 's rule, we have,
F(x)
F(ZI Eh)
lim 8 G( + (ń)
MT (
22

---

The second derivatives of F can be calculated as follows:
(((()))) - 2V μ + V TTD1 - V Tr (29)
Substituting this, we obtain
F( + {n) 5 [z + {Π)
-V (( + (TD
lim
<1-
(30)
2 (
By the assumed Lipschitz property of the derivatives of μ and D, F()) and Vz are L.- Lipschitz in A,) with small enough > 0. Therefore, for any 0 <
F(r + (
V
G( + {
GV, F - FVRG
G
GV))(z) + Le)e (V + - -V - Le)e
(V - LE)e
Le)
(V Le) + Le) - (V) - G() - Le)]
G(,))))))
(AE(E)) 2
(3G())) Le
$ (1+6 (1)+0-)e)(1 Lc)(1 6 Le) (1 8-1 Le)(1 1 -Le) (1 (6 Le)
1(38-1-11 1) (1+ (6 Lc)
K 2L(35- - 1).
Therefore F(2) + G + {n) is 2L - 1)-Lipschitz as function of & in A.\A. This
implies that there exists > 0 such that for any AAAA, F(x)/2G(x) < 1 - 5/2 and hence A((5(x)) < 0. Therefore (xx) is in the domain of AA, and (1) = A(1((z) S O, where AA, is the infinitesimal operator of the stopped process ter with t inf (1(O,) e). By the local version of Doob's 's maximal inequality, for any <c
Plsup 1s (0.) 2 e'10 x] Plsup 1st (((TT) 2 e'100 = x] (x)x)
120
Equivalently, for any e' < 1/6
Plsup ((0,) 2 e'10- = =x] =x (x) 120
1/8
By taking ((x) to be arbitrarily small, we can bound the left-hand side with any positive value >>>>> 0. This means that A is stochastically attractive for 0, and so is for 0z.
E.2 Proof of Theorem 4.1
The sufficient condition is direct corollary of Theorem 4.2 and thus show the necessary condition here. By assumption, L" ("(x) + 1 Dz(r) < -8 at x with some 5 > 0. It suffices to show that there exists 0 such that the hitting time
=: inf [ E
02 2
finite almost surely for any R with initialization Oo x.
We choose e > 0 small enough so that D(x) > 0 for any x E (-e, e)e,e))()) Then we can consider another stochastic process Y, = 2-122 Jorr (21/2 which follows the SDE below
dY. = 1-22 D-1/2 (C'((0ex) ('((en)) di + dW,.
23

---

for any 0 St T. By further taking e small,
D-1/2 (-C((0nt) -
de 2 >
1 p'(Ontt)
- 10-1/2
de
sup D(x) > 6. Deaeer
2
2
/
Thus, Y, > 8t Yo + We. Therefore.
1 /2 [ (=1/2())do) & i R: 8t+W,
[" D (z)dx Yo}.
T- E : Y
V2 e/2
Since 5 > O. the right-hand side is almost surely finite with any value of Yo. This immediately implies that is not stochastically attractive.
24

---

F The attractivity of the sign invariant set of single scalar neuron
We here investigate when neuron may collapse towards the sign invariant set. Consider neural network consisting of only one neuron with scalar input and output: f(z; W1, w2) = wya (w)z), where z, W1, W2 E I and a represents the activation funetion. We assume that a (z) is a smooth Lipschitz continuous function with (0) = O. Let D { 31)..... (En, Nn)} denote the training data where 24, y/ E R, i [n]. We consider the loss function C given by
1 Σ (1(m; n iE[n]
L(W1, 12, 1)
- Si,t)
where Sa is the label noise which assume be increments of the n-dimensional Wiener process. The gradient flow of the loss is given by
dwi dt
OC
1
Σ
(f(#g W1, w2) ( w2zo (w)xo(
dw
dwz
0C
Σ
(f(ggu1 w2) ())o(w1x;) - (w1xs))G..
dt
Owz
(EIm)
This can be rewritten pair of SDEs.
dwi, =
Σ
(f(#g w1, w2) y W2230 (w
V
Σ
W2r,o' U2r,o'(w1z;)d
n iE|M]
LE|n]
dw2,t
Σ n
1
(f(zgwr, w2) ())o(1r;)dd + iE|1]
Σ iE(x]
olwiz,)dBB..
(31)
where ( >> 0 is a constant representing the amplitude of the gradient noise, v > 0 is the learning rate, and B denotes the n-dimensional standard Wiener process. Note that the rescaling factor 1/ 1/vn in the diffusion term is introduced to avoid the data size influencing the SDE's behavior.
F11 Proof of Theorem 5.1
As preparation, we first show modified version of Theorem E.1 incorporate stopped processes.
Lemma F.1. Let A C Rd denote an affine subset and A. be an open neighborhood around A defined A = {0 E Rd inf A -|0= c}. where >> 0 is a constant. Consider stochastic process {02 Rd :t 0} that follows the SDE Eq. 25 within Ae with all the regularity conditions mentioned in Theorem E. Futhermore, let T = inT?! 2 0: 0g & Aa O int B}, where B = {0 E Rd ||0,]| S CB} with i [d] and cB > 0. If An B is an invariant set of 0., then An B is stochastically attractive invariant set of the stopped process {0 t 0}, provided Eq. and and 27 hold for any x E AnB B.
Proof. Since the drift and diffusion terms in Eq 25 are Lipschitz continuous within Ac, there exists a unique solution for the SDE up to the exit time r. This allows us to consider the unique stopped process lerr, continuous strongly Markov process. Let I(0): Rd Rgo be the Euclidian distance between f E Rd and the invariant set A n and A be the infinitesimal operator of tenr where 0, A.}. The same argument Theorem E.1 leads to A(I* (x)) < 0 for any E A. n B with small enough e >> and > 0. By applying a local version of maximial inequality (Lemma 1 in Kushner (4) we obtain, for any A, and 0 < <'< <<
Te inff( >
(xx
sup
100 2 d 0. =
sup smpl((esTAA)) 120
0o =
This implies that A n B is stochastically attractive invariant of the stopped process Otrr-
Theorem 5.1 (formal). Suppose {w1,, W2,1) E RR 1 0} is stochastic process obeying Eq. 31 in
a neighborhood U of A {(0,0) c IR?, and the activation function o R R. is smooth function such that a (0) The invariant set A stochastically attractive if o [EEEto] a.v.
Eiet #7
25

---

Proof. If (0) = 0, the condition can never be met. Therefore, we only consider the case d' * 0. Without loss of generality, assume that U is an open ball centered (0, Since the drift and diffusion terms of Fq.31 are locally Lipschitz continuous, it has a unique solution up to the exit time of U. We define stochastic processes p. = W1,1 + 24,2, 8 = We,1 W2,2- These two processes obey the following SDEs:
dpc Hp(pe, St dtt + 1 Eie(n) op. St Bi.
vn
ds, = Me (e 81 81)dl
in Dieln) (P, (1) B
vn
where
Hp(p,)) 8) = Eienn wi,1, w2,)) 1))22,0 (wi, ow + (f(r w1,1, w2,1) - (.)o(wi,,a
(p,8) = - Eie(n) [m W1.e, w2,x) 12) (w)1 (f(mn W1.6, w2,)) - v1) (ww,
(p.(p, 8) = 1 (W1, )x + (o,
(2) = 10 (w)), (i(
Define stopping times Tp = ipf 2 e). Ta = inf (t 2 E) where e > 0 satisfies {(p, 8) max([p), [s]) {) C U, and consider the stopped processes (PeNTy, SEAT,) and (peATs SeAT,). We have,
0 0p 1p(0,.)
1 82
Σ
(0,00,
o0(0) ==)((0) n
000)2
Σ 2 <0
4n 0p²
yiu; 2n
iE(n)
iE/ml
iE/nl
1 0²
1 2n
22 p.s (0,0)
n 2 (00)²
x > 0,
n 0p
VE(N)
iE/m]
By the assumption
smoothness,
0 (0, - I 02
D (0,
0 and
Bp
0² Eigin) o2 p2,,00, (),s) > 0 for any s e [-c,c] with small enough > O. Therefore by Lemma F.1
Ap == ((0,8) 8) 8 [], is stochastically attractive for the process ((PEAT 8enT,) 1 2
Similarly,
0 05 (,
4n Ōs²
1 82
Σ
2 (0,0)
o'(0)
Σ
yzz;
2n
1
176
0'()²
Σ iE(ml
27 <0
iE|n]
iE|n]
1 80
(0,0)
1
2
00)
Σ
2 >
n. Og²
2n
iE|n]
iEim]
By Lemma
F.1 A, = (p,0) 2 0).
E [-c, c}} is stochastically attractive for the process
{( StAT.):
By combining these two facts, we see that Ap n A, = [(00,)) is stochastically attractive for the process {W w2 : 2 0}.
26

---

G The permutation invariant set in one hidden layer nonlinear neural network
Consider training two neuron neural network fe on a dataset {(zk, yk) where fa(x) = E w (ww 2)). We assume the data are drawn randomly from standard Gaussian distribution. Consider running gradient descent with learning rate 7/ on the mean square error loss with label noise this set up, the invariant set from the permutation symmetry the w/2 for i = 1,2). Empirically, we observe the stochastic collapse towards this invariant set as shown in 7 We found that the speed of collapsing depends on both the learning rate and the noise scale. Interestingly, the attractivity strengthens with increased learning rate and noise scale up to a certain threshold.
drawn freshly from (0,
affine space {0|w] (8)
1.0
109
10a
10
[w: w21
10
0.0
10-6
-0.5
10
10
10
10
0
2500
5000
7500
10000
2500
5000
7500
10000
Steps
Steps
Steps
0-0.0 0-0.1
0-02
0-0.4
Figure 7: Gradient noises in SGD induces stochastic collapsing to the invariant set associated with identical neurons. Left: The collapsing process of the weights of hidden neurons and 2 be- coming identical over the course of training, demonstrating SGD's bias towards simpler subnetworks. Middle: A visualization of the impact of learning rates on the stochastic attractivity. An increase in the learning rate accelerates the collapsing process until reaching critical threshold, beyond which the collapsing ceases. Right: The influence of noise scale on the collapsing process. The attractivity strengthens with increased noise up a certain threshold.
27

---

H Conditions of sign and permutation invariant sets in the presence of residual connection
Definition H.1 (Residual connection). A hidden layer le in neural network is defined as residually connected from a previous layer 1 if its forward pass follows:
h (z) = (wlle-1)
+ b(b)) + (x),
where h denotes the hidden activations of layer I, w is the weight matrix at layer I, and b(t) is the bias vector layer I.
Proposition H.1 (Sign invariant set for residual neural network). Consider hidden neuron p of layer l2 with a residual connection from layer 4 Let (Win.p 6/2) denote the parameters (weights and bias) directly incoming to this neuron, and Wont.p represent the parameters (weights) directly outgoing from the neuron. If the non-linearity is origin-passing (i.e. (0) = 0), then the coordinate plane = {0 wi.p = 0, b O, Want.p for I 1112) forms an invariant set.
(I+1)
[1+1]
Proposition H.2 (Permutation invariant for residual neural network). Consider two hidden neurons q in the same layer 1 with residual connection from layer 4, Let (win.p
in.p" by (w b()) denote
the parameters directly incoming to the neurons, and W Wour (1+1) (I+1) the parameters directly outgoing
Woutg
from the neurons. The affine space A {0] wn wm.gg (1)) b Wout Wouti for 1 12) is an invariant set.
(1) (1)
(I+1)
omit the proofs as they closely resemble the proofs in App. B. and B.I In the presence of residual connections, the conditions for the sign and permutation invariant sets become stronger. Not only do the weights and biases the neuron(s) at layer l2 need to follow the condition, the neuron(s) which are residually connected from must also must also meet these conditions.
28

---

I Understanding Stochastic Collapse within a Teacher-Student Framework
To investigate the impact of the stochastic collapse on generalization, we will extend the gradient flow analysis of training error [53] and test error [51] for two-layer linear neural networks to stochastic gradient flow.
I.1 The Linear Teacher-Student Framework
We follow the teacher-student framework proposed by Lampinen and Ganguli [5I]
Low-rank teacher network. We consider a low-rank teacher that computes the linear map W R - RRm. W IR mxx is low-rank matrix where k < min(m, n) denotes the rank.
Two-layer student network. We consider two-layer student network, parameterized by W1 E R and W2 E Rmxd that computes the linear map W W2W1 R > R". Here d is the hidden-layer dimension and we let d = min(m, n} such that the student has the capacity to represent all linear maps from RD to R
Noisy training data. The student network trained on data generated by the teacher network through the noisy map.
- - + where 7 = Wæ,
(32)
and (O, of I. We will consider fixed training dataset of p noisy input-output pairs (111.))). (Ip. Up]} generated by the teacher network from imput x; ~ (0, Let e Raxp and Y E Rmxp be the matrices with columns I; and y; respectively. This setup yields the second-order training statistics that guide the student network's learning dynamics,
ΣSs = XXT Σ N
Eys = YX = Σ yoo
(33)
Later in our analysis of the learning dynamics we will assume the input data matrix X is whitened such that Σxx = In implying that only the input-output covariance structure Eye will govern the learning dynamics of the student.
Train and test error. The student network is trained to minimize the mean squared error loss between the prediction Y = W X and the output Y, and evaluated the expected prediction error for new test point sampled from the input distribution. We denote the train and test error as
IWX YI |
E
|| E
wz
Eirain
Exest
(34)
W
where r ~ NO II.m) is new test point sampled from the input distribution, || ||e denotes the Frobenius norm, and E [:] represents the expectation over the training data and the input distribution of the test point.
Singular value structure. The train and test performance of student is determined by the relation- ship between three matrices in R the low rank teacher W, the overparameterized student W, and the noisy training data Eyz. Due to the linear nature of this problem, we will use the Singular Value Decomposition (SVD) for all three matrices to describe their relationship.
W Σ S.U.
W Σ 8,
Eye
(35)
Here, the s; denote non-zero singular values, and u; and i denote the left and right singular vectors for their respective matrices. We will sometimes also find it useful to concatenate this information into a matrix representation where we will use S, and V with the appropriate accents 7, or to denote which matrix they are associated with.
I.2 Exact Solutions to the Student Learning Dynamics
The student network trained by gradient descent with learning rate r to minimize Errain- However, we assume that each gradient evaluation is corrupted with Gaussian label-noise such that the training
29

---

dynamics are given by the coupled update equations,
T w w/(²) (w:")) (w2 w/ Exz -Lyz + 2 (T)
(36)
w = w/ 71 (/2h) W E.p Ene + xT) (w:10)
(37)
where W and w are the parameters after steps of training and Z(0) RMxp is matrix of label-noise associated with step t. Directly studying these dynamics is difficult because the gradients equations are coupled between the weights W1 and W2. However, by introducing some assumptions on the covariance of the inputs and the label-noise, selecting a special initialization of the parameters, and taking the limit as 1) - O, we can decouple the dynamics into a system of scalar non-linear SDEs with exact solutions.
Decoupling the dynamics. We introduce the following assumptions:
A1. Whitened-Inpe We assume that Σzz = I, such that the input data matrix X whitened (this implicitly assumes p 2 n such that we have at least as many observations as features).
A2. Structured Label-Noise. We assume the the gradient noise can be decomposed as 2 diia (2() VTX X where 2(() ~ (²),
A3. Spectral Initialization. We assume that the student network is initialized such that W = OA vT and w-1") U OT where ( and (X0) are diagonal matrices such that (0)) = B(o) A and O e Rddd is a random orthonormal matrix such that OTO = la.
A4. Balanced Initialization. We assume balanced initialization such that w/('w(() w0 T (0)
Given the first three assumptions, the dynamies decouple in the eigenbasis of Eyz. Under the change of variables A = or,!") V and B = UTT.!" O, the dynamics transform to
(t+1) (t) n (B()) (B(0A4 (() 5 diag (2(0))
(38)
B(++1) B(t) 1 (B A 3 + diag (2())
A400)
(39)
which because (0) and (()) are diagonal matrices by assumption, decouples into system of scalar equations. Taking the limit as n - 0 can approximate each of these scalar equations with two-dimensional non-linear SDE,
as b.
by (a,b: - =:)] de + V as (a,b; - 8:)]
by
dBt. a
(40)
where a; and by are the tt diagonal element of A and B respectively. Under these much simpler dynamics, is not difficult to prove by Itô's Lemma that the fourth assumption will be obtained no matter the initialization,
Lemma I.1 (Autobalancing). Under the dynamics in equation 44 the lim (t)) b,())
Proof. Let rx(t) denote the difference rt) = a,(t) (t)2. By ItO's Lemma, (t) is driven by the ODE dr; -ns² rdd which has the temporal solution (t) (0)e "x*,, Because n/52 0. then the lim (t) =
The result of all four assumptions Al A4 is decoupling the learning dynamics into a set of independent quartic slices of the loss landscape, (w;) = (w? - 8 with w; = a; = b, where within each slice the dynamics evolve according the non-linear SDE,
dw, = --- w. d + Vn5² wn,Bw.
(41)
The geometry of these slices is controlled by the strength of the singular mode for the training data 3; and the dynamics are determined by the relationship with the student's singular mode strength & wf. By Itô's Lemma we find that the dynamics of 6; are governed by the non-lineat SDE,
d8, = 28; ((
nT²)
8, dt + 2 m
(42)
2
30

---

which is an equation often used to model population growth in stochastic, crowded environment.
See Oksendal [58]. Chapter problem 5.15 for more details.
Theorem 6.1 Under assumptions Al A4, the dynamics of 8:(1) for 2 0 are governed by the
stochastic process,
(24 B
8,(t) 2 So (2 B. dr (:)-)
(43)
Proof.
Let u4 = =, Then by Itô's Lemma, the dynamics of u, are given by the SDE,
du = (2 (23; 3nz5" )u:) dt n uB,
(44)
Rearranging this SDE we can express it in the standard form of geometric Orstein-UUlenbeed process (sometimes referred to as a modified Manstin-UUleenbeck process)
du, = 0(1 - u,)dt dt + ou,dB,
(45)
where 0 = 28, =-= SDE we make the ansatz that the solution takes the form (t) = ((t)G((t where F(t) and G(tt) solve the respective SDEs,
(5;
33 (-12)_1,00 and 222 To derive a solution to this linear
dF DF + oFdB,
dG = a + (()d
(46)
To determine the unknown coefficients att) and (tt) defining the SDE for G(t), we can apply Itô's product rule
du; = dFG F dFdG
(-OFG a(t) + B(t) (oFG
(a(I)F+ +3))) Du,)d + (ar, + ()d)
Aligning this expression with the original SDE for we see that
(tt) = 0pF
p(t) = O.
(47)
We can now solve for Ftt) and G(t). The SDE defining F(t) is standard geometric Brownian motion and thus has the solution,
F(t) = (o+ta")e+aB,
(48)
where we assumed F(O) 1 and thus G(0) un (0). Using this solution for F(t), then G(t) is given by
G(t) un 0p f' (0= dr.
(49)
Combining our solutions for F(f) and G(t) we get that the solution for (t) is
(t) = u. 1o + 0p ( (=00+)0"1C-)))
(50)
Inverting this expression and rearranging terms, gives the expression for (t()).
(0+
a.(() ((( Bn for (0+ 20 - dr +or+++ 8;(0)
(51)
Plugging in the simplifications 8μ 2, 0 0 = 23; 2 and a -2 n gives the final expression.
Corollary 6.1 In the limit of 8,(0) + 0, for any > 0, (t) is given by
(t) 0.8.
(28, - 1422
(52)
31

---

Proof. Since B, is almost surely continuous, (23, + 2 has maximum in E [0,
and therefore 2 So ( dr is finite almost surely. Therefore, by Theorem 6.1
(28)) (22
e (25, ( )
8: (1()/8:(0) = 23,(0) So (23 112 dt + a.s.
(53)
Corollary 6. has important implications for the dynamics of &:()). It demonstrates that the dynamics observe phase transition determined by the ratio 28:/n(², consistent with Theorem 4.1 When this ratio is less than one, the distribution will exhibit stochastic collapse, converging to a delta distribution at the origin. When this ratio is greater than one. the distribution will converge stationary distribution with a positive expectation. Corollary 6.1 demonstrates that stochastic collapse determined by a signal-to-noise ratio where the signal derives from the noisy teacher and the noise originates from stochastic gradients. While Corollary 6.1 describes the dynamics with vanishing initialization for any 0, it does not describe what happens with finite initialization nor infinite time. To understand these scenarios we consider the stationary solution for the dynamics of $;-
Corollary I.1 (Stationary solution). In the limit of -> co, 8t(t) distributed according to the density,
= " 8,
(54)
where Z = 1000 2 8;
d8; the partition function.
Proof. To check if (8:) ) is the stationary solution we the Fokker-Planck equation for the dynamics of 8,, which is
1722 +
O.p V.
28; 8:
p 2m V (sip)
p(s,, 0) ōs.(0)- (55)
where J is commonly referred to the probability current. The stationary solution pas satisfies
O.pas = O, which because we are considering a one-dimensional SDE, is equivalent to Jas = 0. Plugging in the expression for Pes (&x) and using the simplification
2 (s:p) = -5 S,Pea + 2n 8 (8.)
(56)
- s:peo - 2n (nT²)
+
8,
3
3i
)) (8))
(57)
2 1/52/
1122
28; 8;
8i +
Pas
(58)
2
can verify that J.s 0.
The stationary solution can be interpreted as Gibb's distribution (i.e. Pas e=*((x) with the potential
3nsc -5: log(s:)
v(())) 5;
(59)
and temperature constant x= (n) -1 From this interpretation we see two phase transitions,
arg max (.. 8 C2
1--e-~M(()) if > (60)
(so =
2
8;
80
if "C
When n < 8; the most probable state is 8, 2n where already bias from the stochasticity towards the origin. When n 5t the most probable state at the origin. When
n=² > 3. the partition function Z diverges and the stationary distribution collapses to delta distribution at the origin, which is an invariant set.
2
We now conjecture exact expression for the temporal dynamics of the expectation (())
32

---

Conjecture I.1. Under assumptions Al A4, the expectation [[) for 2 0 is given by the equation,
(3: 7/6 ) Je(28:100 =
[[sa(())
(61)
& m
e(23; n)
(0)
Conjecture 1.1 captures both the phase transition for stochastic collapse from Corollary 6.1 and the student's singular value shrinkage from Corollary [.I] Furthermore, in the limit of 0, it aligns with the dynamies from Theorem 6.T] as would be expected. Additionally, one can show that the dynamics predicted by Conjecture [.1 are equivalent to the deterministic dynamics given by L2
regularized gradient flow with a regularization coefficient of nd²
1.3 An Analytic Theory of Error Dynamics in the High-dimensional Limit
So far we have setup a teacher-student setting and demonstrated that the learning dynamics of the student can be decomposed into..nonnlinarr.. Now will show how this implicit shrinking and collapsing process is an effective form of regularization that improves the performance of the student
Decomposition of the error. Using the the SVD structure of the low-rank teacher {ū, 8,, 0g J. overparameterized student {ū, &,. 0,), and noisy training data {ū; 3t, 0/} we can decompose these error terms
m
m
m
Eerain
(
(,)
(62)
j=1
5=
$ 2
Exeat
83
Σ
Σ
818y Kgg) (7
(63)
f=1
=1
From these expressions we see that understanding the dynamics of the train and test error depends on (I) the alignment of the singular values and vectors of the student network with the SVD structure of the training data and (2) the relationship between the SVD structure of the training data and that of the teacher network. Our analysis of the student network training dynamics covers (I) and as done in Lampinen and Ganguli [5T] we can understand (2) in the nigh-dimensional limit through a random matrix theory analysis.
Random matrix theory analysis. We will work in the limit m - ao with finite rank k ~ O(1) and finite aspect ratio A = m E (0, oo). Without loss of generality, we assume o = n -1 such that the singular values of the noise are O(1) and thus the singular values of the teacher network can understood as signal to noise ratios (SNRs). The relationship between the SVD structures of low-rank matrix (the teacher network {ū, 5, } and noisy perturbation (the training data {ū,, 3,, 6 is well studied in [59], which we summarize here.
In the limit, the top k singular values of Eg converge to values given by the transfer function,
(3) (s) V(+ (1++3) (A+ 52) ifs > 1/4
3($)
(64)
11+A
otherwise
while the bottom m singular values are distributed according to the Marchenko-Pase (MP) distribution,
4A (AA
if3 E VA,1 otherwise
(MP(8)
(65)
Additionally, in the limit, the top Is singular vectors of Eng acquire an alignment with the singular vectors the teacher network given by the relationship (ū, ú.)) (ū,, b.) O where the the overlap function is defined as,
( = 1
A(1+82) 1/2 s(++8)
(A++2) 1/2 8(1482)
(1-
if& > A1/4
(66)
otherwise
33

---

Experiment details
The codes to reproduce the experiments in the main paper can be found at https://gitthb wttpe///ooohtttcccoolooooooe wwf.ccfff///toocsccccccccccocccc collapse Our the experiments were run on the Google Cloud Platform (with 4x or 8x NVIDIA A A100 (40GB) GPU). The initial code development occurred on local cluster equipped with 10x NVIDIA TITAN X GPUs. We carried out all the deep learning experiments with VGG-16 [47 and ResNet- 18 [48], training on the CIFAR-10 and CIFAR-100 datasets respectively [4)] For VGG-16, we did not use Batch Normalization or Dropout. Both VGG-16 and ResNet-11 had their nonlinearity activation functions replaced with GELU, as the potentially larger invariant set associated with ReLU could complicate identification and illustration of the two main invariant sets discussed in the paper. We leave it as future work to study in details the structure of the invariant sets associated with ReLU. For all our training, we applied standard data augmentation and used SGD (with momentum 3 = 0.9 and weight decay of 0.0005) as the optimizer.
com/ htt//ccffccfffocc
Stochastic collapse in quartic loss. (Fig I The left three plot are generated by sampling 50 trajectories of the SDE d0;- = - ( 1&Be dt + +B0ddB with μ = 1.5 and Ç = 0.2.1.2.5. The samples are computed using Euler-Maruyama discretization scheme with step size 7) = 0.0025 for 10 steps. All three plots use the same random Gaussian initializations with mean zero and variance four. The side plot for these three subplots shows the theoretical steady-state distribution described in Sec. 4 The rightmost plot is generated by sweeping 50 linearly spaced values of Ç from S = 0 to <= 3. For each value of we sample 10000 trajectories with the same scheme described previously but for varying number of steps: 0, 5000, 10000, 20000, 40000, 80000. We compute the empirical probability of the final position being within e 1e euclidean distance of the origin.
Stochastic collapse towards a sign invariant set. (Fig 2 We trained 103 different single-neurg
models via simulating SDEs (31) up to time T 100, using the Euler-Marumay: discretization
scheme of step size dt = 10 All the models are trained with the same data set, composed with 32 points {axi i=),,=, 1,. .22 where x; is i.i.d. standard Gaussian random variable, and y; is given by y: Iy + 6; where e; is i.i.d. centered Gaussian random variable with standard deviation of 0.1. All the models are initialized with random weights sampled from i.i.d. centered Gaussian distribution with standard deviation 10-3
Evidence of stochastic collapse towards permutation invariant sets. (Fig.) We trained VGG-16 16 for 10a steps on CIFAR- 10 with a learning rate of 0.1 and ResNet-18 for 10th steps on CIFAR- 100 with learning rate of 0.02. We computed the normalized distance between neurons within the same layer by concatenating corresponding vectors of incoming and outgoing weights, and used this to cluster neurons. Neurons within 0.1 normalized distance from each other were grouped together, which then determines the number of clusters or the independent of neurons denoted as Nind- In Fig. 3 we plotted the pairwise distance of the incoming weights and outgoing weights separately. We carried out Agglomerative Clustering based on the normalized distance of the incoming weights with nind clusters and the corresponding pairwise distance between outgoing weights was visualized with the same ordering.
Larger learning rates and increased label noise intensiffy stochastic collapse. (Fig. 4 and 9 These experiments varied learning rates and the extent of label noise. For the learning rate experiments, trained models without label noise for 10 steps. For the label noise experiments, learning rates of 0.01 and 0.02 for VGG-16 and ResNet- 18 respectively over 10 steps. Here, fresh label noise was introduced with each batch sample. A label noise level of e% meant that c% of the labels were randomly assigned an incorrect label. We averaged the results across four replications with different seeds.
maintained
After training, we calculated and displayed the proportion of independent neurons in each layer of the networks. We first removed 'vanishing neurons', defined those with incoming and outgoing weights less than 10% of the maximum norm for that layer. We then clustered the remaining neurons based on the pairwise normalized distance of the concatenated incoming and outgoing weights. Two neurons were defined as identical if their distance in weight space was less than 10% of their norms. This allowed us to determine the number of independent neurons in each layer.
Effect of batch size on stochastic collapse. (Fig. 10 We trained models with different batch sizes for 10s (CIFAR-10) and 008 (CIFAR-100) steps, while keeping the learning rate constant at 0.02. We averaged the results across four replications with different seeds. The method of calculating the proportion of independent neurons remains the same.
34

---

Demonstrating generalization benefits of stochastic collapse in a teacher-student setting, (Fig. 5) The experiments were conducted accordance with assumptions Al A4. In these experiments, we set both the input and output dimensions to be 64, and the dataset size was set at 1024. We used sparse teacher model with a rank of and signals ranging evenly from 0.5 to The input samples were drawn to satisfy the condition X X T. We then constructed the signal y by introducing random Gaussian noise with a standard deviation of 0.5 to the true output values. The network was trained with an initial learning rate warm-up of 1000 steps, followed by 50000 steps at a learning rate of 3.0. After 50000 steps, we reduced the learning rate to 0.3 and continued the training for another 10000 steps. The experimental results averaged over 256 runs with the same dataset.
In Fig. I1 we present the results obtained when replacing label noise with different batch sizes. In this experiment, we used SGD without label noise and maintained similar experimental setup as before. Additionally, we introduced a learning rate drop at step 53000.
Large learning rates aid generalization via stochastic collapse to simpler subnetworks. (Fig. 6 We carried out initial training of VGG-16 and ResNet- 18 using a larger learning rate of 0.1. We created checkpoints in the training process at steps of 10s, 2 x 10a, and x 10 We resumed training from the checkpoints post but reduced learning rate of 0.01. To track the impact of this large learning rate phase on the emergence of simpler subnetworks, we calculated the layer-wise fraction of independent neurons each of these checkpoints, before the learning rate drop. We averaged the results across eight replications with different seeds.
35

---

K Extended Experiments
Layer with little stochastic col apsing
Layers wich strong stochastic collepsing
CIEAR-16 MATR-0 VGG-16
Probabiity Density
102
1100 LiNE
conviLi (b12)
(512)
Steps
20000 Steps
200
LUNE
60000 Steps
40000 Steps
80000 Steps
100000 Steps
CIEAR- 100 RowNet- 18
Probabilkty Density
1ayer2.0.com 2ocr2.0.cowv2221 conv2
(128)
LUBB 1,mn 1110
leyer1.1.comv (64)
leyert.1.comvl (512)
10c
200000 Steps
400000 Steps
16
600000 Steps
800000 Steps
1000000 Steps
Normalized Distance
Norma lzed Distance
Normalized Distance
Figure 8: Evidence of stochastic collapse towards permutation invariant sets in deep neural networks. We show distribution of normalized distance across all pairs of neurons for the same hidden layer of neurons shown in Fig. 3 We show plots for three different layers in VGG-16 trained on CIFAR-10 10 (top row), and for three different layers in a ResNet- 18 trained on CIFAR-I00 (bottom row). Each plot assesses the distribution at various training stages. In layers with little stochastic collapsing (left column), the distributions rapidly thermalize and concentrate around 2. In layers with strong stochastic collapsing (middle and right columns), the distributions evolve to form a peak around 0, indicating stochastic collapsing towards the permutation invariant set.
Fraction of Independent Nerons
CIFAR-10 VGG-16
CIFAR-100 ResNet-18
1.0
Label Noise
1.0
M
0.8
0 %
20 %
0.6
0.6
%
0.4
%
0.4
0.2
80 %
0.2
0.0
0.0
convi
comvz
comve
convo
canvo
convo
comuy
convo
compo
conv10
conv1l
conv
Ic14
1c19
convi
0.covv
comaz
Locuu
camzz O.conu
connzz
canvo
como
L.cnnz R.ooz
conzz d.concl 0.canuz convl L.conzz
Layer
Laya
Figure 9: Increased label noise intensify stochastic collapse. This figure illustrates how the fraction of independent neurons per layer in VGG-16 trained on CIFAR-10 (left column) and ResNet- 18 trained on CIFAR-100 (right column) varies with label noise. The networks are evaluated training steps of 10". A reduced percentage of independent neurons indicates stronger stochastic collapse.
Faatonn Independent Neurans
CIFAR 10 VGG
CIFAR 100 Restict 18
1.0
1.0
MVMMVV
Batch Size
0.6
32
0.4
0.4
128
02
256
0.2
0.0
0.0
827
825
Layer
Layer
Layer
Layer
Figure 10: Effect of batch size on stochastic collapse. This figure illustrates how the fraction of independent neurons layer in VGG-16 16 trained CIFAR-10 (left column) and ResNet-18 18 trained on CIFAR-10 100 (right column) varies with batch size. lower percentage of independent neurons signifies heightened degree of stochastic collapse. As elaborated in Appendix D modifying batch size is expected to parallel changes in the learning rate since both influence the amplitude term the diffusion matrix. Yet, unlike learning rate effects, alterations in batch size produce complex shifts in stochastic collapse. We conjecture that this subtlety may arise from nonrreplaeement sampling and the randomness introduced by transformations applied to input images.
36

---

x10
x10-5
0.4
Train Loss
6
Test Loss
1.4
0.3
Batch Size
1.0
0.2
0.1
0.6
4
0.0
Full Besch
0.00
0.15
0.30
Steps
x10
Steps
x x10
5*
Figure 11: Demonstrating generalization benefits of stochastic collapse a teacher-student
setting with batch noises. Same as Fig. [5] but show effects of different batch sizes. No label noise was used in this setup. We show the train loss (left) and test loss (middle) during the training of the student with different batch sizes. Dashed lines in the leftmost and middle left panels indicate the steps where learning rate is dropped. Training with smaller batch sizes generalizes better. Right: we show the noisy teacher signals against the learned student signals before dropping the learning rate. We only show the smallest 32 singular values.
CHEAR-10 VGG-16
1.3
1.00
80
West Accuracy
10000
Effective
0.4
UA8 WM
0.3
50
10000
4000
Learning
CIFAR-100 RtsNce-18
Tairiin Loss
Test: Accuraty
Test: Accurary
Erfoctiv nkc
2
1.0
8
0.6
0.2
10000
20000
30000
10000
Figure 12: Large learning rates aid generalization via stochastic collapse simpler subnetworks. Left, Middle: Same as Fig. Right: The effective rank (60] of the weight matrices of independent neurons per layer evaluated at different learning rate drop times (indicated by the color) during the initial high learning rate training phase. To compute the effective rank, we first gather the singular values, denoted by s, of the concatenated weight matrices of the incoming and outgoing weights for specified layer. The effective rank is then computed as p Σ, 8; log 8;, where à; represents the normalized singular value defined by 8 = E5
37


Reviews:
[{"review_id": "u6xgHpTeo4", "reviewer": "Reviewer_ZL9R", "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.", "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.", "limitations": "Limitations have been adequately addressed.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "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\n- Assumptions (A1-A4) in Sec. 6 are really restrictive\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful review and the time you took to identify areas where our work could be improved. We'd like to address the concerns you raised:\n\n**Weakness A**: We appreciate your observation regarding our use of a continuous formulation of SGD and we agree that this is a limitation of our current analysis. As we discussed at the end of our manuscript, “_extending our analytic results from continuous SGF to the discrete SGD updates_” is an important direction for future work. Our decision to adopt a continuous formulation was driven by our desire to leverage existing concepts from stochastic control theory, such as the definition of stochastic attractivity and the use of tools such as Itô's calculus to derive concise conditions for stochastic attractivity. However, we'd like to underline that our paper's main arguments retain their strength for several reasons:\n- Our discussion of the invariance properties of sign-invariant and permutation-invariant sets (refer to Lemma 4.1 in our paper) is not confined to the continuous formulation; these properties remain valid in the original discrete SGD setting as well.\n- Our empirical examinations with CNN models utilizes the conventional discrete SGD algorithm instead of a discretization of the continuous SGF variant discussed in our theoretical framework. The consistency of our empirical findings with our theoretical postulates underscores that our continuous assumption is sufficient to explain important features of the outcome of discrete SGD updates at practical finite learning rates typically used.\n\nIn our updated manuscript we will add a discussion further elaborating on our choice to use  a continuous formulation of SGD and the approach to extend our findings to the discrete scenario.\n\n**Weakness B**: We want to emphasize that assumptions A1, A3, and A4 are standard assumptions used in the canonical papers [Saxe et al. 2014] and [Lampinen & Ganguli 2018], which established the linear-teacher student model as a setting for studying generalization. We added A2 such that we could derive analytical expressions for the theory when incorporating the new ingredient of stochasticity to the analysis. We agree that it is important to check how restrictive these assumptions are and we thank the reviewer for pointing this out. We now include additional empirical results (see Figure S3 in the attached PDF) which verifies that the key observed phenomena hold even without these assumptions. Although we acknowledge that our assumptions in section 6 are restrictive, we believe they are necessary to provide an analytic demonstration of how stochastic collapse influences generalization, though the derived conclusions hold more generally. Indeed remarkably, our theory, despite its restrictions made a **new** prediction that training for longer at larger learning rates helps generalization because it promotes stochastic collapse.  We tested this novel prediction successfully in regimes far beyond the restrictions of our theory  - in ResNets and VGG’s trained on CIFAR10/100 (see Fig. 4).  Thus our simple theory makes powerful, generally applicable predictions. \n\n\n**Answer to question for line 176**: We appreciate your inquiry. Upon revisiting this sentence, we agree that it is somewhat ambiguous and might give the impression that converging to a local maxima is beneficial. Although there should exist certain pathological instances where such convergence is beneficial,  in general this would not be the case. For a high-dimensional setting, such as that considered in section 6, stochastic collapse to saddle points, not local maxima, would be beneficial to generalization. The sentence should read, \"_converge to a local maxima or **saddle-point** of the loss landscape_\".  For example, in our linear student-teacher setting, fixed points where the student has some of its singular values stochastically collapse to zero (corresponding to small data singular values that are not learned), are actually saddle points in the training error loss landscape, and these saddle points have **higher** training error than the global minimum (in which all the data singular values are learned for a full rank student), but **lower** test error than the global minimum.  This is thus a concrete example of how convergence to a higher training error saddle point can, remarkably, help generalization.  We will expand upon this important loss landscape perspective in the revised camera-ready paper.  "}, {"review_id": "E6tTWJ7Jj9", "reviewer": "Reviewer_C15e", "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.", "questions": "- Is it possible to (theoretically) analyze stochastic collapse beyond linear models or one-dimensional settings?", "limitations": "No potential negative societal impact.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-motivated, focusing on the study of how gradient noise directs SGD dynamics toward simpler subnetworks.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your comprehensive review and the time you have taken to suggest improvements for our study. We'd like to respond to the weaknesses and queries you raised regarding our paper:\n\n**Weakness 1**: “_not sufficiently clarifi[ing] the applicability of the attractivity condition to permutation invariant sets in high-dimensional settings._” While we did not provide a direct theoretical treatment of our theorems in section 4 to a high-dimensional scenario of the permutation invariant set (while we did do this for the sign invariant set in the teacher-student analysis in section 6). As we explained on lines 221-223, “_we provide intuitive insights into the attractivity of permutation invariant sets through a toy example of a two-neuron neural network in Appendix D_”.  We attempted to analyze this setting theoretically with high-dimensional input/output, but found it very difficult to derive statements stronger than what the main theorem statements in section 4 already implied. Note that our main Theorem 4.1 provides a sufficient condition for stochastic attractivity to **any** affine invariant set, **including** permutation invariant sets.  To further bolster this result, we devoted the majority of section 5 in our paper to empirically verifying and exploring the applicability of our attractivity condition to permutation invariant sets in high-dimensional neural network settings.  We found, remarkably, strong stochastic collapse to permutation invariant sets in which many neurons have identical incoming and outgoing weights (Fig. 1) in realistic networks (ResNets and VGG-16).  We believe our combination of theory and empirics provides strong applicability of the phenomenon of  attraction to permutation invariant sets.\n\n**Weakness 2**: “_extension of the linear teacher-student framework to non-linear models or more general scenarios._” We appreciate your comment and would like to clarify how we extended the linear teacher-student framework to the more realistic neural network settings.  Directly extending our theoretical analysis to more complex non-linear models is exceptionally challenging, demanding the introduction of novel analysis techniques, and would stand as a distinct and independent research endeavor. Thus, the goal of section 6 was to use a well-studied setting for understanding generalization in deep learning dynamics where we could derive exact results highlighting the influence of stochastic collapse on generalization.  We used this understanding to derive testable predictions about the role of stochastic collapse in non-linear models. As explained in the paragraph following the introduction of the linear-teacher student analysis (lines 312 - 319):\n\n“_The analyses in the linear teacher student network provide valuable insights into how stochastic collapsing can enhance generalization…The key prediction is that a large learning rate induces stronger stochastic collapse, thereby regularizing the model complexity. Furthermore, remaining in a phase of larger learning rates for a prolonged period drives SGD closer to the invariant set. Consequently, when the learning rate is eventually dropped, overfitting in these specific directions is mitigated._”\n\nIn the subsequent paragraph, we proceeded to “_test this predicted mechanism_” on a “_VGG-16 and ResNet-16 [trained] on CIFAR-10 and CIFAR-100, respectively_”. And as demonstrated in figure 4, “_we found, as predicted, that models with a longer initial training phase collapse towards simpler invariant sets with fewer independent neurons_”.\n\nThis combination of deriving a theoretical prediction from analytic insights into a simple model, and then testing this prediction in more complex models, that are currently beyond the reach of theory, is a powerful route forward for advancing our conceptual understanding of the complex processes of deep learning; it combines theory and experiment in creative, powerful and synergistic ways, obtaining results that neither theory nor experiment alone could achieve.  Therefore this connection between the linear student-teacher setting and more complex models should, we believe, be considered a strength of our work rather than a weakness.  \n\nAnswer to the question: Yes, our work **does** indeed analyze stochastic collapse in non-linear or high-dimensional setups. Specifically, Section 4 aims to offer a general theorem for stochastic collapse within a broader framework. Following Section 4, we study stochastic collapse in a **non-linear** but single-neuron model (Theorem 5.1), while our teacher-student framework in Section 6 studies stochastic collapse in a **high-dimensional** linear model. We agree that finding a non-linear **and** high-dimensional setting where studying stochastic collapse is theoretically tractable is an important and challenging step for future work. Nonetheless, our current work fills in this theoretical gap by exploring these settings through extensive empirics.  We verified empirically the prevalence and importance for generalization of  stochastic collapse in realistic networks (ResNets and VGG) (Fig. 1,2, and 4), and demonstrate that stochastic collapse in these realistic settings shows similar properties to those we have found in the simple examples."}, {"review_id": "4Zjen2MmYM", "reviewer": "Reviewer_kjwn", "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.", "questions": "1. The invariant sets for SGD are also invariant sets for GD. Is there evidence that GD doesn't exhibit collapse?\n2. Is there any interplay between different types of invariant sets? E.g., which type of invariance does SGD prefer?\n3. Is it possible to characterize an invariance type by a corresponding norm of the parameters? If not, why? \n4. 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?", "limitations": "The authors have adequately addressed the limitations.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "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. \n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your effort on the insightful review and suggestions. We are going to discuss each weaknesses and questions raised:\n\n**Weakness 1**: We value your insight but respectfully disagree.  In Section 4 of our paper, Theorem 4.1 yields sufficient conditions for stochastic collapse applicable to **any** affine invariant set of **any** dimension in **any** network.  We acknowledge that stronger statements in Sec. 5 & 6 are restricted to simpler settings, partially to provide concrete/intuitive examples to illustrate our general theory.  Deriving sufficient and necessary conditions for stochastic collapse for all neural network architectures lies beyond present theory, as we explain in our revised section 4: \n\n“_To extend the collapsing condition derived in Theorem 4.2 to high-dimensional cases, a natural idea would be to consider all one-dimensional slices of parameter space orthogonal to the invariant set. However, the issue is that some of these slices might satisfy the collapsing condition while others do not. This can result in complex dynamics near the boundary of the invariant set making it difficult to derive a necessary and sufficient condition for attractivity in high-dimensions. Nonetheless, we can derive a sufficient condition._”\n\n**Weakness 2**: We would like to clarify the goal of Sec. 6. It is a fundamental challenge to derive analytic learning dynamics for general deep nets. Sec. 6 provides an instructive setting for understanding generalization in deep learning dynamics where we could derive exact results on how stochastic collapse impacts generalization. We used this understanding to derive **new** testable predictions about the role of stochastic collapse in non-linear models. As explained in the paragraph following the introduction of the linear-teacher student analysis (lines 312 - 319):\n\n“_The key prediction is that a large learning rate induces stronger stochastic collapse, thereby regularizing the model complexity. Furthermore, remaining in a phase of larger learning rates for a prolonged period drives SGD closer to the invariant set. Consequently, when the learning rate is eventually dropped, overfitting in these specific directions is mitigated._”\n\nIn the subsequent paragraph, we proceeded to “_test this predicted mechanism_” on a “_VGG-16 and ResNet-16 [trained] on CIFAR-10 and CIFAR-100, respectively_”. And as demonstrated in figure 4, “_we found, as predicted, that models with a longer initial training phase collapse towards simpler invariant sets with fewer independent neurons_”.\n\nThis combination of deriving a theoretical prediction from analytic insights into a simple model, and then testing this prediction in more complex models, that are currently beyond the reach of theory, is a powerful route forward for advancing our conceptual understanding of the complex processes of deep learning; it combines theory and experiment in creative/synergistic ways, obtaining results that neither theory nor experiment alone could achieve.  Therefore this connection between the linear student-teacher setting and more complex models should, we believe, be considered a strength of our work rather than a weakness.  We never would have tested a connection between duration of large learning rate training and degree of collapse in Fig. 4 without our theory.  We also would never have tested for strong stochastic collapse to permutation invariant sets in Fig. 1 & 2 without our theory.  \n\n**Question 1**: While invariant sets of SGD are also invariant sets of GD, such sets can be attractive for SGD but **not** for GD.  Our surprising result is that if the diffusion of SGD is strong enough, SGD can be attracted to invariant sets that **repel** GD (examples: local maximum of double well potential in Sec. 4, saddle points with zero singular values in linear-student teacher setting in Sec. 6). Theorem 4.1 & 4.2 shows precisely how diffusion is essential for collapse of SGD to maxima/saddles; GD cannot similarly collapse because diffusion is absent; GD can only collapse to minima. \n\n**Question 2**:  Here are factors that we believe influence SGD’s likelihood of collapse:\n- Attractivity Rate: SGD is more inclined towards invariant sets with a higher attractivity rate, a key quantity defined in Theorem 4.2.\n- Quantity of Invariant Sets: SGD favors more numerous invariant sets; e.g. the number of permutation invariant sets within each layer of $n$ neurons is $O(n^2)$ while the number of sign invariant sets is only $O(n)$. This suggests collapse to permutation invariant sets is more common (see Fig. 1). \n\n**Question 3**: We did characterize invariant types by a norm of the parameters:\n- Sign Invariant Set: For a given neuron, its sign invariant set is defined by the vanishing norm of both its incoming and outgoing weight vectors (lines 235 - 236). \n- Permutation Invariant Set: For two neurons, their permutation invariant set is determined by the vanishing norm of the difference between their incoming and outgoing weight vectors (lines 227 - 229 and lines 236 - 239).\n\n**Question 4**: Verifying the sufficient condition of theorem 4.1 in deep nets is technically hard because: (1) we must check equation (7) holds for any unit normal vector and every point $\\theta$ in the invariant set; (2) we must compute the full Hessian matrix and second derivative of the diffusion matrix, within the entire invariant set, which is expensive.  Although we cannot directly verify theorem 4.1 for modern neural nets, we have verified it in simpler examples. We now include additional figures (see Figure S1 and S2 in the attached PDF) for the verification of Theorem 4.2 and 5.1, which thus also verifies Theorem 4.1.  But **most importantly** to circumvent these difficulties, we **directly** demonstrate stochastic collapse in realistic settings of ResNets and VGG (Fig. 1,2,4); this **direct** empirics constitutes a powerful demonstration of the prevalence of stochastic collapse in practice.  \n"}, {"review_id": "xRsqIN24bG", "reviewer": "Reviewer_nG1u", "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.", "questions": "Please see the weaknesses.", "limitations": "I did see any potential negative societal impact.\n\nThe main limitation of this work lies its technical contributions, as I mentioned in the weaknesses.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The discovered stochastic collapse is interesting for understanding optimization dynamics of deep learning. Some counterintutive insights are provided to our community.\n\n- Formal theoretical analysis with detailed proofs are presented for understanding stochastic collapse. It suggests that stochastic collapse can be precisely analyzed under certain assumptions. \n\n- 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?\n\n- 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?\n\n\nReference:\n\n[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.\n\n[2] Gur-Ari, G., Roberts, D. A., & Dyer, E. (2018). Gradient descent happens in a tiny subspace. arXiv preprint arXiv:1812.04754.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and the time you spent suggesting ways to make our work better. We discuss the weaknesses of our paper you raised:\n\n**Weakness 1**: We appreciate the referenced works and will incorporate them into our study. We have updated our related work section to cite these works. We welcome any further suggestions on references or other concerns. Here we will focus specifically on these two works and their relationship to our work:\n\nXie, Z., Sato, I., & Sugiyama, M.: We acknowledge the significance of this prior research, which connects the covariance matrix of SGD with the Hessian around local minima and  demonstrates that the escape rate from minima is linked to the Hessian eigenvalues or flatness. Furthermore, the authors briefly discussed that minima selection primarily occurs in a low-dimensional subspace spanned by the top Hessian eigenvectors. Our works both expands and diverges from this analysis in a couple of key ways:\n\n1. New concepts: We introduce the concept of sign/permutation invariant sets, bridging low-dimensionality of SGD dynamics with the sparsity of the network representations. In the vicinity of the invariant sets, the Heissan along orthogonal directions has near-zero eigenvalues, and hence we believe that the low-dimensional subspace discussed by Xie, Z. et al. is nested in one of these invariant sets we introduced. Verifying this hypothesis is interesting future work.\n\n2. Distinct Mechanism: We reveal the mechanism on why the low-dimensional invariant sets can be attractive. This mechanism is distinct from the previous analyses for two reasons. First, their work is based on the Kramers Escape Problem. Contrarily, in our investigation, the stochastic attractivity comes from the multiplicative nature of the noise and the effective entropy of invariant sets, which means that the escape rate doesn't bear direct relevance in our context. Second, while the paper briefly discussed SGD's biases towards low-dimensional subspaces, their analyses are built on the empirical fact that the Hessian has many near-zero eigenvalues with a few large ones. The theoretical explanation of this empirical fact is not provided. In contrast, our analyses are not tied to specific assumptions about the Hessian spectrum, marking a different analytical perspective.\n\nGur-Ari, G., Roberts, D. A., & Dyer, E: This insightful work shows that the Hessian splits into two slowly varying subspaces during the course of training and the gradient lives in the subspace spanned by the principal eigenvectors. They also highlight the exploration of the \"_transparent meaning_\" of these eigenvectors for future research. In line with this, our investigation suggests a potential correlation between the low-dimensional space and the invariant sets. We consider this connection a promising area for exploration and believe it could contribute to understanding the \"_transparent meaning_\" that the previous work hinted at.\n\n**Weakness 2**: We agree that designing improved algorithms based on our insights into stochastic collapse is a major goal for our future work. Nonetheless, the goal of this work was to highlight a novel, not yet well understood (and indeed undiscovered) implicit bias of SGD, stemming from the interplay between stochasticity and sparsity. This said, we are actively developing learning rate schedules that use the understanding of stochastic collapse defined and presented in this work.  We share some additional thoughts on this. As detailed at the end of Section 6 (lines 312-328) and Fig. 4, we found that a prolonged initial training phase with larger learning rates are beneficial for subsequent generalization due to a bias toward sparser representations. Inspired by our theory, we empirically counted the number of independent neurons as a measure of sparsity. We conjecture that this is a good indicator for when to drop down the learning rates. The prolonged training with larger learning rates is beneficial as long as it continuously reduces the network sparsity. One idea for the algorithm is to measure the number of independent neurons (can be done sporadically to reduce computational costs) and drop down the learning rates when this number plateaus. One technical challenge here is to define a single variable to quantify the entire network's simplicity rather than layer-by-layer. This challenge, along with benchmarking the algorithm across various datasets, demands substantial effort and extends beyond this paper's scope. Nevertheless, we recognize the significance of such algorithmic design and embrace it as an essential future endeavor, made possible by our theory/empirics! \n\nBesides, we would like to point out that our work contributes novel and quantitative insights into the effect of stochastic collapse on generalization, explored within a linear teacher-student setting.  Our analysis highlights an interesting tradeoff between maintaining enough gradient noise to induce stochastic collapse in directions with weak signal-to-noise ratios, while also reducing this noise to facilitate annealing in directions with robust signal-to-noise ratios. Also, it reveals the surprising finding that SGD can get attracted to **higher** training error saddle points (student with some zero singular values) that nevertheless have **lower** test error than local minima (student with all nonzero singular values) with lower training error.\n\nIn summary, we believe our paper provides both novel theoretical insights and a foundation for future practical applications. While we recognize the value and necessity of translating these insights into algorithmic designs, the work presented here primarily serves as a theoretical/scientific experiment work. Your feedback has been instrumental in refining our focus, and we hope this response successfully addresses your queries.\n"}, {"review_id": "1wVXqVzzx8", "reviewer": "Reviewer_DqpF", "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. ", "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  ?\n\nb) 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. ", "limitations": "The limitations are already addressed in the sections above. ", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "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. \n\nb) 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. \n\nb) 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. \n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and the suggestions to make our work better. We will discuss the weaknesses and questions of our paper you pointed out:\n\n**Weakness A**: We appreciate your valuable observation of the confusions that arose from our interchangeable use of noise introduced by minibatch SGD and label noise gradient descent (LNGD). This is important and we are committed to enhancing our paper in this regard. We firmly believe that the improvement we discuss below in each section will make our claims more solid:\n- Sec. 3: The existence of sign/permutation invariant sets is valid for both SGD and LNGD. Hence we can discuss these in the same scheme without any problem.\n- Sec. 4: Here, we argue that SGD can be seen as GD with a noisy term, leading us to study the stochastic gradient flow (SGF) as defined in Equation (4). Our theoretical analyses in Sec. 4 are firmly based on SGF. SGF can describe both SGD and LNGD, while their diffusion terms may have different structure.\n- Sec. 5,6: In these sections, we apply the theoretical results in Sec. 4 on SGF. We recognize that these sections need to be improved since we restrict ourselves to the case where the noise of SGF comes from the label noise, and hence we lose the generality of SGF. We used the label noise instead of the minibatch noise for the reason that label noise strength can be directly linked to the noise amplitude in SGF. However, in both these sections the analysis does not change if we let $\\zeta$ be the mini-batch gradient noise amplitude. The other source of confusion might come from our empirical results, which use label noise (for the same reason discussed above). We now also include a new empirical figure (Fig. S3 in the attached PDF), in which we use SGD minibatch noises to validate the outcomes previously demonstrated solely with label noise (Fig. 3). Note that Fig. 1, the top of Fig. 2 and Fig. 4 do **not** use label noise, indicating strong stochastic collapse for minibatch SGD.\n- Appendix: In the appendix of our revised manuscript, we will add a section on the similarities and differences (both theoretically and empirically) between SGD and LNGD.\n\nWe believe these improvements will greatly enhance our treatment of the subject and we hope they will positively influence your evaluation of our work.\n\n**Weakness B**: We thank the reviewer for pointing out the similarity between our analysis and the previous work by Andriushchenko et al., which we cited and discussed (lines 71-73) in our paper. Given the similarity to our work, we will add more discussions in our paper to address the substantial differences. The novelty of their work can be summarized as: 1. They revealed that large-step SGD dynamics have effective slow dynamics with multiplicative noise during loss stabilization. This allowed the authors to model the dynamics as a specific SDE, justified with empirics. 2. They theoretically analyzed diagonal linear networks, showing that the SDE has implicit bias toward sparser representations. 3. They conjectured that deep nonlinear networks also show this phenomenon, which is supported by the empirical results. 4. They argued that SGD first learns sparse features and then fits the data after the step size is decreased. With this foundation laid, we would like to emphasize two fundamental aspects that set our work apart:\n\n1. Our analysis goes far beyond the basic understanding of diagonal linear networks, offering a broader perspective by introducing the invariant sets with a more general theorem (Theorem 4.1).  In their work, the authors conjectured an implicit bias towards sparser features beyond their diagonal linear model, leaving this exploration open as an \"_exciting avenue for future work_\". We believe that our paper makes substantial progress along this avenue by understanding the implicit bias towards sparse representations as stochastic collapse to the invariant sets.  Furthermore, our theory is applicable to general deep neural nets. We also provide a theoretical framework that predicts the quantitative conditions for this attractiveness in a general setting. While applying this theorem to arbitrary neural nets is a challenging future work, our work utilized this theorem to quantitatively analyze the collapsing threshold of simple models. \n2. Our results shed light on the collapsing phenomena of weight vectors, a perspective that contrasts with their emphasis on neuron similarity based on activation patterns. Given that weight vector similarities often lead to similarity in activation patterns, our findings suggest that neural networks adhere to a stronger condition.\n\n**Question A**: We would like to point out that A1, A3, and A4 are standard assumptions used in the canonical papers [Saxe et al. 2014] and [Lampinen & Ganguli 2018], which establish the linear-teacher student model as a setting for studying generalization. As the reviewer rightly points out, the assumption that gradient noise aligns with input-output covariance (A2) is the new assumption we added such that we could derive analytical expressions for the theory when incorporating the new ingredient of stochasticity to the analysis. The technical difficulties for relaxing assumption A2 come from the problem that stochastic noises can perturb the eigenbasis, which adds an extra degree of complexity to the problem. Although we agree that relaxing those assumptions is theoretically interesting and important, we want to point out that our theory with those assumptions captures the key behavior of the general SGD dynamics. We now include additional empirical results (Fig. S3 in the attached PDF) which verifies that the key observed phenomena hold even without those assumptions and on minibatch SGD.\n\n**Question B**: The answer to this question is addressed in our response to Weakness A. In summary, the conclusion that large step sizes improve generalization does apply to mini-batch SGD without label noise (see Fig. 4 and Fig. S3 in the attached pdf)."}]

Author rebuttal:
We would like to thank all reviewers for their careful and detailed comments. We here attach a single-page pdf with the following three figures to address reviewers’ questions. 

1. Figure S1-  Empirics for our illustrative example. This shows that the quantitative prediction via Theorem 4.2 agrees with empirical results.
2. Figure S2 - Empirics for the single neuron model in Section 5. This shows that the quantitative prediction via Theorem 5.1 agrees with empirics results
3. Figure S3 - Empirics for the teacher-student setup in Section 6, where we verify that the key observed phenomena shown in Figure 3 hold even without the assumptions proposed in Section 6 and on SGD minibatch noises.

Please also see our individual responses to each reviewer for more information!

Lastly, it's worth pointing out that a crucial conceptual advance within our paper, though present, hasn't been prominently emphasized. This insightful perspective is as follows: **Anisotropic position dependent diffusion can cause SGD in deep learning to get attracted to higher training error saddle points that nevertheless have lower test error than local minima with lower training error.**  This statement is very surprising to many people we have discussed with, yet it is nevertheless true and our paper explains why.  For example, in our linear student-teacher setting, fixed points where the student has some of its singular values stochastically collapse to zero (corresponding to small data singular values that are not learned), are actually saddle points in the training error loss landscape, and these saddle points have **higher** training error than the global minimum (in which all the data singular values are learned for a full rank student), but **lower** test error than the global minimum (also, more generally, any network at a learning fixed point with identical neurons typically is at a saddle; lower training error can be achieved by specializing the neurons).  This is thus a concrete example of how convergence to a higher training error saddle point can, remarkably, help generalization.  We believe this may become an impactful,  well-known conceptual advance in deep learning due to our paper. We will expand upon this important loss landscape perspective in the revised camera-ready paper.  

We greatly appreciate the time and effort you spent reviewing our paper, which we are certain will make our work better. We hope you will consider our paper an important contribution to the NeurIPS community. We look forward to hearing from you!


Meta-review: This work characterizes the attractivity of stochastic gradient flow to some invariant sets. The paper considers different types of invariant sets such as sign invariant sets and permutation invariant sets, which correspond to sparse or low-rank sub-networks. They show that SGF ends up in these invariant sets during training and that these invariant sets are connected to better generalization.

The reviewers raised a number of concerns, most of which have been successfully addressed in the rebuttal, including the restriction to SGF (SGF is often used as a model to SGD) or the generality of the results (stronger statements in Sec. 5 & 6 are restricted to simpler settings). Another important concern raised by some reviewers is the lack of discussion of prior work. For instance, there is a very large body of literature on double-well potentials or Kramer's escape problems. Perhaps a good reference to start with is https://hal.science/hal-00604399v2/document. Finally, there is a lot of recent work that studies the diffusion properties of SGF, including some references mentioned by the reviewers. I would encourage the authors to expand the related work section to connect their work to these prior results.

Overall, the paper makes some nice contributions but it suffers a bit from a lack of clarity in the writing, which perhaps explains why the scores it receives are on the borderline accept side. I would encourage the authors to try to distill their main messages in a better way if possible. That said, I think this does not require a major revision and would therefore happily accept this paper.
Decision: Accept (poster)