Symmetric Neural-Collapse Representations with Supervised Contrastive Loss: The Impact of ReLU and Batching

Thank you for constructive feedback and insightful questions. Please find our response below in addition to the global response:

Batch size: Thanks for highlighting the role of batch size. Indeed, for contrastive loss and in the analysis of our model, batch size plays an important role in defining the set of global solutions. This is the case as the batching scheme specifies the interactions among samples during training. Our theoretical framework outlines the necessary and sufficient conditions for the batching to achieve a unique minimizer for contrastive loss. Following your suggestion, we conducted additional experiments for varying batch sizes to empirically explore its impact on the final learned embeddings. The detailed results of these new experiments are presented in Section E.4.3 in the appendix of the revised draft. In these experiments, we examine the effect of batch size in two scenarios, depending on whether ReLU is applied on the final embedding layer or not. We compare the learned features with OF when ReLU is present at the final layer, and with ETF (as predicted by Graf et. al (2021)) when ReLU is removed. Our observations indicate that the presence of ReLU reduces the reliance on batch size, leading to consistent convergence to OF. In contrast, in the absence of ReLU, large batches are necessary to achieve the global optimizer ETF. We also note that Khosla et.al (2020) argue that increased batch sizes lead to improved performance. In particular, they typically use batch sizes extending up to 6144. Our findings suggest that with the addition of ReLU, one can potentially achieve comparable performance with smaller batch sizes, which offers venues for reducing the training computational cost.

Activation functions:

Thank you for bringing light to this subtle point. We considered ReLU as the primary focus of our study in terms of the activation function due to its simplicity and frequent use in practice. While it is interesting to study different activation functions, our motivation was to identify that a simple and commonly used activation function can significantly alter the implicit geometry of features.

That said, we agree that this is indeed an interesting question. The theoretical analysis only assumes that the post-activation features are non-negative and span the region $[0, \infty]$. It is reasonable to anticipate that activation functions satisfying this requirement lead to the same result. However, the non-linearity of the activation function is likely to play a role in the optimization dynamics and needs further investigation. We have planned to conduct specific experiments and report our findings in the paper.

Dear Reviewer,

Thank you for your time and effort in reviewing our work. We appreciate your useful inputs, and have made sincere efforts to address your comments. Thank you for the constructive suggestions of batch-size ablation studies and discussion on other activation functions. These discussions will be useful additions to our paper. On the test accuracy discussion, we would like to refer you to our global response (including additional results on gains in worst-class accuracy reported in Sec. E.8 in the revised appendix). If you have more questions, we’re happy to provide further information.

We would appreciate it if you would consider re-evaluating your rating.

Thank you for your appreciation of our messages and analysis, and for your constructive feedback. Please find our response below in addition to the global response:

Additional experiments on transformer architecture: Thank you for the suggestion of extending the experiments to transformer-based architectures. Motivated by the question, for the revised version of our paper, we have included additional experiments on vision transformers (see Fig.22). Based on the observations, our hypothesis of OF geometry holds here as well, thus solidifying the claim of the generality of our results across various architectures.

Questions:

The minimum of the loss is not zero and varies with the number of samples in different classes:

That is correct. Due to fixing the norm of the features to a constant value (say 1), the loss does not go to 0. The minimum value is exactly given by the lower bounds in Eq (3) and (4), respectively, for full-batch and mini-batch optimizations. For full batch SCL, we show the loss evolution in a DNN experiment, where the loss is seen to converge exactly to our lower bounds in Theorem 1. This prediction is made possible due to our theoretical analysis.

To demonstrate the effectiveness of our Theorem 1, in Fig 3, we show how closely the empirical loss in a full-batch DNN experiment converges to the predicted lower bound.

Thank you for your response! Please allow a couple of remarks:

1. First, because of neural collapse (i.e. features of each class collapse to their class means), the features lie on a k-dimensional subspace. In other words, and following the reviewer’s notation for convenience (i.e. n for the feature dimension, and, say, N for the train-set size and k for the number of classes) the rank of the (N x n) feature matrix H is k. Thus, there is no “shrinking of the representation space.”

2. Second, the above holds irrespective of the inclusion or not of ReLU. In fact, even if ReLU is not present and one considers a simplex ETF geometry for which features are maximally separated, then the cosine of the angle between the features is -1/(k-1). Thus, the cosine actually approaches 0 (same as the angles of OF), as the number of classes k increases.

3. Third, let’s even consider the case of CE optimization on balanced data, i.e. the setting of the original work by [Papyan et al. (2020)] and a long list of follow-up papers. In that setting, again, the features of the last layer usually undergo a ReLU nonlinearity in most common architectures (eg. ResNet). Hence, your concern would still apply to that setting, which is however arguably well-established by now in the community. For the above reasons, there are no concerns with regards to the “sufficiency of the method for learning from large and complex datasets.”

With regards to additional experiments: we kindly note that our paper identifies a new phenomenon, which we not only theoretically justify, but also experimentally confirm on four different datasets (MNIST, CIFAR10, CIFAR100, TinyImagenet) of both STEP and LT imbalance trained on four different architectures (MNIST, VGG, ResNet, ViT). Given the range of scales and complexities of these datasets/architectures, we believe this forms sufficient evidence in support of the identified phenomenon. We also kindly note that the breadth of the experiments is consistent with the rest of the literature on neural collapse phenomena (e.g. Graf et al. (2020), Ji et al. (2022), Sukenik et al. (2023), Thrampoulidis et al. (2022), Yaras et al. (2023), Zhu et al. (2021), Zhou et al. (2022a), Zhou et al. (2022b) etc). Please note that our experimental evidence is more extensive than several prior works. As a matter of fact, and to the best of our knowledge, we are the first to demonstrate neural-collapse geometry convergence for a transformer architecture (ViT). (Thanks again for your suggestion!)

PS: By “UAM”, we believe the reviewer is referring to UFM (Unconstrained Features Model)

Thank you for your supportive comments and useful feedback. We reply to the questions below in addition to the global response:

UFM assumption: We would like first to emphasize that our work addresses the problem for highly overparameterized models and training beyond the zero-training error. This is a setting considered in the majority of previous works on neural collapse [Mixon et al. (2020); Fang et al. (2021); Graf et al. (2021);Lu & Steinerberger (2020)]. The UFM-inspired theoretical analysis in the literature is also motivated by this overparameterization assumption. In fact, UFM assumes that overparameterized models are capable of representing any function mapping the input to the embedding space. As pointed out by the reviewer, this indeed is a strong assumption, although one that has proved to be very effective. Our findings alongside the results in a series of previous related works confirm this model's efficacy in predicting the properties of the embeddings learned at the final layer for large enough models. In other words, the high expressivity of overparameterized models makes the UFM a useful tool for identifying the global minimizer feature configuration, so that the impact of the loss function itself can be isolated from the role of the network architecture. Our work shows that the UFM remains a useful mathematical model even with the non-linear activation function of ReLU.

On the optimization aspect, we believe that studying the optimization effects on the UFM has scope for future work. Our experiments provide evidence for consistent convergence towards the global minimizer (OF geometry) despite the non-convexity of the network optimization.

Empirical: We appreciate the reviewer's recommendation regarding experiments on large-scale datasets and consider adding these experiments to the final version of the paper. Along these lines, we are excited to report new experimental findings in support of our hypothesis of OF geometry on transformer architecture. We have included additional experiments on vision transformers (see Fig.22). Based on our current observations, our hypothesis of OF geometry holds here as well, thus solidifying the claim of the generality of our results across various architectures.

In summary, we would like to note that our current findings are backed up by experiments spanning various architectures and datasets, including MLP, ResNet, and DenseNet applied to MNIST, CIFAR10, CIFAR100, and TinyImageNet. Across this range of tasks, from simpler and smaller ones like MNIST (10 classes) to more complex and larger tasks like TinyImageNet (200 classes), we have not observed significant changes in the convergence behavior or the communication of our key message (see Fig 2) as long as the network and dimension of output features are large enough. The observations from all of the above experiments provides us with the confidence to believe that the core message of our paper remains largely unchanged. This is the case for both our main findings: convergence to OF and the efficacy of the new batch-binding scheme.

Missing Citation: Thank you for pointing this out. We agree that Wang, Isola (2020) is a relevant work. It appears that not mentioning the same was an editing oversight on our part. We have included the reference and commented on the relationship.

Empirical differences: At the moment the main empirical differences observed are the following : the test accuracies (Tab. 1) and new analysis on worst-class accuracies between with and without ReLU. We discuss this in the global response as well. We are not aware of differences in the batching aspects. In fact, the observation in Fig. 21 suggests that the batch-binding scheme is beneficial even when training without ReLU, for convergence to the final feature geometry.

Thank you for your time and effort in reviewing our paper. On the motivation, we would kindly refer you to our comments in the global response. At the outset, the explicit goal was not to symmetrize the feature geometry, but to uncover and understand the same. This is in line with the research goal of the rapidly growing interest in neural collapse and its implications for DNN training with various objective functions. Please note that there is no prior work that characterizes the feature geometry and thus the behavior of the loss function in inducing the same, for the SCL under imbalances. Prior to our work, the understanding of the intricate role of mini-batch optimization was limited as well. We have identified a setting where the geometry is predictable, incidentally observing that it is symmetric. Based on the works in literature on related loss functions under class imbalance, this result remains surprising. Additionally, we show that the convergence quality is excellent, in practical experiments, as compared to experimental results on geometry under imbalances in literature. Thus, we believe that our discovery, supported by concrete theoretical results, paves the way to important future questions in understanding the loss functions that are used across datasets and model architectures, and identifying characteristic properties of various loss objectives.

Questions:

Question on CE: It turns out that adding ReLU at the final embedding layer does not guarantee that the classifiers $w_k$ learned by CE are positive. To see this, take a simple binary setting. In the binary setting, we can assume that we only search for a single classifier, as at the optimal solution $w_2=-w_1$. Now, suppose the embeddings $h_i, i=1:n$ are fixed to an OF: $h_i=e_{y_i}$, and further suppose the norm of the classifier is fixed (since only the direction matters for the final decision rule). Then, the max-margin classifier that optimizes the objective is $w\propto[1,-1]$. In short, the optimal classifier can have negative entries despite the embeddings being constrained to the positive orthant.

Our empirical observations also support this claim. First, we note that ResNet networks incorporate ReLU in the final embedding layer (so they fit in this setup). In our ResNet experiments, we have compared the embeddings learned by CE and SCL. First, in Fig. 7, we report the convergence of the embeddings to the OF, and we observe that CE learns embeddings that do not converge to the OF (despite the positivity constraints). Second, in Fig. 1 and 10, we plot the heatmap of the learned embeddings with CE for different imbalance ratios. These results indicate that merely introducing ReLU to the embedding layer does not ensure the orthogonality of features learned by CE. Moreover, the experiments in Thrampoulidis et al. (2022) on networks that have ReLU at the hidden layer show that the embedding geometry alters with imbalance ratio, further suggesting that simply constraining the embeddings to be positive for the CE loss does not guarantee orthogonality of the features. In general, solving for the global minimizer feature-classifier configuration with CE loss is non-trivial. Several works have tackled this problem with different assumptions [Zhu et al. (2021), Graf et al. (2021), Fang et al.(2021), Thrampoulidis et al. (2022), etc]. In light of these observations, we believe exploring the impact of the final layer adjustments on CE deserves a separate study.

We appreciate the interesting questions, and are happy to provide any further clarifications.

We thank you for their detailed review and questions. We are pleased to see your recognition of our efforts in presentation. Please find our detailed response below, in addition to the comments in the global response.

For example, consider the simple case of binary classification...: Thank you for prompting this discussion. The statement regarding antipodal features for binary SCL without ReLU is exactly correct. In fact, this holds in general for balanced multi-class setup as well: When the dataset is balanced Graf et al. (2020) shows that the embedding geometry will form an ETF in the absence of ReLU. In an ETF geometry, we have $\frac{\mu_c^\top\mu_{c'}}{\|\mu_c\|\|\mu_{c'}\|}=\frac{-1}{k-1}$ for the mean-embeddings of any two classes. As opposed to OF where $\frac{\mu_c^\top\mu_{c'}}{\|\mu_c\|\|\mu_{c'}\|}=0$.

Therefore, for ETF, $\|\mu_c-\mu_{c'}\|^2 = \frac{2k}{\tau(k-1)}$ and in the case of OF $\|\mu_c-\mu_{c'}\|^2 = \frac{2}{\tau}$. $\tau$ is the temperature of the SCL loss. While the distance in the ETF geometry is larger for a fixed temperature parameter, we note that by simply scaling the temperature parameter $\tau$ by $\frac{k-1}{k}$, we can achieve an equivalent absolute distance in the learned geometry in the presence of ReLU. Thus, in terms of this measure, the addition of ReLU does not compromise the trained model.

On the contrary, as shown in Fig. 1, in the imbalanced setup, without ReLU, the minority classes are drawn closer to each other, yielding a reduced absolute distance between classes. The introduction of ReLU, however, maintains an equal pairwise distance between classes. Hence, in terms of this measure, the addition of ReLU is in fact helpful.

Finally, as noted in Section D.1.2, centering the OF by $\frac{1}{k}\sum_c\mu_c$ results in an ETF. So, in summary, the geometries learned in balanced or binary scenarios, with or without ReLU, can be considered equivalent, accounting for a shift and scaling.

Novelty and significance

Thank you for appreciating the originality of our work and ideas. We believe that solving for the global minimizer feature-classifier configuration with the SCL loss under is non-trivial and novel. We elaborate below:

1. Symmetric feature geometry: the first surprising observation that we support with theory is the fact that ReLU symmetrizes the feature geometry regardless of imbalances. To the best of our knowledge, such an outcome is unique to the SCL and stands out from several prior works considering class imbalance in CE and other losses. For e.g., Thrampoulidis et al. (2022), Fang et al. (2021) discuss that the feature and classifier geometry depends critically on the imbalance. We further identify the appropriate theoretical model for the observation, and solve for the exact geometry of the global optimizers and the optimal value of the loss function.
2. Batching matters: the first study of global optimizers of the SCL, Graf et al. (2020) only considers the setting of balanced data. In their words, the theoretical analysis is based on “a combinatorial argument which hinges on the balanced class assumption”. The analysis therein requires the specific batching case of all combinations of examples in the dataset of a given batch-size. On the other hand, our mini-batch analysis is applicable to arbitrary batching schemes. We discover the batching conditions under which a unique minimizer geometry exists- this condition allows for a large class of batching schemes, and includes the specific case of Graf et al. (2020) as a special case. We leverage our theoretical results to concretely explain geometry observations made in DNN experiments as a function of batching schemes.

We discuss the question of possible extension to CE below, within the response to the specific question.

Minor issues

overuse of abbreviations: Thanks for the feedback. We have added a table defining all the important abbreviations used in the paper in one place at the beginning of the appendix.

eq(1) and eq(2): Eq (1) and (2) are equivalent, in terms of the loss function. The apparent difference is only in presentation, in that in Eq (2), the sum over $l$ includes $j$, so that when $l = j$, the innermost exponential term becomes $\exp(\mathbf{h}_i^\top\mathbf{h}_l - \mathbf{h}_i^\top\mathbf{h}_j) = \exp(0) = 1$.

******** Continued below *********

Dear Reviewer,

Thank you for your time and effort in reviewing our work. We have made diligent efforts to address your comments and we hope that our responses have addressed your questions. If you have more questions, we’re happy to provide further information.

We would appreciate it if you would consider re-evaluating your rating.

Finally, we outline the main additions to the revised manuscript below:

• Section E.4.3: In response to reviewer ScWC's suggestion, we provide additional experiments examining the impact of batch size on learned embeddings during training. Our findings indicate that when ReLU is present in the final layer, the learned embeddings are less sensitive to the choice of batch size.
• Section E.8: To further compare the performance of the models with and without ReLU, we report the worst-class test errors for different imbalance ratios in each of the two scenarios.
• Section E.8.2: Addressing reviewer q7qd's recommendation, we present preliminary empirical results on vision transformers to showcase the applicability of our findings on more general architectures.