Towards Understanding Neural Collapse: The Effects of Batch Normalization and Weight Decay

We sincerely appreciate your recognition of the strengths of our paper, especially in its contribution to understanding neural collapse and the accessibility of our logical argumentation.

Response to Weaknesses

1. As with prior works, we adapt the layer-peeled model due to the difficulty of analyzing the effect of multiple layers of non-linear activation functions and the variable neural network architecture possibilities of the lower layers. Suppose the activations functions of the Neural Network are linear (e.g. the identity function). In that case, our results can be relatively easily generalized to features of lower layers since the layers can be combined into a single linear transformation. This generalization to lower layers is much more difficult for non-linear activation functions and would likely require more intricate mathematical tools. This direction has recently been pursued by [1][2] for the MSE loss in the optimal case. However, the corresponding generalization for CE loss in the near-optimal case is more involved and likely to require significant future work.

2. We thank the reviewer for pointing out this point of potential misunderstanding. We would like to clarify our choice of using cosine similarity as a measure:

   1. The cosine similarity measure is a widely used adopted similarity measure to compare between features of different samples in both practical feature learning and machine learning theory for its geometric interpretability. Quantitatively analyzing this measure allows our theorems to have more direct implications for researchers in such relevant fields
   2. The cosine similarity measure allows direct value comparison between NC1 and NC2, corresponding to inter-class and intra-class cosine similarity respectively, which is a benefit over prior works. For example, in the original NC paper [1], the authors adopted 2 different measures, SNR ratio and feature mean cosine similarity, for NC1 and NC2, which cannot be directly compared since they are 2 different measures.
   3. Connected to the previous point, while the SNR ratio compares intra-class variability and inter-class distance, a small SNR ratio does not distinguish whether both are small or both are large. In particular, the SNR ratio alone does not ensure proximity to the ETF structure. Explicitly measuring the cosine similarity for both intra-class and inter-class allows separate guarantees using the same measure that implies NC1 and NC2 (without norm equality) when combined.
   4. Cosine similarity also allows relatively simple quantitative analysis.

   We would like to note that we do not intend to decrease the value of the SNR ratio as a measure, which we’re convinced to have significant intuition and insights in other fields that are not captured by cosine similarity. However, the above reasons make us believe that cosine similarity is an appropriate measure for our purpose of quantitative analysis of Neural Collapse.

3. We thank multiple reviewers for making this suggestion and will adjust our plots accordingly. The questions are answered in the Questions section

4. Due to the intricacy of deep neural network training, it is hard to make certain assertions on the result of training complex networks. However, we think that the following reasons may be influential in the current results for the real-world networks and datasets:

   1. Achieving near-1 intra-class cosine similarity is a strong requirement on the features within the same class when the vector dimension is large.
   2. We used the minimum intra-class cosine similarity and maximum inter-class cosine similarity, which again is a strong requirement for the structure of feature vectors. This measure can be affected by outliers since there are $C(C-1)$ pairs of classes for inter-class cosine similarity

   To address the second point, we’re currently running experiments to also record the average intra-class and inter-class cosine similarity for the real-world models and dataset combinations to resolve the sensitivity to outliers. We will update the paper when they’re completed.

Questions

1. Yes, our theorems does assume that $d\geq C$, which is required for forming a equiangular tight frame. We thank the reviewer for pointing this out and will make this assumption explicit to avoid confusion.
2. The answer to both questions is No. We thank the reviewers for pointing out our reuse of variables that can cause confusion and will use different notations for different parameters in our revised version.
3. Yes, we would adjust the term to “BN with zero bias” to better clarify our model.
4. In section 3.1, the number of classes for the synthetic dataset is 4 and the dimension of the penultimate layer is 200. In section 3.2, the dimension is dependent on the specific model architecture used. For VGG models, the dimension is 4096 and for ResNet, the dimension is 512.

Thanks for the clarifications, I appreciate your answers and I think that revising the paper according to them will improve the work. The main concerns, however, -- unclear picture for large architectures, unclear how the results generalize to other models than the layer-peeled model -- remain. After reading the other reviewers' comments, I have thus decided to not change my score. Although the situation is more clear to me know, given that my knowledge in the field of NC is limited, I cannot reasonably champion the paper.

We thank the reviewer for acknowledging the clear presentation of our work and pointing out the effect and batch normalization and weight decay in NC. Below are our responses to the points in the weakness section.

1. We thank multiple reviewers for pointing out the potentially misleading expression of our contributions on the cosine similarity. We are aware that cosine similarity has been widely used in previous literature. Our main intent is to claim that we are novel in applying cosine similarity to the specific phenomenon of analyzing Neural Collapse, which has not been done before to the best of our knowledge at the time of writing. We will clarify our contributions with regard to this statement and discuss the related use of the cosine similarity measure for similar purposes in previous literature.
2. As mentioned in section 2.2, we adapt cosine similarity as a quantitative measure for NC due to the straightforward geometric meaning it provides compared to previous measures. In particular, cosine similarity is a widely adapted measure for similarity between features of different inputs, both in practical feature learning and in machine learning theory. Therefore, we believe that our measures already capture the majority of the implications of NC. The simplistic mathematical definition of the cosine similarity measure also allows for relatively straightforward quantitative analysis in the near-optimal regime that can be adopted by future works. Our main novelty lies in the quantitative analysis of the influence of various factors on the guaranteed proximity to the NC structure in the near-optimal regime. The extension of this analysis to the full predictions of NC, including norm equality, requires future work.
3. We think that the main doubts of the reviewer on the contribution of our theorem can be summarized as a) the asymptotic nature of our theorems, b) the exponential dependency of the theorem, and c) limited theoretical contributions. We will respond to these points separately:

   1. We would like to note that our big-O notation is not intended to indicate that our bound only works when $\epsilon\rightarrow 0$. We addressed this in detail in the global response.

   2. For theorem 3.1, we believe that the exponential dependence on $\alpha\beta$ does not make the theorem vacuous or meaningless. Under the presence of batch normalization and high weight decay, the value of $\alpha\beta$ are not typically large. In fact, the dependence of the bound on the weight and feature norms is the exact point made by the theorem.

      For theorem 3.2, the exponential dependence on $\frac{1}{\lambda}$ may indeed be vacuous in a large number of settings. We have carefully examined our proof and found that our results can be improved to remove this exponential dependence. The newly derived bound has been noted in the global response.

   3. In analyzing the near-optimal scenario, we quantitatively investigate the effects of weight decay parameters and the dependence on the loss value, offering a novel perspective on the understanding of weight decay and training loss value. Moreover, our result—focusing on the near-optimal regime with batch normalization and weight decay—more closely mirrors realistic settings, in contrast to the less realistic assumptions made in other papers analyzing NC under CE loss (e.g. [1][2]) such as exact optimality, strictly normalized feature vectors, or direct optimization of last-layer features.

[1] Jinxin 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. arXiv preprint arXiv:2203.01238, 2022.

[2] Dustin G. Mixon, Hans Parshall, and Jianzong Pi. Neural collapse with unconstrained features, 2020.

Minor

1. We thank the reviewer for pointing out this potentially misleading expression. To avoid confusion, we will change the expression to “neural networks without the bias term”.
2. We will fix this in our updated manuscript.

We thank the reviewers for their insightful reviews!

Theoretical Side:

1. Response to Comment only NC1 and NC2:
   • It is true that our current theorems only provide guarantees on NC1 and NC2. Compared with the full NC phenomenon, the original NC paper showed that NC4 is implied by NC1, NC2 and NC3, therefore the main missing part in our work is NC3. We did not officially include the bounds for NC3 in our main theorem since properties of class-feature vectors are typically of more practical significance than last-layer weight vectors. However, the corresponding bound for NC3 is available as an intermediate result of our proof (page 23, bottom) with the same asymptotic bound as inter-class cosine similarity. We will gladly include this as a separate result if it strengthens the contribution of the paper.
2. Response to Limitation to Last-Layer Feature Vectors:
   • Indeed, the scope of our work, like previous studies on NC, is confined to the last-layer feature vectors. This approach aligns with the original NC claim in [1].
3. Response to Affine Parameter in Theorem 3.2:
   • We believe the reference is to theorem 2.2. Here, the the affine parameter $\gamma$ is regularized by weight decay, as shown in $L_{\text{reg}} = L + \frac{\lambda}{2}(\|\gamma\|^2 + \|\boldsymbol{W}\|^2_F)$ in section 2.1. The theorem's conclusions are a result of the combined effect of batch normalization and weight decay, with the influence of batch normalization being implicit in the $\gamma$ parameter's effect. We will clarify this in the revised manuscript.
4. Response to Ignoring the Bias Term of BN:
   • We appreciate the reviewer highlighting the significance of the bias term in BN, which can indeed impact vector norms and consequently the cosine similarity measures. Our current definitions of intra-class and inter-class cosine similarity assume globally centered feature vectors. To address networks with biased batch normalization, a more general cosine similarity measure for these similarities where the feature vectors are globally centered should be used which centers the feature vectors by subtracting the mean. The globally centered feature vectors is also used in the original proposal of NC in [1]. This revised definition will be included in the updated version of our paper.

Experimental Side:

1. Response to Analysis in Sec 3.:
   • We thank the reviewer for suggesting the inclusion of average values for intra-class and inter-class cosine similarities. This will be incorporated in the revised manuscript
2. Response to Experimental Results in Fig 2.:
   • The phenomenon observed by the authors with the conic hull dataset and 3-layer MLP is indeed interesting. However, our theoretical results suggest that higher weight decay, coupled with batch normalization, leads to a greater degree of Neural Collapse at near-optimal loss. The experiments aim to demonstrate this relationship, particularly under the influence of BN and show that BN models have higher neural collapse than non-BN models at high weight decay. The comparison between models with and without BN across smaller weight decay values is beyond the scope of our current theorems since the theoretical guarantees provided by BN under very small weight decay are mostly vacuous as suggested by our theorems.
3. Response to Inclusion of ResNet Results in Fig 3.:
   • Our objective was to illustrate that Neural Collapse levels increase with weight decay in the presence of BN, as is shown and justified in the ResNet figures even without a non-BN baseline. Training a ResNet model without BN layers to achieve near-optimal loss is challenging, which is why a direct baseline comparison was not included.

Suggestions on Fig 2. and 3.:

• We acknowledge the feedback regarding the figures. We will address these points by generating higher-resolution images and reorganizing the figures for clearer comparison in the revised manuscript.

Response to the Question:

There are 2 main benefits of analyzing the near-optimal regime

1. The exact optimal loss is rarely achieved in practical training. Therefore, while prior works provide insights into the ideal case, our theorems extend their results into more practical final losses.

2. Another more important reason is that analyzing the near-optimal regime allows quantitative analysis of the effect of different parameters on the proximity to NC. For example, while prior works have shown that feature normalization/regularization and optimal loss guarantee the exact NC structure, our theoretical results show the quantitative relationship between the guaranteed proximity to the NC structure and weight decay, which is not possible by only analyzing the optimal case.

[1] Papyan, V., Han, X. Y., & Donoho, D. L. (2020). Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117(40), 24652-24663.

Thank you for your response, but I will keep my score.

1. For the experiments, in Fig. 3, there seems to be no big gap between model with and without BN.
2. In Fig. 2, the experiments cannot be fully predicted by the theorems proposed. The application of the theorems is limited.
3. Still, for theorem 2.2, the effect of $\gamma$ is implicit while a theorem involving $gamma$ explicitly is encouraged.

Also, after proofreading the other reviewers' comment, I decide to retain my score.

[1] Zhihui Zhu, Tianyu Ding, Jinxin Zhou, Xiao Li, Chong You, Jeremias Sulam, and Qing Qu. A geometric analysis of neural collapse with unconstrained features. arXiv preprint arXiv:2105.02375, 2021.

[2] Jianfeng Lu and Stefan Steinerberger. Neural collapse under cross-entropy loss. Applied and Computational Harmonic Analysis, 59:224–241, 2022.

[3] Can Yaras, Peng Wang, Zhihui Zhu, Laura Balzano, and Qing Qu. Neural collapse with normalized features: A geometric analysis over the Riemannian manifold.

General

1. Minimal set of conditions for NC.
   1. Motivation. Past work has demonstrated the emergence of Neural Collapse under various assumptions, such as minimal loss combined with regularization[1], and minimal loss with feature normalization [2][3]. However, there is a significant gap between theory and practice. Empirically, we rarely achieve the optimal loss exactly, nor do we apply regularization and normalization precisely as assumed in these papers. Therefore, the motivation of our paper is to bridge this gap. Specifically, we aim to theoretically identify a minimal set of conditions closely related to practical applications that guarantee the emergence of Neural Collapse.
   2. Batch normalization. The reviewer’s claim that 'The authors argue to add batch normalization to this list but that contradicts minimality' does not accurately reflect our position. Firstly, previous studies have also incorporated various normalization assumptions to facilitate Neural Collapse. Compared to these normalizations, our use of batch normalization and weight decay is more prevalent in practical applications. Therefore, employing batch normalization should not be viewed as an 'additional' assumption, but rather as an alternative or practical implementation of the norm constraint assumption. Secondly, our experiments demonstrate that without batch normalization, Neural Collapse is not significantly observed in some datasets (see Figures 2 and 3). This underscores the necessity of batch normalization from another perspective. Below we outline the role of BN and WD in our theorems in detail:
      1. To illustrate the above point, our first theorem shows the case where “norm constraint on the features” and near-optimal loss implies Neural Collapse. In this theorem, our theoretical result does not rely on Batch Normalization and only requires the feature norm to be bounded.
      2. Our second theorem replaces the feature and weight norm constraint with a batch-normalization layer and weight decay. Our results show that replacing the strict norm constraint with a norm penalty is also sufficient for Neural Collapse. Therefore, the BN assumption is not added as an additional assumption to the bounded norm constraint assumption but as a replacement that better resembles practical training.
   3. Sub-optimal Loss. Another area where our work contributes is in generalizing from the optimal scenario to the sub-optimal one. The theorems in previous works assume a strictly optimal loss, which is a much more stringent requirement than a near-optimal loss, the latter being more representative of practical scenarios. Therefore, in this regard, our assumptions are more 'minimal' compared to those in prior theorems. As we clarified previously, our big-O notation is used to simplify the formula by removing constants and higher order terms does not imply that our theorem only holds when $\epsilon\rightarrow 0$.
   4. Weight decay. Furthermore, we have provided a quantitative analysis of how various parameters influence the bound, specifically the gap from optimal loss and the weight decay parameter, aspects not explicitly analyzed in prior works. Weight decay, while widely used in practice, has a sparse theoretical understanding. Our work contributes to understanding weight decay from the perspective of NC.
2. We thank multiple reviewers for pointing out the potentially misleading expression of our contributions on the cosine similarity. We are aware that cosine similarity has been widely used in previous literature. Our main claim is that we’re novel in applying quantitative analysis of the cosine similarity measure in the theoretical analysis of Neural Collapse. We will clarify our contributions with regard to this statement and discuss the related use of the cosine similarity measure for similar purposes in previous literature.

We also thank the reviewer for pointing out minor errors in our manuscript and will fix them in our updated version.

Questions

1. A simplex ETF structure can be arranged to satisfy batch normalization regardless of the $\gamma>0$ value. To avoid confusion, we would like to clarify that we consider “scaled simplex ETF”, which can have vector norms to be any value (that is the same across different classes) and is also the case in the original NC literature. Therefore, we do not believe that this affects the validity of our theorems
2. As we mentioned in the global response, the big-O notation is mainly intended to omit constants and higher-order terms instead of implying that our bounds only hold when $\epsilon\rightarrow 0$. However, this assumption is still required for our theorem to hold true.

Experiments

• Our bounds provide the worst-case scenario bound on the intra-class and inter-class cosine similarity. Since we proved our theorems in the appendix, it is not necessary to verify our theorems experimentally. I believe the reviewer is referring to verifying how close the actual cosine similarity is to the predicted bound. However, we would like to note that the exact intra-class and inter-class cosine similarity resulting in a training process is heavily dependent on the training dynamics and the underlying model architecture, therefore predicting the exact cosine similarity value or its scaling is much more involved and is not within the scope of the paper. The main insights from the theorems are that 1) Batch normalization and high weight decay imply neural collapse under near-optimal loss and 2) higher weight decay has stronger guarantees of proximity to the Neural Collapse structure. The purpose of our experiments is to investigate these two insights to see whether they transfer into practice.
• As clarified in the previous point, the purpose of the experiments is not to verify whether the worst-case bounds in our theorems provide a good prediction of the actual intra-class and inter-class cosine similarities, which is heavily dependent on the training dynamics and architecture. Rather, the experiments substantiate the insights gained from the theoretical results, namely batch normalization is an important factor in Neural Collapse Emergence, and higher weight decay results in stronger guarantees of Neural Collapse.

  1. We thank the reviewer for pointing out that applying batch normalization for every layer is suboptimal. We have rerun our experiments with only last-layer BN along with other changes and the results have indeed shown improved Neural Collapse under our measures.

  2. We thank the reviewer for pointing out that including training losses will make our plots more informative. We would first like to note that, our theorems do not rely on the unconstrained features assumption so long as near-optimal loss is obtained. On the doubts on the reported intra-class and inter-class results, we think that the following reasons may be influential in the current results for the real-world networks and datasets:

     1. Achieving near-1 intra-class cosine similarity is a strong requirement on the features within the same class when the vector dimension is large.
     2. We used the minimum intra-class cosine similarity and maximum inter-class cosine similarity, which again is a strong requirement for the structure of feature vectors. This measure can be affected by outliers since there are $C(C-1)$ pairs of classes for inter-class cosine similarity

     To address the second point, we would report the average intra-class and inter-class cosine similarity to provide additional measures that are less sensitive to outliers. For the synthetic data experiments, we also adjusted the model architecture to only include batch normalization in the final layer and this has indeed improved the measured results.

  3. We believe that the above 2 replies address this point.

  4. We thank the reviewer for pointing out this result that seems contradictory to the insights from our theorems. We would acknowledge this phenomenon in our updated version.

We thank the reviewer for the meticulous review of the robustness of our theorems and experiments and for providing valuable and constructive advice on our experimental setup. Below are some of our responses to a few points that the reviewer has made in the review.

Theoretical Side

1. While the bound in theorems is formulated using asymptotic notation, the use of this notation is mainly intended to remove the constants and higher-order terms to simplify the formula in the theorem. It is possible to use our result to generate predictions for given values of $\epsilon$ as long as they are small enough for higher order terms of $\epsilon$ (e.g. $\epsilon^2$) to be safely ignored, which is the case for most successful training of over-parameterized neural networks. We have placed the formula for the asymptotic bound in the global response

   We would like to further note that, the main insight gained from the bounds in the theoretical result is its dependency relationship on the parameters rather than the exact value. Therefore, we believe that using the asymptotic notation is sufficient in the main paper as all the dependency relationships are captured with such notation.

2. We thank the reviewer for the meticulous review of the robustness of our second theorem. We acknowledge that the loss function of our second theorem is dependent on $\lambda$. However, we’re currently not convinced that the reviewer’s argument challenges the robustness of our theorem and would appreciate further elaborations on the argument, as outlined below:

   1. In the example given by the authors, the objective function is strongly convex with respect to the bounded variable (i.e. x), and the convexity of the function with respect to x is directly dependent on $\lambda$. This resulted in the trivial conclusion that a larger $\lambda$ value provides stronger guarantees on the value of $x$. However, in our theorems, $\lambda$ is the coefficient of the norm of the weights and the feature vectors, and thus the loss function is strongly convex with respect to the weight and feature norms. However, the intra-class and inter-class cosine similarity that our theorem provides guarantees on is not directly affected by the norm of the feature vectors since proportionally scaling all features does not affect the cosine similarity values. Therefore, the claim that “larger weight decay results in stronger guarantees under the same $\epsilon$ value from the optimal loss” is non-trivial in our theorems, unlike the example provided by the authors to illustrate their point.
   2. If $\epsilon$ is replaced with $\epsilon\lambda$ as the reviewers suggested, the main conclusions of our work are still upheld. Notably, the bounds in theorem 3.2 are still inversely correlated to $\lambda$ if $\epsilon$ is replaced with $\epsilon\lambda$

Reply to Concern on Batch Norm: It is true that Batch Normalization provides stronger guarantees than what is used in theorems and may lead to tighter guarantees. However, the constraint we used is a necessary implication of Batch Normalization and additional constraints of BN thus do not influence the validity of our theorem. It may be possible that additional constraints by BN may lead to better results, but this is not within the scope of this paper. On your second point where “the theorems assume that the loss minimizer is a simplex ETF in the presence of batch normalization”, our theorems do not rely on this particular assumption. Instead, this statement can be derived from our theorem by setting $\epsilon=0$. Could you explain more on where in our proof is this assumption necessary so that we can either provide more explanation or clarification in our proof?