SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Daniel Kunin*
Atsushi Yamamura (山村築志)*
Surya Ganguli
{fengc, kunin, atsushi3, sganguli}@stanford.edu
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.
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:
1. We introduce invariant sets as subsets of parameter space that, once entered, trap SGD. We characterize two such sets that correspond to simpler subnetworks and appear extensively in modern architectures: one for vanishing neurons and the other for identical neurons (Sec. 3).
*Equal contribution. Ordered alphabetically.
2. We formulate a sufficient condition for stochastic attractivity — a process attracting SGD dynamics towards invariant sets — that reveals a competition between the loss landscape's curvature around an invariant set and the noise introduced by stochastic gradients (Sec. 4).
3. We apply the attractivity condition to determine when neurons with origin-passing activation functions collapse their parameters to zero, effectively removing the neuron. We empirically show that the parameters of permutable neurons within the same hidden-layer can collapse towards invariant sets corresponding to identical neurons (Sec. 5).
4. 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).
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<sup>2</sup>,
$$
f _ {\theta} (x) = w ^ {(m)} \sigma \left(w ^ {(m - 1)} \dots \sigma \left(w ^ {(1)} x + b ^ {(1)}\right) \dots + b ^ {(m - 1)}\right) + b ^ {(m)}, \tag {1}
$$
parameterized by  $\theta = (w^{(1)},b^{(1)},\ldots ,w^{(m - 1)},b^{(m - 1)},w^{(m)},b^{(m)})\in \mathbb{R}^d$ , where  $w$  and  $b$  denote weights and biases respectively, the superscript indexes the  $m$  layers, and  $\sigma (\cdot)$  is an activation function. The network is trained by stochastic gradient descent (SGD) over a training dataset  $\{(x_{1},y_{1}),\dots ,(x_{n},y_{n})\}$  of size  $n$ , where  $x_{i}\in \mathbb{R}^{p}$  and  $y_{i}\in \mathbb{R}^{c}$ . The network parameters are randomly initialized at  $\theta^{(0)}$  and iteratively updated according to
$$
\theta^ {(t + 1)} = \theta^ {(t)} - \frac {\eta}{\beta} \sum_ {i \in \mathcal {B} ^ {(t)}} \nabla_ {\theta} \ell \left(\theta^ {(t)}; x _ {i}, y _ {i}\right), \tag {2}
$$
where  $\eta > 0$  is the learning rate,  $\mathcal{B}^{(t)} \subseteq [n]$  is a random $^3$  index set (mini-batch) of size  $\beta$  at step  $t$ , and  $\ell(\theta; x_i, y_i)$  is a loss function. We let  $\mathcal{L}(\theta) = \frac{1}{n} \sum_{i \in [n]} \ell(\theta; 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 \subseteq \mathbb{R}^d$  is an invariant set of SGD if given any initialization  $\theta^{(0)} \in A$ , all future iterates of SGD  $\theta^{(t)}$  for  $t \geq 0$  are contained within  $A$  almost surely, for any batch size  $\beta \in \mathbb{N}$ , learning rate  $\eta > 0$ , and mini-batches  $\{\mathcal{B}^{(t)} \subseteq [n] : t \in \mathbb{N}\}$ .
From this definition, it is immediately clear that all of parameter space  $A = \mathbb{R}^d$  and any interpolating point  $A = \{\theta_*\}$  such that  $\mathcal{L}(\theta_*) = 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_{in,p}^{(l)}, b_p^{(l)})$  and  $w_{out,p}^{(l+1)}$  denote the parameters directly incoming and outgoing from the neuron respectively<sup>4</sup>. If the non-linearity  $\sigma$  is origin-passing  $(\sigma(0) = 0)$ , then the axial subspace  $A = \{\theta \in \mathbb{R}^d | w_{in,p}^{(l)} = 0, b_p^{(l)} = 0, w_{out,p}^{(l+1)} = 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_{in,p}^{(l)}, b_p^{(l)})$ ,  $(w_{in,q}^{(l)}, b_q^{(l)})$  denote the parameters directly incoming to the neurons, and  $w_{out,p}^{(l+1)}, w_{out,q}^{(l+1)}$  the parameters directly outgoing from the neurons. The affine subspace  $A = \{\theta \in \mathbb{R}^d | w_{in,p}^{(l)} = w_{in,q}^{(l)}, b_p^{(l)} = b_q^{(l)}, w_{out,p}^{(l+1)} = w_{out,q}^{(l+1)}\}$  is an invariant set.
The subspace corresponding to a permutation invariant set represents the parameter space of a low-rank 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)\binom{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
2This notation encompasses fully-connected and convolutional networks, excluding more complex architectures such as transformers for simplicity, although much of our theory would directly apply.
3For simplicity we assume sampling with replacement such that there is no dependency between batches.
4The incoming and outgoing weights are related by:  $(w_{\mathrm{in},p}^{(l)})_q = (w_{\mathrm{out},q}^{(l - 1)})_p$
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 \in \mathbb{R}^{d \times d}$  be a symmetric  $(Q = Q^{\top})$  or orthogonal matrix  $(QQ^{\top} = I_d)$ . The affine subspace  $A = \{\theta \in \mathbb{R}^d | Q\theta = \theta\}$  is an invariant set if the loss function  $\ell(\theta; x_i, y_i)$ , for any  $i \in [n]$ , is approximately  $Q$ -symmetric around  $A$ , i.e., for any  $\epsilon > 0$ , there exists  $\delta > 0$  such that  $|\ell(Q\theta; x_i, y_i) - \ell(\theta; x_i, y_i)| < \epsilon d(\theta, A)$  for any  $\theta \in \mathbb{R}^d$  satisfying  $d(\theta, A) < \delta$  where  $d(\theta, A)$  is the Euclidean distance between  $\theta$  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 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  $\xi_{i}(\theta)$  represent the per-sample gradient noise at  $\theta$ . It is easy to check that  $\mathbb{E}[\xi_i(\theta)] = 0$  for all  $\theta$ . This motivates studying the SDE, which we refer to as stochastic gradient flow (SGF),
$$
d \theta_ {t} = - \nabla \mathcal {L} (\theta_ {t}) d t + \sqrt {\frac {\eta}{\beta}} \Sigma (\theta_ {t}) d B _ {t}, \quad \theta_ {0} = \theta^ {(0)}. \tag {3}
$$
Here  $B_{t}$  denotes a standard  $d$ -dimensional Wiener process and  $\Sigma (\theta_t) = \sqrt{\mathbb{E}[\xi_i(\theta)\xi_i(\theta)\tau]}$ . The drift term of this SDE is the negative full-batch gradient  $-\nabla \mathcal{L}(\theta_t)$ , while the diffusion is determined by a spatially-dependent diffusion matrix  $D(\theta_t) = \frac{\eta}{2\beta}\Sigma (\theta_t)\Sigma (\theta_t)^{\intercal}$ . 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$ .<sup>7</sup>
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 \subset \mathbb{R}^d$  of a stochastic process  $\{\theta_t \in \mathbb{R}^d : t \geq 0\}$  is stochastically attractive if for any  $\rho > 0$  and  $\epsilon > 0$ , there exists  $\delta > 0$  such that for any  $\theta \in \mathbb{R}^d$  with the Euclidean distance  $d(\theta, A) < \delta$ ,
$$
\underline {{\mathbb {P}}} \left[ \sup  _ {t \geq 0} d \left(\theta_ {t}, A\right) \geq \epsilon \mid \theta_ {0} = \theta \right] \leq \rho . \tag {4}
$$
5The Euclidian distance between  $\theta$  and  $A$  is defined as  $d(\theta ,A)\coloneqq \min_{a\in A}\| \theta -a\|$
As discussed in Li et al. [39], Eq. 3 is an order 1 weak approximation of Eq. 2.
See App. C for the definition of invariant sets for continuous processes and the formal statement.
This property is called stochastic stability in [44].
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 let's consider a canonical example in one-dimension.
Geometric Brownian motion. Geometric Brownian Motion (GBM) is a linear stochastic process given by the SDE,  $d\theta_{t} = \mu \theta_{t}dt + \zeta \theta_{t}dB_{t}$ , where  $\mu \in \mathbb{R}$  represents the drift rate,  $\zeta \in \mathbb{R}$  the volatility rate. Denote  $\theta_0 \in \mathbb{R}$  as the initialization. The set  $A = \{0\}$  is an invariant set of GBM, which, depending on the relationship between  $\mu$  and  $\zeta$ , can be stochastically attractive. As the trajectory of GBM approaches the invariant set, its movement diminishes. To account for this effect, we take the logarithmic transformation  $z_{t} = \log (\theta_{t})$ , where we assume without loss of generality that  $\theta_0 > 0$ . The stochastic process  $z_{t}$  obeys the transformed SDE,  $dz_{t} = \left(\mu -\zeta^{2} / 2\right)dt + \zeta dB_{t}$ , which has an additional drift term  $-\zeta^2 /2$  from Ito's lemma that attracts the process towards the invariant set. The total drift is now determined by a competition between the original drift rate  $\mu$  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  $\mu < \zeta^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:
$$
d \theta_ {t} \approx - \mathcal {L} ^ {\prime \prime} (0) \theta_ {t} d t + \sqrt {D ^ {\prime \prime} (0)} \theta_ {t} d B _ {t}. \tag {5}
$$
Here we use  $\mathcal{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\} \in \mathbb{R}$  be an invariant set of  $\{\theta_t \in \mathbb{R} : t \geq 0\}$  obeying Eq. 3. Suppose  $\mathcal{L}'(\theta) : \mathbb{R} \to \mathbb{R}, D(\theta) : \mathbb{R} \to \mathbb{R}$  are  $C^2$  functions such that  $D''(0) > 0$ . Define rate of attractivity  $\alpha = \mathcal{L}''(0) + \frac{1}{2} D''(0)$ . A is stochastically attractive if  $\alpha > 0$ , while it is not stochastically attractive if  $\alpha < 0$ .<sup>9</sup>
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(\theta) \geq 0$  by definition, and  $D(0) = 0$  as  $\{0\}$  is an invariant set. As a result,  $D'(0) = 0$  and  $D''(0) \geq 0$  by the continuity assumption of  $D$ . The collapsing condition, therefore crucially depends on the curvature of the loss function  $\mathcal{L}$  at  $\theta = 0$ . When  $D''(0)$  is strictly positive, the collapsing condition can still be satisfied with negative curvature provided that  $\mathcal{L}''(0) > -\frac{1}{2} D''(0)$ . Given that  $D$  is proportional to  $\eta/\beta$ , the learning rate to batch size ratio determines the maximum attainable steepness for an invariant set to be attractive.