SlowGuess/NeurIPS_2023_OCR huggingface dataset:
An illustrative example. Consider SGF in a double-well potential  $\mathcal{L}(\theta) = \frac{1}{4} (\theta^2 -\mu)^2$  with multiplicative diffusion<sup>10</sup>, such that the dynamics are  $d\theta_{t} = -(\theta_{t}^{3} - \mu \theta_{t})dt + \zeta \theta_{t}dB_{t}$  where  $\mu ,\zeta >0$ . The minima of the potential are located at  $\theta = \pm \sqrt{\mu}$ , while  $\theta = 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}(\theta)$ , 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  $\theta_0 > 0$  and consider the logarithm process  $z_{t} = \log (\theta_{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}(\theta)$  is given by a Gibbs distribution  $p_{ss}(\theta)\propto e^{-\kappa \Psi (\theta)}$  for a modified potential<sup>11</sup>  $\Psi (\theta) = \theta^2 /2 - (2\mu /\zeta^2 -2)\log (\theta)$  with constant  $\kappa = 2 / \zeta^2$  and partition function  $Z = \int_{-\infty}^{\infty}e^{-\kappa \Psi (\theta)}d\theta$ . When  $\mu \leq \zeta^2 /2$ , the partition function diverges and  $p_{ss}(\theta)$  collapses to a Dirac delta distribution at  $\theta = 0$ . This transition agrees with the collapsing condition from Theorem 4.1.
<sup>9</sup>When  $\alpha = 0$ , higher-order derivatives of the loss and diffusion determine the attractivity of  $A$ .
Here we build on a substantial literature, originating from the seminal work by Kramers [45], that investigates the escape rate between minima of a particle in a double-well potential with constant diffusion [46].
<sup>11</sup>The interpretation of a modified potential determined by the gradient noise that drives the learning dynamics was explored in [42] and [5]. Stochastic collapse is when the partition function for this modified loss diverges.
Figure 1: Stochastic collapse in a quartic loss. The left three plots show sample trajectories with the same, random initializations driven by the SDE  $d\theta_t = -(\theta_t^3 - \mu \theta_t)dt + \zeta \theta_t dB_t$  for different values of  $\zeta$ . The analytic stationary solution is plotted on the side plot. Leftmost: for small  $\zeta$  the steady-state distribution is two bumps around the minima  $\theta = \pm \sqrt{\mu}$ . Middle left: with increasing  $\zeta$  the distribution spreads out and is biased towards  $\theta = 0$ . Middle right: when  $\zeta$  surpasses the collapsing condition  $\sqrt{2\mu}$  the steady-state distribution collapses to a dirac delta distribution at  $\theta = 0$ . Rightmost: the empirical probability of the sample trajectories at different steps being within  $\epsilon$  of the origin as a function of increasing gradient noise. As we see empirically, there is a sudden phase transition which aligns with the collapsing condition from Theorem 4.1.
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 \subset \mathbb{R}^d$  be a  $d_A$ -dimensional affine subset, and a stochastic process  $\{\theta_t \in \mathbb{R}^d : t \geq 0\}$  obey Eq. 3 in  $A_c$ , open  $c$ -neighborhood of  $A$  with some  $c > 0$ . Suppose  $\mathcal{L}: A_c \to \mathbb{R}$  is a  $C^3$ -function whose first and second-order derivatives are  $L$ -Lipschitz continuous.  $D: A_c \to \mathbb{R}^{d \times 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  $\sqrt{D}: A_c \to \mathbb{R}^{d \times d}$  are  $L$ -Lipschitz continuous. Let  $D_{\perp} = P_{\perp}DP_{\perp}$  where  $P_{\perp}: \mathbb{R}^d \to \mathbb{R}^d$  projects to the normal space of  $A$ . If there exists  $\delta > 0$  such that  $\nabla_n^2\hat{n}^\intercal D\hat{n}^\intercal > \delta$  and
$$
\nabla_ {\hat {n}} ^ {2} \left(\mathcal {L} - \frac {1}{2} \operatorname {T r} D _ {\perp} + (1 - \delta) \hat {n} ^ {\intercal} D \hat {n}\right) > 0, \tag {6}
$$
for any unit normal vector  $\hat{n} \in \mathbb{R}^d$  perpendicular to  $A$  and  $\theta \in A$ , then  $A$  is stochastically attractive.
To prove Theorem 4.2 (see App. E), we consider a family of distance measures between  $\theta_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 \sigma(w_1 x)$  with  $w_1, w_2, x \in \mathbb{R}$ . Suppose this model is trained via label-noise gradient descent and using the  $L^2$ -loss  $\mathcal{L}(w_1, w_2) = \frac{1}{n} \sum_{i \in [n]} (y_i - f(x_i; w_1, w_2))^2$  on a non-empty data
Figure 2: Stochastic collapse of a single neuron. The SDE dynamics of a single-neuron model are simulated with various noise levels  $\zeta$ . Left: A scatter map of trained weights  $(w_{1}, w_{2})$ . The colors of the dots represent the noise level, corresponding to the vertical lines in the right panel. For each noise level, we plot 15 trained models with different noise realizations. The background heatmap shows the loss landscape.
Right: The colored lines represent the empirical probability that  $\sqrt{w_1^2 + w_2^2} < \epsilon = 10^{-2}$  (from  $10^{3}$  samples) after a given number of update steps. We observe a sudden transition that aligns with the collapsing condition from Theorem 5.1.
Figure 3: Evidence of stochastic collapse towards permutation invariant sets in deep neural networks. Each pair of plots shows the normalized distance matrix between incoming parameters  $(w_{\mathrm{in}})$  and outgoing parameters  $(w_{\mathrm{out}})$  of neurons within the same hidden layer. We show two pairs of plots with little stochastic collapsing and four pairs of plots with strong stochastic collapsing: three pairs for three different layers in VGG-16 trained for  $10^{5}$  steps on CIFAR-10 (top row), and three pairs for three different layers in a ResNet-18 trained for  $10^{6}$  steps on CIFAR-100 (bottom row). For each pair of plots, we sort the neurons by hierarchical clustering according to normalized distances between incoming parameters only (left plot in each pair), and then show the distance matrix between outgoing parameters using the same sorting of neurons (right plot in each pair). The similarity between each plot in a pair indicates that clusters of neurons with similar incoming parameters also have similar outgoing parameters. See App. J for experimental details.
set  $\{(x_1, y_1), \dots, (x_n, y_n)\}$  such that  $\sum_{i \in [n]} x_i^2 \neq 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}) \in \mathbb{R}^2 : t \geq 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  $\eta > 0$  and noise amplitude  $\zeta > 0$ . Suppose the activation function  $\sigma: \mathbb{R} \to \mathbb{R}$  is smooth satisfying  $\sigma(0) = 0$ ,  $\sigma'(0) > 0$ . The invariant set  $A$  is stochastically attractive if  $\frac{\sigma'(0)\eta\zeta^2}{2} > \frac{\left|\sum_{i \in [n]} x_i y_i\right|}{\sum_{i \in [n]} x_i^2}$ .
The term  $\frac{\left|\sum_{i\in[n]}x_iy_i\right|}{\sum_{i\in[n]}x_i^2}$  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 proof 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  $\mathrm{dist}(w_i^{(l)}, w_j^{(l)}) \equiv \frac{2\|w_i^{(l)} - w_j^{(l)}\|_2^2}{\|w_i^{(l)}\|_2^2 + \|w_j^{(l)}\|_2^2}$ . Strikingly, hierarchical clustering based on this distance reveals 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
Figure 4: Larger learning rates intensify stochastic collapse. This figure illustrates how the fraction of independent neurons per layer in VGG-16 trained on CIFAR-10 (left column) and ResNet-18 trained on CIFAR-100 (right column) varies with changes in learning rates. The networks are evaluated at training steps of  $10^{5}$ . A reduced percentage of independent neurons indicates stronger stochastic collapse. See App. J for further details.
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 high-dimensionality of the weight vectors.
We found that increased learning rates typically intensify the stochastic attraction to invariant sets by reducing the number of independent neurons (Fig. 4). A large reduction in the fraction of independent neurons was observed in VGG-16 between 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  $\bar{W}$  generates a training dataset by providing noise corrupted outputs to a set of random Gaussian inputs. The input-output data covariance matrix  $\Sigma_{yx}$  then drives the learning dynamics of a two-layer student with composite weight matrix  $\hat{W} = W_2W_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 input-output covariance singular vectors; (A4) balanced initialization with equal power in the two weight layers of the student. Under these assumptions, the learning dynamics of the student weight matrix  $\hat{W} = W_2W_1$  can be decomposed into independent learning dynamics for each singular value  $\hat{s}_i$  of  $\hat{W}$ . An associated singular value  $\tilde{s}_i$  of  $\Sigma_{yx}$  drives the dynamics of  $\hat{s}_i$  via the SDE,
$$
d \hat {s} _ {i} = 2 \hat {s} _ {i} \left(\left(\tilde {s} _ {i} + \eta \zeta^ {2} / 2\right) - \hat {s} _ {i}\right) d t + 2 \sqrt {\eta \zeta^ {2}} \hat {s} _ {i} d B _ {t}, \tag {7}
$$
where  $\zeta > 0$  is the amplitude of the gradient noise. In the small noise limit of  $\zeta \rightarrow 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
Figure 5: Demonstrating generalization benefits of stochastic collapse in a teacher-student setting. We show the train loss (leftmost) and test loss (middle left) during the training of the student with different gradient noises. Dashed lines in the leftmost and middle left panels indicate the step where learning rate is dropped. Training with larger gradient noises  $\zeta$  generalizes better. Middle right: we show the noisy teacher signals against the learned student signals before dropping the learning rate. Larger noises have a stronger implicit bias towards zero. Rightmost: same as the middle right panel but we show the learned signals (for the brightest green in the middle right panel) at different training steps after the learning rate drop. All the results are averaged over 256 replicates.
overfitting to noise in the training data [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  $\hat{s}_i$  evolve according to the process:
Theorem 6.1 (Stochastic Student Dynamics). Under assumptions A1 - A4, the dynamics of  $\hat{s}_i(t)$  given by Eq. 7 are governed by the stochastic process,
$$
\hat {s} _ {i} (t) = e ^ {\left(2 \bar {s} _ {i} - \eta \zeta^ {2}\right) t + 2 \sqrt {\eta \zeta^ {2}} B _ {t}} \left(2 \int_ {0} ^ {t} e ^ {\left(2 \bar {s} _ {i} - \eta \zeta^ {2}\right) \tau + 2 \sqrt {\eta \zeta^ {2}} B _ {\tau}} d \tau + \hat {s} _ {i} (0) ^ {- 1}\right) ^ {- 1}. \tag {8}
$$
In terms of the per layer balanced student weight  $w_{i} = \sqrt{\hat{s}_{i}}$ , the above process can be thought of as originating from the quartic loss landscape  $\ell_{i}(w_{i}) = \frac{1}{4} (w_{i}^{2} - \tilde{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  $\hat{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  $\tilde{s}_i$ . To confirm stochastic collapse to this invariant set, we derive:
Corollary 6.1. In the limit  $\hat{s}_i(0) \to 0$ , for any  $t > 0$ , then  $\hat{s}_i(t) / \hat{s}_i(0) \underset{a.s.}{\rightarrow} e^{(2\tilde{s}_i - \eta \zeta^2)t + 2\sqrt{\eta\zeta^2} B_t}$ .
Corollary 6.1 has important implications. First, if  $\tilde{s}_i < \frac{\eta\zeta^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  $\eta$  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  $\eta$ , because  $\tilde{s}_i < \frac{\eta\zeta^2}{2}$ , and then learning only the subset with large  $\tilde{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
Figure 6: Large learning rates aid generalization via stochastic collapse to simpler subnetworks. Left: The training loss (grey) and test accuracy (blue) during an initial training phase using high learning rate  $(\mathrm{lr} = 0.1)$ . Vertical dashed lines highlight the training steps at which we will drop the learning rate. Middle: Test accuracy preceding and following a learning rate drop to  $\mathrm{lr} = 0.01$ . Colors indicate when the learning rate drop shown in the left plot occurred. Despite plateauing of the training loss and test accuracy, a later learning rate drop yields higher final test accuracy. Right: The fraction of independent neurons per layer evaluated at different learning rate drop times (indicated by the color) during the initial high learning rate training phase. Prolonged training with a large learning rate induces implicit regularization by reducing the number of independent neurons. All curves represent the average across eight replicates. See App. J for experimental details.
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).
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.
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.