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

TL;DR

Stochastic Collapse reveals that SGD with gradient noise is drawn toward invariant sets created by reflection symmetries in neural networks, biasing training toward simpler subnetworks. By formulating a stochastic gradient flow framework and a concrete attractivity criterion $\alpha = \mathcal{L}''(0) + \frac{1}{2} D''(0) > 0$, the authors show how noise and curvature compete to attract SGD toward sign and permutation invariant sets, including a quantitative mechanism for collapse in one and high dimensions. They provide both theoretical results and empirical evidence in deep nets that stochastic collapse leads to vanishing or identical neurons, acting as an implicit low-rank regularizer that can improve generalization, especially when coupled with large initial learning rates. Through a linear teacher-student model, they connect this stochastic collapse to generalization benefits and explain the practical value of early high learning-rate schedules. Overall, the work links gradient-noise geometry to architectural symmetries, offering a principled account of SGD’s implicit bias toward simpler subnetworks and guiding future optimization design.

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.

Figures (12)

• 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 \ref{['theorem:collapsing-condition-1dim-LD-notation']}.
• 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 \ref{['theorem:collapsing-condition-sign-invariant-informal-version']}.
• 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_{\text{in}}$) and outgoing parameters ($w_{\text{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. \ref{['app:exp_details']} for experimental details.
• 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. \ref{['app:exp_details']} for further details.
• 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.

Theorems & Definitions (31)

• Definition 3.1: Invariant Set of SGD
• Proposition 3.1: Sign Invariant Sets
• Proposition 3.2: Permutation Invariant Sets
• Theorem 3.1: Symmetry Induced Invariant Sets
• Proposition 4.1: Informal
• Definition 4.1: Stochastic Attractivity
• Theorem 4.1: Necessary/Sufficient Condition for Stochastic Attraction in One-Dimension
• Theorem 4.2: A Sufficient Condition for Stochastic Attraction in High-Dimensions
• Theorem 5.1: Informal
• Theorem 6.1: Stochastic Student Dynamics