SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 7: Gradient noises in SGD induces stochastic collapsing to the invariant set associated with identical neurons. Left: The collapsing process of the weights of hidden neurons 1 and 2 becoming identical over the course of training, demonstrating SGD's bias towards simpler subnetworks. Middle: A visualization of the impact of learning rates on the stochastic attractivity. An increase in the learning rate accelerates the collapsing process until reaching a critical threshold, beyond which the collapsing ceases. Right: The influence of noise scale on the collapsing process. The attractivity strengthens with increased noise up to a certain threshold.
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:
$$
h ^ {(l _ {2})} (x) = \sigma \left(w ^ {(l _ {2} - 1)} h ^ {(l _ {2} - 1)} (x) + b ^ {(l _ {2})}\right) + h ^ {(l _ {1})} (x),
$$
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_{in,p}^{(l)}, b_{p}^{(l)})$  denote the parameters (weights and bias) directly incoming to this neuron, and  $w_{out,p}^{(l+1)}$  represent the parameters (weights) directly outgoing from the neuron. If the non-linearity  $\sigma$  is origin-passing (i.e.,  $\sigma(0) = 0$ ), then the coordinate plane  $A = \{\theta | w_{in,p}^{(l)} = 0, b_{p}^{(l)} = 0, w_{out,p}^{(l+1)} = 0 \text{ 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_{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 space  $A = \{\theta | w_{in,p}^{(l)} = w_{in,q}^{(l)}, b_{p}^{(l)} = b_{q}^{(l)}, w_{out,p}^{(l+1)} = w_{out,q}^{(l+1)} \text{ 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  $\bar{W}:\mathbb{R}^n\to \mathbb{R}^m$ .  $\bar{W}\in \mathbb{R}^{m\times n}$  is a low-rank matrix where  $k\ll \min \{m,n\}$  denotes the rank.
Two-layer student network. We consider a two-layer student network, parameterized by  $W_{1} \in \mathbb{R}^{d \times n}$  and  $W_{2} \in \mathbb{R}^{m \times d}$ , that computes the linear map  $\hat{W} = W_{2}W_{1} : \mathbb{R}^{n} \to \mathbb{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  $\mathbb{R}^{n}$  to  $\mathbb{R}^{m}$ .
Noisy training data. The student network is trained on data generated by the teacher network through the noisy map,
$$
y = \bar {y} + \epsilon \quad \text {w h e r e} \quad \bar {y} = \bar {W} x, \tag {32}
$$
and  $\epsilon \sim \mathcal{N}(0,\sigma_{\epsilon}^{2}I_{m})$ . We will consider a fixed training dataset of  $p$  noisy input-output pairs  $\{(x_1,y_1),\ldots ,(x_p,y_p)\}$  generated by the teacher network from input  $x_{i}\sim \mathcal{N}(0,\sigma_{x}^{2}I_{n})$ . Let  $X\in \mathbb{R}^{n\times p}$  and  $Y\in \mathbb{R}^{m\times 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,
$$
\Sigma_ {x x} = X X ^ {\intercal} = \sum_ {i = 1} ^ {p} x _ {i} x _ {i} ^ {\intercal}, \quad \Sigma_ {y x} = Y X ^ {\intercal} = \sum_ {i = 1} ^ {p} y _ {i} x _ {i} ^ {\intercal}. \tag {33}
$$
Later in our analysis of the learning dynamics we will assume the input data matrix  $X$  is whitened such that  $\Sigma_{xx} = I_n$  implying that only the input-output covariance structure  $\Sigma_{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  $\hat{Y} = \hat{W} X$  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
$$
\mathcal {E} _ {\text {t r a i n}} = \frac {\| \hat {W} X - Y \| _ {F} ^ {2}}{\| Y \| _ {F} ^ {2}}, \quad \mathcal {E} _ {\text {t e s t}} = \frac {\mathbb {E} \left[ \| \hat {W} x - \bar {W} x \| _ {F} ^ {2} \right]}{\mathbb {E} \left[ \| \bar {W} x \| _ {F} ^ {2} \right]}, \tag {34}
$$
where  $x \sim \mathcal{N}(0, \sigma_x^2 I_n)$  is a new test point sampled from the input distribution,  $\| \cdot \|_F$  denotes the Frobenius norm, and  $\mathbb{E}[\cdot]$  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  $\mathbb{R}^{m\times n}$ : the low rank teacher  $\bar{W}$ , the overparameterized student  $\hat{W}$ , and the noisy training data  $\Sigma_{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,
$$
\bar {W} = \sum_ {i = 1} ^ {k} \bar {s} _ {i} \bar {u} _ {i} \bar {v} _ {i} ^ {\mathsf {T}}, \quad \hat {W} = \sum_ {i = 1} ^ {d} \hat {s} _ {i} \hat {u} _ {i} \hat {v} _ {i} ^ {\mathsf {T}}, \quad \Sigma_ {y x} = \sum_ {i = 1} ^ {d} \tilde {s} _ {i} \tilde {u} _ {i} \tilde {v} _ {i} ^ {\mathsf {T}}. \tag {35}
$$
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  $\tilde{\cdot}$ ,  $\hat{\cdot}$ , or  $\tilde{\cdot}$  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  $\eta$  to minimize  $\mathcal{E}_{\mathrm{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,
$$
W _ {1} ^ {(t + 1)} = W _ {1} ^ {(t)} - \eta \left(W _ {2} ^ {(t)}\right) ^ {\mathsf {T}} \left(W _ {2} ^ {(t)} W _ {1} ^ {(t)} \Sigma_ {x x} - \Sigma_ {y x} + Z ^ {(t)} X ^ {\mathsf {T}}\right), \tag {36}
$$
$$
W _ {2} ^ {(t + 1)} = W _ {2} ^ {(t)} - \eta \left(W _ {2} ^ {(t)} W _ {1} ^ {(t)} \Sigma_ {x x} - \Sigma_ {y x} + Z ^ {(t)} X ^ {\intercal}\right) \left(W _ {1} ^ {(t)}\right) ^ {\intercal}, \tag {37}
$$
where  $W_1^{(t)}$  and  $W_2^{(t)}$  are the parameters after  $t$  steps of training and  $Z^{(t)} \in \mathbb{R}^{m \times 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  $\eta \to 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  $\Sigma_{xx} = I_n$  such that the input data matrix  $X$  is whitened (this implicitly assumes  $p \geq n$  such that we have at least as many observations as features).
A2. Structured Label-Noise. We assume the gradient noise can be decomposed as  $Z^{(t)} = \tilde{U} \mathrm{diag}\left(z^{(t)}\right) \tilde{V}^{\mathsf{T}}X$  where  $z^{(t)} \sim \mathcal{N}(0, \zeta^2)$ .
A3. Spectral Initialization. We assume that the student network is initialized such that  $W_{1}^{(0)} = OA^{(0)}\tilde{V}^{\intercal}$  and  $W_{2}^{(0)} = \tilde{U} B^{(0)}O^{\intercal}$  where  $A^{(0)}$  and  $B^{(0)}$  are diagonal matrices such that  $\hat{S}(0) = B^{(0)}A^{(0)}$  and  $O \in \mathbb{R}^{d \times d}$  is a random orthonormal matrix such that  $O^{\intercal}O = I_{d}$ .
A4. Balanced Initialization. We assume a balanced initialization such that  $W_1^{(0)}W_1^{(0)\top} = W_2^{(0)\top}W_2^{(0)}$ .
Given the first three assumptions, the dynamics decouple in the eigenbasis of  $\Sigma_{yx}$ . Under the change of variables  $A^{(t)} = O^{\intercal}W_1^{(t)}\tilde{V}$  and  $B^{(t)} = \tilde{U}^{\intercal}W_2^{(t)}O$ , the dynamics transform to
$$
A ^ {(t + 1)} = A ^ {(t)} - \eta (B ^ {(t)}) ^ {\intercal} (B ^ {(t)} A ^ {(t)} - \tilde {S} + \operatorname {d i a g} (z ^ {(t)})), \tag {38}
$$
$$
B ^ {(t + 1)} = B ^ {(t)} - \eta \left(B ^ {(t)} A ^ {(t)} - \tilde {S} + \operatorname {d i a g} \left(z ^ {(t)}\right)\right) \left(A ^ {(t)}\right) ^ {\intercal}, \tag {39}
$$
which because  $A^{(0)}$  and  $B^{(0)}$  are diagonal matrices by assumption, decouples into a system of scalar equations. Taking the limit as  $\eta \to 0$  we can approximate each of these scalar equations with a two-dimensional non-linear SDE,
$$
d \left[ \begin{array}{l} a _ {i} \\ b _ {i} \end{array} \right] = - \left[ \begin{array}{l} b _ {i} \left(a _ {i} b _ {i} - \tilde {s} _ {i}\right) \\ a _ {i} \left(a _ {i} b _ {i} - \tilde {s} _ {i}\right) \end{array} \right] d t + \sqrt {\eta \zeta^ {2}} \left[ \begin{array}{l} b _ {i} \\ a _ {i} \end{array} \right] d B _ {t}, \tag {40}
$$
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 Ito'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\to \infty}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 Ito's Lemma,  $r_i(t)$  is driven by the ODE  $dr_i = -\eta \zeta^2 r_i dt$ , which has the temporal solution  $r_i(t) = r_i(0)e^{-\eta \zeta^2 t}$ . Because  $\eta \zeta^2 > 0$ , then the  $\lim_{t \to \infty} 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,  $\ell_i(w_i) = (w_i^2 -\tilde{s}_i)^2$  with  $w_{i} = a_{i} = b_{i}$ , where within each slice the dynamics evolve according to the non-linear SDE,
$$
d w _ {i} = - w _ {i} \left(w _ {i} ^ {2} - \tilde {s} _ {i}\right) d t + \sqrt {\eta \zeta^ {2}} w _ {i} d B _ {t}. \tag {41}
$$
The geometry of these slices is controlled by the strength of the singular mode for the training data  $\tilde{s}_i$  and the dynamics are determined by the relationship with the student's singular mode strength  $\hat{s}_i = w_i^2$ . By Ito's Lemma we find that the dynamics of  $\hat{s}_i$  are governed by the non-linear SDE,
$$
d \hat {s} _ {i} = 2 \hat {s} _ {i} \left(\left(\tilde {s} _ {i} + \frac {\eta \zeta^ {2}}{2}\right) - \hat {s} _ {i}\right) d t + 2 \sqrt {\eta \zeta^ {2}} \hat {s} _ {i} d B _ {t}, \tag {42}
$$
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  $\hat{s}_i(t)$  for  $t\geq 0$  are governed by the stochastic process,
$$
\hat {s} _ {i} (t) = \frac {e ^ {\left(2 \tilde {s} _ {i} - \eta \zeta^ {2}\right) t + 2 \sqrt {\eta \zeta^ {2}} B _ {t}}}{2 \int_ {0} ^ {t} e ^ {\left(2 \tilde {s} _ {i} - \eta \zeta^ {2}\right) \tau + 2 \sqrt {\eta \zeta^ {2}} B _ {\tau}} d \tau + \hat {s} _ {i} (0) ^ {- 1}}. \tag {43}
$$
Proof. Let  $u_{i} = \hat{s}_{i}^{-1}$ . Then by Itô's Lemma, the dynamics of  $u_{i}$  are given by the SDE,
$$
d u _ {i} = \left(2 - \left(2 \tilde {s} _ {i} - 3 \eta \zeta^ {2}\right) u _ {i}\right) d t - 2 \sqrt {\eta \zeta^ {2}} u _ {i} d B _ {t} \tag {44}
$$
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)
$$
d u _ {i} = \theta (\mu - u _ {i}) d t + \sigma u _ {i} d B _ {t} \tag {45}
$$
where  $\theta = 2\tilde{s}_i - 3\eta \zeta^2$ ,  $\mu = (\tilde{s}_i - \frac{3}{2}\eta \zeta^2)^{-1}$ , and  $\sigma = -2\sqrt{\eta\zeta^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,
$$
d F = - \theta F d t + \sigma F d B _ {t}, \quad d G = \alpha (t) d t + \beta (t) d B _ {t}. \tag {46}
$$
To determine the unknown coefficients  $\alpha(t)$  and  $\beta(t)$  defining the SDE for  $G(t)$ , we can apply Itô's product rule
$$
\begin{array}{l} d u _ {i} = d F G + F d G + d F d G \\ = (- \theta F G + \alpha (t) F + \beta (t) \sigma F) d t + (\sigma F G + \beta (t) F) d B _ {t} \\ = (\alpha (t) F + \beta (t) \sigma F - \theta u _ {i}) d t + (\sigma u _ {i} + \beta (t) F) d B _ {t} \\ \end{array}
$$
Aligning this expression with the original SDE for  $u_{i}$  we see that
$$
\alpha (t) = \theta \mu F ^ {- 1}, \quad \beta (t) = 0. \tag {47}
$$
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,
$$
F (t) = e ^ {- \left(\theta + \frac {1}{2} \sigma^ {2}\right) t + \sigma B _ {t}} \tag {48}
$$
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
$$
G (t) = u _ {0} + \theta \mu \int_ {0} ^ {t} e ^ {\left(\theta + \frac {1}{2} \sigma^ {2}\right) \tau - \sigma B _ {s}} d \tau . \tag {49}
$$
Combining our solutions for  $F(t)$  and  $G(t)$  we get that the solution for  $u_{i}(t)$  is
$$
u _ {i} (t) = u _ {i} (0) e ^ {- \left(\theta + \frac {1}{2} \sigma^ {2}\right) t + \sigma B _ {t}} + \theta \mu \int_ {0} ^ {t} e ^ {- \left(\theta + \frac {1}{2} \sigma^ {2}\right) (t - \tau) + \sigma \left(B _ {t} - B _ {\tau}\right)} d \tau . \tag {50}
$$
Inverting this expression and rearranging terms, gives the expression for  $\hat{s}_i(t)$
$$
\hat {s} _ {i} (t) = u _ {i} (t) ^ {- 1} = \frac {e ^ {\left(\theta + \frac {1}{2} \sigma^ {2}\right) t - \sigma B _ {t}}}{\theta \mu \int_ {0} ^ {t} e ^ {\left(\theta + \frac {1}{2} \sigma^ {2}\right) \tau - \sigma B _ {\tau}} d \tau + \hat {s} _ {i} (0) ^ {- 1}} \tag {51}
$$
Plugging in the simplifications  $\theta \mu = 2$ ,  $\theta + \frac{1}{2}\sigma^2 = 2\tilde{s}_i - \eta \zeta^2$ , and  $\sigma = -2\sqrt{\eta\zeta^2}$  gives the final expression.
Corollary 6.1. In the limit of  $\hat{s}_i(0) \to 0$ , for any  $t > 0$ ,  $\hat{s}_i(t)$  is given by
$$
\hat {s} _ {i} (t) / \hat {s} _ {i} (0) \underset {a. s.} {\rightarrow} e ^ {\left(2 \tilde {s} _ {i} - \eta \zeta^ {2}\right) t + 2 \sqrt {\eta \zeta^ {2}} B _ {t}}. \tag {52}
$$
Proof. Since  $B_{t}$  is almost surely continuous,  $(2\tilde{s}_{i} - \eta \zeta^{2})\tau + 2\sqrt{\eta\zeta^{2}} B_{\tau}$  has a maximum in  $\tau \in [0,t]$  and therefore  $2\int_0^t e^{(2\tilde{s}_i - \eta \zeta^2)\tau + 2\sqrt{\eta\zeta^2} B_\tau}d\tau$  is finite almost surely. Therefore, by Theorem 6.1,
$$
\hat {s} _ {i} (t) / \hat {s} _ {i} (0) = \frac {e ^ {\left(2 \tilde {s} _ {i} - \eta \zeta^ {2}\right) t + 2 \sqrt {\eta \zeta^ {2}} B _ {t}}}{2 \hat {s} _ {i} (0) \int_ {0} ^ {t} e ^ {\left(2 \tilde {s} _ {i} - \eta \zeta^ {2}\right) \tau + 2 \sqrt {\eta \zeta^ {2}} B _ {\tau}} d \tau + 1} \rightarrow {} ^ {a. s.} e ^ {\left(2 \tilde {s} _ {i} - \eta \zeta^ {2}\right) t + 2 \sqrt {\eta \zeta^ {2}} B _ {t}}. \tag {53}
$$
Corollary 6.1 has important implications for the dynamics of  $\hat{s}_i(t)$ . It demonstrates that the dynamics observe a phase transition determined by the ratio  $2\tilde{s}_i / \eta \zeta^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  $\hat{s}_i$ .
Corollary I.1 (Stationary solution). In the limit of  $t \to \infty$ ,  $\hat{s}_i(t)$  is distributed according to the density,
$$
p _ {s s} \left(\hat {s} _ {i}\right) = Z ^ {- 1} e ^ {- \frac {\hat {s} _ {i}}{\eta \zeta^ {2}}} \hat {s} _ {i} ^ {- \left(\frac {3}{2} - \frac {\bar {s} _ {i}}{\eta \zeta^ {2}}\right)} \tag {54}
$$
where  $Z = \int_0^{+\infty}e^{-\frac{\hat{s}_i}{\eta\zeta^2}}\hat{s}_i^{-\left(\frac{3}{2} -\frac{\bar{s}_i}{\eta\zeta^2}\right)}d\hat{s}_i$  is the partition function.
Proof. To check if  $p_{ss}(\hat{s}_i)$  is the stationary solution we use the Fokker-Planck equation for the dynamics of  $\hat{s}_i$ , which is
$$
\partial_ {t} p = \nabla \cdot \underbrace {\left(2 \hat {s} _ {i} \left(\hat {s} _ {i} - \left(\tilde {s} _ {i} + \frac {\eta \zeta^ {2}}{2}\right)\right) p + 2 \eta \zeta^ {2} \nabla \cdot \left(\hat {s} _ {i} ^ {2} p\right)\right)} _ {- J}, \quad p \left(\hat {s} _ {i}, 0\right) = \delta_ {\hat {s} _ {i} (0)}, \tag {55}
$$
where  $J$  is commonly referred to as the probability current. The stationary solution  $p_{ss}$  satisfies  $\partial_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}(\hat{s}_i)$  and using the simplification
$$
\begin{array}{l} 2 \eta \zeta^ {2} \nabla \cdot \left(\hat {s} _ {i} ^ {2} p\right) = 4 \zeta^ {2} \hat {s} _ {i} p _ {s s} + 2 \eta \zeta^ {2} \hat {s} _ {i} ^ {2} \nabla p _ {s s} (\hat {s} _ {i}) (56) \\ = 4 \zeta^ {2} \hat {s} _ {i} p _ {s s} - 2 \eta \zeta^ {2} \hat {s} _ {i} ^ {2} \left(\left(\eta \zeta^ {2}\right) ^ {- 1} + \hat {s} _ {i} ^ {- 1} \left(\frac {3}{2} - \frac {\tilde {s} _ {i}}{\eta \zeta^ {2}}\right)\right) p _ {s s} (\hat {s} _ {i}) (57) \\ = - 2 \hat {s} _ {i} \left(\hat {s} _ {i} - \left(\tilde {s} _ {i} + \frac {\eta \zeta^ {2}}{2}\right)\right) p _ {s s} (58) \\ \end{array}
$$
we can verify that  $J_{ss} = 0$
The stationary solution can be interpreted as a Gibb's distribution (i.e.  $p_{ss} \propto e^{-\kappa \Psi(x)}$ ) with the potential
$$
\Psi \left(\hat {s} _ {i}\right) = \hat {s} _ {i} + \left(\frac {3 \eta \zeta^ {2}}{2} - \tilde {s} _ {i}\right) \log \left(\hat {s} _ {i}\right) \tag {59}
$$
and temperature constant  $\kappa = \left(\eta \zeta^{2}\right)^{-1}$ . From this interpretation we see two phase transitions,
$$
\arg \max  _ {s s} \left(\hat {s} _ {i}\right) = \left\{ \begin{array}{l l} \tilde {s} _ {i} - \frac {3}{2} \eta \zeta^ {2} & \text {i f} \tilde {s} _ {i} > \frac {3}{2} \eta \zeta^ {2} \\ 0 & \text {i f} \tilde {s} _ {i} \leq \frac {3}{2} \eta \zeta^ {2} \end{array} \quad p _ {s s} (\hat {s} _ {i}) = \left\{ \begin{array}{l l} Z ^ {- 1} e ^ {- \kappa \Psi (w)} & \text {i f} \tilde {s} _ {i} > \frac {\eta \zeta^ {2}}{2} \\ \delta_ {0} & \text {i f} \tilde {s} _ {i} \leq \frac {\eta \zeta^ {2}}{2} \end{array} \right. \right. \tag {60}
$$
When  $\frac{3}{2}\eta \zeta^2 \leq \tilde{s}_i$ , the most probable state is  $\tilde{s}_i - \frac{3}{2}\eta \zeta^2$ , where we already see a bias from the stochasticity towards the origin. When  $\frac{3}{2}\eta \zeta^2 > \tilde{s}_i$ , the most probable state is at the origin. When  $\frac{\eta \zeta^2}{2} \geq \tilde{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  $\mathbb{E}[\hat{s}_i(t)]$ :
Conjecture I.1. Under assumptions A1 - A4, the expectation  $\mathbb{E}[\hat{s}_i(t)]$  for  $t\geq 0$  is given by the equation,
$$
\mathbb {E} \left[ \hat {s} _ {i} (t) \right] = \frac {\left(\tilde {s} _ {i} - \frac {\eta \zeta^ {2}}{2}\right) e ^ {\left(2 \tilde {s} _ {i} - \eta \zeta^ {2}\right) t}}{e ^ {\left(2 \tilde {s} _ {i} - \eta \zeta^ {2}\right) t} - 1 + \frac {\tilde {s} _ {i} - \frac {\eta \zeta^ {2}}{2}}{\hat {s} _ {i} (0)}}. \tag {61}
$$
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  $\zeta \to 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  $\lambda = \frac{\eta\zeta^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 SVD structure of the low-rank teacher  $\{\bar{u}_i,\bar{s}_i,\bar{v}_i\}$ , overparameterized student  $\{\hat{u}_i,\hat{s}_i,\hat{v}_i\}$ , and noisy training data  $\{\tilde{u}_i,\tilde{s}_i,\tilde{v}_i\}$  we can decompose these error terms as
$$
\mathcal {E} _ {\text {t r a i n}} = \left(\sum_ {i = 1} ^ {m} \tilde {s} _ {i} ^ {2}\right) ^ {- 1} \left(\sum_ {i = 1} ^ {m} \tilde {s} _ {i} ^ {2} + \sum_ {j = 1} ^ {d} \hat {s} _ {j} ^ {2} - 2 \sum_ {i = 1} ^ {m} \sum_ {j = 1} ^ {d} \tilde {s} _ {i} \hat {s} _ {j} \langle \tilde {u} _ {i}, \hat {u} _ {j} \rangle \langle \tilde {v} _ {i}, \hat {v} _ {j} \rangle\right), \tag {62}
$$
$$
\mathcal {E} _ {\text {t e s t}} = \left(\sum_ {i = 1} ^ {k} \bar {s} _ {i} ^ {2}\right) ^ {- 1} \left(\sum_ {i = 1} ^ {k} \bar {s} _ {i} ^ {2} + \sum_ {j = 1} ^ {d} \hat {s} _ {j} ^ {2} - 2 \sum_ {i = 1} ^ {k} \sum_ {j = 1} ^ {d} \bar {s} _ {i} \hat {s} _ {j} \langle \bar {u} _ {i}, \hat {u} _ {j} \rangle \langle \bar {v} _ {i}, \hat {v} _ {j} \rangle\right). \tag {63}
$$
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.