Kernel vs. Kernel: Exploring How the Data Structure Affects Neural Collapse
Abstract
Recently, a vast amount of literature has focused on the “Neural Collapse” (NC) phenomenon, which emerges when training neural network (NN) classifiers beyond the zero training error point. The core component of NC is the decrease in the within-class variability of the network’s deepest features, dubbed as NC1. The theoretical works that study NC are typically based on simplified unconstrained features models (UFMs) that mask any effect of the data on the extent of collapse. In this paper, we provide a kernel-based analysis that does not suffer from this limitation. First, given a kernel function, we establish expressions for the traces of the within- and between-class covariance matrices of the samples’ features (and consequently an NC1 metric). Then, we turn to focus on kernels associated with shallow NNs. First, we consider the NN Gaussian Process kernel (NNGP), associated with the network at initialization, and the complement Neural Tangent Kernel (NTK), associated with its training in the “lazy regime”. Interestingly, we show that the NTK does not represent more collapsed features than the NNGP for prototypical data models. As NC emerges from training, we then consider an alternative to NTK: the recently proposed adaptive kernel, which generalizes NNGP to model the feature map** learned from the training data. Contrasting our NC1 analysis for these two kernels enables gaining insights into the effect of data distribution on the extent of collapse, which are empirically aligned with the behavior observed with practical training of NNs.
1 Introduction
Deep Neural Network classifiers are often trained beyond the zero training error point [1, 2]. In this regime, a phenomenon dubbed “Neural Collapse” (NC) emerges [3]. NC is typically described by the following components: (NC1) the networks’ deepest features exhibit a significant decrease in the variability of within-class samples, (NC2) the mean features of different classes approach a certain symmetric structure, and (NC3) the last layer’s weights become more aligned with the penultimate layer features’ means. This behavior has been observed both when using the cross-entropy (CE) loss [3] and the mean squared error (MSE) loss [4].
Recently, a vast amount of literature has been dedicated to exploring NC (as surveyed in [5]), studying the effect of imbalanced data [6, 7], depthwise evolution [8, 9, 10, 11], fine-grained structures [12, 13, 14], and implications [15, 16, 17]. Note that without a sufficient decrease in the features’ within-class variability around the class means, measured by NC1 metrics, one may not gain valuable insights from the structure of the means of different classes. Therefore, oftentimes, NC papers focus specifically on NC1 rather than on other components of NC [12, 13, 14, 16, 11, 18]. Notably, most of the works that attempt to theoretically analyze the NC behavior [19, 6, 15, 4, 20, 8, 21, 7, 22, 12, 23, 10, 24] are based on variants of the unconstrained features model (UFM) [19], which treats the deepest features of the training samples as free optimization variables. Therefore, these analyses cannot predict the effect of the data structure on the extent of collapse.
Since theoretically analyzing the behavior of (deep) NNs is challenging, simplifying approaches that are based on kernel methods [25] have gained massive popularity [26, 27, 28, 29, 30, 31]. Prominent examples include the NN Gaussian Process kernel (NNGP) [26, 27, 32], associated with the infinitely wide NNs at initialization, and the Neural Tangent Kernel (NTK) [28], associated with their training in the “lazy regime”, where the learning rate is sufficiently small. These approaches, and in particular NTK, were used to provide mathematical reasoning for deep learning phenomena such as achieving zero training loss [28, 30, 31], faster learning of lower frequencies [33], benefits of ResNets over fully connected networks [34, 35, 36], usefulness of positional encoding in coordinated-based NNs [37], and more. More recently, finite width variants of the aforementioned kernels have been studied [38, 39, 40], aiming to mitigate the gap between practical NNs behavior and infinite width analyses [41, 42, 43, 44, 45].
In this paper, we provide a kernel-based analysis for NC1, the core component of NC, which does not suffer from the limitations of UFM-based analysis. Thus, it allows us to explore how the data structure affects the collapse. Since kernels provide fixed feature map**, we propose a “kernel vs. kernel” analysis — that is, gaining insights by comparing the properties across NN-related kernels.
Our main contributions can be summarized as follows:
-
•
Given an arbitrary kernel function, we establish expressions for the traces of the within- and between-class covariance matrices of the samples’ features — and consequently, an NC1 metric that depends on the features only through the kernel function.
-
•
We specialize our kernel-based NC1 to kernels associated with shallow NNs. We analyze it for NNGP (NN at initialization), and NTK (NN trained in the “lazy regime”) and show that, perhaps surprisingly, the NTK does not represent more collapsed features than the NNGP for prototypical data models.
-
•
As NC emerges from training, we consider an alternative to NTK: the recently proposed adaptive kernel [40], which generalizes NNGP to model the feature map** learned from the training data. Contrasting our NC1 analysis for these two kernels enables gaining deeper insights into the effect of data distribution on the extent of collapse, which are empirically aligned with the behavior observed with practical training of NNs.
2 Related Work
As mentioned in the previous section, most of the works that attempt to analyze the NC behavior theoretically [19, 6, 15, 4, 20, 8, 21, 7, 22, 23, 10] are based on variants of the UFM [19]. The work in [12] attempts to generalize the UFM by adding a penalty term to the loss that ensures that the features matrix is in the vicinity of a predefined matrix. Yet, the model still lacks an explicit connection between the features and the data. In [46], the authors avoid optimizing the features directly but assume that the model is linear and that the data is nearly orthogonal, which are restrictive assumptions. In [47], the authors claim that having an exact class-wise block structure in the Gram matrix of the empirical NTK on training samples implies NC. Yet, they do not provide reasoning for reaching this collapsed Gram matrix and the analysis is still disconnected from the data. Here, however, we fully depart from the UFM approach and provide analysis that explicitly depends on the data. Furthermore, unlike most works, our analysis is applicable to the less studied case where the data is class imbalanced [6, 17, 23, 7]. Our kernel-based analysis utilizes results on NNGP and NTK from [48, 49, 27, 28, 31]. To simplify the analysis, we theoretically analyze kernels associated with shallow fully connected NNs. Focusing on shallow networks is justified by recent works demonstrating monotonic depthwise evolution of NC1 [8, 9, 10, 11, 12]. Specifically, a data structure that promotes a larger reduction in NC1 for shallow NNs is expected to be more collapsed when using deep NNs. Since NC is related to training, and we show the limitation of NTK to capture it when compared to NNGP, we also utilize the generalization of NNGP that has been proposed in [40]. In this adaptive kernel model, there is an explicit kernel function that depends on the training data. Recently, this richer model has been used to study phase transition behaviors, such as grokking [50], that cannot be captured with the data-independent kernels.
3 Problem Setup
In this section, we outline the notations and the setup.
Data: We consider a dataset , comprising data points of dimension belonging to classes. Each class has size , where and . The dataset is represented in an “organized” matrix form as , where represents the data point of the class. Specific assumptions on the data distribution will be presented during the paper together with the related theory or experiments.
Neural Network: Unless stated otherwise, we consider a 2-layer fully connected neural network (2L-FCN) with layer width , and point-wise activation function . Let denote the weight and bias parameters of the layer. At initialization, the entries are drawn i.i.d from Gaussian distributions of mean and variance , respectively. For an input to the network , we denote the component of the output vector as follows:
| (1) |
Task: We train the network to classify the data points to their respective classes. Let denote the prediction and ground truth label matrices respectively:
| (2) |
We aim to minimize the Mean Squared Error (MSE) between using the following objective:
| (3) |
Note that training deep classifiers with MSE loss has been shown to be a useful strategy [51, 4], which is commonly considered in NC analyses [4, 8, 21].
Pre- and Post-activation Kernels: For any two inputs , we denote their corresponding pre- and post-activation features as and , respectively. The pre and post-activation kernels corresponding to layer are given by:
| (4) | ||||
4 Within-Class Variability Metric (NC1) for Kernels
In this section, we derive an NC1 metric that depends on the features only through the kernel function. Let be a matrix encapsulating arbitrary feature vectors associated with samples of the classes. We define the within-class covariance and between-class covariance matrices of the features as follows:
| (5) |
where and represent the class mean vectors and the global mean vector, respectively. Additionally, we consider the total covariance and non-centered between-class covariance matrices as follows:
| (6) |
Based on these formulations, we define the variability metric , introduced in [12] and used also in [14, 46, 52], as:
| (7) |
In the following theorem, we formulate the traces and using an arbitrary kernel function that expresses inner product of data samples in feature space.
Theorem 4.1.
For any two data points , let the inner-product of their associated features be given by a kernel as . The traces of covariance matrices and can now be formulated as:
| (8) | ||||
| (9) |
The proof (in Appendix A) leverages matrix trace properties and direct expansions of the covariance matrices in terms of vector outer-products to arrive at the results.
Observe that Theorem 4.1 allows us to replace with any suitable kernel formulation corresponding to the features. In the following sections, we leverage this flexibility to analyze and compare NC1 for kernels that model the behavior of neural networks.
5 Activation Variability in the Lazy Learning Regime
In the UFM-based analysis of NC, the assumption is that the deepest features associated with the training samples are free optimization variables, thus losing any ability to analyze the effect of the training data, apart from the balancedness of its labels . In this section, we address these shortcomings by analyzing the role of data distributions on in the infinite width regime, where the NN behavior can be well modeled by NNGP and NTK. In particular, we focus on the case where data is sampled from a mixture of Gaussians and understand the limits of reduction.
5.1 Limiting NNGP Kernel
Under the NN model and initialization stated in Section 3, as the hidden layer width , we can characterize the pre-activation kernel in terms of the GP limit [26] (commonly referred to as the NNGP limit [27]) as follows:
| (10) |
In this limit, the post-activation kernel can have a closed form representation depending on the choice of activation function [27, 31]. The expression for Erf activation [48] is:
| (11) | ||||
The formulation for ReLU-based kernel [49] is presented in the Appendix B. Given a kernel function and samples , we can formulate the kernel Gram matrix as:
| (12) |
Considering the NNGP kernel in (11), we illustrate an example matrix in Figure 1 to visualize the sub-matrices based on the imbalance/balance of class sizes.
5.2 Limiting NTK
In the infinite width limit, we also analyze the NTK [28] to understand the effect of optimization on the NN’s features in the “lazy regime”. Specifically, a well-known result is that in the infinite width limits (with initialization as per Section 3), the deepest feature map** of the NN is fixed during gradient descent optimization with small enough learning rate, and is characterized by the NTK.
Formally, the recursive relationship between the NTK and NNGP [28, 31] can be given as follows:
| (13) |
Here, can be defined using the recursive formulation [27, 28]:
| (14) | ||||
Similar to the activation function specific formulations of in (11), we define the Erf based derivative kernel as follows:
| (15) | ||||
The formulation for ReLU-based kernel is presented in the Appendix B.
Remark on the limiting NTK. Denote by the parameters and the activation function of an -layer NN . The limiting NNGP is defined directly on the product of neurons, , and provides kernel expressions for the inner product of features at different layers of the network (via recursion similar to (11)). On the other hand, the limiting NTK is defined as the inner product of the output’s gradients . The NTK theory shows that this can model the inner product only of the deepest features (i.e the output of the penultimate layer)111Note that we denote NNGP and NTK () with different superscripts to be consistent with the literature. Yet both are associated with the output of the single hidden layer. .
5.3 1-D Gaussian Data with 2 Classes
Notice that even for shallow NNs, it is challenging to theoretically analyze NNGP and NTK for general data. Therefore, we consider a simplified setting to analyze the NC1 properties of these kernels. Formally, consider a -dimensional Gaussian dataset (i.e., ) with classes. The data points belonging to class are independently sampled from and have the labels . Similarly, the data points belonging to class are independently sampled from and have the labels .
Assumption 1: For , let be small enough such that , and almost surely.
Assumption 2: The dataset consists of large enough samples .
Assumption 3: The 2L-FCN has output layer dimension and .
These assumptions present a scenario where the samples of the two classes are sufficiently far from the origin in opposite directions. Thus, the simplest prediction rule pertains to the sign of a sample.
Theorem 5.1 (ReLU Activation).
Under Assumptions 1-3, let be the ReLU activation. Denote by the features associated with NNGP and NTK , respectively. Then:
| (16) |
where is a term that vanishes as increase.
Appendix D presents the proof by calculating the expected values of and employing Theorem 4.1. For a better understanding of the result, consider the balanced class scenario with . This gives us: , which intuitively captures the sum of the within-class variance of in the numerator and the between-class variance in the denominator of the first term.
Erf Activation. We present a similar analysis with the NNGP and NTK with Erf activation in Appendix E (as the terms involved in the formulation are relatively complex than the ReLU case). In the case of ReLU with balanced data, observe that the numerator (corresponding to ) solely depends on , however, for Erf activation, our analysis shows a dependence on terms i.e, it depends on inverse of higher powers of class means as well (see (92) in Appendix E). Similar analysis for with Erf in (100) shows a dependence on terms . Importantly, under Assumptions 1-3, we show similar values of the expected NC1 metric for NNGP and NTK even for the Erf activation, which are smaller than the ReLU case.
Remark. These results reveal the effect of the activation function on the NC1 metric when . Especially under Assumptions 1-3, the Erf-based kernels reflect a larger extent of ‘variability collapse’ (NC1) of the hidden layer post-activations of our 2L-FCN (both at initialization via NNGP and during training via NTK). Additionally, they indicate that the expected of the NTK closely approximates the NNGP counterparts on our 1-D Gaussian dataset. Perhaps surprisingly, this shows that NTK does not represent more collapsed features than NNGP, despite being associated with NN gradient-based optimization. Namely, we have established another result that shows that training in the lazy regime provably deviates from the practical feature learning of NNs [41, 42, 43, 44, 45].
5.4 Experiments with High-Dimensional Gaussian Data
Setup: We conduct experiments on datasets with varying sample sizes and input dimensions to verify our theoretical results and show that insights generalize (e.g., beyond ). For , a dataset size chosen from , and input dimension chosen from , we create the data vector and label pairs as follows:
| (17) | ||||
The vectors and labels from the dataset can then be arranged into the matrix form (as described in the setup) for analysis. The sampling procedure is repeated times for each 222The code is available at: https://github.com/kvignesh1420/shallow_nc1 .
Observations: Figure 2 illustrates the mean and standard deviation (std) of for the post-activation NNGP kernel and NTK for Erf and ReLU activations. For the low-dimensional case of and Erf activation, observe from Figure 2(a) that has a small value of . On the contrary, Figure 2(c) illustrates that for ReLU, is more than an order of magnitude larger () than the former. Furthermore, Figures 2(b) and 2(d) corresponding to NTK (with Erf and ReLU respectively) do not exhibit significantly different values from the NNGP counterparts. These observations empirically verify our theoretical results.
As increases, increases at similar rates for and . With ReLU activation, remains almost constant for and exhibits an increasing trend for — implying less collapse for NTK. All these observations corroborate our theory on the limitation of analyzing NC with NTK. We present additional experimental results for imbalanced datasets in Appendix H and show that for a given , the trends of for increasing can vary based on the imbalance ratio of classes (i.e., ). Nevertheless, the trends in for NNGP with increasing , resemble that of a trained 2L-FCN with a large hidden layer width (see Figure 3(b) corresponding to the Erf activation). These observations provide zero-order reasoning for NN behavior where the feature map** is learned based on data properties such as dimension.
6 Activation Variability in the Feature Learning Regime
The explicit kernel formulations in the infinite width limit have allowed us to go beyond the unconstrained features assumption and preserve a link between the features and the data. Yet, the NC phenomenon relates to NN training, while NNGP relates to NN at initialization and NTK has been found unsuitable for NC analysis. Thus, we wish to contrast NNGP with a kernel that is an alternative to NTK and takes into account both optimization and data. To this end, we transition to a large but finite width () and large sample () setting and analyze the recently introduced ‘adaptive kernels’ approach by Seroussi et al. [40] for fully connected networks.
6.1 Equations of State (EoS)
A transition from the infinite to finite width regime can introduce various corrections to the pre-and post-activations of a -layer FCN. In this context, Seroussi et al. [40] have observed the following dominant corrections: (1) The mean and covariance of the pre-activations deviate from that of a random FCN and, (2) the collective effect of activations from the and layers determine the covariance of activations in the layer. Based on these observations, they employ a Variational Gaussian Approximation (VGA) approach to propose the following system of equations for the pre and post-activation kernels respectively. We formally define the EoS for a 2-layer FCN as follows (based on specializing the generic -layer formulation in equation 5 in [40] to , as done in equation 95 in their arxiv extended version):
Definition 6.1.
The “Equations of State” (EoS) for pre and post-activation kernels of a -layer FCN with Erf activation, no bias, and are given by:
| (18) | ||||
Here, models the statistical covariance of a row of , initialized with , , is the regularization parameter, and corresponds to the prediction of the 2-layer FCN (governed by the EoS). Additionally, are the kernel matrices associated with kernel functions .
Relationship with NNGP: At initialization, we set , which implies that . This resulting exactly matches the kernel matrix for the pre-activation GP kernel (10) as . Similarly, the matrix corresponds to the kernel function defined in (11). The EoS provides a mechanism for transitioning from NNGP kernels to finite-width-based kernels that adapt to the data. Intuitively, observe that the predictions are formulated based on kernel ridge regression with and . The matrix along with and the initial predictions are then used to update the weight covariance matrix . Notice that every entry involves a trace operation on a matrix product, resulting in a weighted sum across entries of (i.e the pairs of data samples).
Finite width corrections of EoS on NC1: We solve the EoS (initialized with ) and obtain the stable state using the Newton-Krylov method with an annealing schedule, as originally proposed by [40]. (see Appendix G for details). At initialization, is exactly described by the limiting NNGP kernel matrix, which has been analyzed in the previous section. Now, by solving the EoS with the final annealing factors as and (which correspond to a 2L-FCN with hidden layer widths respectively), we illustrate the NC1 metrics of in Figure 3 for the running example of a balanced 2 class dataset (17). Notice that for , the values for in Figure 3(a) closely resemble the plots for the limiting NNGP kernel in Figure 2(a) and the NTK in Figure 2(b). However, we observe noticeable changes in the metrics when . Especially, for and , the values for the EoS in Figure 3(c) exhibit a noticeable reduction compared to in Figure 2(a) and in Figure 2(b). Based on this ‘kernel vs. kernel’ analysis, the reduction in values for EoS reflect the departure of from the initial NNGP state to a feature learning state (based on finite widths).
6.2 Activation Variability with Adaptive Kernels (EoS) and 2L-FCN
Setup: We train a 2L-FCN with and Erf activation using (vanilla) Gradient Descent with a learning rate of and weight-decay for epochs on datasets described below. For EoS, we employ the same setup described above, with the final annealing factor of and (as per the critical scaling value [40]). Similar to the formulation of for , we formulate for as follows:
| (19) | ||||
of EoS provides a good approximation of FCN. Let us consider the running example with the balanced 2 class dataset for varying as per (17). The EoS primarily aims to capture the finite-width corrections (as discussed above) depending on the scaling of , and . To this end, observe from Figure 3(c) that for , the trends of for EoS vary depending on the scale of . Figure 3(d) illustrates the the actual 2L-FCN behavior. For , although is scaled to larger values, the of EoS results in similar values and resemble the 2L-FCN case. However, as the input dimension increases to , the larger value of tends to deviate from the 2L-FCN behaviour. Similar deviations were observed even after choosing higher/lower learning rates for 2L-FCN (i.e ) and weight-decays (ex: ). Nonetheless, for the higher dimension of notice that the larger values of are required for the EoS to match the 2L-FCN. Additionally, by training deeper FCN’s with layers and hidden layer widths on , we observed that the EoS trends (which correspond to 2L-FCN) can also be used to estimate the trends of reduction in the penultimate layers of these networks (see Figure 14 in Appendix H).
The trends in for EoS and 2L-FCN hold even for datasets with . Observe from Figure 4 that the values for the kernels , EoS and for 2L-FCN are consistently lower than the case as shown in Figures 2(a), 2(b), 3(c), 3(d), even when comparing them at . We believe this is because is constructed by adding new classes to whose distance between the means is much larger than the within class co-variances.
On the effects of class imbalances: One of the key conditions for the EoS to approximate the behavior of a 2L-FCN is the presence of sufficient data samples [40] (depending on the scale of the input dimension as shown above). Thus, extreme class imbalances can lead to biased trends in EoS. Consider a collection of imbalanced datasets based on where is split into two classes as follows: Case 1: , Case 2: , Case 3: . We observe that for Case 3 where the imbalance ratio is relatively large, the values in EoS are slightly larger than 2L-FCN (Figure 9 in Appendix H). On the other hand, the imbalance ratios in Case 1 and Case 2 lead to good approximations with 2L-FCN. A similar observation can be made for the dataset with and from Figure 10.
Implications: Our analysis showcases the dependence of on the activation functions and indicates that an increase in the data complexity by increasing the dimension of the data typically leads to larger values. Furthermore, the values depend heavily on the relative positions of the data points. In particular, when the means are well separated with smaller , even a large number of classes can lead to smaller values. These observations explain the empirical results by Papyan et al. [3] (Figure 6) where extensive experiments on complex datasets (e.g., ImageNet) led to relatively less collapse (i.e larger ) than simpler datasets such as MNIST. We present a broader discussion on the limitations of our work and future efforts in Appendix J, and a discussion on formulating NC1 metrics that consider the structure of the data in Appendix F.
7 Conclusion
In this paper, we presented a kernel-based approach to understanding the role of data in the emergence of the Neural Collapse (NC) phenomenon. By considering a general kernel function, we first formulated the trace expressions for the variability collapse (NC1) of the features of the data samples. By leveraging these results, we provided theoretical and empirical results to showcase that the NTK does not represent more collapsed features than the NNGP for various Gaussian datasets. Next, to capture the feature-learning aspects of finite-width neural networks, we switched to an ‘adaptive kernel’ approach whose state equations (EoS) facilitate the transition of the post-activation kernel beyond the GP limit. Through this “kernel vs. kernel” approach for limiting NNGP, NTK, and adaptive kernels, we showcased a promising direction to analyze the properties of data for which the NC1 behavior of an actual FCN can be understood. Thus addressing the limitations of the unconstrained features based analysis of NC and explaining the empirical observations of [3] on datasets of varying complexity. We believe that future work on analyzing the EoS for multi-layer FCN and convolutional networks [40] can provide further insights into the depthwise reduction of NC1 for deeper networks and provide a framework for analyzing datasets with more complex distributions.
Acknowledgments and Disclosure of Funding
The authors would like to thank Zhengdao Chen for helpful discussions during the preparation of this manuscript. The work of Tom Tirer is supported by the ISF grant No. 1940/23.
References
- Hoffer et al. [2017] Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the generalization gap in large batch training of neural networks. Advances in neural information processing systems, 30, 2017.
- Ma et al. [2018] Siyuan Ma, Raef Bassily, and Mikhail Belkin. The power of interpolation: Understanding the effectiveness of sgd in modern over-parametrized learning. In International Conference on Machine Learning, pages 3325–3334. PMLR, 2018.
- Papyan et al. [2020] Vardan Papyan, XY Han, and David L Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117(40):24652–24663, 2020.
- Han et al. [2022] XY Han, Vardan Papyan, and David L Donoho. Neural collapse under mse loss: Proximity to and dynamics on the central path. In International Conference on Learning Representations, 2022.
- Kothapalli [2023] Vignesh Kothapalli. Neural collapse: A review on modelling principles and generalization. Transactions on Machine Learning Research, 2023.
- Fang et al. [2021] Cong Fang, Hangfeng He, Qi Long, and Weijie J Su. Exploring deep neural networks via layer-peeled model: Minority collapse in imbalanced training. Proceedings of the National Academy of Sciences, 118(43):e2103091118, 2021.
- Thrampoulidis et al. [2022] Christos Thrampoulidis, Ganesh Ramachandra Kini, Vala Vakilian, and Tina Behnia. Imbalance trouble: Revisiting neural-collapse geometry. Advances in Neural Information Processing Systems, 35:27225–27238, 2022.
- Tirer and Bruna [2022] Tom Tirer and Joan Bruna. Extended unconstrained features model for exploring deep neural collapse. In International Conference on Machine Learning, pages 21478–21505. PMLR, 2022.
- Rangamani et al. [2023] Akshay Rangamani, Marius Lindegaard, Tomer Galanti, and Tomaso A Poggio. Feature learning in deep classifiers through intermediate neural collapse. In International Conference on Machine Learning, pages 28729–28745. PMLR, 2023.
- Súken\́mathbf{missing}ik et al. [2023] Peter Súken\́mathbf{i}k, Marco Mondelli, and Christoph H Lampert. Deep neural collapse is provably optimal for the deep unconstrained features model. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
- He and Su [2023] Hangfeng He and Weijie J Su. A law of data separation in deep learning. Proceedings of the National Academy of Sciences, 120(36):e2221704120, 2023.
- Tirer et al. [2023] Tom Tirer, Haoxiang Huang, and Jonathan Niles-Weed. Perturbation analysis of neural collapse. In International Conference on Machine Learning, pages 34301–34329. PMLR, 2023.
- Yang et al. [2023] Yongyi Yang, Jacob Steinhardt, and Wei Hu. Are neurons actually collapsed? on the fine-grained structure in neural representations. In International Conference on Machine Learning, pages 39453–39487. PMLR, 2023.
- Kothapalli et al. [2023] Vignesh Kothapalli, Tom Tirer, and Joan Bruna. A neural collapse perspective on feature evolution in graph neural networks. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
- Zhu et al. [2021] Zhihui Zhu, Tianyu Ding, **xin Zhou, Xiao Li, Chong You, Jeremias Sulam, and Qing Qu. A geometric analysis of neural collapse with unconstrained features. Advances in Neural Information Processing Systems, 34:29820–29834, 2021.
- Galanti et al. [2022] Tomer Galanti, András György, and Marcus Hutter. On the role of neural collapse in transfer learning. In International Conference on Learning Representations, 2022.
- Yang et al. [2022] Yibo Yang, Shixiang Chen, Xiangtai Li, Liang Xie, Zhouchen Lin, and Dacheng Tao. Inducing neural collapse in imbalanced learning: Do we really need a learnable classifier at the end of deep neural network? Advances in Neural Information Processing Systems, 35:37991–38002, 2022.
- Xu and Liu [2023] **g Xu and Haoxiong Liu. Quantifying the variability collapse of neural networks. In International Conference on Machine Learning, pages 38535–38550. PMLR, 2023.
- Mixon et al. [2020] Dustin G Mixon, Hans Parshall, and Jianzong Pi. Neural collapse with unconstrained features. arXiv preprint arXiv:2011.11619, 2020.
- Ji et al. [2022] Wenlong Ji, Yi** Lu, Yiliang Zhang, Zhun Deng, and Weijie J Su. An unconstrained layer-peeled perspective on neural collapse. In International Conference on Learning Representations, 2022.
- Zhou et al. [2022] **xin Zhou, Xiao Li, Tianyu Ding, Chong You, Qing Qu, and Zhihui Zhu. On the optimization landscape of neural collapse under mse loss: Global optimality with unconstrained features. In International Conference on Machine Learning, pages 27179–27202. PMLR, 2022.
- Yaras et al. [2022] Can Yaras, Peng Wang, Zhihui Zhu, Laura Balzano, and Qing Qu. Neural collapse with normalized features: A geometric analysis over the riemannian manifold. Advances in neural information processing systems, 35:11547–11560, 2022.
- Dang et al. [2023] Hien Dang, Tan Minh Nguyen, Tho Tran, Hung The Tran, Hung Tran, and Nhat Ho. Neural collapse in deep linear networks: From balanced to imbalanced data. In International Conference on Machine Learning, 2023.
- Wojtowytsch et al. [2020] Stephan Wojtowytsch et al. On the emergence of simplex symmetry in the final and penultimate layers of neural network classifiers. arXiv preprint arXiv:2012.05420, 2020.
- Schölkopf et al. [2002] Bernhard Schölkopf, Alexander J Smola, Francis Bach, et al. Learning with kernels: support vector machines, regularization, optimization, and beyond. 2002.
- Neal [1995] Radford M Neal. Bayesian learning for neural networks, volume 118. Springer Science & Business Media, 1995.
- Lee et al. [2018] Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks as gaussian processes. In International Conference on Learning Representations, 2018.
- Jacot et al. [2018] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
- Chizat et al. [2019] Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019.
- Arora et al. [2019] Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In Advances in Neural Information Processing Systems, pages 8141–8150, 2019.
- Lee et al. [2019] Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
- Matthews et al. [2018] Alexander G de G Matthews, Jiri Hron, Mark Rowland, Richard E Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks. In International Conference on Learning Representations, 2018.
- Ronen et al. [2019] Basri Ronen, David Jacobs, Yoni Kasten, and Shira Kritchman. The convergence rate of neural networks for learned functions of different frequencies. Advances in Neural Information Processing Systems, 32, 2019.
- Huang et al. [2020] Kaixuan Huang, Yuqing Wang, Molei Tao, and Tuo Zhao. Why do deep residual networks generalize better than deep feedforward networks?—a neural tangent kernel perspective. Advances in neural information processing systems, 33:2698–2709, 2020.
- Tirer et al. [2022] Tom Tirer, Joan Bruna, and Raja Giryes. Kernel-based smoothness analysis of residual networks. In Mathematical and Scientific Machine Learning, pages 921–954. PMLR, 2022.
- Barzilai et al. [2022] Daniel Barzilai, Amnon Geifman, Meirav Galun, and Ronen Basri. A kernel perspective of skip connections in convolutional networks. arXiv preprint arXiv:2211.14810, 2022.
- Tancik et al. [2020] Matthew Tancik, Pratul Srinivasan, Ben Mildenhall, Sara Fridovich-Keil, Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan Barron, and Ren Ng. Fourier features let networks learn high frequency functions in low dimensional domains. Advances in neural information processing systems, 33:7537–7547, 2020.
- Hanin and Nica [2019] Boris Hanin and Mihai Nica. Finite depth and width corrections to the neural tangent kernel. arXiv preprint arXiv:1909.05989, 2019.
- Lee et al. [2020] Jaehoon Lee, Samuel Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study. Advances in Neural Information Processing Systems, 33:15156–15172, 2020.
- Seroussi et al. [2023] Inbar Seroussi, Gadi Naveh, and Zohar Ringel. Separation of scales and a thermodynamic description of feature learning in some cnns. Nature Communications, 14(1):908, 2023.
- Woodworth et al. [2020] Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, pages 3635–3673. PMLR, 2020.
- Ghorbani et al. [2020] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When do neural networks outperform kernel methods? Advances in Neural Information Processing Systems, 33:14820–14830, 2020.
- Wei et al. [2019] Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets vs their induced kernel. Advances in Neural Information Processing Systems, 32, 2019.
- Yehudai and Shamir [2019] Gilad Yehudai and Ohad Shamir. On the power and limitations of random features for understanding neural networks. Advances in Neural Information Processing Systems, 32, 2019.
- Li et al. [2020] Yuanzhi Li, Tengyu Ma, and Hongyang R Zhang. Learning over-parametrized two-layer neural networks beyond ntk. In Conference on learning theory, pages 2613–2682. PMLR, 2020.
- Wang et al. [2023] Peng Wang, Xiao Li, Can Yaras, Zhihui Zhu, Laura Balzano, Wei Hu, and Qing Qu. Understanding deep representation learning via layerwise feature compression and discrimination. arXiv preprint arXiv:2311.02960, 2023.
- Seleznova et al. [2023] Mariia Seleznova, Dana Weitzner, Raja Giryes, Gitta Kutyniok, and Hung-Hsu Chou. Neural (tangent kernel) collapse. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
- Williams [1996] Christopher Williams. Computing with infinite networks. Advances in neural information processing systems, 9, 1996.
- Cho and Saul [2009] Youngmin Cho and Lawrence Saul. Kernel methods for deep learning. Advances in neural information processing systems, 22, 2009.
- Rubin et al. [2024] Noa Rubin, Inbar Seroussi, and Zohar Ringel. Grokking as a first order phase transition in two layer networks. In The Twelfth International Conference on Learning Representations, 2024.
- Hui and Belkin [2021] Like Hui and Mikhail Belkin. Evaluation of neural architectures trained with square loss vs cross-entropy in classification tasks. In The Ninth International Conference on Learning Representations (ICLR), 2021.
- Yaras et al. [2023] Can Yaras, Peng Wang, Wei Hu, Zhihui Zhu, Laura Balzano, and Qing Qu. The law of parsimony in gradient descent for learning deep linear networks. arXiv preprint arXiv:2306.01154, 2023.
- Seltman [2012] Howard Seltman. Approximations for mean and variance of a ratio. unpublished note, 2012.
- Vershynin [2012] Roman Vershynin. How close is the sample covariance matrix to the actual covariance matrix? Journal of Theoretical Probability, 25(3):655–686, 2012.
Appendix A Proof of Theorem 4.1
To obtain the NC1 formulation corresponding to an arbitrary feature matrix , we start with a simple relationship between as follows:
| (20) | ||||
Similarly, by considering , we get:
| (21) | ||||
Formulating : Expanding into individual outer-products of vectors and leveraging the trace properties leads to the following:
Formulating : Similar to the above analysis, we can reformulate the trace of non-centered between-class covariance matrix as:
Formulating : Reformulation of can be approached along the same lines:
By using these intermediate results, we can formulate as:
Hence, proving the theorem.
Appendix B Limiting NNGP and NTK for ReLU
Consider the GP limit characterization of the pre-activation kernel as follows:
| (22) |
Observe that is independent of the activation function. Now, the closed form representation of the post-activation NNGP kernel for the ReLU activation is given by:
| (23) | ||||
Next, we define the ReLU based derivative kernel as follows:
| (24) | ||||
Finally, the NTK can be formulated as follows:
| (25) |
Here, can be defined using the recursive formulation:
| (26) | ||||
Appendix C General Results for NC1 with Kernels
In this section, we present some general results to calculate the expected value of for any given kernel function that is associated with the features . To begin with, we consider a generic formulation of the three cases for :
| (27) | ||||
Lemma C.1.
Given the cases for the expected values of a kernel function as per (27), the is given by:
| (28) |
Proof.
Lemma C.2.
Given the cases for the expected values of a kernel function as per (27), the is given by:
| (30) |
Lemma C.3.
Given the cases for the expected values of a kernel function as per (27), the is given by:
| (32) |
Proof.
Note that the expectation of the ratios can be given as:
| (33) | |||
| (34) |
Here, corresponds to higher order terms given by [53]:
| (35) |
where, based on the well-studied concentration of sample covariance matrices around the statistical covariance [54], tend to for large values.
∎
Lemma C.4.
For a random variable which represents the sample of class (as per notation in Section 5.3), the expected value is given by:
| (36) |
Proof.
Based on the standard result on the expectation of ratios [53], we get:
| (37) | ||||
| (38) |
Based on the results from the moment-generating function, we know that:
| (39) |
which gives us:
| (40) | ||||
| (41) |
Hence proving the lemma.
∎
Appendix D Proof of Theorem 5.1
D.1 NC1 of limiting NNGP with ReLU activation
In the limit , we leverage the kernels in the GP limit as per (10), (23). Observe that for any two data points , the value of can be given as:
Since , the value of simplifies to:
| (42) |
which follows from and almost surely. This leads to:
| (43) | ||||
| (44) |
For the case, the value of the kernel boils down to the product of norms of independent random variables drawn from the same distribution. Since we assume if , the equation 44 can be rewritten as:
| (45) |
Additionally, since are random variables, the expected value of the kernel can be formulated as:
| (46) |
As and , Lemma C.3 gives us:
| (48) | ||||
D.2 NC1 of limiting NTK with ReLU activation
The recursive relationship between the NTK and NNGP [31, 35] can be given as follows (13):
| (49) |
Here, can be defined using the following recursive formulation:
| (50) | ||||
Based on (24), the derivative can be given as follows:
| (51) | ||||
We build on the results from the NNGP analysis (with ) for computing the variability collapse with the limiting NTK. First, note that:
| (52) |
Appendix E Results for NC1 with Erf activation
E.1 NC1 of Limiting NNGP with Erf activation
Considering :
| (60) | ||||
| (61) |
where the last equality comes from:
| (62) |
For notational simplicity, consider:
| (63) |
and represent as:
| (64) |
Based on Assumption 1, we know that almost surely. This leads to:
| (65) |
E.1.1 Calculating
For , we consider the expansion of to obtain:
| (66) |
To this end, based on Assumption 1 of large enough , we approximate the expectation with only the first term and denote to capture the effects of the higher order terms. Notice that since for finite , the effects of are finite but decay rapidly compared to the first term. To this end, we get:
| (67) |
Calculating the expectation can now be split based on .
Case :
| (68) | ||||
| (69) | ||||
| (70) |
The last equality is based on Lemma C.4 which gives the expanded version of .
Finally, the value of can be given as:
| (71) | ||||
| (72) |
Notice that even in this simple case, the expressions are non-trivial to fully expand. Nonetheless, along with Assumption 1, we consider large enough such that:
| (73) |
Thus, based on the expansion of , we obtain the following cleaner approximation of:
| (74) |
Here captures all the higher order terms corresponding to and as denoted above.
Case :
In the case of , the expectations on the square roots do not have a particular closed form. To this, end we leverage Assumption 1 to obtain the following approximation:
| (75) | ||||
| (76) | ||||
| (77) |
Observe that the inner terms in the expectations are scaled versions of the above case. To this end, we approximate as:
| (78) | ||||
| (79) |
Similar to the assumption that led to (74), we get:
| (80) | ||||
Case
A similar analysis as above applies in this case:
| (81) | ||||
| (82) | ||||
| (83) |
Observe that the inner terms in the expectations are similar to the above case. To this end, we approximate as:
| (84) | ||||
| (85) |
Similar to the assumption that led to (74), we get:
| (86) | ||||
Here are the collective higher order terms that tend to as increases relative to smaller values of . These cases can now be plugged into our generic formulation of expected values of a kernel function (i.e ) as per (27) in Appendix C. Thus, based on Lemma C.3 for sufficiently large we get :
| (88) |
Numerator in the balanced class setting.
To better understand the result, let’s consider the balanced class scenario with , for which the numerator simplifies to:
| (89) | ||||
| (90) |
If we were to ignore the effects of the higher order terms, then observe that the numerator primarily depends on , which can be given based on Lemma C.4 as:
| (91) |
Thus, showcasing the dependence on in determining the extent of collapse. For sufficiently large , we can approximate this value to:
| (92) |
Denominator in the balanced class setting.
Similar to the numerator analysis, observe that when , the denominator can be given as:
| (93) | |||
| (94) | |||
| (95) | |||
| (96) | |||
| (97) | |||
| (98) |
Observe that the term represents:
| (99) |
and for sufficiently large , essentially represents:
| (100) |
E.2 NC1 of Limiting NTK with Erf activation
Recall from (13) that the recursive relationship between the NTK and NNGP can be given as follows:
| (101) |
where:
| (102) | ||||
| (103) | ||||
| (104) | ||||
Considering , (as per the setting and assumptions), we get:
| (105) |
| (106) |
| (107) | ||||
This gives us:
| (108) | ||||
| (109) | ||||
| (110) |
For notational simplicity, consider:
| (111) |
which simplifies the kernel formulation to:
| (112) |
E.2.1 Calculating
Similar to the NNGP analysis, we break down the calculation of into three cases.
Case
| (113) | ||||
| (114) | ||||
| (115) |
This gives us:
| (116) | ||||
Based on the results from the moment-generating function, we know that:
| (117) |
which can be used along with Lemma C.4 to obtain:
| (118) | ||||
For notational simplicity, we define a helper function as follows:
| (119) | ||||
which gives us:
| (120) |
Finally, the value of can be given as:
| (121) | ||||
| (122) |
Notice that even in this simple case, the expressions are non-trivial to fully expand. Nonetheless, along with Assumption 1, we consider large enough such that:
| (123) |
Thus, based on the expansion of , we obtain the following cleaner approximation of:
| (124) |
Case :
| (125) | ||||
| (126) | ||||
| (127) | ||||
| (128) |
Case :
| (135) | ||||
| (136) |
Thus, based on Lemma C.4, we get:
| (137) | ||||
| (138) | ||||
| (139) |
This gives us:
| (140) | ||||
| (141) | ||||
| (142) |
Finally, the cases for the expected value of the kernel can be given as:
| (143) |
From (87), we know that:
| (144) | ||||
To simplify the presentation, we can ignore the higher-order terms and obtain:
| (145) | |||
| (146) |
Observe that the order of the terms involved here resemble that of the NNGP scenario in (87). Thus, we can make similar conclusions regarding the role of the order of in determining the value of .
Appendix F Activation Variability Relative to Data
In this section, we introduce a relative measure of activation variability collapse with respect to the data. First, we begin by defining the within-class and between-class data covariance matrices for the data samples as:
| (147) |
where and represent the data class mean vectors and the data global mean vector respectively.
Definition F.1.
Set a small . The variability collapse relative to the data is given by:
| (148) |
The constant prevents numerical instabilities. Through this approach, we capture the extent of variability collapse of activation features relative to the variability collapse of the data samples itself.
Corollary F.2.
Under Assumptions 1-3 (as per Section 5.3), let be the ReLU activation, and the limiting NNGP kernel be , then:
| (149) |
Proof.
To keep the derivation similar to those for the kernel formulation in equation 45, we consider a simplified kernel on (identity feature map):
| (150) |
Additionally, since are 1-d random variables, the expected value of the kernel is given by:
| (151) |
We use Lemma C.3 with cases , and to obtain:
| (152) | ||||
To better understand the result, let us consider the balanced class scenario where . This results in a ratio of . Furthermore, if (so ), then the ratio . Thus, it emphasizes the interplay between class imbalance/balance and the values of expected class means on the relative variability collapse.
Addressing misleading values: Consider the case where . Then Theorem 5.1 for indicates that (considering smaller fluctuations from ) in the balanced class setting. Such an observation can be misleading if one were to ignore . For instance, such an empirical result while training deep neural networks fails to differentiate between settings where the network learned meaningful features and learned to classify complex datasets or was simply able to leverage the already collapsed data vectors. This applies to Erf activation as well. We justify this argument with the following experiment. For a sample size chosen from , and input dimension chosen from , we sample the vectors for class and for class as our dataset. From Figure 5(a), 5(b), observe that values for can be orders of magnitude larger than , and for high-dimensions . Essentially, the raw data is ‘more’ collapsed than the activations in these settings. Similar observations can be made for the NTK in Figure 5(c), 5(d).
Appendix G Numerical solutions of EoS
We solve the EoS using the Newton-Krylov method with an annealing schedule (as originally proposed by [40]) using the scipy.optimize.newton_krylov python API. We initialize with the GP limit value of and choose a large annealing factor (ex: ) as the value for . The result of optimizing with newton_krylov is a new , which in addition to a lower annealing factor is used as an input for the next newton_krylov function call. This loop is repeated until the end of an annealing schedule. For instance, to analyze the EoS corresponding to , we choose the following list of step-wise annealing factors:
| (156) |
Similarly, for a choice of , we select the slice of the above list up to . Selecting the schedule is a manual operation and can be treated as a hyper-parameter. In our experiments, we observed that this schedule is sufficient to obtain insights on the NC1 metrics of . Thus, we leave the exploration of various annealing strategies as future work.
Comparing the spectrum of weight covariance matrices: Since is subject to change while obtaining the stable state of the EoS, we analyze its initial and final (normalized) spectra for two different datasets of dimension and . Dataset 1: , . Dataset 2: The first dataset is our running example, and the second is pure random noise data. Surprisingly, we observed that the EoS solution captures correlations in the data for both datasets, which is reflected in its final spectrum. In particular, the singular values shift from being constant at initialization to exhibiting a decay in their values (Figure 7). Such a shift does not exactly match the case of L-FCN because of (1) the difference in the dynamics of GD and Newton-Krylov with annealing, and (2) we start with a GP-based initial value for in EoS. A rigorous analysis of the EoS dynamics is an open research direction (as also highlighted by the Seroussi et al. [40]). Nonetheless, the EoS offers a richer data-dependent setup to analyze the activations and weights, than the UFM.
Appendix H Additional Experiments
Compute Resources: All the experiments in this paper were executed on a machine with GB of host memory and CPU cores. Experiments with the EoS on datasets of varying dimensions and sample sizes took the longest time hour to finish.
Appendix I Impact Statement
This paper aims to address the limitations of the unconstrained features model to understand the role of data on the Neural Collapse phenomenon. Our work does not have any direct negative societal impact. On the contrary, our work takes a step forward in understanding the characteristics of datasets that govern the performance of deep neural network classifiers. Thus, laying the groundwork for theoretically informed training practices for wide societal use.
Appendix J Limitations and Future Work
In certain cases, we have observed that none of the kernel methods approximate the 2L-FCN reasonably. One such instance is the following, where we sample for class and for class of our dataset. Essentially, these are scenarios where there is a significant overlap between samples of the two classes. First, we note that we had to increase the learning rate of our 2L-FCN from to and run GD for epochs for convergence. For dimensions , the EoS reasonably approximates the 2L-FCN but for , the values for 2L-FCN turned out to be almost twice as large as the EoS (see Figure 13). To this end, we leave modifications to the EoS for handling such noisy data cases and different activation functions as future work.
Additionally, we highlight the difficulties in the theoretical/empirical analysis of NC1 with EoS. The primary bottleneck is a lack of rigorous study on the existence and uniqueness of solutions (As also highlighted by [40]). Since we deviate from the lazy regime and deal with kernels in the feature learning setup, we cannot expect simpler closed-form solutions like the limiting NNGP/NTK for the EoS. However, analytical solutions to the EoS can sometimes be time-consuming and require a manual selection of the annealing schedule. This is a tradeoff that can be improved with future research. Furthermore, the role of scaling on NC1 is yet to be fully understood and we hope that our analysis lays the groundwork for such efforts.
Finally, we point the reader to Appendix F for a discussion on a relative NC1 metric that explicitly incorporates the variability collapse of the data vectors into the NC1 metric. In particular, we aim to differentiate between settings where the neural network learned meaningful features and learned to classify complex datasets or was simply able to leverage the already collapsed data vectors. Our results showcase that in higher dimensions, the data vectors are ‘more’ collapsed than the activations themselves. Thus showcasing the limitations of the current NC1 metrics and encouraging the reader to explore much richer variants.