SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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}(\theta)$ . 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  $\zeta^2$  Gaussian label noise per training label, denoted as  $\xi_{i}$ , it results in an update rule we refer to as label-noise gradient descent (LNGD),
$$
\theta^ {(t + 1)} = \theta^ {(t)} - \eta \nabla \mathcal {L} \left(\theta^ {(t)}\right) - \sqrt {\eta} \frac {\sqrt {\eta}}{N} \sum_ {i \in [ N ]} \nabla f _ {\theta^ {(t)}} \left(x _ {i}\right) \xi_ {i}. \tag {20}
$$
This is a discretization of the following continuous-time stochastic differential equation, which we denote as label-noise gradient flow (LNGF),
$$
d \theta_ {t} = - \eta \nabla \mathcal {L} \left(\theta_ {t}\right) + \sqrt {2 D _ {L N} \left(\theta_ {t}\right)} d B _ {t}, \quad \theta_ {0} = \theta^ {(0)}. \tag {21}
$$
where the diffusion matrix of LNGF can be expressed as,
$$
D _ {L N} \left(\theta_ {t}\right) = \frac {\eta \zeta^ {2}}{2} \frac {1}{N} \sum_ {i \in [ N ]} \nabla f _ {\theta_ {t}} \left(x _ {i}\right) \nabla f _ {\theta_ {t}} \left(x _ {i}\right) ^ {T}, \tag {22}
$$
which can also be decomposed into a magnitude term and a shape term,
$$
D _ {m} = \frac {\eta \zeta^ {2}}{2} \tag {23}
$$
$$
D _ {s} = \frac {1}{N} \sum_ {i \in [ N ]} \nabla f _ {\theta_ {t}} \left(x _ {i}\right) \nabla f _ {\theta_ {t}} \left(x _ {i}\right) ^ {T} \tag {24}
$$
Compared with SGF, the magnitude term here depends on the learning rate and the variance of the label noise. Compared with SGF, the shape term here 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 \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 the following SDE in  $A_c$ , open  $c$ -neighborhood of  $A$  with some  $c > 0$ :
$$
d \theta_ {t} = \mu (\theta_ {t}) d t + \sqrt {2 D (\theta_ {t})} d W _ {t}, \tag {25}
$$
where  $\{W_t : t \geq 0\}$  is the  $d$ -dimensional standard Wiener process. Here  $\mu : A_c \to \mathbb{R}$  is a  $L$ -Lipschitz  $C^2$ -function with its first-order derivatives being  $L$ -Lipschitz as well.  $D : A_c \to \mathbb{R}^{m \times 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  $\sqrt{D(\theta)}$  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$ . Suppose  $A$  is an invariant set of  $\theta_t$ .  $A$  is stochastically attractive if there exists  $\delta > 0$  such that for any unit normal vector  $\hat{n} \in \mathbb{R}^d$  perpendicular to  $A$  and any  $\theta \in A$ ,
$$
\nabla_ {\hat {n}} ^ {2} \hat {n} ^ {\intercal} D \hat {n} > \delta \tag {26}
$$
$$
\nabla_ {\hat {n}} \hat {n} ^ {\intercal} \mu + \nabla_ {\hat {n}} ^ {2} \left(- \frac {1}{2} \operatorname {T r} D _ {\perp} + (1 - \delta) \hat {n} ^ {\intercal} D \hat {n}\right) > 0. \tag {27}
$$
Proof. We exploit a local version of Doob's maximal inequality (Lemma 1 in [44]). By the Lipschitz property of  $\mu$  and  $\sqrt{D}$ , the SDE Eq. 25 has a unique strong solution up to the exit time  $\tau_{c} = \inf \{t\geq 0:\theta_{t}\notin A_{c}\}$ . The stopped process  $\{\tilde{\theta}_t = \theta_{t\wedge \tau_c}:t\geq 0\}$  is a continuous strong Markov process. It suffices to show that  $A$  is stochastically attractive for  $\tilde{\theta}_t$ . We first show that  $\tilde{\mathcal{A}} l^{\lambda}(x)\leq 0$  for any  $x\in A_{\epsilon}\backslash A$  with some  $0 < \epsilon < c$ , and  $\lambda >0$ , where  $\tilde{\mathcal{A}}$  is the infinitesimal operator of the stochastic process  $\tilde{\theta}_t$  and  $A_{\epsilon}$  denotes the open neighborhood of  $A$  within distance  $\epsilon >0$ . We consider the Euclidean distance  $l(\tilde{\theta}_t)$  between the process  $\tilde{\theta}_t$  and the affine invariant set  $A$ .
The SDE for  $l^{\lambda}(\tilde{\theta}_t)$  is given by,
$$
\begin{array}{l} d l ^ {\lambda} \left(\tilde {\theta} _ {t}\right) = \left(- \lambda l ^ {\lambda - 2} \left(\tilde {\theta} _ {t}\right) \theta_ {t} ^ {\intercal} P _ {\perp} \mu + \lambda l ^ {\lambda - 2} \left(\tilde {\theta} _ {t}\right) \operatorname {T r} D _ {\perp} + \lambda (\lambda - 2) l ^ {\lambda - 4} \left(\tilde {\theta} _ {t}\right) \left(\theta_ {t} ^ {\intercal} D _ {\perp} \tilde {\theta} _ {t}\right)\right) d t \\ + \lambda l ^ {\lambda - 2} (\tilde {\theta} _ {t}) \tilde {\theta} _ {t} ^ {\intercal} \sqrt {2 D _ {\perp}} d W _ {t}. \\ \end{array}
$$
This means that for any  $x \in A_c$ ,
$$
\tilde {\mathcal {A}} l ^ {\lambda} (x) = - \lambda l ^ {\lambda - 2} (x) x ^ {\intercal} P _ {\perp} \mu + \lambda l ^ {\lambda - 2} (x) \operatorname {T r} D _ {\perp} + \lambda (\lambda - 2) l ^ {\lambda - 4} (x) \left(x ^ {\intercal} D _ {\perp} x\right). \tag {28}
$$
By the fact that  $A$  is an invariant set,  $\tilde{\mathcal{A}}l^2(x) = 0$  holds for any  $x \in A$ . Therefore, we obtain  $\operatorname{Tr} D_{\perp}(x) = 0$ . Since  $D_{\perp}$  is positive semi-definite,  $D_{\perp}(x) = 0$  for any  $x \in A$ . Furthermore, since  $\tilde{\mathcal{A}}f(x) = 0$  with  $f(\theta_t) = P_{\perp}\theta_t$  for any  $x \in A$ ,  $P_{\perp}\mu(x) = 0$  for any  $x \in A$ .
Going back to Eq.28, we obtain
$$
\tilde {\mathcal {A}} l ^ {\lambda} (x) = 2 \lambda l ^ {\lambda - 2} (x) G (x) \left(\frac {F (x)}{2 G (x)} - 1 + \lambda / 2\right),
$$
$$
F (x) := (- x ^ {\intercal} P _ {\perp} \mu (x) + \operatorname {T r} D _ {\perp} (x)), \qquad G (x) := \left(\frac {P _ {\perp} x}{\| P _ {\perp} x \|}\right) ^ {\intercal} D (x) \frac {P _ {\perp} x}{\| P _ {\perp} x \|}.
$$
Notice that  $G(x) > 0$  for any  $x \in A_{\epsilon} \backslash A$  with small enough  $\epsilon > 0$  because Eq. 26 holds and the second derivatives of  $D(x)$  are  $L$ -Lipschitz. We want to show that there exists  $\lambda > 0$  such that for any  $x \in A_{\epsilon} \backslash A$ ,  $\frac{F(x)}{2G(x)} - 1 + \lambda / 2 < 0$  with some  $\lambda > 0$ . Parameterize  $x$  by  $x(\xi) = x_{\parallel} + \xi \hat{n}$ , where  $x_{\parallel} \in A$  and  $\hat{n}$  is the unit normal vector perpendicular to  $A$ . From L'Hôpital's rule, we have,
$$
\lim _ {\xi \to 0} \frac {F (x _ {\parallel} + \xi \hat {n})}{G (x _ {\parallel} + \xi \hat {n})} = \frac {\nabla_ {\hat {n}} ^ {2} F (x _ {\parallel})}{\hat {n} ^ {\intercal} \nabla_ {\hat {n}} ^ {2} D (x _ {\parallel}) \hat {n}}.
$$
The second derivatives of  $F$  can be calculated as follows:
$$
\nabla_ {\hat {n}} ^ {2} F \left(x _ {\parallel}\right) = - x _ {\parallel} ^ {\intercal} P _ {\perp} \nabla_ {\hat {n}} ^ {2} \mu - 2 \nabla_ {\hat {n}} \hat {n} ^ {\intercal} \mu + \nabla_ {\hat {n}} ^ {2} \operatorname {T r} D _ {\perp} = - 2 \nabla_ {\hat {n}} \hat {n} ^ {\intercal} \mu + \nabla_ {\hat {n}} ^ {2} \operatorname {T r} D _ {\perp}. \tag {29}
$$
Substituting this, we obtain
$$
\lim  _ {\xi \rightarrow 0} \frac {F \left(x _ {\parallel} + \xi \hat {n}\right)}{2 G \left(x _ {\parallel} + \xi \hat {n}\right)} = \frac {- 2 \nabla_ {\hat {n}} \hat {n} ^ {\intercal} \mu \left(x _ {\parallel}\right) + \nabla_ {\hat {n}} ^ {2} \operatorname {T r} D _ {\perp} \left(x _ {\parallel}\right)}{2 \hat {n} ^ {\intercal} \nabla_ {\hat {n}} ^ {2} D \left(x _ {\parallel}\right) \hat {n}} <   1 - \delta . \tag {30}
$$
By the assumed Lipschitz property of the derivatives of  $\mu$  and  $D$ ,  $\nabla_{\hat{n}}^{2}F(x)$  and  $\nabla_{\hat{n}}^{2}G(x)$  are  $L$ -Lipschitz in  $A_{\epsilon}\backslash A$  with small enough  $\epsilon >0$ . Therefore, for any  $0 < \xi < \epsilon$ ,
$$
\begin{array}{l} \left| \nabla_ {\hat {n}} \frac {F (x _ {\parallel} + \xi \hat {n})}{G (x _ {\parallel} + \xi \hat {n})} \right| \\ = \left| \frac {G \nabla_ {\hat {n}} F - F \nabla_ {\hat {n}} G}{G ^ {2}} \right| \\ \leq \frac {\left| \left(\nabla_ {\hat {n}} ^ {2} G (x) + L \epsilon\right) \epsilon^ {2} \left(\nabla_ {\hat {n}} ^ {2} F \left(x _ {\parallel}\right) + L \epsilon\right) \epsilon - \left(\nabla_ {\hat {n}} ^ {2} F \left(x _ {\parallel}\right) - L \epsilon\right) \epsilon^ {2} \left(\nabla_ {\hat {n}} ^ {2} G \left(x _ {\parallel}\right) - L \epsilon\right) \epsilon \right|}{\left(\left(\nabla_ {\hat {n}} ^ {2} G \left(x _ {\parallel}\right) - L \epsilon\right) \epsilon^ {2}\right) ^ {2}} \\ \leq \frac {\left| \left(\nabla_ {\hat {n}} ^ {2} G \left(x _ {\parallel}\right) + L \epsilon\right) \left(\nabla_ {\hat {n}} ^ {2} F \left(x _ {\parallel}\right) + L \epsilon\right) - \left(\nabla_ {\hat {n}} ^ {2} F \left(x _ {\parallel}\right) - L \epsilon\right) \left(\nabla_ {\hat {n}} ^ {2} G \left(x _ {\parallel}\right) - L \epsilon\right) \right|}{\epsilon \left(\nabla_ {\hat {n}} ^ {2} G \left(x _ {\parallel}\right)\right) ^ {2}} \\ \times \left(\frac {\nabla_ {\hat {n}} ^ {2} G (x _ {\parallel})}{\nabla_ {\hat {n}} ^ {2} G (x _ {\parallel}) - L \epsilon}\right) ^ {2} \\ \leq \epsilon^ {- 1} | (1 + \delta^ {- 1} L \epsilon) (1 - \delta + \delta^ {- 1} L \epsilon) - (1 - \delta - \delta^ {- 1} L \epsilon) (1 - \delta^ {- 1} L \epsilon) | \left(1 + (\delta - L \epsilon) ^ {- 1} L \epsilon\right) ^ {2} \\ \leq L \left(3 \delta^ {- 1} - 1\right) \left(1 + (\delta - L \epsilon) ^ {- 1} L \epsilon\right) ^ {2} \\ \leq 2 L \left(3 \delta^ {- 1} - 1\right). \\ \end{array}
$$
Therefore  $F(x_{\parallel} + \xi \hat{n}) / G(x_{\parallel} + \xi \hat{n})$  is  $2L(3\delta^{-1} - 1)$ -Lipschitz as a function of  $\xi$  in  $A_{\epsilon} \backslash A$ . This implies that there exists  $\epsilon > 0$  such that for any  $x \in A_{\epsilon} \backslash A$ ,  $F(x) / 2G(x) < 1 - \delta / 2$  and hence  $\tilde{\mathcal{A}}(l^{\delta}(x)) \leq 0$ . Therefore  $l^{\delta}(x)$  is in the domain of  $\tilde{\mathcal{A}}_{A_{\epsilon}}$  and  $\tilde{\mathcal{A}}_{A_{\epsilon}}(l^{\delta}(x)) = \tilde{\mathcal{A}}(l^{\delta}(x)) \leq 0$ , where  $\tilde{\mathcal{A}}_{A_{\epsilon}}$  is the infinitesimal operator of the stopped process  $\tilde{\theta}_{t \wedge \tau}$  with  $\tau = \inf_{t \geq 0} \{ l(\theta_t) \geq \epsilon \}$ . By the local version of Doob's maximal inequality, for any  $\epsilon' < \epsilon$ ,
$$
\mathbb {P} _ {t \geq 0} ^ {[ \sup  l ^ {\delta} (\tilde {\theta} _ {t}) \geq \epsilon^ {\prime} | \tilde {\theta} _ {0} = x ]} = \mathbb {P} _ {t \geq 0} ^ {[ \sup  l ^ {\delta} (\tilde {\theta} _ {t \wedge \tau}) \geq \epsilon^ {\prime} | \tilde {\theta} _ {0} = x ]} \leq \frac {l ^ {\delta} (x)}{\epsilon^ {\prime}}.
$$
Equivalently, for any  $\epsilon' < \epsilon^{1/\delta}$ ,
$$
\mathbb {P} [ \sup  _ {t \geq 0} l (\tilde {\theta} _ {t}) \geq \epsilon^ {\prime} | \theta_ {0} = x ] \leq \frac {l ^ {\delta} (x)}{\epsilon^ {\prime 1 / \delta}}.
$$
By taking  $l(x)$  to be arbitrarily small, we can bound the left-hand side with any positive value  $\rho > 0$ . This means that  $A$  is stochastically attractive for  $\tilde{\theta}_t$  and so is for  $\theta_t$ .
The sufficient condition is a direct corollary of Theorem 4.2, and thus we show the necessary condition here. By assumption,  $\mathcal{L}''(x) + \frac{1}{2} D''(x) < -\delta$  at  $x = 0$  with some  $\delta > 0$ . It suffices to show that there exists  $\epsilon > 0$  such that the hitting time
$$
\tau := \inf  \left\{t \in \mathbb {R}: \theta_ {t} \geq \epsilon \right\}
$$
is finite almost surely for any  $x\in \mathbb{R}$  with initialization  $\theta_0 = x$
We choose  $\epsilon > 0$  small enough so that  $D(x) > 0$  for any  $x \in (-\epsilon, \epsilon) \setminus \{0\}$ . Then we can consider another stochastic process  $Y_{t} = 2^{-1/2} \int_{\epsilon/2}^{\theta_{t \wedge \tau}} D^{-1/2}(x) dx$ , which follows the SDE below
$$
d Y _ {t} = D ^ {- 1 / 2} \left(- \mathcal {L} ^ {\prime} \left(\theta_ {t \wedge \tau}\right) - \frac {1}{2} D ^ {\prime} \left(\theta_ {t \wedge \tau}\right)\right) d t + d W _ {t}.
$$
for any  $0 \leq t < \tau$ . By further taking  $\epsilon$  small,
$$
D ^ {- 1 / 2} \left(- \mathcal {L} ^ {\prime} \left(\theta_ {t \wedge \tau}\right) - \frac {1}{2} D ^ {\prime} \left(\theta_ {t \wedge \tau}\right)\right) > \frac {\delta \epsilon}{2} D ^ {- 1 / 2} \geq \frac {\delta \epsilon}{2} / \sqrt {\sup  _ {0 <   x <   \epsilon} D (x)} > \delta .
$$
Thus,  $Y_{t} > \delta t + Y_{0} + W_{t}$ . Therefore,
$$
\tau = \inf \{t \in \mathbb {R}: Y _ {t} \geq \frac {1}{\sqrt {2}} \int_ {\epsilon / 2} ^ {\epsilon} D ^ {- 1 / 2} (x) d x \} \leq \inf \{t \in \mathbb {R}: \delta t + W _ {t} \geq \frac {1}{\sqrt {2}} \int_ {\epsilon / 2} ^ {\epsilon} D ^ {- 1 / 2} (x) d x - Y _ {0} \}.
$$
Since  $\delta > 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 \sigma(w_1 x)$ , where  $x, w_1, w_2 \in \mathbb{R}$  and  $\sigma$  represents the activation function. We assume that  $\sigma(x)$  is a smooth Lipschitz continuous function with  $\sigma(0) = 0$ . Let  $\mathcal{D} = \{(x_1, y_1), \dots, (x_n, y_n)\}$  denote the training data set where  $x_i, y_i \in \mathbb{R}, i \in [n]$ . We consider the loss function  $\mathcal{L}$  given by
$$
\mathcal {L} (w _ {1}, w _ {2}, t) = \frac {1}{n} \sum_ {i \in [ n ]} \big (f (x _ {i}; w _ {1}, w _ {2}) - y _ {i} - \zeta_ {i, t} \big) ^ {2},
$$
where  $\zeta_{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
$$
\frac {d w _ {1}}{d t} = - \frac {\partial \mathcal {L}}{\partial w _ {1}} = - \frac {1}{n} \sum_ {i \in [ n ]} (f (x _ {i}; w _ {1}, w _ {2}) - y _ {i}) w _ {2} x \sigma^ {\prime} (w _ {1} x _ {i}) - w _ {2} x \sigma^ {\prime} (w _ {1} x _ {i}) \zeta_ {i, t}
$$
$$
\frac {d w _ {2}}{d t} = - \frac {\partial \mathcal {L}}{\partial w _ {2}} = - \frac {1}{n} \sum_ {i \in [ n ]} (f (x _ {i}; w _ {1}, w _ {2}) - y _ {i}) \sigma (w _ {1} x _ {i}) - \sigma (w _ {1} x _ {i}) \zeta_ {i, t}.
$$
This can be rewritten as a pair of SDEs.
$$
{d w _ {1, t}} = {- \frac {1}{n} \sum_ {i \in [ n ]} (f (x _ {i}; w _ {1}, w _ {2}) - y _ {i}) w _ {2} x _ {i} \sigma^ {\prime} (w _ {1} x _ {i}) d t + \frac {\sqrt {\eta} \zeta}{\sqrt {n}} \sum_ {i \in [ n ]} w _ {2} x _ {i} \sigma^ {\prime} (w _ {1} x _ {i}) d B _ {i, t}}
$$
$$
d w _ {2, t} = - \frac {1}{n} \sum_ {i \in [ n ]} \left(f \left(x _ {i}; w _ {1}, w _ {2}\right) - y _ {i}\right) \sigma \left(w _ {1} x _ {i}\right) d t + \frac {\sqrt {\eta} \zeta}{\sqrt {n}} \sum_ {i \in [ n ]} \sigma \left(w _ {1} x _ {i}\right) d B _ {i, t}, \tag {31}
$$
where  $\zeta > 0$  is a constant representing the amplitude of the gradient noise,  $\eta > 0$  is the learning rate, and  $B_{i,t}$  denotes the  $n$ -dimensional standard Wiener process. Note that the rescaling factor  $1 / \sqrt{n}$  in the diffusion term is introduced to avoid the data size  $n$  influencing the SDE's behavior.
As a preparation, we first show a modified version of Theorem E.1 to incorporate stopped processes.
Lemma F.1. Let  $A \subset \mathbb{R}^d$  denote an affine subset and  $A_c$  be an open neighborhood around  $A$  defined as  $A_c = \{\theta \in \mathbb{R}^d : \inf_{x \in A} \| \theta - x \| < c\}$ , where  $c > 0$  is a constant. Consider a stochastic process  $\{\theta_t \in \mathbb{R}^d : t \geq 0\}$  that follows the SDE Eq. 25 within  $A_c$  with all the regularity conditions mentioned in Theorem E.1. Furthermore, let  $\tau = \inf \{t \geq 0 : \theta_t \notin A_c \cap \operatorname{int} B\}$ , where  $B = \{\theta \in \mathbb{R}^d : \| \theta_i \| \leq c_B\}$  with  $i \in [d]$  and  $c_B > 0$ . If  $A \cap B$  is an invariant set of  $\theta_t$ , then  $A \cap B$  is a stochastically attractive invariant set of the stopped process  $\{\theta_{t \wedge \tau} : t \geq 0\}$ , provided Eq. 26 and 27 hold for any  $x \in A \cap 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  $\tau$ . This allows us to consider the unique stopped process  $\theta_{t\wedge \tau}$ , a continuous strongly Markov process. Let  $l(\theta):\mathbb{R}^d\to \mathbb{R}_{\geq 0}$  be the Euclidean distance between  $\theta \in \mathbb{R}^d$  and the invariant set  $A\cap B$ , and  $\tilde{\mathcal{A}}$  be the infinitesimal operator of  $\theta_{t\wedge \tau}$  where  $\tau_{\epsilon} = \inf \{t\geq 0:\theta_t\notin A_\epsilon \}$ . The same argument in Theorem E.1 leads to  $\tilde{\mathcal{A}} (l^{\lambda}(x))\leq 0$  for any  $x\in A_{\epsilon}\cap B$  with small enough  $\epsilon >0$  and  $\lambda >0$ . By applying a local version of maximal inequality (Lemma 1 in Kushner [44]), we obtain, for any  $x\in A_{\epsilon}$  and  $0 < \epsilon^{\prime} < \epsilon$
$$
\mathbb {P} \left[ \sup _ {t \geq 0} l (\theta_ {t \wedge \tau}) \geq \epsilon^ {\prime} \Big | \theta_ {0} = x \right] = \mathbb {P} \left[ \sup _ {t \geq 0} l (\theta_ {t \wedge \tau \wedge \tau_ {\epsilon}}) \geq \epsilon^ {\prime} \Big | \theta_ {0} = x \right] \leq \frac {l ^ {\lambda} (x)}{\epsilon^ {\prime}}.
$$
This implies that  $A \cap B$  is a stochastically attractive invariant set of the stopped process  $\theta_{t \wedge \tau}$ .
Theorem 5.1 (formal). Suppose  $\{(w_{1,t}, w_{2,t}) \in \mathbb{R}^2 : t \geq 0\}$  is a stochastic process obeying Eq. 31 in a neighborhood  $U$  of  $A = \{(0,0)\} \subset \mathbb{R}^2$ , and the activation function  $\sigma : \mathbb{R} \to \mathbb{R}$  is a smooth function such that  $\sigma(0) = 0$ . The invariant set  $A$  is stochastically attractive if  $\frac{\eta \zeta^2 |\sigma'(0)|}{2} > \frac{\left| \sum_{i \in [n]} x_i y_i \right|}{\sum_{i \in [n]} x_i^2}$ .
Proof. If  $\sigma'(0) = 0$ , the condition can never be met. Therefore, we only consider the case  $\sigma'(0) \neq 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:
$$
\left\{ \begin{array}{l} d p _ {t} = \mu_ {p} (p _ {t}, s _ {t}) d t + \frac {1}{\sqrt {n}} \sum_ {i \in [ n ]} \sigma_ {p, i} (p _ {t}, s _ {t}) d B _ {i, t} \\ d s _ {t} = \mu_ {s} (p _ {t}, s _ {t}) d t + \frac {1}{\sqrt {n}} \sum_ {i \in [ n ]} \sigma_ {s, i} (p _ {t}, s _ {t}) d B _ {i, t} \end{array} \right.
$$
$$
\left\{ \begin{array}{l} \mu_ {p} (p, s) = - \frac {1}{n} \sum_ {i \in [ n ]} \left[ (f (x _ {i}; w _ {1, t}, w _ {2, t}) - y _ {i}) w _ {2, t} \sigma^ {\prime} (w _ {1, t} x _ {i}) x _ {i} + (f (x _ {i}; w _ {1, t}, w _ {2, t}) - y _ {i}) \sigma (w _ {1, t} x _ {i}) \right] \\ \mu_ {s} (p, s) = - \frac {1}{n} \sum_ {i \in [ n ]} \left[ (f (x _ {i}; w _ {1, t}, w _ {2, t}) - y _ {i}) w _ {2, t} \sigma^ {\prime} (w _ {1, t} x _ {i}) x _ {i} - (f (x _ {i}; w _ {1, t}, w _ {2, t}) - y _ {i}) \sigma (w _ {1, t} x _ {i}) \right] \\ \sigma_ {p, i} (p, s) = \sqrt {\eta} \zeta w _ {2, t} \sigma^ {\prime} (w _ {1, t} x _ {i}) x _ {i} + \sqrt {\eta} \zeta \sigma (w _ {1, t} x _ {i}) \\ \sigma_ {s, i} (p, s) = \sqrt {\eta} \zeta w _ {2, t} \sigma^ {\prime} (w _ {1, t} x _ {i}) x _ {i} - \sqrt {\eta} \zeta \sigma (w _ {1, t} x _ {i}) \end{array} \right.
$$
Define stopping times  $\tau_{p} = \inf_{t\geq 0}\{p_{t}\geq \epsilon \}$ ,  $\tau_{s} = \inf_{t\geq 0}\{s_{t}\geq \epsilon \}$  where  $\epsilon >0$  satisfies  $\{(p,s):$ $\max (|p|,|s|) < \epsilon \} \subset U$ , and consider the stopped processes  $(p_{t\wedge \tau_p},s_{t\wedge \tau_p})$  and  $(p_{t\wedge \tau_s},s_{t\wedge \tau_s})$ . We have,
$$
\begin{array}{l} \frac {\partial}{\partial p} \mu_ {p} (0, 0) - \frac {1}{4 n} \frac {\partial^ {2}}{\partial p ^ {2}} \sum_ {i \in [ n ]} \sigma_ {p, i} ^ {2} (0, 0) = - \frac {\sigma^ {\prime} (0)}{n} \sum_ {i \in [ n ]} y _ {i} x _ {i} - \frac {1}{2 n} \eta \zeta^ {2} \sigma^ {\prime} (0) ^ {2} \sum_ {i \in [ n ]} x _ {i} ^ {2} <   0 \\ \frac {1}{n} \frac {\partial^ {2}}{\partial p ^ {2}} \sum_ {i \in [ n ]} \sigma_ {p, i} ^ {2} (0, 0) = \frac {1}{2 n} \eta \zeta^ {2} \sigma^ {\prime} (0) ^ {2} \sum_ {i \in [ n ]} x _ {i} ^ {2} > 0, \\ \end{array}
$$
By the assumption on smoothness,  $\frac{\partial}{\partial p}\mu_p(0,s) - \frac{1}{4}\frac{\partial^2}{\partial p^2}\sum_{i\in [n]}\sigma_{p,i}^2 (0,s) < 0$  and  $\frac{\partial^2}{\partial p^2}\sum_{i\in [n]}\sigma_{p,i}^2 (0,s) > 0$  for any  $s\in [-\epsilon ,\epsilon ]$  with small enough  $\epsilon >0$ . Therefore by Lemma F.1,  $A_{p}\coloneqq \{(0,s):s\in [-\epsilon ,\epsilon ]\}$  is stochastically attractive for the process  $\{(p_{t\wedge \tau_p},s_{t\wedge \tau_p}):t\geq 0\}$ . Similarly,
$$
\begin{array}{l} \frac {\partial}{\partial s} \mu_ {s} (0, 0) - \frac {1}{4 n} \frac {\partial^ {2}}{\partial s ^ {2}} \sum_ {i \in [ n ]} \sigma_ {s, i} ^ {2} (0, 0) = \frac {\sigma^ {\prime} (0)}{n} \sum_ {i \in [ n ]} y _ {i} x _ {i} - \frac {1}{2 n} \eta \zeta^ {2} \sigma^ {\prime} (0) ^ {2} \sum_ {i \in [ n ]} x _ {i} ^ {2} <   0 \\ \frac {1}{n} \frac {\partial^ {2}}{\partial s ^ {2}} \sum_ {i \in [ n ]} \sigma_ {s, i} ^ {2} (0, 0) = \frac {1}{2 n} \eta \xi^ {2} \sigma^ {\prime} (0) ^ {2} \sum_ {i \in [ n ]} x _ {i} ^ {2} > 0. \\ \end{array}
$$
By Lemma F.1,  $A_{s} \coloneqq \{(p,0) : p \in [-\epsilon, \epsilon]\}$  is stochastically attractive for the process  $\{(p_{t \wedge \tau_{s}}, s_{t \wedge \tau_{s}}) : t \geq 0\}$ .
By combining these two facts, we see that  $A_{p} \cap A_{s} = \{(0,0)\}$  is stochastically attractive for the process  $\{(w_{1,t}, w_{2,t}) : t \geq 0\}$ .
G The permutation invariant set in one hidden layer nonlinear neural network
Consider training a two neuron neural network  $f_{\theta}$  on a dataset  $\{(x_k, y_k)\}_{k=1}^N$ , where  $f_{\theta}(x) = \sum_{i=1}^{2} w_i^{(2)} \sigma(w_i^{(1)} x)$ . We assume the data are drawn randomly from standard Gaussian distribution. Consider running gradient descent with learning rate  $\eta$  on the mean square error loss with label noise drawn freshly from  $\mathcal{N}(0, \sigma^2)$ . In this set up, the invariant set from the permutation symmetry is the affine space  $\{\theta | w_1^{(i)} = w_2^{(i)} \text{ 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.