SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\left\{ \begin{array}{l l} \dot {\mathbf {H}} = & - \mathbf {W} ^ {\top} \left[ \left(\kappa_ {d} - \kappa_ {c}\right) \mathbf {R} + \left(\kappa_ {c} - \kappa_ {n}\right) m \mathbf {R} _ {\text {c l a s s}} + \kappa_ {n} N \mathbf {R} _ {\text {g l o b a l}} \right] \\ \dot {\mathbf {W}} = & - \mathbf {R} \mathbf {H} ^ {\top} \\ \dot {\mathbf {b}} = & - \mathbf {R} _ {\text {g l o b a l}} \mathbf {1} _ {N}. \end{array} \right. \tag {19}
$$
We note that at the end of training, where  $\mathbf{R}_{\mathrm{class}}$  and  $\mathbf{R}_{\mathrm{global}}$  are zero, the system (19) reduces to
$$
\dot {\mathbf {H}} = - \left(\kappa_ {d} - \kappa_ {c}\right) \nabla_ {\mathbf {H}} \tilde {\mathcal {L}}, \quad \dot {\mathbf {W}} = - \nabla_ {\mathbf {W}} \tilde {\mathcal {L}}, \quad \tilde {\mathcal {L}} (\mathbf {W}, \mathbf {H}) := \frac {1}{2} \| \mathbf {W} \mathbf {H} + \mathbf {b} \mathbf {1} _ {N} ^ {\top} - \mathbf {Y} \| _ {F} ^ {2}, \tag {20}
$$
and  $\dot{\mathbf{b}} = 0$ . This system differs from the unconstrained features dynamics only by a factor of  $\kappa_{d} - \kappa_{c}$  before  $\mathbf{H}$ . Moreover, such a form of the loss function also appears in the literature of implicit regularization [4, 6, 11], where the authors show that  $\mathbf{WH}$  converges to a low rank matrix.
5 NTK Alignment Drives Neural Collapse
The main goal of this work is to demonstrate how NC results from the NTK block structure. To this end, in Section 5.1 we further analyze the dynamics presented in Theorem 4.1, in Section 5.2 we derive the invariant of this training dynamics, and in Section 5.3 we finally derive NC.
5.1 Features Decomposition
We first decompose the features dynamics presented in Theorem 4.1 into two parts:  $\mathbf{H}_1$ , which lies in the subspace of the labels  $\mathbf{Y}$ , and  $\mathbf{H}_2$ , which is orthogonal to the labels and eventually vanishes. To achieve this, note that the SVD of  $\mathbf{Y}$  has the following form:
$$
\mathbf {P} ^ {\top} \mathbf {Y} \mathbf {Q} = \left[ \sqrt {m} \mathbb {I} _ {C}, \mathbb {O} \right], \tag {21}
$$
where  $\mathbb{O} \in \mathbb{R}^{C \times (N - C)}$  is a matrix of zeros, and  $\mathbf{P} \in \mathbb{R}^{C \times C}$  and  $\mathbf{Q} \in \mathbb{R}^{N \times N}$  are orthogonal matrices. Moreover, we can choose  $\mathbf{P}$  and  $\mathbf{Q}$  such that  $\mathbf{P} = \mathbb{I}_C$  and
$$
\mathbf {Q} = \left[ \mathbf {Q} _ {1}, \mathbf {Q} _ {2} \right], \quad \mathbf {Q} _ {1} = \frac {1}{\sqrt {m}} \mathbb {I} _ {C} \otimes \mathbf {1} _ {m} \in \mathbb {R} ^ {N \times C}, \quad \mathbf {Q} _ {2} = \mathbb {I} _ {C} \otimes \tilde {\mathbf {Q}} _ {2} \in \mathbb {R} ^ {N \times (N - C)}, \tag {22}
$$
where  $\otimes$  is the Kronecker product. Note that by orthogonality,  $\tilde{\mathbf{Q}}_2\in \mathbb{R}^{m\times (m - 1)}$  has full rank and  $\mathbf{1}_m^\top \tilde{\mathbf{Q}}_2 = \mathbb{O}$ . We can now decompose  $\mathbf{HQ}$  into two components as follows:
$$
\mathbf {H Q} = \sqrt {m} \left[ \mathbf {H} _ {1}, \mathbf {H} _ {2} \right], \quad \mathbf {H} _ {1} = \frac {1}{\sqrt {m}} \mathbf {H Q} _ {1}, \quad \mathbf {H} _ {2} = \frac {1}{\sqrt {m}} \mathbf {H Q} _ {2}. \tag {23}
$$
The following equations reveal the meaning of these two components:
$$
\mathbf {H} _ {1} = \left[ \langle h \rangle_ {1}, \dots , \langle h \rangle_ {C} \right], \quad \mathbf {H} _ {2} = \frac {1}{\sqrt {m}} \left[ \mathbf {H} ^ {(1)} \tilde {\mathbf {Q}} _ {2}, \dots , \mathbf {H} ^ {(C)} \tilde {\mathbf {Q}} _ {2} \right], \tag {24}
$$
where  $\langle h\rangle_c\in \mathbb{R}^n$  is the mean of  $h$  over inputs  $x_{i}^{c}$  from class  $c\in [C]$ , and  $\mathbf{H}^{(c)}\in \mathbb{R}^{n\times m}$  is the submatrix of  $\mathbf{H}$  corresponding to samples of class  $c$ , i.e.,  $\mathbf{H} = [\mathbf{H}^{(1)},\dots ,\mathbf{H}^{(C)}]$ . We see that  $\mathbf{H}_1$
is simply the matrix of the last-layer features' class means, which is prominent in the NC literature. We also see that the columns of  $\mathbf{H}^{(c)}\tilde{\mathbf{Q}}_2$  are  $m - 1$  different linear combinations of  $m$  vectors  $h(x_i^c)$ ,  $i\in [m]$ . Moreover, the coefficients of each of these linear combinations sum to zero by the choice of  $\tilde{\mathbf{Q}}_2$ . Therefore,  $\mathbf{H}_2$  must reduce to zero in case of variability collapse (NC1), when all the feature vectors within the same class become equal. We prove that  $\mathbf{H}_2$  indeed vanishes in DNNs with block-structured NTK as part of our main result (Theorem 5.2).
5.2 Invariant
We now use the former decomposition of the last-layer features to further simplify the dynamics and deduce a training invariant in Theorem 5.1. The proof is given in Appendix B.2.
Theorem 5.1. Suppose Assumption 3.2 holds. Define  $\mathbf{H}_1$  and  $\mathbf{H}_2$  as in (23). Then the class-means of the residuals (defined in (17)) are given by  $\mathbf{R}_1 = \mathbf{W}\mathbf{H}_1 + \mathbf{b}\mathbf{1}_C^\top -\mathbb{I}_C$ , and the training dynamics of the DNN can be written as
$$
\left\{ \begin{array}{l l} \dot {\mathbf {H}} _ {1} & = - \mathbf {W} ^ {\top} \mathbf {R} _ {1} \left(\mu_ {\text {c l a s s}} \mathbb {I} _ {C} + \kappa_ {n} m \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}\right) \\ \dot {\mathbf {H}} _ {2} & = - \mu_ {\text {s i n g l e}} \mathbf {W} ^ {\top} \mathbf {W} \mathbf {H} _ {2} \\ \dot {\mathbf {W}} & = - m \left(\mathbf {R} _ {1} \mathbf {H} _ {1} ^ {\top} + \mathbf {W} \mathbf {H} _ {2} \mathbf {H} _ {2} ^ {\top}\right) \\ \dot {\mathbf {b}} & = - m \mathbf {R} _ {1} \mathbf {1} _ {C}, \end{array} \right. \tag {25}
$$
where  $\mu_{\mathrm{single}} \coloneqq \kappa_d - \kappa_c$  and  $\mu_{\mathrm{class}} \coloneqq \mu_{\mathrm{single}} + m(\kappa_c - \kappa_n)$  are the two smallest eigenvalues of the kernel  $\Theta_{k,k}^h(X)$  for any  $k \in [n]$ . Moreover, the quantity
$$
\mathbf {E} := \frac {1}{m} \mathbf {W} ^ {\top} \mathbf {W} - \frac {1}{\mu_ {\text {c l a s s}}} \mathbf {H} _ {1} \left(\mathbb {I} _ {C} - \alpha \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}\right) \mathbf {H} _ {1} ^ {\top} - \frac {1}{\mu_ {\text {s i n g l e}}} \mathbf {H} _ {2} \mathbf {H} _ {2} ^ {\top} \tag {26}
$$
is invariant in time. Here  $\alpha := \frac{\kappa_n m}{\mu_{\mathrm{class}} + C \kappa_n m}$ .
We note that the invariant  $\mathbf{E}$  derived here resembles the conservation laws of hyperbolic dynamics that take the form  $\mathbf{E}_{\mathrm{hyp}} := a^2 - b^2 = const$  for time-dependent quantities  $a$  and  $b$ . Such dynamics arise when gradient flow is applied to a loss function of the form  $\mathcal{L}(a, b) := (ab - q)^2$  for some  $q$ . Since the solutions of such minimization problems, given by  $ab = q$ , exhibit symmetry under scaling  $a \to \gamma a, b \to b / \gamma$ , the value of the invariant  $\mathbf{E}_{\mathrm{hyp}}$  uniquely specifies the hyperbola followed by the solution. In machine learning theory, hyperbolic dynamics arise as the gradient flow dynamics of linear DNNs [42], or in matrix factorization problems [3, 15]. Moreover, the end of training dynamics defined in (20) has a hyperbolic invariant given by
$$
\mathbf {E} _ {\mathrm {e o t}} := \mathbf {W} ^ {\top} \mathbf {W} - \frac {1}{\mu_ {\text {s i n g l e}}} \mathbf {H} \mathbf {H} ^ {\top}. \tag {27}
$$
Therefore, the final phase of training exhibits a typical behavior for the hyperbolic dynamics, which is also characteristic for the unconstrained features models [21, 38]. Namely, "scaling"  $\mathbf{W}$  and  $\mathbf{H}$  by an invertible matrix does not affect the loss value but changes the dynamic's invariant. On the other hand, minimizing the invariant  $\mathbf{E}_{\mathrm{ot}}$  has the same effect as joint regularization of  $\mathbf{W}$  and  $\mathbf{H}$  [48].
However, we also note that our invariant  $\mathbf{E}$  provides a new, more comprehensive look at the DNNs' dynamics. While unconstrained features models effectively make assumptions on the end-of-training invariant  $\mathbf{E}_{\mathrm{eot}}$  to derive NC [21, 38, 48], our dynamics control the value of  $\mathbf{E}_{\mathrm{eot}}$  through the more general invariant  $\mathbf{E}$ . This way we connect the properties of end-of-training hyperbolic dynamics with the previous stages of training.
5.3 Neural Collapse
We are finally ready to state and prove our main result in Theorem 5.2 about the emergence of NC in DNNs with NTK alignment. We include the proof in Appendix B.3.
Theorem 5.2. Assume that the NTK has a block structure as defined in Assumption 3.2. Then the DNN's training dynamics are given by the system of equations in (25). Assume further that the last-layer features are centralized, i.e  $\langle h\rangle = 0$ , and the dynamics invariant (26) is zero, i.e.,  $\mathbf{E} = \mathbb{O}$ . Then the DNN's dynamics exhibit neural collapse as defined in (NC1)-(NC4).
Below we provide several important remarks and discuss the implications of this result:
(1) Zero invariant assumption: We assume that the invariant (26) is zero in Theorem 5.2 for simplicity and consistency with the literature. Indeed, similar assumptions arise in matrix decomposition papers, where zero invariant guarantees "balance" of the problem [3, 15]. However, our proofs in fact only require a weaker assumption that the invariant terms containing features  $\mathbf{H}$  are aligned with the weights  $\mathbf{W}^{\top}\mathbf{W}$ , i.e.
$$
\mathbf {W} ^ {\top} \mathbf {W} \propto \frac {1}{\mu_ {\text {c l a s s}}} \mathbf {H} _ {1} \mathbf {H} _ {1} ^ {\top} - \frac {1}{\mu_ {\text {s i n g l e}}} \mathbf {H} _ {2} \mathbf {H} _ {2} ^ {\top}, \tag {28}
$$
where we have taken into account our assumption on the zero global mean  $\langle h\rangle = 0$
(2) Necessity of the invariant assumption: The relaxed assumption on the invariant (28) is necessary for the emergence of NC in DNNs with block-structured NTK. Indeed, NC1 implies  $\mathbf{H}_2 = \mathbb{O}$ , and NC3 implies  $\mathbf{H}_1\mathbf{H}_1^\top \propto \mathbf{W}^\top \mathbf{W}$ . Therefore, DNNs that do not satisfy this assumption do not display NC. Our numerical experiments described in Section 6 strongly support this insight (see Figure 2, panes a-e). Thus, we believe that the invariant derived in this work characterizes the difference between models that do and do not exhibit NC.
(3) Zero global mean assumption: We note that the zero global mean assumption  $\langle h\rangle = 0$  in Theorem 5.2 ensures that the biases are equal to  $\mathbf{b} = \frac{1}{C}\mathbf{1}_C$  at the end of training. This assumption is common in the NC literature [21, 38] and is well-supported by our numerical experiments (see figures in Appendix C, pane i). Indeed, modern DNNs typically include certain normalization (e.g. through batch normalization layers) to improve numerical stability, and closeness of the global mean to zero is a by-product of such normalization.
(4) General biases case: Discarding the zero global mean assumption allows the biases  $\mathbf{b}$  to take an arbitrary form. In this general case, the following holds for the matrix of weights:
$$
\left(\mathbf {W} \mathbf {W} ^ {\top}\right) ^ {2} = \frac {m}{\mu_ {\text {c l a s s}}} \left(\mathbb {I} _ {C} - \alpha \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top} + (1 - \alpha C) \left(C \mathbf {b b} ^ {\top} - \mathbf {b 1} _ {C} ^ {\top} - \mathbf {1} _ {C} \mathbf {b} ^ {\top}\right)\right). \tag {29}
$$
For optimal biases  $\mathbf{b} = \frac{1}{C}\mathbf{1}_C$ , this reduces to the ETF structure that emerges in NC. Moreover, if biases are all equal, i.e.  $\mathbf{b} = \beta \mathbf{1}_C$  for some  $\beta \in \mathbb{R}$ , the centralized class means still form an ETF (i.e., NC2 holds), and the weights exhibit a certain symmetric structure given by
$$
\mathbf {W} \mathbf {W} ^ {\top} \propto \left(\mathbb {I} _ {C} - \gamma \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}\right), \quad \mathbf {M} ^ {\top} \mathbf {M} \propto \left(\mathbb {I} _ {C} - \frac {1}{C} \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}\right), \tag {30}
$$
where  $\gamma := \frac{1}{C} (1 - |1 - \beta C| \sqrt{1 - \alpha C}) < \frac{1}{C}$ . The proof and a discussion of this result are given in Appendix B.4. In general, the angles of these two frames are different, and thus NC3 does not hold. This insight leads us to believe that normalization is an important factor in the emergence of NC.
(5) Partial NC: Our proofs and the discussion suggest that all the four phenomena that form NC do not have to always coincide. In particular, our proof of NC1 only requires the block-structured NTK and the invariant to be P.S.D, which is much weaker than the total set of assumptions in Theorem 5.2. Therefore, variability collapse can occur in models that do not exhibit the ETF structure of the class-means or the duality of the weights and the class means. Moreover, as shown above, NC2 can occur when NC3 does not, i.e., the ETF structure of the class means does not imply duality.
We conducted large-scale numerical experiments to support our theory. While we only showcase our results on a single dataset-architecture pair in the main text (see Figure 2) and refer the rest to the appendix, the following discussion covers all our experiments.
Datasets and models. Following the seminal NC paper [39], we use three canonical DNN architectures: VGG [46], ResNet [24] and DenseNet [26]. Our datasets are MNIST [35], FashionMNIST [51] and CIFAR10 [34]. We choose VGG11 for MNIST and FashionMNIST, and VGG16 for CIFAR10. We add batch normalization after every layer in the VGG architecture, set dropout to zero and choose the dimensions of the two fully-connected layers on the top of the network as 512 and 256. We use ResNet20 architecture described in the original ResNet paper [24], and DenseNet40 with bottleneck layers, growth  $k = 12$ , and zero dropout for all the datasets.
Figure 2: ResNet20 trained on MNIST with three initialization settings and varying learning rates (see Section 6 for details). We chose a model that exhibits NC (red lines, filled markers) and a model that does not exhibit NC (blue lines, empty markers) for each initialization. The vertical lines indicate the epoch when the training accuracy reaches  $99.9\%$  (over the last 10 batches). a) Frobenius norm of the invariant  $\| \mathbf{E}\| _F$ . b) Alignment of the invariant terms as defined in (28). c) NC1: standard deviation of  $h(x_i^c)$  averaged over classes. d) NC2:  $\| \mathbf{M}^\top \mathbf{M} / \| \mathbf{M}^\top \mathbf{M}\| _F - \Phi \| _F$ , where  $\Phi$  is an ETF. e) NC3:  $\| \mathbf{W}^\top /\| \mathbf{W}\| _F - \mathbf{M} / \| \mathbf{M}\| _F\| _F$ . The legend displays the test accuracy achieved by each model and the last-layer features kernel alignment given by  $\langle \Theta^h /\| \Theta^h\| _F,\mathbf{Y}^\top \mathbf{Y} / \| \mathbf{Y}^\top \mathbf{Y}\| _F\rangle_F$ . The curves in panes a-e are smoothed by Savitzky-Golay filter with polynomial degree 1 over window of size 10. Panes f, g and h show the NC metrics and the test accuracy as functions of the learning rate.
Optimization and initialization. We use SGD with Nesterov momentum 0.9 and weight decay  $5 \times 10^{-4}$ . Every model is trained for 400 epochs with batches of size 120. To be consistent with the theory, we balance the batches exactly. We train every model with a set of initial learning rates spaced logarithmically in the range  $\eta \in [10^{-4}, 10^{0.25}]$ . The learning rate is divided by 10 every 120 epochs. On top of the varying learning rates, we try three different initialization settings for every model: (a) LeCun normal initialization (default in Flax), (b) uniform initialization on  $[- \sqrt{k}, \sqrt{k}]$ , where  $k = 1 / n_{\ell - 1}$  for a linear layer, and  $k = 1 / (Kn_{\ell - 1})$  for a convolutional layer, where  $K$  is the convolutional kernel size (default in PyTorch), (c) He normal initialization in fan_out mode.
Results. Our experiments confirm the validity of our assumptions and the emergence of NC as their result. Specifically, we make the following observations:
- While most of the DNNs that achieve high test performance exhibit NC, we are able to identify DNNs with comparable performance that do not exhibit NC (see Figure 2, panes f-h). We note that such models still achieve near-zero error on the training set in our setup.
- Comparing DNNs that do and do not exhibit NC, we find that our assumption on the invariant (see Theorem 5.2 and (28)) holds only for the models with NC (see Figure 2, panes a-e). This confirms our reasoning about the necessity of the invariant assumption for NC emergence.
- The kernels  $\Theta$  and  $\Theta^h$  are strongly aligned with the labels  $\mathbf{Y}^\top \mathbf{Y}$  in the models with the best performance, which is in agreement with the NTK alignment literature and justifies our assumption on the NTK block structure.
We include the full range of experiments along with the implementation details and the discussion of required computational resources in Appendix C. Specifically, we present a figure analogous to Figure 2 for every considered dataset-architecture pair. Additionally, we report the norms of matrices  $\mathbf{H}_1\mathbf{H}_1^\top$ ,  $\mathbf{H}_2\mathbf{H}_2^\top$ , and  $\langle h\rangle \langle h\rangle^\top$ , as well as the alignment of both the NTK  $\Theta$  and the last-layer features kernel  $\Theta^h$  in the end of training, to further justify our assumptions.
7 Conclusions and Broad Impact
This work establishes the connection between NTK alignment and NC, and thus provides a mechanistic explanation for the emergence of NC within realistic DNNs' training dynamics. It also contributes to the underexplored line of research connecting NC and local elasticity of DNNs' training dynamics.
The primary implication of this research is that it exposes the potential to study NC through the lens of NTK alignment. Indeed, previous works on NC focus on the top-down approach (layer-peeled models) [18, 21, 38, 48], and fundamentally cannot explain how NC develops through earlier layers of a DNN and what are the effects of depth. On the other hand, NTK alignment literature focuses on the alignment of individual layers [7], and recent theoretical results even quantify the role of each hidden layer in the final alignment [37]. Therefore, we believe that the connection between NTK alignment and NC established in this work provides a conceptually new method to study NC.
Moreover, this work introduces a novel approach to facilitate theoretical analysis of DNNs' training dynamics. While most theoretical works consider the NTK in the infinite-width limit to simplify the dynamics [1, 20, 28, 49], our analysis shows that making reasonable assumptions on the empirical NTK can also lead to tractable dynamics equations and new theoretical results. Thus, we believe that the analysis of DNNs' training dynamics based on the properties of the empirical NTK is a promising approach also beyond NC research.
The main limitation of this work is the simplifying Assumption 3.2 on the kernel structure. While the NTK of well-trained DNNs indeed has an approximate block structure (as we discuss in detail in Section 3), the NTK values also tend to display high variance in real DNNs [22, 44]. Thus, we believe that adding stochasticity to the dynamics considered in this paper is a promising direction for the future work. Moreover, the empirical NTK exhibits so-called specialization, i.e., the kernel matrix corresponding to a certain output neurons aligns more with the labels of the corresponding class [45]. In block-structured kernels, specialization implies different values in blocks corresponding to different classes. Thus, generalizing our theory to block-structured kernels with specialization is another promising short-term research goal. In addition, our theory relies on the assumption that the dataset (or the training batch) is balanced, i.e., all the classes have the same number of samples. Accounting for the effects of non-balanced datasets within the dynamics of DNNs with block-structured NTK is another possible future work direction.
More generally, we believe that empirical observations are essential to demistify the DNNs' training dynamics, and there are still many unknown and interesting connections between seemingly unrelated empirical phenomena. Establishing new theoretical connections between such phenomena is an important objective, since it provides a more coherent picture of the deep learning theory as a whole.
R. Giryes and G. Kutyniok acknowledge support from the LMU-TAU - International Key Cooperation Tel Aviv University 2023. R. Giryes is also grateful for partial support by ERC-StG SPADE grant no. 757497. G. Kutyniok is grateful for partial support by the Konrad Zuse School of Excellence in Reliable AI (DAAD), the Munich Center for Machine Learning (BMBF) as well as the German Research Foundation under Grants DFG-SPP-2298, KU 1446/31-1 and KU 1446/32-1 and under Grant DFG-SFB/TR 109, Project C09 and the Federal Ministry of Education and Research under Grant MaGriDo.
[45] Haozhe Shan and Blake Bordelon. A theory of neural tangent kernel alignment and its influence on training. arXiv preprint arXiv:2105.14301, 2021.
[46] Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. In 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
[47] Yifan Sun, Changmao Cheng, Yuhan Zhang, Chi Zhang, Liang Zheng, Zhongdao Wang, and Yichen Wei. Circle loss: A unified perspective of pair similarity optimization. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 6398-6407, 2020.
[48] Tom Tirer and Joan Bruna. Extended unconstrained features model for exploring deep neural collapse. In Proceedings of the 39th International Conference on Machine Learning, volume 162, pages 21478-21505, 2022.
[49] Tom Tirer, Joan Bruna, and Raja Giryes. Kernel-based smoothness analysis of residual networks. In Mathematical and Scientific Machine Learning, volume 145, pages 921-954, 2021.
[50] Tom Tirer, Haoxiang Huang, and Jonathan Niles-Weed. Perturbation analysis of neural collapse. arXiv preprint arXiv:2210.16658, 2022.
[51] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.07747, 2017.
[52] Greg Yang. Tensor programs II: neural tangent kernel for any architecture. arXiv preprint arXiv:2006.14548, 2020.
[53] Jiayao Zhang, Hua Wang, and Weijie J. Su. Imitating deep learning dynamics via locally elastic stochastic differential equations. In Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pages 6392-6403, 2021.
NC with MSE loss. NC was first introduced for DNNs with cross-entropy (CE) loss, which is commonly used in classification problems [39]. Since then, numerous papers discussed NC with MSE loss, which provides more opportunities for theoretical analysis, especially after the MSE loss was shown to perform on par with CE loss for classification tasks [14, 29].
Most previous works on MSE-NC adopt the so-called unconstrained features model [21, 38, 48]. In this model, the last-layer features  $\mathbf{H}$  are free variables that are directly optimized during training, i.e., the features do not depend on the input data or the DNN's trainable parameters. Fang et al. [18] also introduced a generalization of this approach called  $N$ -layer-peeled model, where features of the  $N$ -th-to-last layer are free variables, and studied the 1-layer-peeled model (equivalent to the unconstrained features model) with CE loss as a special case.