SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Random matrix theory analysis. We will work in the limit  $n, m \to \infty$  with a finite rank  $k \sim O(1)$  and finite aspect ratio  $\mathcal{A} = \frac{m}{n} \in (0, \infty)$ . Without loss of generality, we assume  $\sigma_{\epsilon}^{2} = n^{-1}$  such that the singular values of the noise are  $O(1)$  and thus the singular values of the teacher network can be understood as signal to noise ratios (SNRs). The relationship between the SVD structures of a low-rank matrix (the teacher network  $\{\bar{u}_i, \bar{s}_i, \bar{v}_i\}$ ) and a noisy perturbation (the training data  $\{\tilde{u}_i, \tilde{s}_i, \tilde{v}_i\}$ ) is well studied in [59], which we summarize here.
In the limit, the top  $k$  singular values of  $\Sigma_{yx}$  converge to values given by the transfer function,
$$
\tilde {s} (\bar {s}) = \left\{ \begin{array}{l l} (\bar {s}) ^ {- 1} \sqrt {\left(1 + \bar {s} ^ {2}\right) \left(\mathcal {A} + \bar {s} ^ {2}\right)} & \text {i f} \bar {s} > \mathcal {A} ^ {1 / 4} \\ 1 + \mathcal {A} & \text {o t h e r w i s e} \end{array} \right. \tag {64}
$$
while the bottom  $m - k$  singular values are distributed according to the Marchenko-Pastur (MP) distribution,
$$
p _ {M P} (\tilde {s}) = \left\{ \begin{array}{l l} \frac {\sqrt {4 \mathcal {A} - (\tilde {s} ^ {2} - (1 + \mathcal {A})) ^ {2}}}{\pi \mathcal {A} \tilde {s}} & \text {i f} \tilde {s} \in [ 1 - \sqrt {\mathcal {A}}, 1 + \sqrt {\mathcal {A}} ] \\ 0 & \text {o t h e r w i s e} \end{array} \right. \tag {65}
$$
Additionally, in the limit, the top  $k$  singular vectors of  $\Sigma_{yx}$  acquire an alignment with the  $k$  singular vectors the teacher network given by the relationship  $|\langle \tilde{u}_i,\bar{u}_i\rangle ||\langle \tilde{v}_i,\bar{v}_i\rangle | = \mathcal{O}(\bar{s}_i)$  where the overlap function is defined as,
$$
\mathcal {O} \left(\bar {s} _ {i}\right) = \left\{ \begin{array}{l l} \left(1 - \frac {\mathcal {A} \left(1 + \bar {s} ^ {2}\right)}{\bar {s} \left(\mathcal {A} + \bar {s} ^ {2}\right)}\right) ^ {1 / 2} \left(1 - \frac {\left(\mathcal {A} + \bar {s} ^ {2}\right)}{\bar {s} \left(1 + \bar {s} ^ {2}\right)}\right) ^ {1 / 2} & \text {i f} \bar {s} > \mathcal {A} ^ {1 / 4} \\ 0 & \text {o t h e r w i s e} \end{array} \right. \tag {66}
$$
J Experiment details
The codes to reproduce the experiments in the main paper can be found at https://github.com/ccffccffcc/stochastic-collapse. Our the experiments were run on the Google Cloud Platform (with  $4 \times$  or  $8 \times$  NVIDIA A100 (40GB) GPU). The initial code development occurred on a local cluster equipped with  $10 \times$  NVIDIA TITAN X GPUs. We carried out all the deep learning experiments with VGG-16 [47] and ResNet-18 [48], training on the CIFAR-10 and CIFAR-100 datasets respectively [49]. For VGG-16, we did not use Batch Normalization or Dropout. Both VGG-16 and ResNet-18 had their nonlinearity activation functions replaced with GELU, as the potentially larger invariant set associated with ReLU could complicate identification and illustration of the two main invariant sets discussed in the paper. We leave it as future work to study in details the structure of the invariant sets associated with ReLU. For all our training, we applied standard data augmentation and used SGD (with momentum  $\beta = 0.9$  and weight decay of 0.0005) as the optimizer.
Stochastic collapse in a quartic loss. (Fig.1) The left three plot are generated by sampling 50 trajectories of the SDE  $d\theta_{t} = -(\theta_{t}^{3} - \mu \theta_{t})dt + \zeta \theta_{t}dB_{t}$  with  $\mu = 1.5$  and  $\zeta = 0.1, 1, 2.5$ . The samples are computed using an Euler-Maruyama discretization scheme with step size  $\eta = 0.0025$  for  $10^{4}$  steps. All three plots use the same random Gaussian initializations with mean zero and variance four. The side plot for these three subplots shows the theoretical steady-state distribution described in Sec. 4. The rightmost plot is generated by sweeping 50 linearly spaced values of  $\zeta$  from  $\zeta = 0$  to  $\zeta = 3$ . For each value of  $\zeta$  we sample 10000 trajectories with the same scheme described previously but for a varying number of steps: 0, 5000, 10000, 20000, 40000, 80000. We compute the empirical probability of the final position being within  $\epsilon = 1e^{-15}$  euclidean distance of the origin.
Stochastic collapse towards a sign invariant set. (Fig.2) We trained  $10^{3}$  different single-neuron models via simulating SDEs (31) up to time  $T = 100$ , using the Euler-Marumaya discretization scheme of step size  $dt = 10^{-2}$ . All the models are trained with the same data set, composed with 32 points  $\{(x_{i},y_{i})\}_{i = 1,\dots ,32}$  where  $x_{i}$  is an i.i.d. standard Gaussian random variable, and  $y_{i}$  is given by  $y_{i} = x_{i} + \epsilon_{i}$  where  $\epsilon_{i}$  is i.i.d. centered Gaussian random variable with standard deviation of 0.1. All the models are initialized with random weights sampled from i.i.d. centered Gaussian distribution with standard deviation  $10^{-3}$ .
Evidence of stochastic collapse towards permutation invariant sets. (Fig. 3) We trained VGG-16 for  $10^{5}$  steps on CIFAR-10 with a learning rate of 0.1 and ResNet-18 for  $10^{6}$  steps on CIFAR-100 with a learning rate of 0.02. We computed the normalized distance between neurons within the same layer by concatenating corresponding vectors of incoming and outgoing weights, and used this to cluster neurons. Neurons within 0.1 normalized distance from each other were grouped together, which then determines the number of clusters or the independent of neurons denoted as  $n_{ind}$ . In Fig. 3, we plotted the pairwise distance of the incoming weights and outgoing weights separately. We carried out Agglomerative Clustering based on the normalized distance of the incoming weights with  $n_{ind}$  clusters and the corresponding pairwise distance between outgoing weights was visualized with the same ordering.
Larger learning rates and increased label noise intensify stochastic collapse. (Fig. 4 and 9) These experiments varied learning rates and the extent of label noise. For the learning rate experiments, we trained models without label noise for  $10^{5}$  steps. For the label noise experiments, we maintained learning rates of 0.01 and 0.02 for VGG-16 and ResNet-18 respectively over  $10^{6}$  steps. Here, fresh label noise was introduced with each batch sample. A label noise level of  $\epsilon \%$  meant that  $\epsilon \%$  of the labels were randomly assigned an incorrect label. We averaged the results across four replications with different seeds.
After training, we calculated and displayed the proportion of independent neurons in each layer of the networks. We first removed 'vanishing neurons', defined as those with incoming and outgoing weights less than  $10\%$  of the maximum norm for that layer. We then clustered the remaining neurons based on the pairwise normalized distance of the concatenated incoming and outgoing weights. Two neurons were defined as identical if their distance in weight space was less than  $10\%$  of their norms. This allowed us to determine the number of independent neurons in each layer.
Effect of batch size on stochastic collapse. (Fig. 10) We trained models with different batch sizes for  $10^{5}$  (CIFAR-10) and  $10^{6}$  (CIFAR-100) steps, while keeping the learning rate constant at 0.02. We averaged the results across four replications with different seeds. The method of calculating the proportion of independent neurons remains the same.
Demonstrating generalization benefits of stochastic collapse in a teacher-student setting. (Fig. 5) The experiments were conducted in accordance with assumptions A1 - A4. In these experiments, we set both the input and output dimensions to be 64, and the dataset size was set at 1024. We used a sparse teacher model with a rank of 8 and signals ranging evenly from 0.5 to 1. The input samples were drawn to satisfy the condition  $X^T X = I$ . We then constructed the signal  $y$  by introducing a random Gaussian noise with a standard deviation of 0.5 to the true output values. The network was trained with an initial learning rate warm-up of 1000 steps, followed by 50000 steps at a learning rate of 3.0. After 50000 steps, we reduced the learning rate to 0.3 and continued the training for another 10000 steps. The experimental results are averaged over 256 runs with the same dataset.
In Fig. 11, we present the results obtained when replacing label noise with different batch sizes. In this experiment, we used SGD without label noise and maintained a similar experimental setup as before. Additionally, we introduced a learning rate drop at step 53000.
Large learning rates aid generalization via stochastic collapse to simpler subnetworks. (Fig. 6) We carried out initial training of VGG-16 and ResNet-18 using a larger learning rate of 0.1. We created checkpoints in the training process at steps of  $10^{4}$ ,  $2 \times 10^{4}$ , and  $4 \times 10^{4}$ . We resumed training from the checkpoints post but at a reduced learning rate of 0.01. To track the impact of this large learning rate phase on the emergence of simpler subnetworks, we calculated the layer-wise fraction of independent neurons at each of these checkpoints, before the learning rate drop. We averaged the results across eight replications with different seeds.
K Extended Experiments
Figure 8: Evidence of stochastic collapse towards permutation invariant sets in deep neural networks. We show the distribution of normalized distance across all pairs of neurons for the same hidden layer of neurons as shown in Fig. 3. We show plots for three different layers in VGG-16 trained on CIFAR-10 (top row), and for three different layers in a ResNet-18 trained on CIFAR-100 (bottom row). Each plot assesses the distribution at various training stages. In layers with little stochastic collapsing (left column), the distributions rapidly thermalize and concentrate around 2. In layers with strong stochastic collapsing (middle and right columns), the distributions evolve to form a peak around 0, indicating stochastic collapsing towards the permutation invariant set.
Figure 9: Increased label noise 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 label noise. The networks are evaluated at training steps of  $10^{6}$ . A reduced percentage of independent neurons indicates stronger stochastic collapse.
Figure 10: Effect of batch size on 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 batch size. A lower percentage of independent neurons signifies a heightened degree of stochastic collapse. As elaborated in Appendix D, modifying batch size is expected to parallel changes in the learning rate since both influence the amplitude term in the diffusion matrix. Yet, unlike learning rate effects, alterations in batch size produce complex shifts in stochastic collapse. We conjecture that this subtlety may arise from non-replacement sampling and the randomness introduced by transformations applied to input images.
Figure 11: Demonstrating generalization benefits of stochastic collapse in a teacher-student setting with batch noises. Same as Fig. 5 but we show effects of different batch sizes. No label noise was used in this setup. We show the train loss (left) and test loss (middle) during the training of the student with different batch sizes. Dashed lines in the leftmost and middle left panels indicate the steps where learning rate is dropped. Training with smaller batch sizes generalizes better. Right: we show the noisy teacher signals against the learned student signals before dropping the learning rate. We only show the smallest 32 singular values.
Figure 12: Large learning rates aid generalization via stochastic collapse to simpler subnetworks. Left, Middle: Same as Fig. 6. Right: The effective rank [60] of the weight matrices of independent neurons per layer evaluated at different learning rate drop times (indicated by the color) during the initial high learning rate training phase. To compute the effective rank, we first gather the singular values, denoted by  $s_i$ , of the concatenated weight matrices of the incoming and outgoing weights for a specified layer. The effective rank is then computed as  $\rho = -\sum_{i} \hat{s}_i \log \hat{s}_i$ , where  $\hat{s}_i$  represents the normalized singular value defined by  $\hat{s}_i = \frac{s_i}{\sum_j s_j}$ .