Title: Stochastic Collapse: How Gradient Noise Attracts SGD Dynamics Towards Simpler Subnetworks
Venue: NeurIPS 2023
Accepted: yes
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
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

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.

1 Introduction

The remarkable performance of modern deep learning systems relies on a complex interplay between a training dataset, a network's architecture, and an optimization strategy. Contrary to traditional statistical learning theory, these highly expressive models exhibit impressive generalization capabilities on natural tasks, even without explicit regularization, despite having the capacity to memorize random data [1]. 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. It is widely believed that a central source of this implicit bias is the randomness introduced by stochastic gradient descent (SGD) [2]. In this work we 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 as 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).

We formulate a sufficient condition for stochastic attractivity -aprocess 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).

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

We demonstrate how this process of stochastic collapse influences generalization in a linear teacher-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 [3] 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-dependent 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 a 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 on a smoother version of the original loss allowing it to escape sharp local minima. Zhu et al. [11] further demonstrated how the anisotropic structure of SGD noise helps SGD escape efficiently. Xie et 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 or 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 [3-5], minimizes a measure of curvature among equivalent minima [7-9], or escapes from sharp local minima [10-13]. 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.

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 to sparse solutions via an implicit max-margin process. Woodworth et al. [15] demonstrated how an implicit L 1 penalty occurs for diagonal linear networks trained with gradient flow in the limit of small initialization. Nacson et 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 et al. [20] and Ziyin et al. [21] showed how weight decay induce a 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 a few top eigenvectors of the Hessian. Ziyin et al. [27] demonstrated how data augmentation can promote dimensional collapse of representations in self-supervised learning. 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 a clear underlying mechanism. In our work, we focus on a fundamental characteristic of neural networks trained with SGD that leads to novel empirical observations.

Our analysis 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 a fixed point [29]. In App. A we further discuss these works and other relevant works.

3 Invariant Sets of SGD Generated by Reflection Symmetries

Throughout this work, we consider a feed-forward network 2 ,

parameterized by θ = ( w (1) , b (1) , . . . , w ( m -1) , b ( m -1) , w ( m ) , b ( m ) ) ∈ R d , where w and b denote weights and biases respectively, the superscript indexes the m layers, and σ ( · ) is an activation function. The network is trained by stochastic gradient descent (SGD) over a training dataset { ( x 1 , y 1 ) , . . . , ( x n , y n ) } of size n , where x i ∈ R p and y i ∈ R c . The network parameters are randomly initialized at θ (0) and iteratively updated according to

where η > 0 is the learning rate, B ( t ) ⊆ [ n ] is a random 3 index set (mini-batch) of size β at step t , and ℓ ( θ ; x i , y i ) is a loss function. We let L ( θ ) = 1 n ∑ i ∈ [ n ] ℓ ( θ ; x i , y i ) 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-dimensional, non-convex loss landscapes to find generalizing solutions is still a 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-measurable set A ⊆ R d is an invariant set of SGD if given any initialization θ (0) ∈ A , all future iterates of SGD θ ( t ) for t ≥ 0 are contained within A almost surely, for any batch size β ∈ N , learning rate η > 0 , and mini-batches {B ( t ) ⊆ [ n ] : t ∈ N } .

From this definition, it is immediately clear that all of parameter space A = R d and any interpolating point A = { θ ∗ } such that L ( θ ∗ ) = 0 are invariant sets of dimension d and 0 respectively. However, the highly over-parameterized 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-rank) subnetworks.

Proposition 3.1 (Sign Invariant Sets) . Consider a hidden neuron p within layer l of a feed-forward neural network. Let ( w ( l ) in ,p , b ( l ) p ) and w ( l +1) out ,p denote the parameters directly incoming and outgoing from the neuron respectively 4 . If the non-linearity σ is origin-passing ( σ (0) = 0 ), then the axial subspace A = { θ ∈ R d | w ( l ) in ,p = 0 , b ( l ) p = 0 , w ( l +1) out ,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 (ReLU) [30], Leaky ReLU [31], Exponential Linear Unit (ELU) [32], 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 m 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 l of a feed-forward neural network. Let ( w ( l ) in ,p , b ( l ) p ) , ( w ( l ) in ,q , b ( l ) q ) denote the parameters directly incoming to the neurons, and w ( l +1) out ,p , w ( l +1) out ,q the parameters directly outgoing from the neurons. The affine subspace A = { θ ∈ R d | w ( l ) in ,p = w ( l ) in ,q , b ( l ) p = b ( l ) q , w ( l +1) out ,p = w ( l +1) out ,q } is an invariant set.

The subspace corresponding to a permutation invariant set represents the parameter space of a lowrank subnetwork obtained by constraining two neurons to be identical. All feed-forward networks possess this permutation invariance among the neurons in each hidden-layer. For a fully-connected network of depth m and width k there are ( m -1) ( k 2 ) 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

context, a symmetry is defined as any transformation of the parameter that preserves the network's input-output mapping, regardless of the input. As stated in the following theorem, any approximate linear symmetry defined by a symmetric or orthogonal matrix generates an invariant set:

Theorem 3.1 (Symmetry Induced Invariant Sets) . Let Q ∈ R d × d be a symmetric ( Q = Q ⊺ ) or orthogonal matrix ( QQ ⊺ = I d ). The affine subspace A = { θ ∈ R d | Qθ = θ } is an invariant set if the loss function ℓ ( θ ; x i , y i ) , for any i ∈ [ n ] , is approximately Q -symmetric around A , i.e., for any ϵ > 0 , there exists δ > 0 such that | ℓ ( Qθ ; x i , y i ) -ℓ ( θ ; x i , y i ) | < ϵd ( θ, A ) for any θ ∈ R d satisfying d ( θ, A ) < δ where d ( θ, A ) is the Euclidian distance between θ and A 5 .

Remarkably, the sign and permutation invariant sets defined earlier can be understood through this symmetry perspective. In both cases, the symmetry is defined by a symmetric and orthogonal matrix. This class of matrices describes the generalized notion of a reflection in d -dimensional space and the invariant sets are the axes of reflection, which can be at most ( d -1) -dimensional. 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 an 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 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.1 implies that the axial subspace A , where input and output weights are exactly zero, is 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 a further discussion on the connection between symmetry and invariant sets. We also discuss how our definition of invariant sets relates to other works [35-38] where symmetry and invariance are used to explore geometric properties of the loss landscape and the presence of critical points.

4 Gradient Noise Attracts SGD Dynamics towards Invariant Sets

To study the dynamics of SGD we will approximate 6 its trajectory with a stochastic differential equation (SDE), a common analysis technique applied in many works [5, 39-43]. SGD can be interpreted as full-batch gradient descent plus a per-step gradient noise term introduced by the random minibatch. Let ξ i ( θ ) represent the per-sample gradient noise at θ . It is easy to check that E [ ξ i ( θ )] = 0 for all θ . This motivates studying the SDE, which we refer to as stochastic gradient flow (SGF) ,

Here B t denotes a standard d -dimensional Wiener process and Σ( θ t ) = √ E [ ξ i ( θ ) ξ i ( θ ) ⊺ ] . The drift term of this SDE is the negative full-batch gradient -∇L ( θ t ) , while the diffusion is determined by a spatially-dependent diffusion matrix D ( θ t ) = η 2 β Σ( θ t )Σ( θ t ) ⊺ . 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 D for further discussions on the relationship between SGD and SGF. Additionally, SGF preserves all affine invariant sets of SGD:

Proposition 4.1 (Informal) . All affine invariant sets of SGD are also affine invariant sets of SGF . 7

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 set A ⊂ R d of a stochastic process { θ t ∈ R d : t ≥ 0 } is stochastically attractive 8 if for any ρ > 0 and ϵ > 0 , there exists δ > 0 such that for any θ ∈ R d with the Euclidian distance d ( θ, A ) < δ ,

Intuitively, a set is stochastically attractive if there always exists a close initial condition that guarantees all future iterates remain close to the set with high probability. Understanding when the dynamics are attracted by a particular invariant set becomes fundamental for understanding the implicit bias of SGD. To gain insight into this process lets consider a canonical example in one-dimension.

Geometric Brownian motion. Geometric Brownian Motion (GBM) is a linear stochastic process given by the SDE, dθ t = µθ t dt + ζθ t dB t , where µ ∈ R represents the drift rate, ζ ∈ R the volatility rate. Denote θ 0 ∈ R 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 GBMapproaches the invariant set, its movement diminishes. To account for this effect, we take the logarithmic transformation z t = log( θ t ) , where we assume without loss of generality that θ 0 > 0 . The stochastic process z t obeys the transformed SDE, dz t = ( µ -ζ 2 / 2 ) dt + ζdB t , which has an additional drift term -ζ 2 / 2 from Itô's lemma that attracts the process towards the invariant set. The total drift is now determined by a competition between the original drift rate µ 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 µ < ζ 2 / 2 the invariant set is attractive).

Stochastic attraction in one-dimension. Stochastic attractivity is a local property and thus we 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:

Here we use L ′ (0) = D (0) = D ′ (0) = 0 , which is a consequence of { 0 } being an invariant set. From Eq. 5 we can formulate a necessary/sufficient condition for stochastic attractivity in one-dimension:

Theorem 4.1 (Necessary/Sufficient Condition for Stochastic Attraction in One-Dimension) . Let A = { 0 } ∈ R be an invariant set of { θ t ∈ R : t ≥ 0 } obeying Eq. 3. Suppose L ′ ( θ ) : R → R , D ( θ ) : R → R are C 2 functions such that D ′′ (0) > 0 . Define rate of attractivity α = L ′′ (0) + 1 2 D ′′ (0) . A is stochastically attractive if α > 0 , while it is not stochastically attractive if α < 0 . 9

One of the most surprising implications of Theorem 4.1 is that SGF can potentially converge to a saddle-point or local maxima of a loss landscape, an observation previously made by Ziyin et al. [28]. To see this, recall that D ( θ ) ≥ 0 by definition, and D (0) = 0 as { 0 } is an invariant set. As a result, D ′ (0) = 0 and D ′′ (0) ≥ 0 by the continuity assumption of D . The collapsing condition, therefore crucially depends on the curvature of the loss function L at θ = 0 . When D ′′ (0) is strictly positive, the collapsing condition can still be satisfied with negative curvature provided that L ′′ (0) > -1 2 D ′′ (0) . Given that D is proportional to η/β , 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 double-well potential L ( θ ) = 1 4 ( θ 2 -µ ) 2 with multiplicative diffusion 10 , such that the dynamics are dθ t = -( θ 3 t -µθ t ) dt + ζθ t dB t where µ, ζ > 0 . The minima of the potential are located at θ = ± √ µ , while θ = 0 is a local maximum and forms an invariant set. This example is special as the dynamics have an analytical steady-state distribution p ss ( θ ) , 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 θ 0 > 0 and consider the logarithm process z t = log( θ t ) . 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 p ss ( θ ) is given by a Gibbs distribution p ss ( θ ) ∝ e -κ Ψ( θ ) for a modified potential 11 Ψ( θ ) = θ 2 / 2 -(2 µ/ζ 2 -2) log( θ ) with constant κ = 2 /ζ 2 and partition function Z = ∫ ∞ -∞ e -κ Ψ( θ ) dθ . When µ ≤ ζ 2 / 2 , the partition function diverges and p ss ( θ ) collapses to a Dirac delta distribution at θ = 0 . This transition agrees with the collapsing condition from Theorem 4.1.

1.00

Theorem 4.1

+√

+√

+√μ

3

0.75

Steps

=√2u

0

0

0

0.50

0

10000

0.25

20000

口

40000

0.00

80000

0.00 0.25 0.50 0.75 1.00

0.00 0.25 0.50 0.75 1.00

0.00 0.25 0.50 0.75 1.00

0.0

0.8

1.6

2.4

Time (t)

1e4

Time (t)

1e4

Time (t)

1e4

2

Stochastic attraction in high-dimensions. To extend the collapsing condition derived in Theorem 4.1 to high-dimensional cases, a natural idea would be to 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 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:

Theorem 4.2 (A Sufficient Condition for Stochastic Attraction in High-Dimensions) . Let A ⊂ R d be a d A -dimensional affine subset, and a stochastic process { θ t ∈ R d : t ≥ 0 } obey Eq. 3 in A c , open c -neighborhood of A with some c > 0 . Suppose L : A c → R is a C 3 -function whose first and second-order derivatives are L -Lipschitz continuous. D : A c → R d × d 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 √ D : A c → R d × d are L -Lipschitz continous. Let D ⊥ = P ⊥ DP ⊥ where P ⊥ : R d → R d projects to the normal space of A . If there exists δ > 0 such that ∇ 2 ˆ n ˆ n ⊺ D ˆ n ⊺ > δ and

for any unit normal vector ˆ n ∈ R d perpendicular to A and θ ∈ A , then A is stochastically attractive.

To prove Theorem 4.2 (see App. E), we consider a family of distance measures between θ t and the invariant set A . Eq. 6 ensures that one of the distance measures is super-martingale in a neighborhood of the invariant set, allowing us to apply a version of Doob's maximal inequality (Lemma 1 in [44]).

5 Attractivity of Invariant Sets in Deep Neural Networks

In this section we explore theoretically and empirically the stochastic attractivity of invariant sets in deep neural networks.

Sign invariant sets. Sign invariant sets correspond to the parameter space of a sparse subnetwork obtained by removing a hidden neuron. Here, we demonstrate the stochastic collapse to a sign invariant set in a minimal model of neural networks-a scalar single neuron model: f ( x ; w 1 , w 2 ) = w 2 σ ( w 1 x ) with w 1 , w 2 , x ∈ R . Suppose this model is trained via label-noise gradient descent and using the L 2 -loss L ( w 1 , w 2 ) = 1 n ∑ i ∈ [ n ] ( y i -f ( x i ; w 1 , w 2 )) 2 on a non-empty data

Steps

CIFAR-100  ResNet-18

CIFAR-10  VGG-16

Layer with little stochastic collapsing

Layers with strong stochastic collapsing

Win:conv5(256)

Wout:conv6(256)

Win:conv11(512)

Wout:conv12(512)

Win: fc14 (512)

Wout:fc15(512)

2.

2.

1.

1.

0.5

0.

2.

2.

1.

1.

0.5

0.

̸

set { ( x 1 , y 1 ) , · · · , ( x n , y n ) } such that ∑ i ∈ [ n ] x 2 i = 0 . Exploiting Theorem 4.1, we can determine the attractivity of the sign invariant set A = { (0 , 0) } for this network.

Theorem 5.1 (Informal) . Let { ( w 1 ,t , w 2 ,t ) ∈ R : t ≥ 0 } be a process obeying label-noise gradient flow with L 2 -loss of a single scalar neuron in a neighborhood U of A = { (0 , 0) } , with learning rate η > 0 and noise amplitude ζ > 0 . Suppose the activation function σ : R → R is smooth satisfying

2 .

The term | ∑ i ∈ [ n ] x i y i | ∑ i ∈ [ n ] x 2 i can be thought of as the signal of the dataset. Thus, Theorem 5.1 states that stochastic attractivity is determined by balancing the signal and noise - an idea we will see later in Sec. 6 as the origin of generalization benefits from stochastic collapse. We leave the poof in App. F.

Permutation invariant sets. Permutation invariant sets correspond to subnetworks with fewer unique hidden neurons. An important question is to what extent are permutation invariant sets attractive - as this would indicate an implicit bias towards a 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 can be attractive in deep neural networks via experiments using VGG-16 [47] 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 a permutation invariant set, we compute the normalized pairwise distance between parameters for neurons from the same layer, defined by

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 ResNet-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 non-vanishing neurons by clustering the concatenated vector of incoming

Layer 1

Layer 2

Layer 3

Layer 4

and outgoing parameters based on normalized distance. Two neurons were considered identical if their distance in parameter space was 10% of their norms, a stringent criterion given the highdimensionality 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 conv7 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 in Proposition H.2 and H.1. We believe, this stricter constraint is the source of layer-to-layer 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 a further discussion.

6 How Stochastic Collapse Can Improve Generalization

To understand how stochastic collapse can benefit generalization, we will theoretically analyze learning dynamics in a teacher-student model for linear networks, a well established framework for studying generalization [50-52]. This analysis not only provably connects the phenomenon of stochastic collapse to generalization benefits, but also suggests a 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 a two-layer linear neural network. We build on 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 a student-teacher setting. In our 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. I.

A teacher neural network with a 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 Σ yx then drives the learning dynamics of a two-layer student with composite weight matrix ˆ W = W 2 W 1 . We analyze the learning dynamics of the student under a set of four assumptions precisely stated in App. I. 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 inputoutput 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 = W 2 W 1 can be decomposed into independent learning dynamics for each singular value ˆ s i of ˆ W . An associated singular value ˜ s i of Σ yx drives the dynamics of ˆ s i via the SDE,

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 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 [53] and larger data singular modes, when learned, help the student generalize, while smaller singular modes, when learned, impede generalization by

ζ

ζ

ζ

ζ

ζ

4

4

overfitting to noise in the training data [51]. Both these insights will apply in our analysis of SGF. However, we will also uncover a 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. I reveals that the student singular values ˆ s i evolve according to the process:

Theorem 6.1 (Stochastic Student Dynamics) . Under assumptions A1 - A4, the dynamics of ˆ s i ( t ) given by Eq. 7 are governed by the stochastic process,

In terms of the per layer balanced student weight w i = √ ˆ s i , the above process can be thought of as originating from the quartic loss landscape ℓ i ( w i ) = 1 4 ( w 2 i -˜ s i ) 2 , equivalent to the example presented in Sec. 4. Thus, there exists an invariant set for each student singular value corresponding to the affine space ˆ s i = w i = 0 where both drift and diffusion terms vanish. Stochastic collapse to this invariant set would thus induce a 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 ˜ s i . To confirm stochastic collapse to this invariant set, we derive:

Corollary 6.1 has important implications. First, if ˜ s i < ηζ 2 2 , the distribution will exhibit stochastic collapse, converging to a delta distribution at the origin, consistent with Theorem 4.1 (Fig. 5: 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 η will continue to drive stochastic collapse of small singular modes. After a subsequent learning rate drop, some of these modes will cease to obey the collapsing condition and as a consequence, the associated invariant set will repulse 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 η , because ˜ s i < ηζ 2 2 , and then learning only the subset with large ˜ s 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-noise ratios, resulting in improved generalization after the learning rate drop as shown in Fig. 5: middle left panel. In summary, the linear teacher-student model highlights that with an appropriate learning rate schedule, a noisy student can effectively avoid overfitting through stochastic collapse and achieve superior generalization compared to a noiseless student.

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 a long period of time (even after training and test loss have plateaued) helps generalization after a subsequent learning rate drop. The key prediction is that a large learning rate induces stronger

CIFAR-100  ResNet-18

CIFAR-10  VGG-16

Fraction of Independent Nerons

Fraction of Independent Nerons

Layer 1

Layer 2

Layer 3

Layer 4

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.

To test this predicted mechanism, we trained VGG-16 and ResNet-16 on CIFAR-10 and CIFAR-100, respectively. The training loss and test accuracy had already plateaued at 10 4 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 an 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. 6: right column).

7 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 a 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 explain the importance of large learning rates early in the training process.

Limitations and future directions. Activation function with additional continuous symmetries, such as ReLU, might have a 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 of general affine invariant sets is an important goal for future work. Extending our analytic results from the continuous SGF to the discrete SGD updates is an interesting theoretical direction. Lastly, designing new optimization algorithms based on our insights into stochastic collapse is a major goal for our future work.

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 Fellowship for support. S.G. thanks the James S. McDonnell and Simons Foundations, NTT Research, and an 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 to a saddle-point introduced by symmetry. A.Y. 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

Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Commun. ACM , 64(3):107-115, feb 2021. ISSN 0001-0782. doi: 10.1145/3446776. URL https://doi.org/10.1145/3446776 .

Fengxiang He, Tongliang Liu, and Dacheng Tao. Control batch size and learning rate to generalize well: Theoretical and empirical evidence. 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, Inc., 2019. URL https://proceedings.neurips.cc/ paper_files/paper/2019/file/dc6a70712a252123c40d2adba6a11d84-Paper.pdf .

David Barrett and Benoit Dherin. Implicit gradient regularization. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id= 3q5IqUrkcF .

Jonas Geiping, Micah Goldblum, Phil Pope, Michael Moeller, and Tom Goldstein. Stochastic training is not necessary for generalization. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id=ZBESeIUB5k .

Daniel Kunin, Javier Sagastuy-Brena, Lauren Gillespie, Eshed Margalit, Hidenori Tanaka, Surya Ganguli, and Daniel LK Yamins. Limiting dynamics of sgd: Modified loss, phase space oscillations, and anomalous diffusion. arXiv preprint arXiv:2107.09133 , 2021.

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 Thirty Fourth Conference on Learning Theory , volume 134 of Proceedings of Machine Learning Research , pages 2315-2357. PMLR, 15-19 Aug 2021. URL https: //proceedings.mlr.press/v134/haochen21a.html .

Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an ornstein-uhlenbeck like process. In Jacob Abernethy and 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://proceedings.mlr.press/v125/blanc20a.html .

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 J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , volume 34, pages 27449-27461. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/paper_files/ paper/2021/file/e6af401c28c1790eaef7d55c92ab6ab6-Paper.pdf .

Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss? -a mathematical framework. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id=siCt4xZn5Ve .

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 on Machine Learning , volume 80 of Proceedings of Machine Learning Research , pages 2698-2707. PMLR, 10-15 Jul 2018. URL https://proceedings.mlr.press/v80/ kleinberg18a.html .

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 on Machine Learning , volume 97 of Proceedings of Machine Learning Research , pages 7654-7663. PMLR, 09-15 Jun 2019. URL https://proceedings.mlr. press/v97/zhu19e.html .

Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. arXiv preprint arXiv:2002.03495 , 2020.

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:1609.04836 , 2016.

Daniel Kunin, Atsushi Yamamura, Chao Ma, and Surya Ganguli. The asymmetric maximum margin bias of quasi-homogeneous neural networks. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum? id=IM4xp7kGI5V .

Blake Woodworth, Suriya Gunasekar, Jason D. Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized 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://proceedings.mlr.press/v125/woodworth20a. html .

Mor Shpigel Nacson, Kavya Ravichandran, Nathan Srebro, and Daniel Soudry. Implicit bias of the step size in linear diagonal neural networks. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 16270-16295. PMLR, 17-23 Jul 2022. URL https://proceedings.mlr.press/v162/nacson22a.html .

Scott Pesme, Loucas Pillaud-Vivien, and Nicolas Flammarion. Implicit bias of sgd for diagonal linear networks: a provable benefit of stochasticity. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , volume 34, pages 29218-29230. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/paper_files/paper/2021/file/ f4661398cb1a3abd3ffe58600bf11322-Paper.pdf .

Maksym Andriushchenko, Aditya Varre, Loucas Pillaud-Vivien, and Nicolas Flammarion. Sgd with large step sizes learns sparse features. arXiv preprint arXiv:2210.05337 , 2022.

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 on Learning Theory , volume 178 of Proceedings of Machine Learning Research , pages 2127-2159. PMLR, 02-05 Jul 2022. URL https://proceedings.mlr.press/v178/vivien22a.html .

Daniel Kunin, Jonathan Bloom, Aleksandrina Goeva, and Cotton Seed. Loss landscapes of regularized linear autoencoders. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors,

Proceedings of the 36th International Conference on Machine Learning , volume 97 of Proceedings of Machine Learning Research , pages 3560-3569. PMLR, 09-15 Jun 2019. URL https://proceedings.mlr.press/v97/kunin19a.html .

Liu Ziyin, Botao Li, and Xiangming Meng. Exact solutions of a deep linear network. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. URL https://openreview.net/forum?id= X6bp8ri8dV .

Tomer Galanti, Zachary S Siegel, Aparna Gupte, and Tomaso Poggio. Sgd and weight decay provably induce a low-rank bias in neural networks. arXiv preprint arxiv.org/abs/2206.05794 , 2022.

Zihan Wang and Arthur Jacot. Implicit bias of sgd in l 2 -regularized linear dnns: One-way jumps from high to low rank. arXiv preprint arXiv:2305.16038 , 2023.

Arthur Jacot. Implicit bias of large depth networks: a notion of rank for nonlinear functions. arXiv preprint arXiv:2209.15055 , 2022.

Minyoung Huh, Hossein Mobahi, Richard Zhang, Brian Cheung, Pulkit Agrawal, and Phillip Isola. The low-rank simplicity bias in deep networks. arXiv preprint arXiv:2103.10427 , 2021.

Guy Gur-Ari, Daniel A. Roberts, and Ethan Dyer. Gradient descent happens in a tiny subspace, 2018. URL https://arxiv.org/abs/1812.04754 .

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 https://openreview.net/forum?id=3zSn48RUO8M .

Liu Ziyin, Botao Li, James B Simon, and Masahito Ueda. Sgd with a constant large learning rate can converge to local maxima. arXiv preprint arXiv:2107.11774 , 2021.

Liu Ziyin, Botao Li, Tomer Galanti, and Masahito Ueda. The probabilistic stability of stochastic gradient descent. arXiv preprint arXiv:2303.13093 , 2023.

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

Andrew L Maas, Awni Y Hannun, Andrew Y Ng, et al. Rectifier nonlinearities improve neural network acoustic models. In Proc. icml , volume 30, page 3. Atlanta, Georgia, USA, 2013.

Djork-Arné Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). arXiv preprint arXiv:1511.07289 , 2015.

Quoc V. Le Prajit Ramachandran, Barret Zoph. Searching for activation functions, 2018. URL https://openreview.net/forum?id=SkBYYyZRZ .

Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08415 , 2016.

Daniel Kunin, Javier Sagastuy-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 Representations , 2021. URL https://openreview. net/forum?id=q8qLAbQBupm .

Johanni Brea, Berfin Simsek, Bernd Illing, and Wulfram Gerstner. Weight-space symmetry in deep networks gives rise to permutation saddles, connected by equal-loss valleys across the loss landscape. arXiv preprint arXiv:1907.02911 , 2019.

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 9722-9732. PMLR, 18-24 Jul 2021. URL https:// proceedings.mlr.press/v139/simsek21a.html .

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-843, mar 2008. ISSN 0899-7667. doi: 10.1162/neco.2007.12-06-414. URL https://doi.org/10.1162/ neco.2007.12-06-414 .

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 https://proceedings.mlr. press/v70/li17f.html .

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 https://proceedings.mlr.press/v48/mandt16.html .

Stanisław Jastrz˛ ebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in sgd. arXiv preprint arXiv:1711.04623 , 2017.

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: 10.1109/ITA.2018.8503224.

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://proceedings. mlr.press/v119/ali20a.html .

Harold J Kushner. Stochastic stability and control . Academic Press, 1967.

Hendrik Anthony Kramers. Brownian motion in a field of force and the diffusion model of chemical reactions. Physica , 7(4):284-304, 1940.

Nils Berglund. Kramers' law: Validity, derivations and generalisations. Markov Processes And Related Fields , 19(3):459-490, October 2013. URL https://hal.science/hal-00604399 . 26 pages.

Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.1556 , 2014.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , June 2016.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. University of Toronto , 2009.

Madhu S. Advani, Andrew M. Saxe, and Haim Sompolinsky. High-dimensional dynamics of generalization error in neural networks. Neural Networks , 132:428-446, 2020. ISSN 0893-6080. doi: https://doi.org/10.1016/j.neunet.2020.08.022. URL https://www.sciencedirect. com/science/article/pii/S0893608020303117 .

Andrew K. Lampinen and Surya Ganguli. An analytic theory of generalization dynamics and transfer learning in deep linear networks. In International Conference on Learning Representations , 2019. URL https://openreview.net/forum?id=ryfMLoCqtQ .

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, Inc., 2019. URL https://proceedings.neurips.cc/paper_files/paper/2019/file/ cab070d53bd0d200746fb852a922064a-Paper.pdf .

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:1312.6120 , 2013.

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: https://doi.org/10.1016/0893-6080(89)90014-2. URL https://www.sciencedirect. com/science/article/pii/0893608089900142 .

Kenji Fukumizu and Shun-ichi Amari. Local minima and plateaus in hierarchical structures of multilayer perceptrons. Neural networks , 13(3):317-327, 2000.

Shun-ichi Amari, Hyeyoung Park, and Tomoko Ozeki. Singularities affect dynamics of learning in neuromanifolds. Neural computation , 18(5):1007-1065, 2006.

Christopher DiMattina and Kechen Zhang. How to modify a neural network gradually without changing its input-output functionality. Neural computation , 22(1):1-47, 2010.

Bernt Oksendal. Stochastic differential equations: an introduction with applications . Springer Science & Business Media, 2013.

Florent Benaych-Georges 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: https://doi.org/10.1016/j.jmva.2012.04.019. URL https://www.sciencedirect.com/science/article/pii/S0047259X12001108 .

Olivier Roy and Martin Vetterli. The effective rank: A measure of effective dimensionality. In 2007 15th European Signal Processing Conference , pages 606-610, 2007.

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. In particular, their work conjectured that the parameter dependent noise introduced by SGD has an implicit bias effect'toward parameters where the noise covariance is smaller'. They studied this effect in two-layer quadratically-parameterized model introduced by [15]. Our work generalizes this effect to a broad range of models by introducing invariant sets and demonstrates that this implicit bias towards these regions is 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 a fixed point [29]. Ziyin et al. [28] studied the behavior of SGD in one-dimensional quadratic and quartic losses and showed that SGD can converge to local maxima, escape from saddle-points slowly, and prefer sharp minima due to the gradient noise. A strength of their analysis is 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 it to a simplicity bias that can benefit generalization. Concurrent to our work, Ziyin et al. [29] proposes a new notion of stability for SGD near a 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.

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 4.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 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. Second, 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.

B Further Discussion on Invariant Sets of SGD

In our work, we 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 of 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-57] 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 ∈ A , thus it suffices to show that ∂ℓ ∂θ ∈ A and is independent of the training data. Denote the activation at layer l as h ( l ) = σ ( w ( l ) T h ( l -1) + b ( l ) ) . The gradients associated with the outgoing weights are,

If θ ∈ A , ∂ℓ ∂w ( l +1) out ,p = 0 . The gradients associated with the incoming weights are,

If θ ∈ A , ∂ℓ ∂w ( l ) in ,p = 0 . The gradients associated with the biases are,

Proof. We can directly see from Eq. 9-11 that if θ ∈ A , ∂ℓ ∂w ( l +1) out ,p = ∂ℓ ∂w ( l +1) out ,q , ∂ℓ ∂w ( l ) in ,p = ∂ℓ ∂w ( l ) in ,q and ∂ℓ ∂b ( l ) p = ∂ℓ ∂b ( l ) q . Therefore, ∂ℓ ∂θ ∈ A .

B.2 Proof of Theorem 3.1

Proof. It suffices to show that ∇ ℓ ( θ ; x, y ) ∈ A for all θ ∈ A and training data. In the following, we use the notation of ℓ ( θ ) as a shorthand for ℓ ( θ ; x, y ) . First, we show ∇ ℓ Q ( θ ) = ∇ ℓ ( θ ) for θ ∈ A , where ℓ Q ( θ ) = ℓ ( Qθ ) . The derivative of ℓ Q along an arbitrary unit vector n is given by

If n is a tangent vector of A , the second term is 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

∇ ℓ Q ( θ ) = ∇ ℓ ( θ ) . Then, for any θ ∈ A , ∇ ℓ ( θ ) = ∇ ℓ Q ( θ ) = Q ⊺ ∇ ℓ ( Qθ ) . Since Q is symmetric ( Q = Q ⊺ ) or orthogonal ( QQ ⊺ = I d ), we obtain

which implies ∇ ℓ ( θ ; x, y ) ∈ A . Hence, regardless of the sequence of mini-batches, starting from θ (0) ∈ A , the SGD trajectory remains within A , demonstrating A is an invariant set.

Theorem 3.1 provides an alternative proof of Proposition 3.1 and Proposition 3.2, simply identify a 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 a softmax layer, the columns of the weight matrix and the vector of biases immediately preceding this layer can be shifted by a constant without changing the outputs. As a result, the dynamics for these parameters are constrained to an affine space with a constant sum determined by their initialization. This subspace is an invariant set. Similarly, if the training data has a 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 of the network. Again, this constrains the learning dynamics to an affine space, which is an invariant set. In general, any translational symmetry should generate invariant sets.

C Invariant sets of SGF

C.1 Definition of invariant sets for continuous processes

Here, we provide the formal definition of invariant sets for a continuous process and a formal statement of Proposition 4.1.

Definition C.1 (Invariant sets of continuous processes) . For a given stochastic process { θ t ∈ R d : t ≥ 0 } , a Borel-measurable set A ⊂ R d is defined as invariant if for every initial point θ 0 in A , the probability that the process stays in A for all t ≥ 0 is 1, i. e.,

C.2 Proof of Proposition 4.1

Proposition 4.1 (Formal) . Consider an SGD process { θ ( t ) ∈ R d : t ∈ N } given by Eq. 2, where the gradients ∇ ℓ ( · ; x i , y i ) : R d → R d in Eq. 2 are Lipschitz continuous and bounded for any i ∈ [ n ] . If A ⊆ R d is an affine invariant set of SGD { θ ( t ) : t ∈ N } , then A is also an invariant set of an SGF process { θ t : t ≥ 0 } given by Eq. 3.

proof. By assumption, ∇ ℓ ( θ ; x i , y i ) is L -Lipschitz in θ and ∥∇ ℓ ( θ ; x i , y i ) ∥ ≤ L with some constant L > 0 . Then we can see that ∇L is also Lipschitz continuous and bounded as follows:

Let ¯ ℓ ( θ, x i , y i ) = ℓ ( θ, x i , y i ) -L ( θ, x i , y i ) . ∇ ¯ ℓ ( · , x i , y i ) is clearly 2 L -Lipschitz continuous and is bounded by 2 L for any i ∈ [ n ] . Then,

This implies that ∥ D ( θ 1 ) -D ( θ 2 ) ∥ ∞ ≤ 4 L 2 ∥ θ 1 -θ 2 ∥ 2 , where ∥ · ∥ ∞ denotes the operator norm. On the other hand, the Powers-Stormer inequality gives,

where ∥ · ∥ F and ∥ · ∥ 1 denote the Frobenius norm and Schatten 1-norm, respectively. Thus, √ D ( θ ) is Lipschitz continuous with respect to the Frobenius norm. This establishes the existence of a unique strong solution to the SDE in Eq. 3. It now suffices to show that the SDE has a solution { θ t : t ≥ 0 } satisfying Eq. 12. Let P : R d → R d be the projection operator to the affine subset A , and we construct a projected SDE:

By construction, this projected SDE also has a unique solution ˜ θ t , which clearly satisfies the condition Eq. 12. Moreover, given that A is an invariant set of the original SGD process, we have P ∇ ℓ ( θ ; x i , y i ) = ∇ ℓ ( θ ; x i , y i ) for all i ∈ B and x ∈ A . This means that { ˜ θ t : t ≥ 0 } satisfies Eq.12 almost surely. Hence, we conclude that A is also an invariant set of the solution to Eq.3.

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.

We factorize the right-hand side into the negative full-batch gradient and the deviation from the full-batch gradient,

Let ξ B ( θ ) represent the gradient noise at θ . Assuming we sample with replacement, the gradient noise can be disentangled into a sum of i.i.d. per-sample gradient noises: ξ B ( θ ) = √ η β ∑ i ∈B ξ i where ξ i = ∇ ℓ ( θ ; x i , y i ) -∇L ( θ ) . It is easy to check that E [ ξ i ( θ )] = 0 for all θ . Assuming that the gradient noises are gaussian random variables, Eq. 14 turns into

where ∆ W ∼ N (0 , ηI ) and the diffusion matrix D ( θ ( t ) ) = 1 2 Var ( ξ B ) = η 2 β Var ( ξ i ) . Eq. 15 is a Euler-Maruyama discretization of the following continuous-time stochastic differential equation, which we denote as stochastic gradient flow (SGF) ,

The diffusion matrix can be expressed as:

The diffusion matrix can be decomposed into a product of a parameter-independent magnitude term D m and a parameter-dependent shape term D s ( θ ) such that D ( θ ) = D m D s ( θ ) , where

The optimization hyperparameters learning rate and batch size affect the magnitude term D m , while the network architecture, the training data, and the loss function affect the shape term D s ( θ ) . For generic network architectures, determining the exact form of the shape term is quite complicated. For the purpose of analytical assessments, we sometimes pivot to 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, we 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 on this, what happens 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 ζ 2 Gaussian label noise per training label, denoted as ξ i , it results in an update rule we refer to as label-noise gradient descent (LNGD), √

This is a discretization of the following continuous-time stochastic differential equation, which we denote as label-noise gradient flow (LNGF) ,

where the diffusion matrix of LNGF can be expressed as,

which can also be decomposed into a magnitude term and a shape term,

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 is much simpler, which allows us to get analytical results in Sec. 5 and 6. 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 a more comprehensive grasp of the more complex dynamics inherent to mini-batch noise.

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. Let A ⊂ R d be a d A -dimensional affine subset, and a stochastic process { θ t ∈ R d : t ≥ 0 } obey the following SDE in A c , open c -neighborhood of A with some c > 0 :

where { W t : t ≥ 0 } is the d -dimensional standard Wiener process. Here µ : A c → R is a L -Lipschitz C 2 -function with its first-order derivatives being L -Lipschitz as well. D : A c → R m × m is the diffusion matrix such that the second-order derivatives of its elements are L -Lipschitz continuous. We further assume that all the elements of √ D ( θ ) are L -Lipschitz continuous. Let D ⊥ = P ⊥ DP ⊥ where P ⊥ : R d → R d projects to the normal space of A . Suppose A is an invariant set of θ t . A is stochastically attractive if there exists δ > 0 such that for any unit normal vector ˆ n ∈ R d perpendicular to A and any θ ∈ A ,

Proof. We exploit a local version of Doob's maximal inequality (Lemma. 1 in [44]). By the Lipschitz property of µ and √ D , the SDE Eq. 25 has a unique strong solution up to the exit time τ c = inf { t ≥ 0 : θ t / ∈ A c } . The stopped process { ˜ θ t = θ t ∧ τ c : t ≥ 0 } is a continuous strong Markov process. It suffices to show that A is stochastically attractive for ˜ θ t . We first show that ˜ A l λ ( x ) ≤ 0 for any x ∈ A ϵ \ A with some 0 < ϵ < c , and λ > 0 , where ˜ A is the infinitesimal operator of the stochastic process ˜ θ t and A ϵ denotes the open neighborhood of A within distance ϵ > 0 . We consider the Euclidian distance l ( ˜ θ t ) between the process ˜ θ t and the affine invariant set A .

The SDE for l λ ( ˜ θ t ) is given by,

This means that for any x ∈ A c ,

By the fact that A is an invariant set, ˜ A l 2 ( x ) = 0 holds for any x ∈ A . Therefore, we obtain Tr D ⊥ ( x ) = 0 . Since D ⊥ is positive semi-definite, D ⊥ ( x ) = 0 for any x ∈ A . Furthermore, since ˜ A f ( x ) = 0 with f ( θ t ) = P ⊥ θ t for any x ∈ A , P ⊥ µ ( x ) = 0 for any x ∈ A .

Going back to Eq.28, we obtain

where

Notice that G ( x ) > 0 for any x ∈ A ϵ \ A with small enough ϵ > 0 because Eq. 26 holds and the second derivatives of D ( x ) are L -Lipschitz. We want to show that there exists λ > 0 such that for any x ∈ A ϵ \ A , F ( x ) 2 G ( x ) -1 + λ/ 2 < 0 with some λ > 0 . Parameterize x by x ( ξ ) = x ∥ + ξ ˆ n , where x ∥ ∈ A and ˆ n is the unit normal vector perpendicular to A . From L'Hôpital's rule, we have,

The second derivatives of F can be calculated as follows:

Substituting this, we obtain

By the assumed Lipschitz property of the derivatives of µ and D , ∇ 2 ˆ n F ( x ) and ∇ 2 ˆ n G ( x ) are L -Lipschitz in A ϵ \ A with small enough ϵ > 0 . Therefore, for any 0 < ξ < ϵ ,

Therefore F ( x ∥ + ξ ˆ n ) /G ( x ∥ + ξ ˆ n ) is 2 L (3 δ -1 -1) -Lipschitz as a function of ξ in A ϵ \ A . This implies that there exists ϵ > 0 such that for any x ∈ A ϵ \ A , F ( x ) / 2 G ( x ) < 1 -δ/ 2 and hence ˜ A ( l δ ( x )) ≤ 0 . Therefore l δ ( x ) is in the domain of ˜ A A ϵ and ˜ A A ϵ ( l δ ( x )) = ˜ A ( l δ ( x )) ≤ 0 , where ˜ A A ϵ is the infinitesimal operator of the stopped process ˜ θ t ∧ τ with τ = inf t ≥ 0 { l ( θ t ) ≥ ϵ } . By the local version of Doob's maximal inequality, for any ϵ ′ < ϵ ,

Equivalently, for any ϵ ′ < ϵ 1 /δ ,

By taking l ( 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 ˜ θ t and so is for θ t .

E.2 Proof of Theorem 4.1

The sufficient condition is a direct corollary of Theorem 4.2, and thus we show the necessary condition here. By assumption, L ′′ ( x ) + 1 2 D ′′ ( x ) < -δ at x = 0 with some δ > 0 . It suffices to show that there exists ϵ > 0 such that the hitting time

is finite almost surely for any x ∈ R with initialization θ 0 = x .

We choose ϵ > 0 small enough so that D ( x ) > 0 for any x ∈ ( -ϵ, ϵ ) \{ 0 } . Then we can consider another stochastic process Y t = 2 -1 / 2 ∫ θ t ∧ τ ϵ/ 2 D -1 / 2 ( x ) dx , which follows the SDE below

for any 0 ≤ t < τ . By further taking ϵ small,

Thus, Y t > δt + Y 0 + W t . Therefore,

Since δ > 0 , the right-hand side is almost surely finite with any value of Y 0 . This immediately implies that A is not stochastically attractive.

F The attractivity of the sign invariant set of single scalar neuron

We here investigate when a neuron may collapse towards the sign invariant set. Consider a neural network consisting of only one neuron with scalar input and output: f ( x ; w 1 , w 2 ) = w 2 σ ( w 1 x ) , where x, w 1 , w 2 ∈ R and σ represents the activation function. We assume that σ ( x ) is a smooth Lipschitz continuous function with σ (0) = 0 . Let D = { ( x 1 , y 1 ) , · · · , ( x n , y n ) } denote the training data set where x i , y i ∈ R , i ∈ [ n ] . We consider the loss function L given by

where ζ i,t is the label noise which we assume to be increments of the n -dimensional Wiener process. The gradient flow of the loss is given by

This can be rewritten as a pair of SDEs.

where ζ > 0 is a constant representing the amplitude of the gradient noise, η > 0 is the learning rate, and B i,t denotes the n -dimensional standard Wiener process. Note that the rescaling factor 1 / √ n in the diffusion term is introduced to avoid the data size n influencing the SDE's behavior.

F.1 Proof of Theorem 5.1

As a preparation, we first show a modified version of Theorem E.1 to incorporate stopped processes.

Lemma F.1. Let A ⊂ R d denote an affine subset and A c be an open neighborhood around A defined as A c = { θ ∈ R d : inf x ∈ A ∥ θ -x ∥ < c } , where c > 0 is a constant. Consider a stochastic process { θ t ∈ R d : t ≥ 0 } that follows the SDE Eq. 25 within A c with all the regularity conditions mentioned in Theorem E.1. Futhermore, let τ = inf { t ≥ 0 : θ t / ∈ A c ∩ int B } , where B = { θ ∈ R d : ∥ θ i ∥ ≤ c B } with i ∈ [ d ] and c B > 0 . If A ∩ B is an invariant set of θ t , then A ∩ B is a stochastically attractive invariant set of the stopped process { θ t ∧ τ : t ≥ 0 } , provided Eq. 26 and and 27 hold for any x ∈ A ∩ B .

Proof. Since the drift and diffusion terms in Eq.25 are Lipschitz continuous within A c , there exists a unique solution for the SDE up to the exit time τ . This allows us to consider the unique stopped process θ t ∧ τ , a continuous strongly Markov process. Let l ( θ ) : R d → R ≥ 0 be the Euclidian distance between θ ∈ R d and the invariant set A ∩ B , and ˜ A be the infinitesimal operator of θ t ∧ τ where τ ϵ = inf { t ≥ 0 : θ t / ∈ A ϵ } . The same argument in Theorem E.1 leads to ˜ A ( l λ ( x )) ≤ 0 for any x ∈ A ϵ ∩ B with small enough ϵ > 0 and λ > 0 . By applying a local version of maximial inequality (Lemma 1 in Kushner [44]), we obtain, for any x ∈ A ϵ and 0 < ϵ ′ < ϵ ,

This implies that A ∩ B is a stochastically attractive invariant set of the stopped process θ t ∧ τ .

Theorem 5.1 (formal) . Suppose { ( w 1 ,t , w 2 ,t ) ∈ R 2 : t ≥ 0 } is a stochastic process obeying Eq. 31 in a neighborhood U of A = { (0 , 0) } ⊂ R 2 , and the activation function σ : R → R is a smooth function such that σ (0) = 0 . The invariant set A is stochastically attractive if ηζ 2 | σ ′ (0) | 2 > | ∑ i ∈ [ n ] x i y i | ∑ i ∈ [ n ] x 2 i .

where

Define stopping times τ p = inf t ≥ 0 { p t ≥ ϵ } , τ s = inf t ≥ 0 { s t ≥ ϵ } where ϵ > 0 satisfies { ( p, s ) : max( | p | , | s | ) < ϵ } ⊂ U , and consider the stopped processes ( p t ∧ τ p , s t ∧ τ p ) and ( p t ∧ τ s , s t ∧ τ s ) . We have,

By the assumption on smoothness, ∂ ∂p µ p (0 , s ) -1 4 ∂ 2 ∂p 2 ∑ i ∈ [ n ] σ 2 p,i (0 , s ) < 0 and ∂ 2 ∂p 2 ∑ i ∈ [ n ] σ 2 p,i (0 , s ) > 0 for any s ∈ [ -ϵ, ϵ ] with small enough ϵ > 0 . Therefore by Lemma F.1, A p := { (0 , s ) : s ∈ [ -ϵ, ϵ ] } is stochastically attractive for the process { ( p t ∧ τ p , s t ∧ τ p ) : t ≥ 0 } . Similarly,

By Lemma F.1, A s := { ( p, 0) : p ∈ [ -ϵ, ϵ ] } is stochastically attractive for the process { ( p t ∧ τ s , s t ∧ τ s ) : t ≥ 0 } .

By combining these two facts, we see that A p ∩ A s = { (0 , 0) } is stochastically attractive for the process { ( w 1 ,t , w 2 ,t ) : t ≥ 0 } .

̸

Proof. If σ ′ (0) = 0 , the condition can never be met. Therefore, we only consider the case σ ′ (0) = 0 . Without loss of generality, assume that U is an open ball centered at (0 , 0) . Since the drift and diffusion terms of Eq.31 are locally Lipschitz continuous, it has a unique solution up to the exit time of U . We define stochastic processes p t = w t, 1 + w t, 2 , s = w t, 1 -w t, 2 . These two processes obey the following SDEs:

G The permutation invariant set in one hidden layer nonlinear neural network

Consider training a two neuron neural network f θ on a dataset { ( x k , y k ) } N k =1 , where f θ ( x ) = ∑ 2 i =1 w (2) i σ ( w (1) i x ) . We assume the data are drawn randomly from standard Gaussian distribution. Consider running gradient descent with learning rate η on the mean square error loss with label noise drawn freshly from N (0 , σ 2 ) . In this set up, the invariant set from the permutation symmetry is the affine space { θ | w ( i ) 1 = w ( i ) 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.

1.0-

100

100

22

22

0.5

2

10-3

2

10-3

W

1

0.0

M

10-6

1

10-6

-0.5

10-9

10-9

101

103

0

25005

5000

750010000

0

2500

5000

750010000

Steps

Steps

Steps

w1)

W2)

W2)

W2)

-lr=0.1—lr=0.2—lr=0.3-lr=0.4

一Ir=0.5

0=0.0-0=0.1-0=0.2

=0.4

0=0.8

H Conditions of sign and permutation invariant sets in the presence of residual connection

Definition H.1 (Residual connection) . A hidden layer l 2 in a neural network is defined as residually connected from a previous layer l 1 if its forward pass follows:

where h ( l ) denotes the hidden activations of layer l , w ( l ) is the weight matrix at layer l , and b ( l ) is the bias vector at layer l .

Proposition H.1 (Sign invariant set for residual neural network) . Consider a hidden neuron p of layer l 2 with a residual connection from layer l 1 . Let ( w ( l ) in ,p , b ( l ) p ) denote the parameters (weights and bias) directly incoming to this neuron, and w ( l +1) out ,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 A = { θ | w ( l ) in ,p = 0 , b ( l ) p = 0 , w ( l +1) out ,p = 0 for l = l 1 , l 2 } forms an invariant set.

Proposition H.2 (Permutation invariant set for residual neural network) . Consider two hidden neurons p, q in the same layer l 1 with a residual connection from layer l 1 . Let ( w ( l ) in ,p , b ( l ) p ) , ( w ( l ) in ,q , b ( l ) q ) denote the parameters directly incoming to the neurons, and w ( l +1) out ,p , w ( l +1) out ,q the parameters directly outgoing from the neurons. The affine space A = { θ | w ( l ) in ,p = w ( l ) in ,q , b ( l ) p = b ( l ) q , w ( l +1) out ,p = w ( l +1) out ,q for l = l 1 , l 2 } is an invariant set.

We omit the proofs as they closely resemble the proofs in App. B.1 and B.1. In the presence of residual connections, the conditions for the sign and permutation invariant sets become stronger. Not only do the weights and biases of the neuron(s) at layer l 2 need to follow the condition, the neuron(s) which are residually connected from must also must also meet these conditions.

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

Low-rank teacher network. We consider a low-rank teacher that computes the linear map ¯ W : R n → R m . ¯ W ∈ R m × n is a low-rank matrix where k ≪ min { m,n } denotes the rank.

Two-layer student network. Weconsider a two-layer student network, parameterized by W 1 ∈ R d × n and W 2 ∈ R m × d , that computes the linear map ˆ W = W 2 W 1 : R n → R m . 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 R n to R m .

Noisy training data. The student network is trained on data generated by the teacher network through the noisy map,

and ϵ ∼ N (0 , σ 2 ϵ I m ) . We will consider a fixed training dataset of p noisy input-output pairs { ( x 1 , y 1 ) , . . . , ( x p , y p ) } generated by the teacher network from input x i ∼ N (0 , σ 2 x I n ) . Let X ∈ R n × p and Y ∈ R m × p be the matrices with columns x i and y i respectively. This setup yields the second-order training statistics that guide the student network's learning dynamics,

Later in our analysis of the learning dynamics we will assume the input data matrix X is whitened such that Σ xx = I n implying that only the input-output covariance structure Σ yx 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 = ˆ WX and the output Y , and evaluated on the expected prediction error for a new test point sampled from the input distribution. We denote the train and test error as

where x ∼ N (0 , σ 2 x I n ) is a new test point sampled from the input distribution, ∥ · ∥ F 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 a student is determined by the relationship between three matrices in R m × n : the low rank teacher ¯ W , the overparameterized student ˆ W , and the noisy training data Σ yx . Due to the linear nature of this problem, we will use the Singular Value Decomposition (SVD) for all three matrices to describe their relationship,

Here, the s i denote non-zero singular values, and u i and v 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 U , S , and V with the appropriate accents ¯ · , ˆ · , or ˜ · to denote which matrix they are associated with.

I.2 Exact Solutions to the Student Learning Dynamics

The student network is trained by gradient descent with learning rate η to minimize E train . However, we assume that each gradient evaluation is corrupted with Gaussian label-noise such that the training

dynamics are given by the coupled update equations,

where W ( t ) 1 and W ( t ) 2 are the parameters after t steps of training and Z ( t ) ∈ R m × p is a matrix of label-noise associated with step t . Directly studying these dynamics is difficult because the gradients equations are coupled between the weights W 1 and W 2 . 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 η → 0 , 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-Input. We assume that Σ xx = I n such that the input data matrix X is whitened (this implicitly assumes p ≥ 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 Z ( t ) = ˜ U diag ( z ( t ) ) ˜ V ⊺ X where z ( t ) ∼ N (0 , ζ 2 ) .

A3. Spectral Initialization. We assume that the student network is initialized such that W (0) 1 = OA (0) ˜ V ⊺ and W (0) 2 = ˜ UB (0) O ⊺ where A (0) and B (0) are diagonal matrices such that ˆ S (0) = B (0) A (0) and O ∈ R d × d is a random orthonormal matrix such that O ⊺ O = I d .

A4. Balanced Initialization. We assume a balanced initialization such that W (0) 1 W (0) 1 ⊺ = W (0) 2 ⊺ W (0) 2 .

Given the first three assumptions, the dynamics decouple in the eigenbasis of Σ yx . Under the change of variables A ( t ) = O ⊺ W ( t ) 1 ˜ V and B ( t ) = ˜ U ⊺ W ( t ) 2 O , the dynamics transform to

which because A (0) and B (0) are diagonal matrices by assumption, decouples into a system of scalar equations. Taking the limit as η → 0 we can approximate each of these scalar equations with a two-dimensional non-linear SDE,

where a i and b i are the i th diagonal element of A and B respectively. Under these much simpler dynamics, it 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 (40), the lim t →∞ a i ( t ) 2 -b i ( t ) 2 = 0 .

Proof. Let r i ( t ) denote the difference r i ( t ) = a i ( t ) 2 -b i ( t ) 2 . By Itô's Lemma, r i ( t ) is driven by the ODE dr i = -ηζ 2 r i dt , which has the temporal solution r i ( t ) = r i (0) e -ηζ 2 t . Because ηζ 2 > 0 , then the lim t →∞ r i ( t ) = 0 .

The result of all four assumptions A1 - A4 is a decoupling the learning dynamics into a set of independent quartic slices of the loss landscape, ℓ i ( w i ) = ( w 2 i -˜ s i ) 2 with w i = a i = b i , where within each slice the dynamics evolve according to the non-linear SDE,

The geometry of these slices is controlled by the strength of the singular mode for the training data ˜ s i and the dynamics are determined by the relationship with the student's singular mode strength ˆ s i = w 2 i . By Itô's Lemma we find that the dynamics of ˆ s i are governed by the non-linear SDE,

which is an equation often used to model population growth in a stochastic, crowded environment. See Oksendal [58], Chapter 5 problem 5.15 for more details.

Theorem 6.1. Under assumptions A1 - A4, the dynamics of ˆ s i ( t ) for t ≥ 0 are governed by the stochastic process,

Proof. Let u i = ˆ s -1 i . Then by Itô's Lemma, the dynamics of u i are given by the SDE,

Rearranging this SDE we can express it in the standard form of a geometric Ornstein-Uhlenbeck process (sometimes referred to as a modified Ornstein-Uhlenbeck process)

where θ = 2˜ s i -3 ηζ 2 , µ = (˜ s i -3 2 ηζ 2 ) -1 , and σ = -2 √ ηζ 2 . To derive a solution to this linear SDE we make the ansatz that the solution takes the form u i ( t ) = F ( t ) G ( t ) where F ( t ) and G ( t ) solve the respective SDEs,

To determine the unknown coefficients α ( t ) and β ( t ) defining the SDE for G ( t ) , we can apply Itô's product rule

Aligning this expression with the original SDE for u i we see that

We can now solve for F ( t ) and G ( t ) . The SDE defining F ( t ) is standard geometric Brownian motion and thus has the solution,

where we assumed F (0) = 1 and thus G (0) = u i (0) . Using this solution for F ( t ) , then G ( t ) is given by

Combining our solutions for F ( t ) and G ( t ) we get that the solution for u i ( t ) is

Inverting this expression and rearranging terms, gives the expression for ˆ s i ( t ) ,

Plugging in the simplifications θµ = 2 , θ + 1 2 σ 2 = 2˜ s i -ηζ 2 , and σ = -2 √ ηζ 2 gives the final expression.

Corollary 6.1. In the limit of ˆ s i (0) → 0 , for any t > 0 , ˆ s i ( t ) is given by

Proof. Since B t is almost surely continuous, (2˜ s i -ηζ 2 ) τ +2 √ ηζ 2 B τ has a maximum in τ ∈ [0 , t ] and therefore 2 ∫ t 0 e (2˜ s i -ηζ 2 ) τ +2 √ ηζ 2 B τ dτ is finite almost surely. Therefore, by Theorem 6.1,

Corollary 6.1 has important implications for the dynamics of ˆ s i ( t ) . It demonstrates that the dynamics observe a phase transition determined by the ratio 2˜ s i /ηζ 2 , 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 to a stationary distribution with a positive expectation. Corollary 6.1 demonstrates that stochastic collapse is 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 t > 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 ˆ s i .

Corollary I.1 (Stationary solution) . In the limit of t → ∞ , ˆ s i ( t ) is distributed according to the density,

where Z = ∫ + ∞ 0 e -ˆ s i ηζ 2 ˆ s -( 3 2 -˜ s i ηζ 2 ) i d ˆ s i is the partition function.

Proof. To check if p ss (ˆ s i ) is the stationary solution we use the Fokker-Planck equation for the dynamics of ˆ s i , which is

where J is commonly referred to as the probability current . The stationary solution p ss satisfies ∂ t p ss = 0 , which because we are considering a one-dimensional SDE, is equivalent to J ss = 0 . Plugging in the expression for p ss (ˆ s i ) and using the simplification

we can verify that J ss = 0 .

The stationary solution can be interpreted as a Gibb's distribution (i.e. p ss ∝ e -κ Ψ( x ) ) with the potential

and temperature constant κ = ( ηζ 2 ) -1 . From this interpretation we see two phase transitions,

When 3 2 ηζ 2 ≤ ˜ s i , the most probable state is ˜ s i -3 2 ηζ 2 , where we already see a bias from the stochasticity towards the origin. When 3 2 ηζ 2 > ˜ s i , the most probable state is at the origin. When ηζ 2 2 ≥ ˜ s i , the partition function Z diverges and the stationary distribution collapses to a delta distribution at the origin, which is an invariant set.

We now conjecture an exact expression for the temporal dynamics of the expectation E [ˆ s i ( t )] :

Conjecture I.1. Under assumptions A1 - A4, the expectation E [ˆ s i ( t )] for t ≥ 0 is given by the equation,

Conjecture I.1 captures both the phase transition for stochastic collapse from Corollary 6.1 and the student's singular value shrinkage from Corollary I.1. Furthermore, in the limit of ζ → 0 , it aligns with the dynamics from Theorem 6.1, as would be expected. Additionally, one can show that the dynamics predicted by Conjecture I.1 are equivalent to the deterministic dynamics given by L 2 regularized gradient flow with a regularization coefficient of λ = ηζ 2 2 .

I.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...nonlinear... 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 { ¯ u i , ¯ s i , ¯ v i } , overparameterized student { ˆ u i , ˆ s i , ˆ v i } , and noisy training data { ˜ u i , ˜ s i , ˜ v i } we can decompose these error terms as

From these expressions we see that understanding the dynamics of the train and test error depends on (1) 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 (1) and as done in Lampinen and Ganguli [51] we can understand (2) in the high-dimensional limit through a random matrix theory analysis.

Random matrix theory analysis. We will work in the limit n, m →∞ with a finite rank k ∼ O (1) and finite aspect ratio A = m n ∈ (0 , ∞ ) . Without loss of generality, we assume σ 2 ϵ = n -1 such that the singular values of the noise are O (1) and thus the singular values of the teacher network can be understood as signal to noise ratios (SNRs). The relationship between the SVD structures of a low-rank matrix (the teacher network { ¯ u i , ¯ s i , ¯ v i } ) and a noisy perturbation (the training data { ˜ u i , ˜ s i , ˜ v i } ) is well studied in [59], which we summarize here.

In the limit, the top k singular values of Σ yx converge to values given by the transfer function,

while the bottom m -k singular values are distributed according to the Marchenko-Pastur (MP) distribution, √

Additionally, in the limit, the top k singular vectors of Σ yx acquire an alignment with the k singular vectors the teacher network given by the relationship |⟨ ˜ u i , ¯ u i ⟩| |⟨ ˜ v i , ¯ v i ⟩| = O (¯ s i ) where the the overlap function is defined as,

J Experiment details

The codes to reproduce the experiments in the main paper can be found at https://github. com/ccffccffcc/stochastic_collapse . Our the experiments were run on the Google Cloud Platform (with 4 × or 8 × NVIDIA A100 (40GB) GPU). The initial code development occurred on a local cluster equipped with 10 × 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 [49]. For VGG-16, we did not use Batch Normalization or Dropout. Both VGG-16 and ResNet-18 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 β = 0 . 9 and weight decay of 0.0005) as the optimizer.

Stochastic collapse in a quartic loss. (Fig.1) The left three plot are generated by sampling 50 trajectories of the SDE dθ t = -( θ 3 t -µθ t ) dt + ζθ t dB t with µ = 1 . 5 and ζ = 0 . 1 , 1 , 2 . 5 . The samples are computed using an Euler-Maruyama discretization scheme with step size η = 0 . 0025 for 10 4 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 ζ = 0 to ζ = 3 . For each value of ζ we sample 10000 trajectories with the same scheme described previously but for a varying number of steps: 0 , 5000 , 10000 , 20000 , 40000 , 80000 . We compute the empirical probability of the final position being within ϵ = 1 e -15 euclidean distance of the origin.

Stochastic collapse towards a sign invariant set. (Fig.2) We trained 10 3 different single-neuron models via simulating SDEs (31) up to time T = 100 , using the Euler-Marumaya discretization scheme of step size dt = 10 -2 . All the models are trained with the same data set, composed with 32 points { ( x i , y i ) } i =1 , ··· , 32 where x i is an i.i.d. standard Gaussian random variable, and y i is given by y i = x i + ϵ i where ϵ i 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. 3) We trained VGG-16 for 10 5 steps on CIFAR-10 with a learning rate of 0.1 and ResNet-18 for 10 6 steps on CIFAR-100 with a 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 n ind . 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 n ind clusters and the corresponding pairwise distance between outgoing weights was visualized with the same ordering.

Larger learning rates and increased label noise intensify stochastic collapse. (Fig. 4 and 9) These experiments varied learning rates and the extent of label noise. For the learning rate experiments, we trained models without label noise for 10 5 steps. For the label noise experiments, we maintained learning rates of 0.01 and 0.02 for VGG-16 and ResNet-18 respectively over 10 6 steps. Here, fresh label noise was introduced with each batch sample. A label noise level of ϵ %meant that ϵ %of the labels were randomly assigned an incorrect label. We averaged the results across four replications with different seeds.

After training, we calculated and displayed the proportion of independent neurons in each layer of the networks. We first removed 'vanishing neurons', defined as 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 10 5 (CIFAR-10) and 10 6 (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.

Demonstrating generalization benefits of stochastic collapse in a teacher-student setting. (Fig. 5) The experiments were conducted in accordance with assumptions A1 - 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 a sparse teacher model with a rank of 8 and signals ranging evenly from 0.5 to 1. The input samples were drawn to satisfy the condition X T X = I . We then constructed the signal y by introducing a 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 are averaged over 256 runs with the same dataset.

In Fig. 11, we present the results obtained when replacing label noise with different batch sizes. In this experiment, we used SGD without label noise and maintained a 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 10 4 , 2 × 10 4 , and 4 × 10 4 . We resumed training from the checkpoints post but at a 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 at each of these checkpoints, before the learning rate drop. We averaged the results across eight replications with different seeds.

K Extended Experiments

Layer with little stochastic collapsing

Layers with strong stochastic collapsing

CIFAR-10 VGG-16

Probability Density

101

conv5 (256)

conv11 (512)

fc14 (512)

-0Steps

20000Steps

40000Steps

100

60000Steps

80000Steps

10-1

LIL

100000 Steps

0

1

2

3

4

0

1

2

3

4

1

2

3

4

CIFAR-100 ResNet-18

layer2.0.conv2(128)

layerl.1.conv1 (64)

layer4.1.conv1 (512)

Probability Density

LLLI

-0Steps

200000Steps

10°

TITTTT

10-1

600000Steps

800000Steps

1000000 Steps

10-2

0

1

2

3

4

0

1

2

3

4

0

1

2

3

4

Normalized Distance

Normalized Distance

Normalized Distance

Nerons

CIFAR-10VGG-16

CIFAR-100ResNet-18

1.0

LabelNoise

1.0

0.8

一

0%

0.8

Fraction

Independent

0.6

20%

0.6

40%

0.4

60%

0.4

0.2

80%

0.2

0.0

0.0

conv1

conv2

conv3

conv4

conv5

conv6

conv7

conv8

conv9

conv10

conv11

conv1

o.conv1

.conv2

.conv1

.conv2

O.

O

.conv1

.conv2

7

.conv1

O

.conv2

.conv1

.conv2

Layer 1

Layer 2

Layer 3

Layer 4

Layer 1

Layer 2

Layer 3

Layer 4

x10-1

x10-5

0.4

Train Loss

S

6

1.4

0.3

BatchSize

1

1.0

SGD

0.2

2

5

Test

4

S>

0.1

8

0.6

16

4

0.0

0

2

4

6

0

2

4

6

0.00

0.15

0.30

Steps

x104

Steps

x104

S

*

CIFAR-10 VGG-16

1.3

90

93

Training Loss

1.0

Steps

TestAccuracy

Test Accuracy

92

EffectiveRank

5

10000

91

0.7

20000

70

4

06

40000

0.4-

60

89

3

0.1

50

88

10000

20000

3000040000

-2000

0

2000

4000

Steps

Steps after Learning Rate Drop

conv3

conv4

conv5

conv6

conv7

conv8

conv9

conv11

CIFAR-100 ResNet-18

1.8

65

76

Training Loss

1.4

Steps

55

TestAccuracy

TestAccuracy

Effective Rank

10000

72

５

1.0-

20000

40000

0.6

35

68

4

0.2

25

64

3

10000200003000040000

0

5000

10000

Steps

Steps afterLearningRateDrop

conv1

Layer 1

1.conv2

.conv2

Layer 2

1.conv

1.conv2

o.conv2

Layer 3

1.conv1

1.conv2

Layer 4

1.conv

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\na) 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.\n\nWEAKNESSES:\na)  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\nQUESTIONS:\na) 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.\n\nLIMITATIONS:\nThe limitations are already addressed in the sections above.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nPlease see the weaknesses.\n\nLIMITATIONS:\nI did see any potential negative societal impact.\n\nThe main limitation of this work lies its technical contributions, as I mentioned in the weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-08"}

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

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

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

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

LIMITATIONS:
The authors have adequately addressed the limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-09

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

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

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

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

LIMITATIONS:
No potential negative societal impact.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-12

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

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

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

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

LIMITATIONS:
Limitations have been adequately addressed.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-30