Symmetric Neural-Collapse Representations with Supervised Contrastive Loss: The Impact of ReLU and Batching
Authors: Ganesh Ramachandra Kini
Vala Vakilian
Tina Behnia
Jaidev Gill
Christos Thrampoulidis
Venue: ICLR 2024 poster (2024)
Keywords: Supervised contrastive learning
neural collapse
implicit bias
class imbalance
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Abstract: Supervised contrastive loss (SCL) is a competitive and often superior alternative to the cross-entropy loss for classification. While prior studies have demonstrated that both losses yield symmetric training representations under balanced data, this symmetry breaks under class imbalances. This paper presents an intriguing discovery: the introduction of a ReLU activation at the final layer effectively restores the symmetry in SCL-learned representations. We arrive at this finding analytically, by establishing that the global minimizers of an unconstrained features model with SCL loss and entry-wise non-negativity constraints form an orthogonal frame. Extensive experiments conducted across various datasets, architectures, and imbalance scenarios corroborate our finding.  Importantly, our experiments reveal that the inclusion of the ReLU activation restores symmetry without compromising test accuracy. This constitutes the first geometry characterization of SCL under imbalances. Additionally, our analysis and experiments underscore the pivotal role of batch selection strategies in representation geometry. By proving necessary and sufficient conditions for mini-batch choices that ensure invariant symmetric representations, we introduce batch-binding as an efficient strategy that guarantees these conditions hold.

Published as a conference paper at ICLR 2024
SYMMETRIC NEURAL-COLLAPEE REPRESENTATIONS WITH SUPERVISED CONTRASTIVE LosS: THE IMPACT OF RELU AND BATCHING
Ganesh Ramachandra Kini', Vala Vakilian", Tina Behnia Jaidev Gill", Christos Thrampoulidis'
University of California, Santa Barbara, USA 'University of British Columbia, Canada
ABSTRACT
Supervised contrastive loss (SCL) is a competitive and often superior alternative to the cross-entropy loss for classification. While prior studies have demonstrated that both losses yield symmetric training representations under balanced data, this symmetry breaks under class imbalances. This paper presents an intriguing discovery: the introduction of ReLU activation at the final layer effectively restores the symmetry in SCL-learned representations. We arrive at this finding analytically, by establishing that the global minimizers of an unconstrained features model with SCL loss and entry-wise non-negativity constraints form an orthogonal frame. Extensive experiments conducted across various datasets, architectures, and imbalance scenarios corroborate our finding. Importantly, our experiments reveal that the inclusion of the ReLU activation restores symmetry without compromising test accuracy. This constitutes the first geometry characterization of SCL under imbalances. Additionally, our analysis and experiments underscore the pivotal role of batch selection strategies in representation geometry. By proving necessary and sufficient conditions for mini-batch choices that ensure invariant symmetric representations, we introduce batch-binding as an efficient strategy that guarantees these conditions hold.
INTRODUCTION
The prevalence of deep-neural networks (DNNs) has led to growing research interest in understand- ing their underlying mechanisms. A recent research thread, focusing on classification tasks, explores whether it possible describe the structure of weights learned by DNNs when trained beyond zero-training error. The specific characteristics of this structure will depend on the DNN being used, the dataset being trained on, and the chosen optimization hyperparameters. Yet, is it possible to identify macroscopic structural characteristics that are common among these possibilities?
In an inspiring study, Papyan al. (2020) demonstrates this is possible for the classifiers and for the embeddings when training with cross-entropy (CE) loss and balanced datasets. Through extensive experiments over multiple architectures and datasets with an equal number of examples per class, they found that the geometries of classifiers and of centered class-meat embeddings (outputs of the last hidden layer) consistently converge during training to a common simplex equiangular tight frame (ETF), structure composed of vectors that have equal norms and equal angles between them, with the angles being the maximum possible. Moreover, they observed neural-collapse (NC), a property where embeddings of individual examples from each class converge to their class mean embedding.
Numerous follow-up studies have delved deeper into explaining this phenomenon and further inves- tigating how the converging geometry changes with class imbalances. The unconstrained features featurss model (UFM), proposed independently by Mixon et al. (2020); Fang et al. (2021): 1); Graf et al. (2021) )); Lu & Steinerberger (2020), plays central role in the majority of these follow-up studies. Specifically,
"This work is supported by an NSERC Discovery Grant, NSF Grant CCF- and by KAUSTKAUST award. JG and CT gratefully acknowledge the support of NSERC Undergraduate Student Research Grant. The authors also acknowledge use of the Sockeye cluster by UBC Advanced Research Computing.

---

Published as a conference paper at ICLR 2024
the UFM serves theoretical abstraction for DNN training, in which the network architecture is viewed as a powerful black- box that generates embeddings without any restrictions in the last (hidden) layer. For CE loss, the UFM minimizes min b,eRR LCE (f (We) (e())) in which both classifiers We for the & classes and embeddings h, for the n training examples are unconstrainedly optimized over RR Zhu et al. (2021); Graf et al. (2021); Fang et al. (2021) have verified that the global minimum of this non-convex problem satisfies NC and follows the ETF geometry observed in DNN experiments by Papyan et al. (2020). Recent works by Thrampoulidis et al. (2022); Fang al. (2021); Behnia et al. (2023) have demonstrated that the global optimum of the UFM changes when classes are imbalanced. Yet, the new solution still predicts the geometry observed in DNN experiments, providing evidence that, despite its oversimplification, the UFM is valuable in predicting structural behaviors.
Expanding beyond the scope of CE minimization, Graf al. (2021) also used the UFM determine whether optimizing with the geometric structure of the learned embeddings.'
loss (SCL) results in any alterations to the
SCL. ReLU
SCL
CF
They proved for balanced datasets that the global solution of min Lscu. ( (hh)) simplex ETF, suggesting that CE and SCL find the same embedding geometries. In Fig. I, we
remains
R=1
investigate the geometry of class- ence of class-ime mass-imbalance. The prediction in Graf al. (2021) only applies to balanced data (SCL. with R = 1 in Fig. 1) and no prior work has explicitly characterized the geometry of SCL with class imbalances. Observing the middle column of the figure, note that the geometry of embeddings changes drastically" as the imbal- ance ratio R increases. This behavior of SCL is consistent with the CE embeddings geome- try, which, as discussed previously, also changes with the imbalance; see last column of Fig. 1.
in pres-
R= 10
100 100
In this paper, we arrive at a surprising finding: simply introducing ReLU activation at the fi- nal layer restores the symmetry in SCL-trained representations. This phenomenon is clearly il- lustrated in the first column of Fig. 1. Below is detailed description of our contributions.
Figure Gram matrices, Gw// maxi, [GM[
[]]]]
of class-means of learned feature embeddings last epoch of training, with ResNet-18 on MNIST. SCL+ReLU: the mean feature embeddings for dif- ferent classes are mutually orthogonal, forming OF, regardless
imbalance (imbalance ratio At 1,10, 100). This invariance does not hold in the ab- sence of ReLU for SCL. Further, CE feature geome- try is also sensitive to imbalance. The label distribu- tion is STEP imbalanced, with the first five classes majorities and rest as minorities. See Sec.E.2 SM for more information.
1.1 SUMMARY CONTRIBUTIONS
We conduct an in-depth examination of the ge- ometry of training representations (embeddings) of SCL,, particularly in relation to varying levels of dataset imbalances. Our study leads us to identify and capture, both analytically and empirically, and for the first time, the impact the representation geometry of: (i) straightforward architec- tural modification, specifically the incorporation of a ReLU activation at the final layer, and (ii) the batch-selection strategy.
Impact of ReLU, We find that the addition of a ReLU activation to the last-layer of the architecture results in orthogonal frame (OF) geometry. That is, a geometry structure composed of class-mest mean embedding vectors that have equal norms and are mutually orthogonal to each other. A preliminary experimental validation of this finding is illustrated in Fig. 1. The experiments correspond to STEP- classes and imbalance ratio R 1, 10, 100.
imbalanced MNIST data with five majority/minoriin Each heatmap represents the pairwise inner products between the class-means of learned feature embeddings. The figure (first column) reveals that with the addition of ReLU at the last layer, SCL
'SCI is designed to train only embeddings and is an extension of the well-known contrastive loss used for unsupervised learning. adapted to supervised datasets (Khosla et al., 2020); see Equation (2). AAnlooouus finding reported by Zhu et al. (2022), although not visualized as done here.
2

---

Published as a conference paper at ICLR 2024
MNIST
CIFARIO
CIFARI00
Tiny ImageNet
10
Step
10
(0
Long-tail
0.6
0.4
0.
200
200
Epoch
Tt. F across epochs where GM M M is the Gram matrix of class-mean mean embeddings. For MNIST and CIFARIO ResNet-18 and for the complex CIFAR 100 and Tiny-ImageNet, batch-binding (see Sec. 5). Regardless of imbalance level R. the embeddings consistently converge to the OF geometry.
Figure 2: Distance of learned embeddings the OF geometry measured by
|
GM
ResNet-34 and train using
learns orthogonal features regardless of class imbalance. This in stark contrast to the varying geometries of vanilla SCL (optimized commonly without ReLU) and CE loss. Extensive additional experiments for Long-tailes (LT) distributions, other datasets and architectures, are presented in Sec. and also Sec. E in the SM. We also present experimental findings that confirm ReLU symmetrises the geometry without compromising test accuracy.
We support this finding by theoretically investigating the global minimizers of extended version, denoted as UFM,, of the original unconstrained-featte for the presence of ReLU activations which impose the last-layer by minimizing SCL over the non-negative orthant. Concretely, we show that all global solutions of minn, LSCu. [(h )e(()) satisfy NC and the corresponding class-mean features form an OF. This explains the convergence of feature geometry an OF consistently observed in our experiments; for example, see Fig. 2.
model (UFM). This refinement accounts
Batching and its implications. Furthermore, we identify the crucial impact of batch selection strategies during SCL training in shaping the learned embedding geometry. Concretely, by analyzing the UFM., we establish straightforward criteria for the batching scheme, ensuring that any global minimizer UFM. forms an OF. These conditions explain the substantial degradation in convergence towards an OF when employing fixed batch partition instead of randomly reshuffling the batches each training epoch. Intriguingly, we demonstrate that any arbitrary batching partition can be transformed into one that fuliill our theoretical conditions by performing so-called batch-binding Through extensive experiments in Sec. and in the SM, we show that batch-binding consistently accelerates the convergence of the learned geometry towards OF and even improves convergence to ETF geometry under balanced training in the absence of ReLU.
SETUP
Notation. For a positive integer [n] (1.,,3,..., For matrix V Ruxx V[i.j] denotes its (ih entry, denotes the j-th column, VT its transpose. We denote | the Frobenius norm of V. We use V X whenever the two matrices are equal up scalar constant. We V 2 to denote the entry-wise non-negativity, i.c.. V[i.j] 2 0, Vi [m],j [n]. Finally, Im denotes the identity matrix of size m.
Consider k-class classification and training set S (xx, ) where x, Rd, yi [k] represent data samples and corresponding class labels. We use ne to denote the number examples in class ee [k], and (-(-) to denote the last-layer ceiture-embedddngg of deep-net parameterized by 0. We compute the SCL on given batch B $ [Nx] of training examples as follows.
(B = Σ
Σ
log 1 Σ exp (bx hho(e)) - ho(x)) P+i.j
hx (oxx)))oxx))))))))))))
(1)
F
BB.
By ded
3

---

Published as a conference paper at ICLR 2024
where 7 is positive scalar temperature hyper-parameter, and B. is the number of examples in batch B belonging to class c = y/. To train a deep-net, we minimize SCL (1) over the network parameters 0 set of batches B chosen from the training set. As introduced by Khosla al. (2020), SCI. requires normalized features, so we assume [ (x;) = 1. Furthermore, we assume that 2 k and ne. B.B 2 2, c [A We let Haxn = [h,,",, h denote an embeddings matrix where each column corresponds to one of the examples in the training set, i.e., h, (x,).. The embeddings geometry or so-called implicil geometry* refers to the norms and pairwise-angles of these vectors & e [n]. Note these quantities correspond exactly to the entries of the Gram matrix H H. To characterize the geometry. we need the following definitions.
Definition (Neural Collapse (NC)). NC occurs if h; = hj, V vi : Vi- Definition 2 (k- (k-Orthogonal Frame (k-OF)) We say that l vectors V = (Vi,..,,vs) E Rdck form k-OF if V Lg. f.e. each pair of (i.j) é [k]. |v, = lvyll and v, 00.
When the within-class variation of embeddings is negligible, i.e., H satisfies NC, it suffices to focus on the class-mean embeddings thc th Eigze ci,CE [k] and the respective matrices Matc
[Mi,""", [R]] and GM = M instead of H and H'H With these, we can formally define the OF geometry for embeddings.
Definition 3 (OF geometry). We say that a feature-embeddiig matrix H. follows an OF geometry if satisfies NC and the class-means form a k-OF, i.e. GM oc Ilg.
3 SCL WITH RELU LEARNS OF GEOMETRIES: EMPIRICAL FINDINGS
Experimental setup. Following the setup of Graf et al. (2021), our models consist of a backbone (ResNet, NenseNet, etc) with a normalizing layer to output features with unit norm. For all geometric convergence results, we apply ReLU on output features before normalization. We employ Stochastic Gradient Descent (SGD) with a learning rate of 0.1, momentum set to 0.9 and with no weight decay. Furthermore, as in Graf et al. (2021); Khosla et al. (2020), we set the temperature parameter = 0.1 as Khosla et al. (2020) have found that it yields optimal performance. In addition, we empirically for values near 0.1, yet the specific choice
find that convergence to OF is not highly dependent affects the speed of convergence.
We study the behavior of models trained with SCL under, 1) R-STEP imbalance having k/2 majority classes with maaj examples per class and k/2 minority classes with nimin "Wa// it examples per class, and 2) R-Long-taiin tailed (LT) imbalance where the number of training datapoints exponentially decreases across classes such that n. n, R- (e-1//1)-,, for c [k:. Regardless of the imbalance ratio R. we ensure that n. 2 2 by adding a vertically flipped version of each image as a method of batch duplication. Unless stated, we do not perform any additional data augmentation on the datasets, and use a batch size of 1024 with random reshuffling.
Finally, when studying the impact of ReLU on generalization, we use a Nearest Center Classifier (NCC) on the output features to evaluate model performance. For such experiments, following the setup of Khosla et al. (2020), we employ a projection head by adding 2 layer non-linear MLP (with or without ReLU at the last layer) and compare the test accuracy under difference imbalance ratios.
Metries. Since in all experiments we observe the within- class variations of the last- embeddings (h,) become negligible towards the end of training (see NC plots in Fig. 8 in SM), we focus here on the geometry of class- mean embeddings M. Thus, to measure the distance of learned embedding to the OF geometry, we compute the Gm
Imbalance Ratio (R)
w/o ReLU
w RcLU
72.17
0.23
72.32 = 0.60
10 (Siep)
56.58 a 0.50
58.16 x 0.76
10 (LT)
37.10 1.04
57.88 # 1.05
100 (Step)
43.49 0.30
43.80 a 0.25
100 (LT)
37.19 2.50
39.71 + 0.09
Is distance metric AGM := I |GM1 r | Ir..
Table 1: Test accuracy comparison for a ResNet- 18 trained using SCL on CIFAR100 with and without ReLU after the projection head, We use NCC for classification. Note that the addition ReI.U does not compro- mise the accuracy. For details, see E.7 in SM. 'When using SCI., it is common practice to add augmented version of the datapoints in each hatch to itself (Khosla et al., 2020; Graf et al., 2021). This practice, referred to as batch duplication, helps ensure that each datapoint in a batch has at least one other example in the same class to compare to during training. "The specific terminology is adopted from Behnia et al. (2023).
Observations Geometry, In Fig. 2, we plot the distance AGM between the learned Teature-embeddee geometry training progresses. We consistently observe
and the OF
4

---

Published as a conference paper at ICLR 2024
that the learned features during training converge to the OF geometry, irrespective of the imbalance level R of the training set and the imbalance pattern (STEP or LT). This suggests that the feature geometry learned by SCL is invariant to the training label distribution.
Interestingly, the invariance that is revealed by our study is distinct for SCL with ReLU as opposed to the case of SCL without ReLU, and that of CE and MSE loss. CE and MSE are known to be sensitive to imbalances (Thrampoulidii et al., 2022; Fang et al., 2021; Liu et al., 2023; Behnia et al., 2023; Dang al., 2023). Additionally, for different values of Rt, there is significant difference between the speed of convergence, unlike the CE loss (Thrampoulidis et al., 2022), where it is empirically observed that the rate worsens with larger imbalance.
Generalization. To ensure that ReLU does not yield any adverse effects on test accuracy, we trained ResNet-18 (with a projection head) with and without ReLU on CIFARIO, CIFAR 100, Tiny Imagenet and evaluated the balanced test accuracy when training under label imbalance. Following (Zhu et al., 2022; Khosla et al., 2020; Chen et al., 2020), we consider the features before a projection head for evaluating the test accuracy of the models. Our results, e.g. in Tab. 1, indicate that the addition of ReLU does not compromise test accuracy. For a detailed discussion on the results see E.7 in SM.
4 THEORETICAL JUSTIFICATION SCL WITH NON-NEGATIVITY CONSTRAINTS
In this section, we analytically justify the convergence of the embeddings learned by SCL to the OF geometry. For our theoretical analysis, we use the Unconstrained Features Model (UFM) (Mixon et al., 2020; Fang et al., 2021: Graf et al., 2021: Ji et al., 2021; Lu & Steinerberger, 2020; Tirer & Bruna, 2022), where we treat the last-layer features H as free variables, removing the dependence to the network parameters 0. However, we refine the UFM, and consider the SCL minimization over the non-negan non-negative orthant to accommodate for the presence of ReLU activations in the last layer.' Following Khosla et al. (2020), we further constrain the embeddings h, to be normalized. We prove, irrespective of the labels distribution, that the global optimizers form an OF. Formally, consider the following refined UFM for SCL, which we call UFM, for convenience:
H E arg min ESCL H
(H) subj. to H and |h,| 2 1, Vie [m].
(UFM,)
Recall H 2 denotes entry-wise non-negativity. We begin our analysis by considering the full-batch loss in Sec. 1. and extend results the SCL minimized mini-batches in Sec. 4.2. We provide specific conditions that full-batch version.
must satisfy in order to have the same global optimum as the
4.1 FULL-BATCH SCL
Here, we focus on the full-batch SCL, where the entire training set is treated as single batch. Specifically, we consider UFM. with LscI(()) being the full-batch SCL defined as,
Cran (H) == Σ
1
Σ
log Σ exp (h h,1))))
:c[a]
ny.
Thm. 1 specifies the optimal cost and optimizers of UFM, with the full-batch SCL (2). Theorem (Full-batch SCL minimizers). Let d 2 k. For any H feasible in UFM,, it holds,
Lrall (H)) 2 Σ log (ne 1+ (n nn)
(3)
re[k]
Moreover, equality is achieved if and only H satisfies NC and the class-means form a k-OF.
We defer the proof details to Sec. A in SM. The bound relies on successive uses of Jensen's inequality and the fact that for feasible H, each pair of training samples i.j E {[n] satisfies 0 s hhTh, s 1. We complete the proof by verifying that equality is attained only if h h, when y = My and h h, = when yi yj- Rather, to achieve the optimal cost, features with similar labels must align (NC) and the class-mean features must form an OF.
TTirer& & Bruna (2022) (and several follow ups, e.g. Súkeník et al. (2023)) has also studied incorporating the ReLU activation in the UFM, albeit focusing MSE loss.
5

---

Published as a conference paper at ICLR 2024
Loss
10-
B
Thm. 1 shows that any optimal embedding geometry learned by the full- batch SCL with ReLU constraints uniquely follows the OF geometry (Defn. 3) and the conclusion is independent of the training label distribution. We have already seen in Sec. 3, that this conclusion is empirically verified by deep-net experiments. To further verify the lower bound on the cost of the SCL loss given by the theorem, we compare it in Fig. with the empirical loss of ResNet-18 model trained with full-batch SCL (2) on R-STEP MNIST dataset and n 1000 total examples. Note the remarkable convergence of the loss to the lower bound (dashed horizontal lines) as training progresses.
9-
10
8
R 120
200
400
Epoch
Figure 3: Full-batch SCL converges to the lower bound (dashed lines) in Thm.
4.2 MINI-BATCH SCL
Thm. I shows that minimizing SCL over the training set as a single batch recovers an OF in the feature space, regardless of the training label distribution. However, in practice, SCL optimization is performed over batches chosen from training set (Khosla et al., 2020; Graf et al., 2021). Specifically, we have a set of batches B and we compute the loss on each batch B é B as in (1). While in the full-batch SCL (2), all pairs of training samples interact with each other, in the mini-batch version, two samples (i.j) interact only if there exists batch B such that i, e B. A natural question is whether the mini-batch construction impacts the embeddings learned. In this section, we study the role of batching on the embeddings geometry more closely.
Similar to the previous section, we consider UFM,, where we directly optimize the SCL over embeddings H. However, this time we study the mini-batch SCL defined as follows,
Lhatch (H) == Σ Σ BeB B
Σ log Σ JEB jeiays-2
exp (h, - h2 h/]h)) E
RB
Eas
Ki: = B, = c3{. Thm. 2 is an extension of Thm. for the mini-batch SCL (4). The proof proceeds similarly to that Thm. 1, where the loss over each batch can be independently bounded from below. Interestingly, due to the orthogonality of the optimal embeddings in every batch, the lower bound for each batch can be achieved by the common minimizer. The detailed proof deferred to Sec. in the SM.
where recall nB.
Theorem 2 (Mini-batch SCL minimizers). Let d k and B be an arbitrary set of batches of examples. For any feasible H in UFM., the following lower bound holds,
Lbathh H 2
Σ Σ
nB,cloog (NB, - 1 + (nn- - (nB=n nB.c)e
(5)
BeB c=[R]
Equality holds if and only if for every B E B and every pair of samples i,j B, h2 = 0 if Uj and h, = h, if yt y5-
We remark that the only previous study of SCL geometry with batches by Graf et al. (2021) requires significantly more restrictive conditions on the batch set B to guarantee a common geometry among all batches B e B. Specifically, Graf et al. (2021) assumes a batch set that includes all possible combinations of examples. Instead, thanks to the addition of ReLU, our requirements are significantly relaxed. This is detailed in the following remark. See also Sec. B.2 in SM.
analysis of Graf et al. (2021)). We elaborate on the comparison to Graf el al. (2021) with an illustrative example below. When optimizing the loss decomposed over a set of mini-batches it is crucial to consider whether the individual batch minimizers are also overall optimal solutions. Consider the following setting: k= d== = 3, y1 =1,92 = 2, Y3 = 3, B (1,2), B2 = (1,3), B3 =2,3), , [B1, B2, B3). Let us identify the optimal embeddings for the two scenarios: (i) with and (īi) without non-negativity constraints on the embedding coordinates. (i) Without non-negativity (UFM), the optimal embedding configurations can be shown to be ETF with 2 vectors for every batch, which is simply antipodal structure. In other words, the batch-wise optimal solutions are h1 = -hg h -h and h2 = -hg. for the batches B1, B2 and By, respectively. However, the batch-wise optimal solutions are infeasible due to their contradictory nature. From Graf et al. (2021), ). it can be deduced that the overall optimal configuration is instead an ETF with 3 vectors. This example underscores the difficulty in optimizing the loss over every batch separately in
Remark 1 (Comparison

---

Published as a conference paper at ICLR 2024
case of SCL without non-negativity. See (a) in Fig. 6 in SM.
(ii) With non-negativity (UFM.), our results imply that the optimal embeddings form a 2-OF for every batch, i.e., h1 L h2, h1 h2 and h2 h3, for the batches B1, B2 and B3. respectively. The three conditions are compatible with each other and the fact that the overall optimal solution is a 3- OF in R Therefore, were able break down the overall optimization into individual batches. See (b) in Fig, 6 in SM.
Global minimizer geometry may not be unique. It is easy to verify that any H following the OF geometry achieves the lower-bound of Thm. 2 and it is also a global optimizer of the mini-batch SCL. However, depending on the choice of B, the lower bound can possibly be attained by other optimal embeddings that do not necessarily satisfy NC or orthogonality. Hence the global optimizer may not have unique geometry.
14
To highlight the importance of B in the characterization of the optimal geometry of UFM. when using the batch-loss, consider the toy exam- ple in Fig. 4. Suppose we have le = 3 classes and we want to find the optimal normalized and non-negative features in R3 by minimizing (4) for B- 1 B2)- Although the OF satisfies the global optimality condition in Thm. 2, the theorem requires milder conditions for the optimal At: hú and hú need to be aligned (y; = yj) or orthogonal (y: yj) only if they interact within one of the selected batches. Fig. shows one such optimal geometry that does not satisfy either of NC
Figure 4: non-OF geome- try minimizing mini-batch SCL with B (Bi, B2}. Samples with different labels marked with different colors.
or orthogonality conditions: Firstly, the samples i 1 and 3 have significantly different features despite belonging to the same class. Second, hgg h4 align with hs, h1 although they have different labels, and h7, hg are not orthogonal to samples h, h2. With this example, we are now ready to discuss the role of batching more formally in the next section.
BATCHING MATTERS
Since the OF may not be the only solution for the UFM. with mini-batch SCL (4), here we study the role of mini-batches to avoid ambiguous solutions with different geometries. In Sec. 5.1, we identify necessary and sufficient conditions for batching strategy to yield a unique global solution geometry. By this result, in Sec. 5.2, we propose a simple yet effective scheme that provably succeeds in improving the convergence to an OF.
5.1 WHEN IS OF GEOMETRY THE UNIQUE MINIMIZER?
As discussed in Sec. 4.2 the uniqueness of the global minimizer's 's geometry when considering mini- batches depends on the interaction of samples in the loss function, or, in other words, the choice of B. Before specifying for which B the minimizer is unique, we need to define the Batch Interaction Graph, graph that captures the pairwise interactions of samples within the batches.
Definition 4 (Batch Interaction Graph). Consider an undirected graph G (V, E) where V [n]] We define the Batch Interaction Graph for a given set of batches B as follows: vertices u, u V are connected if and only if there exists a batch B B such that u, U E B. Moreover, G. denotes the induced subgraph of G with vertices V. = (u: : Ya = c).
We state necessary and sufficient conditions on B for the minimizer of mini-batch SCL to be unique:
Corollary 2.1. Consider the Batch Interaction Graph G corresponding to B. The global optimizer geometry of UFM. with mini-batch SCL (4) is unique and forms OF if and only if G satisfies the following conditions: (1) For every class e E [k], Ge is a connected graph. (2) For every pair C1,C2 E [k], there exists at least one edge between Ge, and Gez.
Corollary 2.1 serves as a check for whether a given batching scheme yields a unique global minimizer geometry not. It also provides guidance for designing mini-batches. We elaborate this below.
7

---

Published as a conference paper at ICLR 2024
No Batch Shuffling Batch-Binding
No Batch Shuffling
Batch Shuffling
1.2
0.5
0.4
110 230 3.03
LX]
2001
500
130 2.0 2001
Epoch
Figure 5: Convergence the OF geometry for various batching schemes including the analysis-ingt inspired scheme "No Batch Shuffling Batch- Baich-Bindingg. ResNet-18 trained CIFAR-10 with small batch size of 128, see text. See also Fig. 18 in SM
5.2 Batch-binding ENSURES UNIQUE OF MINIMIZER AND IMPROVES CONVERGENCE
Corollary 2.1 shows that an arbitrary set of mini batcheatcessssse is not guaranteed to have OF geometry as its minimizing configuration in the embedding space. To enforce the OF geometry as the unique minimizer, we introduce a simple scheme with low computational overhead as follows. We define binding examples as a set of examples including exactly one example from every class. This set has, thus, & elements. The process of batch-binding is simply adding the above set of binding examples to every mini-batch in the given original set. See Sec. D of SM for an algorithm and Fig. 16 SM for visual example.
Batch- Batch-binding adds a constant to the size of every mini-batch. A is typically smaller than the batch sizes used to train SCL which could range from 2048 6144 (Khosla et al., 2020). While adding small computational overhead, this method guarantees that B' satisfies the conditions of Cor. 2.1 and SCL learns OF geometry. We conduct series of experiments illustrate the impact of batch binding on convergence to OF.
To verify the role of batch selection the learned features, train ResNet-18 CIFAR-10 under three batching scenarios: 1) No batch-shuflling: mini-batch set B that partitions the training examples once initialization and is held constant across training epochs. 2) Batch-shuftting: the examples are divided into mini-batch partition, with a random reshuffle of the examples at every epoch. 3) No batch-shuffling batch-binding: we construct a fixed partition of examples into mini-batches; then, we perform batch-binding. In order to focus on the impact of batch selection strategies, we consider a relatively small batch-size of 128.
We remark that Thm. 2 and Cor. 2.1 can explain the behaviors of all three schemes. Moreover, the batch- batch-binding strategy guided by our analytical results effective at ensuring fast and predictable convergence of deep-nets to unique OF geometry. Fig. shows the distance of learned embeddings to the OF geometry for the three batching scenarios mentioned above. In case of the fixed partition batching (left), the features do not converge OF, especially with large imbalance ratios. This behavior is consistent with the conditions identified by Thm. 2, because in a typical random partition, necessarily connected. On the other hand, when the
the corresponding induced subgraphs batching is performed with a randomly reshuffled data ordering (center), the geometry converges more closely to the OF, albeit with certain epochs deviating significantly from the convergence direction. We hypothesize that random reshuffling creates an opportunity for different examples to interact sufficiently, thus, when trained long enough, to converge close to an OF geometry. Finally, with the third batching strategy of patch-binding (right), even without shuffling the samples at every epoch, we observe a consistently fast convergence to the global optimizer, OF. These observations provide strong evidence supporting our analysis of batching predicting the behavior of SCL trained deep-nets.
While the above studies the impact of batching itself, we perform number of additional experiments in order to illustrate the impact binding examples. particular, we observe that batch binding helps improve convergence to OF geometry when training with less powerful models (see Fig. 19 for DenseNet experiments in SM), more complex dataset (see Fig. 20 in SM and Fig. 2 for 100 and TinylmageNet results) and even convergence ETF in the absence of ReLU under balanced training (See Fig. 21 in SM). Such results emphasize the importance of batch binding and encourage further analysis of batching under contrastive learning.
8

---

Published as a conference paper at ICLR 2024
6
MORE RELATED WORK
The simple concept behind SCL, introduced as an extension of the contrastive loss (e.g., Chen et al., 2020; Tian et al., 2019) to the fully supervised setting by Khosla et al. (2020), involves pulling together the normalized features of examples belonging to the same class while pushing them away from examples of other classes. Despite its simplicity, SCL offers a generalization of existing loss functions like the triplet loss (Weinberger et al., 2005) and the N-pair loss (Sohn, 2016), and surpasses the performance of the CE loss on standard vision classification datasets, while also heing more robust natural corruptions in the data and less sensitive to hyperparameters (Khosla et al., 2020). The feature geometry of the self-supervised constrastive loss was considered by (Wang & Isola, 2020), where the loss function was shown to optimize uniformity and alignment of features, as the number of negative examples approaches infinity. Instead, we study the supervised loss, providing exact non -symmtottii solutions for the corresponding UFM and the impact of batching in detail.
In recent years, there has been a notable interest regarding the properties of the embedding space and weights learned through unsupervised//sppriised contrastive losses, as well as CE. Of particular importance is the endeavor to uncover the fundamental distinctions between these losses to effectively harness each of their strengths and formalize principled ways to combine them. Specifically, starting increasing series of recent works have focused on uncovering the
by Papyan et al. (2020), embedding geometry of CE, Mean-Squared Error (MSE) losses, and their variants. We highlighted some of these works in Sec. 1 while numerous other works contribute to these investigation (e.g., Zhou et al., 2022a;b; Yaras et al., 2022; Gao et al., 2023; Han et al., 2021; Súkeník et al., 2023). In addition to the scientific curiosity surrounding such studies, this exploration has paved the way for novel CE-based approaches to training DNNs on imbalanced data (Behnia et al., 2023; Liang & Davis, 2023; Xie et al., 2023; Sharma et al., 2023; Dang et al., 2023; Yang et al., 2022).
Less attention has been devoted in the existing literature to a comparable set of results for the SCL. Graf al. (2021) the first to study the geometry of SCL for balanced data and comparing it to CE. To tackle imbalances, and inspired by Graf et al. (2021) that suggests balanced data is crucial for obtaining symmetric embeddings, Zhu et al. (2022) have recently proposed and investigated a modification of the SCL that boosts performance under class imbalances. Our work extends and complements the two studies by showing that SCL with ReLU can actually learn symmetric embedding structure even in the presence of imbalances. While preparing the manuscript, we became aware of a recent work (Cho et al., 2023), where the authors consider the mini-batch optimization of the unsupervised contrastive loss and propose a choice of mini-batches to speed up convergence.
7 CONCLUDING REMARKS
We have shown consistent empirical and theoretical evidence that SCL learns training embeddings that converge to the OF geometry irrespective of the level of class imbalance. We believe this finding contributes a unique result the growing literature of neural-collapse/ / -eepmitttgeometty phenomena. For balanced data, we extend the contributions of Graf et al. (2021) to the case in presence of ReLU activations. Further, our results identify exact conditions on the mini-batch selection to achieve the OF geometry, allowing for wider set of possibilities in batching than found by Graf et al. (2021) for SCL without ReLU, For imbalanced data, our results are the first explicit characterization of the geometry, concluding that imbalances do not alter the geometry if ReLU activation is added at the last layer. In both cases, these advancements are achieved by analyzing a refined UFM model, UFM,, that closely aligns with experimental conditions by constraining embeddings to the non-negative orthant. This model also leads to new findings about the intricate role of batching in SCL optimization, that may be of independent interest. A future investigation of exact characterization of the geometry with imbalances and without ReLU can help complete the understanding of the implicit-geometry of SCL. Although the majority of implicit geometry characterizations in the literature require d k. is interesting to extend our findings to settings where d < k following the steps of Gao et al. (2023) for the CE loss. Likewise, our finding regarding the crucial role of batching in geometry opens the door to further investigations aimed at devising efficient batch strategies in cases with a large number of classes. Finally, like most studies on neural-collapse phenomena, our results share a common limitation: they only provide insights into the behavior during training. On the other hand, there is a line of research that explores algorithmic modifications to SCL (Gunel et al., 2020; Jitkrittum et al., 2022; Samuel & Chechik, 2021; Li et al., 2022; Kang et al., 2021; Li et al., 2022), often rooted but not explicitly connected to geometric principles, to enhance generalization under imbalances. Ultimately, we envision the two research streams merging through the ongoing exchange of ideas.

---

Published as a conference paper at ICLR 2024
REFERENCES
Tina Behnia, Ganesh Ramachandra Kini, Vala Vakilian, and Christos Thrampoulidis. On the implicit geometry of cross-entropy parameterizations for label-imbalanced data. In International Conference on Artificial Intelligence and Statistics, pp. 10815-10838. PMLR, 2023.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pp. 1597-1607. PMLR, 2020.
Jaewoong Cho, Kartik Sreenivasan, Keon Lee, Kyunghoo Mun, Soheun Yi, Jeong- Lec, Anna Lee, Jy-yong Sohn, Dimitris Papailiopoulos, and Kangwook Lee. Mini-batch optimization of contrastive loss. arXiv preprint arXiv:2307. .0590, 2023.
Hien Dang. Tan Nguyen, Tho Tran, Hung Tran, and Nhat Ho. Neural collapse in deep linear network: From balanced to imbalanced data. arXiv preprint 2tt1::2331.003377 2023.
Alexey Dosovitskiy. Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image worth 16x words: Transformers for image recognition scale, 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), 2021.
Peifeng Gao, Qianqian Xu, Peisong Wen, Huiyang Shao, Zhiyong Yang, and Qingming Huang. study of neural collapse phenomenon: Grassmannian frame, symmetry, generalization. arXiv preprint arXiv:2304.089) 2023.
Florian Graf, Christoph Hofer, Mare Niethammer, and Roland Kwitt. Dissecting supervised con-
strastive learning. In International Conference on Machine Learning. pp. 3821-3830, PMLR. 2021.
Beliz Gunel, Jingfei Du, Alexis Conneau, and Ves Stoyanov. Supervised contrastive learning for pre-trained language model fine-tuning. arXiv preprint arXiv:2011.01403, 2020.
XY Han, Vardan Papyan, and David L Donoho. Neural collapse under mse loss: Proximity to and dynamics on the central path. arXiv preprint arXiv:2106.02017 2021.
Wenlong Ji, Yiping Lu, Yiliang Zhang, Zhun Deng, and Weijie Su. An unconstrained layer-peeled perspective neural collapse. arXiv preprint arXiv:21 2021.
Wittawat Jitkrittum, Aditya Krishna Menon, Ankit Singh Rawat, and Sanjiv Kumar. Elm: Embedding and logit margins for long-tail learning, arXiv preprint arXiv:2204. 13208, 2022.
Bingyi Kang, Yu Li, Sa Xie, Zehuan Yuan, and Jiashi Feng. Exploring balanced feature spaces for representation learning. International Conference on Learning Representations, 2021.
Prannay Khosla. Piotr Teterwak, Chen Wang. Aaron Sarna, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishnan. Supervised contrastive learning. Advances neural information processing systems, 33:1 11111866111866 2020.
Tianhong Li, Peng Cao, Yuan Yuan, Lijie Fan, Yuzhe Yang, Rogerio S Feris, Piotr Indyk, and Dina Katabi. Targeted supervised contrastive learning for long-tailed recognition. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp.
2022.
Tong Liang and Jim Davis. Inducing neural collapse fixed hierarchy-awarg frame for reducing mistake severity. arXiv preprint arXiv:2303. 05689, 2023.
Xuantong Liu, Jianfeng Zhang, Tianyang Hu, He Cao, Yuan Yao, and Lujia Pan. Inducing neural collapse in deep long-tailed learning. In International Conference on Artificial Intelligence and Statistics, pp. 11534-11 1544. PMLR, 2023.
10

---

Published as a conference paper at ICLR 2024
Jianfeng Lu and Stefan Steinerberger. Neural collapse with cross-entropy loss. arXiv preprint arXiv:2012.08465, 2020.
Dustin G Mixon, Hans Parshall, and Jianzong Pi. Neural collapse with unconstrained features. arXiv preprint arXiv:201 1619, 2020.
Vardan Papyan. The full spectrum of deep net hessians at scale: Dynamics with sample size. arXiv preprint arXiv:1811.07002 2018.
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.
Dvir Samuel and Gal Chechik. Distributional robustness loss for long-tail learning. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 9495-9504, 2021.
Saurabh Sharma, Yongqin Xian, Ning Yu, and Ambuj Singh. Learning prototype classifiers for long-tailed recognition. arXiv preprint arXiv:2302.0043 2023.
Kihyuk Sohn. Improved deep metric learning with multi-class n-pair loss objective. Advances in neural information processing systems, 29, 2016.
Peter Súkeník, Marco Mondelli, and Christoph Lampert. Deep neural collapse is provably optimal for the deep unconstrained features model. arXiv preprint arXiv:2305.13166 2023.
Christos Thrampoulidis, Ganesh Kini, Vala Vakilian, and Tina Behnia. Imbalance trouble: Revisit- ing neural-collapse geometry. arXiv preprint arXiv:2208.05512 2022.
Y Tian, D Krishnan, and P Isola. Contrastive multiview coding. arxiv. arXiv: 1X06::990.05500 05849, 2019.
arXiv preprint
Tom Tirer and Joan Bruna. Extended unconstrained features model for exploring deep neural collapse. arXiv preprint #:Xiv:2202.08083 2022.
Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through align- ment and uniformity the hypersphere. In International Conference Machine Learning, pp. 9929-9939. PMLR, 2020.
Kilian Q Weinberger, John Blitzer, and Lawrence Saul. Distance metric learning for large margin nearest neighbor classification. Advances in neural information processing systems, 18, 2005.
Liang Xie, Yibo Yang, Deng Cai, and Xiaofei He. Neural collapse inspired attraction-repulsion balanced loss for imbalanced learning. Neurocomputiing. 2023.
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.
Can Yaras, Peng Wang, Zhihui Zhu, Laura Balzano, and Qing Qu. Neural collapse with normalized features: A geometric analysis over the ricmannian manifold. arXiv preprint arXiv:2209.097) 2,110 2022.
Jinxin Zhou, Xiao Li, Tianyu Ding, Chong You, Qing Qu, and Zhihui Zhu. On the optimization landscape neural collapse under mse loss: Global optimality with unconstrained features. arXiv preprint arXiv:2203.012.3 2022a.
Jinxin Zhou, Chong You, Xiao Li, Kangning Liu, Sheng Liu, Qing Qu, and Zhihui Zhu. Are all losses created equal: A neural collapse perspective. arXiv preprint arX1v:2210.02/19,, 2022b.
Jianggang Zhu, Zheng Wang, Jingjing Chen, Yi-Ping Phoehe Chen, and Yu-Gang Jiang. Balanced contrastive learning for long-tailed visual recognition. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6908-6917, 2022.
Zhihui Zhu, Tianyu Ding, Jinxin 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, 2021.
11

---

Published as a conference paper at ICLR 2024
CONTENTS
Introduction
1.1
Summary of contributions
2 Setup
3.
SCL with ReLU learns OF geometries: Empirical findings
4
Theoretical justification: SCL with non-negativity constraints
4.1
Full-batch SCL
4.2
Mini-batch SCL
5
Batching matters
5.1
When is OF geometry the unique minimizer?
5.2
Batch-binding ensures unique OF minimizer and improves convergence
More Related Work
7 Concluding remarks
A Proof of Thm. 1
13
Proof of Thm. 2
15
B.1 Proof of Cor. 2.1
15
B.2 Batch analysis for UFM, vs UFM
16
Proofs for Section D.1
17
C.1 Proof of Lemma D.1
17
C.2 Proof of Lemma D.2
17
D Additional Discussion
18
D.1 Detailed comparison between UFM and UFM,
18
D.1.1 Centering heuristic
18
D.1.2 Centered OF simplex ETF
18
D.1.3 UFM can fail to predict the true geometry
19
D.2
Comparing implicit Geometries: SCL vs CE summary
19
E
Additional experimental results and discussion
20
E.1 Details on the main experimental setup
20
E.2 Details on Fig. Heatmaps
20
E.3 Additional geometric analysis
20
E.3.1
Neural Collapse
12

---

Published as a conference paper at ICLR 2024
E.3.2
Angular convergence
21
E.3.3
Embedding heatmaps
22
E.3.4
Experiments with MLPs
22
E,4
Optimization dynamics..
22
E.4.1
Loss convergence
................................
22
E.4.2
Effect of T
23
E.4.3
Impact of Batch Size
24
E.5
Complementary results and discussions on batch- batch-binding
24
E.5.1
How batch-binding ensures unique OF geometry
24
E.5.2
Convergence non-ReLU settings
27
E.6
On the convergence of batch- batch-shuflling to OF
28
E.7
ReLU does not compromise Accuracy
30
E.8
ReLU Helps improve worst class accuracy
30
E.8.1
Impact of batch-binding on generalization
31
E.8.2
Experiments with additional architectures
32
Additional Notations. While recalling the notation mentioned in Sec. 1, we note a few more below. denotes Kronecker products. We use 1m to denote m-dimensional vector of all ones. For vectors/matrices with all zero entries, we simply write 0, as dimensions are easily understood from context. V its Moore-Penrose pseudoinverse. VvL [R*** is the gradient of scalar differentiable function f(.) with respect to V.
List of Frequently Used Abbreviations.
CE
Cross Entropy
ETF
Equiangular Tight Frame
LT
Long-Tailed
MLP
Multi-Layer Perceptron
MSE
Mean Squared Error
NC
Neural Collapse
NCC
Nearest Class-Center
OF
Orthogonal Frame
ReLU
Rectified Linear Unit
SCL
Supervised Contrastive Loss
UFM
Unconstrained Features Model
A PROOF OF THM. 1
For c [k], let C. ii y: = c} be the set of examples belonging to class c, and define (H as follows.
(H Σ log Σ exp (h. (b..c--hhh,))
jaC Jei
Be(n],eti
Σ log Σ
exp - h,h,)) 1 + Σ -x,(bt/hu-bbbb)))
(6)
jeCy,
PeCy
13

---

Published as a conference paper at ICLR 2024
Then, rewrite the full-batch SCL in (2) as,
1
Lin H)
Σ [
(H)
Σ (H) scCc
fly
cc[K]
ng
Now for [A]], consider example E Ce, and define
a
he: ReCo.eei
Σ
Ebe. he. b, The EfC.
Note a; is the mean-embedding of all examples in class but i, and b, is the mean-embedding of all examples not in class c. Starting from (6), by applying Jensen's 's inequality to the strictly convex function fi (x) == exx((= h,,,) we have,
(H)
Σ
log (ne exp exp (hfa. h₂ (n n.) exp (h, b. h h,))
jeCejai
Σ log (ex (n ex) ex; exp (- (-hīh,)
jeCojei
Σ (xggaxexpp-hii),)) (-h,hy))
jeCoffsi
( (n, - 1) log a; exp (-h? nr
Σ jeCeijsi
h2)))
= (ne 1)
a:))
-n 1) log (ne 1) (n ne) exp n The
h he EfC.
1 The
h Σ ho)): tcCo,tti
(7)
where in (i). we define == (n, 1) exp (h. a.) + (n ne) exp (h b,), and in (ii), we use the fact that the function f2(x) log(a; exp(x)) an affine function. Now consider all samples Ca. From (7),
CCHH
Σ
log (no- 1) (n no) exp
(
1
h Σ he EfCe
h
he)) LeCo,tti
Σ
ne
IeC.
n2 ne
Nee
ne log (nc-1) + (n - (n)) exp nn nec)
Σ
h
Σ
h hhThe))
(n -1)
icC
seC
EfC.
BeCu..
(8)
In the last line we use Jensen's inequality on fs(x) log((ne 1) (n nx) exp(x)) which is strictly convex since n. 1 and n n, > 0.
By the ReLU constraints we have h, 2 0, i [n], which implies h,h, O, Vi,j E [n]. Further, by Cauchy-Schward inequality h, < |h; |h; = 1, with equality achieved if and only if h, h..
Since fa non-decreasing, we can simplify the bound (8) follows,
(H) 2 ne log (ne - 1) + (a ne)e Iec.
(9)
Ne
Summing over all classes
[k:]. we get the final bound on the full-batch SCL:
Luu ( Σ cE[R]
E (H) RC.
2 Σ log (ne 1 1) (n - -n)e
cc[R]
To achieve the fower-bound from (9), we require h, h, if ys = Dj and h, = otherwise. This requirement also satisfies the equality condition for the Jensen inequalities applied in (7) and (8). Thus, any H achieving the lower-bound follows an OF geometry (Def. 3), which exists as long as d 2 k.
14

---

Published as a conference paper at ICLR 2024
B
PROOF OF THM. 2
Consider the mini-batch SCL in (4), repeated below for convenience:
Dbatth (H) = ΣΣ BeB B RB
Σ
log
Σ exp (b,
6cB Et1
j+iyo-
Denoting the loss over a batch B by
Cn(H) == Σ iB B.
E
log
Σ
exp ee (hThe-h?hh,)
JEB
KB Eo1
we have Launh (H) EeB L (H). Now, for given batch B, we apply Thm. 1 to get the lower bound and conditions for equality:
LB (B(H) 2
DBBBclog (n - 1+ (nB- - -B,c)e cc[R]
Σ
Moreover, equality holds if and only for every pair of samples i.j B: h, 0f y y, and h; hy if yo yy- With this, the overall loss can be bounded from below by summing the individual lower bounds over different batches:
Lauthh (H) 2 Σ BEB [c(
Σ
nB. log (N - 1 (nz - (onB-n,,)e n)).)e
with equality achieved if and only if for every B E B and every pair of samples i, B. h. h, 0 yi yj and h, h, if yı yj. As long as d 2 k, the equality conditions of every individual batch can be satisfied simultancously.
B.1 PROOF OF COR. 2.1
From Thm. 2, the optimal configuration of embeddings UFM, under the mini-batch SCL depends on how batch conditions interact with each other. Specifically, recall that the equality in (5) is achieved if and only if for any batch B e B, and each pair of samples (i.j) E B, h; = h, if y; = yj- and h,h, 0 otherwise. The OF geometry clearly satisfies all these conditions and achieves the minimal cost of UFM However, as discussed in Sec. 4.2, these equality conditions may be satisfied by configurations other than OF for arbitrary batching scheme. So, in Cor. 2.1, specify the requirements of a batching scheme in order for the global optimal of the mini- batch SCL to be unique (up global rotations), and match the optimal the full-batch SCL, which is the OF geometry.
To prove the corollary, we separately address the 'IF' and 'ONLY IF" parts of the Cor. 2.1.
"IF' direction. Assume the batching scheme satisfies both conditions of Cor. 2.1: 1) Ve e [k], the induced subgraph Ge is connected, and 2) Vc é [k], there exists an edge between Gg and Ge We show below that under these two conditions the optimum of UFM, under mini-batch SCL follows an OF geometry. In other words, we show that the optimal embeddings align if they belong to the same class (NC) and are orthogonal if they have different labels (mean-embeddinng follow k-OF). (See Def. 3.)
NC: Consider class c. From our assumption, the induced subgraph G. is connected. Thus, path exists from any node (representing corresponding example) to any other node in G.. Consider three nodes along a path, indexed by i,j, that belong to class c. Let there be an edge between i,j and j each. By Def. 4 of the Batch Interaction Graph, examples i and batch, say B1, and examples j. belong to a batch Bg that is possibly different from Bi. From infer that at the optimal solution, we have h, hj. due equality conditions for batch B. Also, h; = h,, due to equality conditions for batch Bz. Thus, by transitivity, we have h; h, h, Repeating this argument along any path in G. for every [k), we obtain the NC property: h, = h, Vi,j y4 yj e for every class C [k].
present in
Thm. 2,
"The arguments and considerations apply when Ge consists of only two nodes.
15

---

Published as a conference paper at ICLR 2024
k-OF: From our assumption, for any pair of classes e1 and 022 there exists an edge between Ge and Ge, connecting at least one pair of nodes i.j : i GeyrJ e Gez. By definition of the Batch Interaction Graph, we know that examples i and j belong in at least one batch. Thus. by Thm. 2. at the optimal solution, we have h; 1 h since examples yi C1 02 = yj- Now, by NC, h, May h Meg. and thus, Mei 1 Meg. This holds for every pair of classes er * c [R]] Therefore, the matrix M = [,,,..., [4 of class-mean embeddings forms a K-OF.
Combining the two statements above, we conclude that the global optimal solution, the embeddings follow the OF geometry, as desired.
'ONLY IF" direction. For the other direction, it suffices to show that if either of the two conditions in Cor. 2.1 does not hold, then there exists an optimizer that does not follow the OF geometry. Specifically, we show that when either of the two conditions is violated and d 2 k 1, there exists an embeddings matrix H attaining the loss lower-bound that does not satisfy one of the two requirements of the OF geometry: 1) H does not follow NC, 2) the corresponding mean-embeddinng M do not arrange as a k-OF.
Case 1. Suppose for c [k]. the induced subgraph Ge is not connected, and without loss of generality, assume c 1. This means that G1 has at least two separate components. Denote the nodes in each of the two component by Vi and V2² respectively, and recall for e 2, Ve = {i y/4 = c}. 22 As d k + 1, we can choose set of k 1 vectors [M], M1; M2, Pa. [bx such that they form (k 1) OF Define H = [h],... h as follows:
ViE VieVi,
h = m
Vie V?
h b= M
Vie Ve, c [k]] =
Then, H satisfies the equality conditions in Thm. 2 since there is edge between the nodes in V and V? by the assumption. However, the embeddings in class y = do not align, since ç7| is orthogonal to p.. Thus, H optimizes UFM, while it does not satisfy NC and differs from the OF geometry.
Case 2. Suppose there exists # C2 E [k] for which there is no edges between Ge, and Go and without loss of generality, assume 0₁ = 12 = 2. Consider an embedding matrix H that satisfies NC, and the corresponding mean-embedding matrix M = [71. P is such that Vo # c' (2...,,A, k),, T.. c.. 0 and ãi Tá2. Such M exists as d 2 k and all we need is to have [M2,,... Pa be (A: - 1)-OF the non-negative orthant. Since, there is no edge between G and G2, there is no B B that includes samples from classes y= 1 and y == simultaneously. Equivalently, to achieve the lower-bound of Thm. 2, we do not require orthogonality between any pairs of samples e C₁ and j Cz. Therefore, H is an optimal solution of the UFM.. However, the mean-embeddings do not follow K-OF. In fact, this optimal geometry, does not distinguish the samples in classes y and 2 2.
Therefore, the global minimizer will be uniquely an OF if and only if the batching scheme satisfies both the conditions stated in Cor. 2.1.
B.2 BATCH ANALYSIS FOR UFM, VS UFM
Graf al. (2021, Thm. 2) studies the UFM without ReLU constraints and with balanced labels, and proves that the global solution is simplex ETF. The authors note that their proof relies on the batch construction as the set of all combinations of given size. contrast, we have proved that for UFM, the global solution is an OF for a wider range of batching scenarios (specifically. as long as they satisfy the interaction properties characterized in Cor. 2.1). Without ReL.U, as noted by Graf et al. (2021, Fig. 5), the optimal configuration of examples in each batch can have a different geometry. depending on the label distribution of the examples within the batch. However, with ReLU, the optimal configuration of every batch is an OF over the classes whose examples are present in the batch. In particular, there is no contradiction between two mini mani-batches having both overlapping and mutually exclusive classes, since the optimal configuration of one batch does not violate the optimal configuration of another. Furthermore, the overall batch construction can have a unique OF as the optimal configuration provided the conditions in Cor. 2.1 are satisfied. The conditions include the
16

---

Published as a conference paper at ICLR 2024
(iv)
(a) UFM optimal embeddings
(i)
(ii)
(iii)
(iv)
(b) UFM, optimal emheddings
Figure Rem. visualized, (a)(i,ii,ii)) indicate antipodal structure the optimal embeddings under UFM each the 3 mini-batches respectively, whereas the overall optimal geometry is an ETF (a)(iv). This contrasts the optimal embeddings under UFM, where each mini-hatch (b)ii,iiiii) is consistent with the overall optima (b)(iv).
batching assumed by Graf et al. (2021) as a special case, while being applicable less restrictive scenarios. Sec. 5.2 explored the implications of this finding, by studying the criteria for batching scheme to lead to a unique minimizer geometry and further suggesting a simple scheme to convert an arbitrary batching
satisfying these criteria.
C PROOFS FOR SECTION D.1
C. PROOF OF LEMMA D.I
Let [v v&] and V [V, Without loss generality suppose the constant of proportionality is so that V I. Since all v,s orthonormal, this then implies that V Va & l and v These put together gives the desired as follows:
vzJ. By definition of a k-OF, we know that V oc I.
1
V = -oVVvvl)) (V-va1)) V L Vivall (vora) =I 1
C.2 PROOF OF LEMMA D.2
To rule out the optimality of ETF in imbalanced setups, we show ETF does not minimize the loss in k class example. Specifically, consider STEP-imballnced training set with 3 classes of sizes [Rmmin mmin minn] with M'nn 2 and R 2 10. Suppose HETF follows the ETF geometry. Then, from Def. 5 and the feasibility condition, for the mean-embeddingg METF we have MTTPMETF k-1 (II 1.) Thus, the value of the loss function ETF is
Lial HET)) Mminn [loog PRnminn -1+ mm #) 2 2log (Tiin 1 Pemin R 1)+ t)).
Now, consider H [ 1 "Rnm"" - @ T - 8 T. where ç is an arbitrary unit-norm vector.
Since | 1, it follows that H is in the feasible set. With this choice of H. we have,
Laa H Plmin log (Rrmin 1 20 2) log (21min + Ruimine )))
It is easy to check that, for imbalance levels higher than R 2 10.
Lau H < Liuu (HEF).
In other words, ETF does not always attain the optimal cost in UFM.
17

---

Published as a conference paper at ICLR 2024
D ADDITIONAL DISCUSSION
D.1 DETAILED COMPARISON BETWEEN UFM AND UFM,
Within this section, we provide an in-depth analysis of the disparities between the predictions generated by the UFM, commonly utilized in previous studies, and the refined version UFM. that we employ in this work to study the embedding geometry of SCL.
The majority of previous works (e.g., Papyan et al., 2020; Zhu et al., 2021; Fang al., 2021) study the geometry of learned embeddings by investigating proxy unconstrained features model (UFM). Specifically, the UFM drops the dependence of the learned embeddings H on the network parameters 0, and finds optimal H that minimizes the loss function. In case of SCL, as studied in e.g., Graf et al. (2021), 1), the corresponding UFM is as follows,
arg min Lscu(H) subj. to |h,12= 1, Vi [n].
(UFM)
H
As explained in Sec. 4, this paper employs refined version of the UFM, namely UFM,, character- ize the representations learned by SCL. Specifically, UFM. further constrains the embeddings H to be non-negative, this way accounting for the ReL.U activation commonly used several deep-net architectures. Thus the search for the optimal H is restricted to the non-negative orthant.
D.1.1 CENTERING HEURISTIC
Many deep neural network models incorporate ReLU activations, resulting in nonnegative feature embeddings at the last layer. In contrast, the UFM does not impose any constraints on the optimal H. allowing it to contain negative entries. To bridge this gap. previous works (e.g., Papyan et al., 2020; Zhu et al., 2021; Fang al., 2021; Zhou al., 2022a; Thrampoulidis et al., 2022b) apply heuristic approach called global-mean centering on the learned embeddings before comparing their arrangement with the theoretical prediction given by the UFM. Specifically, the centering heuristic is applied as follows: Prior to comparing the geometries, we subtract the global- mean embedding μG 1 class-mean embeddings Pe:
al., 2022; Zhou
k: Ee((k) Me from all meas embeddings μ.c to form the centered
P== MAc PG: M [m
ps
Indeed, this heuristic centering operation effectively ensures that the mean- meea-embeddding matrix M. after the necessary shifting, is centered at the origin, satisfying M1, = 0. This property also holds for the global optimizers of the UFM found in the previous work, which served motivation for the heuristic centering.
Unlike the approach mentioned above, our findings do not necessitate any heuristic embedding processing for comparing the geometries. This is because our model directly provides the geometry of the embeddings in their original form.
D.1.2 CENTERED OF IS SIMPLEX ETF
Focusing on the SCL, Graf et al. (2021) have used the UFM to characterize the learned embeddings and have found that, for balanced classes, the simplex ETF geometry the global optimizer of UFM. In other words, the optimal embeddings H adhere to the NC (Defn. 1), and the corresponding mean embeddings M form a simplex ETF, For the reader's convenience, we recall below the definition of simplex ETF from Papyan et al. (2020).
Definition 5 (Simplex ETF). We say that vectors V =[V1,'> v&] Rack form a simplex-ETF if VTV T 1. i.e., for each pair of (i,j) [k:], |v.| |vyl and v.V,
E-
In this paper, show instead that the global optimizer of UFM. is OF. Remarkably, this result holds true regardless of the label-imbalance profile. However, for the purpose comparison with the findings of Graf et al. (2021), let us specifically consider the scenario of balanced data. In this case, an intriguing question arises:
Specifically for balanced data, how does our discovery that embeddings with ReLU converge an OF compare to the previous finding by Graf et al. 22021) that embeddings without ReLU converge to an ETF?
18

---

Published as a conference paper at ICLR 2024
The answer to this question lies on the centering trick dis- cussed in Section D.I.1. Specifically, stating that embed- dings form OF implies that centered embeddings form an ETF. This straightforward fact is formalized in the following lemma.
GM convergence
1.0
SCL OF
SCL ETF
0.8
CE OF
CE ETF
Lemma D.1. Suppose that set of vectors {V1,.... V&) form
0.6
k-OF in Rd with d k, and their mean-centered counter-
0.4
parts are V. = Ve -G. Vi e [k]. where VG = k. Leejk] Ve. Then, the centered embeddings (V1.....vs, form a simplex ETF spanning k 1-dimensional subspace.
0.2
100
200
300
Lemma D.1 states that based on our finding that embeddings form an OF, we can deduce that centered embeddings form an ETF. This conclusion aligns with the analysis conducted by Graf et al. (2021) on the UFM (See Fig. 7 for repre- sentative comparison). However, the UFM analysis itself does not provide information regarding the necessity of the centering technique, which remains heuristic. Additionally, it is important to note that it is unclear how to establish that (uncentered) embeddings form an OF we initially know that centered embeddings form an ETF, This is due to the unknown vector used for centering, which cannot be determined.
Epoch
Figure 7: Convergence of embeddings geometries for SCL (in green) and CE (in blue) on balanced CIFARIO with ResNet-18. Solid lines track the distance of uncentered embeddings to OF and dashed lines track the distance of cen- tered embeddings CE, the learning rate and weight decay and batch size are set according to Pa- pyan et al. (2020). Note the convergence of SCL (solid green) to its implicit ge- ometry is noticeably better compared to CE (dashed blue).
simplex ETF. For
D.1.3 UFM CAN FAIL TO PREDICT THE TRUE GEOMETRY
In the previous section, we demonstrated that for balanced data, the solution found by the UFM does not predict the true geometry of the embeddings (i.e., OF) in presence of ReLU, but it can predict the geometry of the embeddings after the application of global-mean centering heuristic (i.e., ETF). Naturally this raises the following question:
Does the UFM consistently provide accurate predictions for the geometry of the centered embeddings? In other words, does the OF geometry predicted by our refined UFM, always align with being a global optimizer of the UFM after centering?
The lemma presented below demonstrates that the answer to this question is negative. Specifically, it shows that the global optimizer of UFM is sensitive to the label distribution of the training set, thus ETF is not necessarily& a global optimizer in the presence of imbalances. This is in contrast to UFM., for which we showed that the global optimizer is consistently an OF, irrespective of the labels distribution. In other words, the difference between UFM and UFM, cannot be addressed by only considering the centering of the optimal embeddings in general.
Lemma D.2. If classes in the training set not balanced. the global solution of UFM is not necessarily an ETF. Thus, the solution of UFM, (which according to Thm. I is an OF) is in general different to that of UFM even after centering is applied.
D.2 COMPARING IMPLICIT GEOMETRIES: SCL VS CE - A SUMMARY
We conclude this discussion by providing a summary of the observations on the contrasting implicit geometries of embeddings between the two loss functions: SCL and CE.
The SCL geometry with ReLU robust to label-imbalances. Specifically, the geometry of SCL embeddings exhibits symmetry consistently, whereas imbalances introduce distortions in the implicit geometry of SCL without ReLU and also for CE (Fang et al., 2021; Thrampoulidis et al., 2022). 2. The convergence of SCL to its implicit geometry is notably superior. The measured geometry of SCL exhibits a faster decrease in distance from its analytic formula (i.e., OF) with each epoch, eventually reaching lower values. This holds true not only for label-imbalancer data, where Thram- poulidis et al. (2022) finds that the convergence of embeddings deteriorates as the imbalance ratio increases, but also for balanced data, where the convergence of embeddings with CE is inferior
19

---

Published as a conference paper at ICLR 2024
compared to classifiers (Papyan et al., 2020); see Fig. 7 for visualization.
3. For SCL, batching matters. During CE training, when using any arbitrary batching method, there exists an implicit interaction among all the examples, that is facilitated by the inclusion of the classifier vectors in the loss objective. On the other hand, with SCI., the interactions between examples are explicitly controlled by the mini-batches. Specificallyy batches in Secs. 4.2 and 5.1, along with the batch-binding scheme presented in Sec. 5.2 reveal that the structure of mini-batches critically affects the geometry of the learnt embeddings.
analysis of the role of
E ADDITIONAL. EXPERIMENTAL RESULTS AND DISCUSSION
E. DETAILS ON THE MAIN EXPERIMENTAL SETUP
In our deep-net experiments, focus two common network architectures, ResNet and DenseNet. Specifically, we use ResNet-18, 18, ResNet-34 and DenseNet-40 with approximately 63, 11 million and 200 thousand trainable parameters, respectively. For all models, we replace the last layer linear classifier with normalization layer (normalizing such that |b, | = 1 for E [r1]) of feature dimension d 512. Specifically, following Graf et al. (2021), we directly optimize the normalized features that are then used for inference. (We note this is slightly different compared to Khosla et al. (2020), where the authors train the output features of a projection head, then discard this projection head at inference time). We use this projection head when studying NCC test accuracy. In particular, we add a 2 layer non-linear MLP with with and without ReLU to report the test accuracies provided in Tab. 1, Tab. and Tab. 5. Following Khosla et al. (2020), in addition to the NCC test accuracy on post-prot post-projection feature (final output), we evaluate the accuracies on the pre-projection features (without input to projector) as well. We remark that both ResNet and DenseNet architectures include ReLU activations before the final output, which enforces a non-negativity on the learned embeddings. Resnet- 18 models with CIFAR10, MNIST have been trained for 350 epochs with a constant learning rate of 0.1, no weight decay, batch size of 1024, and SCL temperature parameter T = 0.1 (consistent with the choice of made in Khosla al. (2020); Graf et al. (2021)). For CIFAR100 and Tiay-Im experiments, we train a ResNet-34 under a similar setup, for 500 epochs, (with batch binding in Fig. 2). For all experiments with other models or datasets, we provide experimental details in the following sections. All ResNet-18 18 and DenseNet models have been trained on single Tesla V100 GPU machine and all ResNet-34 models were trained on a single machine with 2-4 Tesla V100 GPU, depending on experiment.
E.2 DETAILS ON FIG. 2 HEATMAPS
Since this particular experiment provides the most insight into the geometric behaviour of features, we have provided further explanation of the experimental setup and method. ResNet- 18 models were trained for 350 epochs MNIST using CE. SCL and SCL + ReLU under difference imbalance ratios. For SCL and SCL + ReLU we extract the features, calculate the class means and plot a heatmap of the gram matrices G 2. For CE, following previous work by Papyan (2018); Thrampoulidis et al. (2022); Zhu et al. (2021) we center the class means by the global mean, ie t Me K then proceed to calculate GM. In order to have a fair comparison, we normalize all gram matrices by dividing by the matrix maximum GM (G
1 ΣK 14 and
Gy No projection layer was implemented for these
results. The heatmaps were included to provide a clear visual, emphasizing the impact of ReLU on achieving symmetric geometry under imbalance. Each cell [i,j] of the Gas matrix represents the inner product 4i μJ. rather a normalized value representing the angle between the two class centers. In addition, as SCL includes a normalization layer, the diagonals G [i, i] would be closer to zero as the features exhibit stronger neural collapse properties.
E.3 ADDITIONAL GEOMETRIC ANALYSIS
In addition to analyzing the convergence of the Gram-matrix of class-mean embeddings Gm to the OF geometry (as provided in Fig. 2), we also keep track of Neural Collapse (Defn. 1) of individual embeddings and orthogonality of their class (Defn. 2) separately. Furthermore, qualitatively study the Gram matrices GM and GH and compare them to corresponding matrices for the CE loss under imbalances.
20

---

Published as a conference paper at ICLR 2024
MNIST
CIFARIO
CIFARI00
Tiny ImageNet
102
R
307
Step
10
See Dnz
R 50
10"
10
1.7
307
101
Long-tail
107
10
Epoch
Figure 8: Neural Collapse metric IvC tr(2ww //k/ for corresponding experiments Fig. >'y
Values usually the order of 10-2 suggesting strong convergence of embeddings to their class means.
E.3.1 NEURAL COLLAPSE
In order to quantify the collapse (NC) of embeddings, h, i [m]. to their class-means, we measure ANC tr(Ew Z))// (Papyan et al., 2020). Here, Σn Eee[k] (#e (G) (AG) is the between class covariance matrix, HG 1. 7 Ees[a] Me the global mean, and Ew Eie[n] Eie (b,
is the within class covariance matrix.
For the experiments shown in Fig. 2 and Fig. 3 in the main body, the corresponding values of BNC can be found in Fig. 8 and Fig. 12(b) respectively.
E.3.2 ANGULAR CONVERGENCE
Having confirmed convergence of embeddings to their respective class-means via NC (h; μe for y/ c), we can now compare the feature geometry the OF geometry by calculating the average between each pair of classes. Fig. 9 plots the average cosine similarity A (c, c') between class means for the same experiment that Fig. The graphs indicate strong convergence orthogonality between feature representations of different classes. Similarly, results angular convergence corresponding Fig. 3 are provided in Fig. 12(c), indicating similar convergence OF for the full-batch experiments.
Osim (c, c') == ///ee//le
MNIST
CIFAR1O
CIFAR100
Tiny ImageNet
Step
Long-tail
100
Epoch
Figure Average cosine similarity between different class means asim (c, (ee))== N/mo//ellleel for corresponding results from Fig. 2. Final values mostly on the order of 10-?, indicating strong orthogonality between class-mes mean embeddings.
21

---

Published as a conference paper at ICLR 2024
E.3.3 EMBEDDING HEATMAPS
As a qualitative measure, have generated heatmaps that visually represent the learned embedding geometries; see Figs. 1, 10, 11. Specifically, we generate heatmaps of the Gram-matrices GM train ResNet-18 with the full MNIST dataset. In Fig. 10
M'M and GH H'H H. In Fig. run on subset of the dataset (n 10000 total examples) with batch size of 1000. Specifically for Fig. 10, when optimizing with CE loss we modify the network (ResNet-18) such that it has normalization layer before the linear classifier head. We consider this additional setting to allow for a comparison where CE features are constrained to the unit sphere akin to our SCL experiments. Lastly, in Fig. 11 plot the learned features Gram-matrix GH for ResNet-18 18 trained on CIFARIO (n = 10000 total examples) with a batch size of 1000. This heatmap qualitatively shows a more complete picture
plotting GH H H rather than Gm, thus simultaneously illustrating both Neural Collapse and convergence to the k-OF structure.
E.3.4 EXPERIMENTS
MLPs
In Fig. 14 we run experiments with a simple layer multilayer perceptron (MLP) to further explore the impact of model complexity on geometric convergence. The MLP includes batch normalization and ReLU activation between each layer. Each layer has 512 hidden units. We train the model with batch size of 1000 with random reshuffling at each epoch. Furthermore, imbalanced MNIST. No batch duplication was used. All other aspects of the implementation are described in Sec. E. As shown in Fig. 14 all metrics AGM. ÂNC, and Ave (sim((,c) indicate strong convergence to the OF geometry, irrespective of imbalance ratio R.
train under R-STEP
R
R 10
R = 100
08
08
105
4
SCI
04
05
(12
CE
Figure 10: GM Gram-matrices eabedddnd various R-STEP imbalances at last epoch (350) training with ResNet-18 on MNIST with n 10000. To allow fair comparison, the CE features are normalized before the classifier head akin to the SCL experiments.
E.4 OPTIMIZATION DYNAMICS
E.4.1 LOSS CONVERGENCE
In order to compute the lower bounds shown in Fig. 3, we use Thm. substituting e1 with 1/r using 0.1 as employed in experiments; this substitution is allowed thanks to Remark 2. Furthermore, we compute the lower bounds and ((H) on a per sample basis, thus we divide by 1000 which corresponds to the number of datapoints in single batch. Lastly, method of
In both Fig. 11 and Fig. 10 batch duplication used descrihed Sec. 3.
22

---

Published as a conference paper at ICLR 2024
R =
R = 10
100
0931
43
JAC
60
330
OEE
64
931
5003
Figure 11: GH Gram-matrices of feature embeddings for various imbalances at last epoch (350) of training SCL with CIFARIO 10 and ResNet.
(a) Loss
(b) NC
(c) Angles
1.0
10
0.5
10-3
0.0
1000
2000
0
1000
2000
1000
2000
Epech
Epoch
Epoch
Figure Full-batch SCL: ResNet-18 trained on a R-STEP imbalanced subset of MNIST of size n 1000.
(a) Loss converges to the lower bound (dashed lines) computed in Thm. 1. (b) Within-class feature variation becomes negligible (NC). (c) The average pairwisc cosine similarity of class- means approaches zero. Each epoch is equivalent iteration of gradient descent.
maintaining numerical stability (as implemented in Khosla et al. (2020)) we apply global scaling of the loss by a factor 7/r., where тか = 0.07 is the base temperature Khosla al. (2020). The complete experiment, conducted over 2000 epochs (with axes limited to 500 epochs for clarity in Figure 3), is available in Figure 12.
E.4.2
EFFECT OF
Remark 2. The normalization of the embeddings in UFM. corresponds SCL with temperature 1. More generally, the normalization becomes |ha]l2 1/T. The conclusion of Thm. 1 is insensitive the choice of thus stated above for T without loss of generality. Although the value of T does not affect the global optimizers of UFM., we have empirically observed that it impacts the speed of convergence during training.
As described in Sec. and Sec. 4.1 the optimality of the OF geometry for the UFM. is invariant to the choice of the temperature parameter т. However, we have found that the speed of convergence to the OF geometry is dependent on the choice of r. Shown in Fig. 13 is full batch SCL experiment on 1000 samples of MNIST trained on ResNet-18 with 7 = 0.1, 1, 10.
clear from Fig. 13 that: (a) the within-class feature variation converges significantly faster for smaller ;; (b) the angles converge to k-OF faster and smoother for smaller J as well (see Fig. 13 (b)). In all cases, values of SNC and Cssi (c,c') continue decrease, and anticipate further convergence if the networks are trained for longer. These results qualitatively agree with the findings of Khosla et al. (2020) which suggest that smaller improves training speed.
23

---

Published as a conference paper at ICLR 2024
(a) NC
(b) Angles
1.00
0.1
10
Aveerd (sim((ee))
0.75
8xc
0.50
10-2
0.25
200
200
Epoch
Fpocc
Figure 13: Full-batch SCT.; ResNet-18 18 trained on a subset of MNIST of size 1000 with different temperature parameters (a) Within-class feature variation (NC). (b) Average pairwise cosine similarity of class-means. Each epoch equivalent to one iteration gradient descent.
(a) AG
(b) 3NC
(c) Avecrd Osim((e.c))
R
1.0
0.6
10-2
0.4
0.5
10-5
0.0
0.0
200
200
0
200
Epoch
Epoch
Epoch
Figure 14: Geometry convergence metries (a) AGM (b) 9sc. and (c) Ave Ouam(e,e for a 6 layer multilayer perceptron (MLP) model with ReLU activations trained with SCL and MNIST under R-STEP imbalance.
E.4.3 IMPACT OF BATCI SIZE
Khosla et al. (2020) have found performance increases with larger batch sizes, and thus we consider the effects of batch size on the convergence of the learned features their corresponding predicted geometry. We see in Fig. 15 that the features learned when ReLU is applied are more robust to variations in the batch size. Without ReLU, despite having ETF as the optimal geometry, as predicted by Graf et al. (2021), convergence is considerably more sensitive to batch size choice. This sensitivity exacerbated when the number of classes increases, see CIFAR100 Fig. 15.
E.5 COMPLEMENTARY RESULTS AND DISCUSSIONS ON BATCIIBIINDING
batch-binding binding help improve the In this section, describe examples where batching methods and batch- convergence speed of embeddings geometries to OF.
E.5.1 HOW BINDING ENSURES UNIQUE OF GEOMETRY
Fig. 16 provides simple illustration demonstrating how adding binding examples can satisfy the requirements of Cor. 2.1. While there alternative approaches satisfy the graph conditions stated Cor. 2.1 for ensuring unique OF geometry, the method of adding the same k examples to each batch is straightforward technique that is often computationally efficient, considering that batch sizes typically exceed the number of classes k.
Geometry convergence for different batching schemes. A comparison of OF convergence with three batch selection schemes was introduced in Sec. 5. We considered three strategies: 1) No batch-shuffling. Batch-shuffling, and 3) No batch-shuffing + batch-binding. We noted that
24

---

Published as a conference paper at ICLR 2024
ReLU
ReLU
Excd
1.4
1.0
1.2
Brac Sin
Exc Sta
Brsck
CIEARIO
0.8-
1.0
0.5
0.6-
0.5
0.4
0.4
02
0.
00
100
200
500
0
100
200
300
Epoch
Egoch
1.4
12-
CIEAR100
1.0-
10-
6.8-
0.5
0.6-
0.6
0.4
6.0
0.0
300
Epoch
Epoch
Figure 15: Convergence of learned features Gram matrices, GM to predicted geometry; OF in when ReLU applicd to features and ETF otherwise. Batch size varied with balanced data.
optimization with batch- binding, even in absence of shuffling is very effective at guiding the solution the unique OF geometry.
Further, in Table 5 in the appendix, we present experimental evidence in suggesting that batch binding does not negatively effect the balanced test accuracy of Nearest Class Center (NCC) classifier. This observation, paired with the impact of binding examples on OF convergence, motivates further analysis of the relationship between OF geometry and test accuracy.
Improving convergence for less powerful models. When conducting our experiments, we noticed that DenseNet-40 converges slower compared to ResNet-18. A reason for this may be related the very different complexities of the two models: ResNet-18 has substantially more trainable parameters compared to DenseNet-40. size for DenseNet experiments to 128 to increase the number of SGD iterations. We also train DenseNet for 500 epochs instead of 350 while reducing the learning rate from 0.1 to 0.01 at epoch 300. Yet, we observe that with much smaller batch sizes, the embedding geometry does not always converge to the OF geometry, especially when training with high imbalance ratios. Specifically when training with randomly reshuffled batches, there is higher chance that the examples do not interact with each other even if trained for long, suggesting that the optimal solution for all batches is not necessarily OF. We hypothesize that this likelihood increases when using smaller batch sizes during training. In order to combat this effect, we ran the DenseNet experiments with the addition of binding examples to every batch. As shown Fig. 19, we find that adding such binding examples significantly improves the convergence to OF, particularly when training on more complex datasets (CIFAR 10 compared to MNIST) and under more severe imbalances (STEP imbalance with It = 50, 100). These results emphasize the impact of batch-binding and provide further evidence regarding our claims in Sec. 4.2. While the convergence values AGM are higher when compared to ResNet-18, we deem them reasonable considering the difference in the number of parameters between the two architectures. Similar to ResNet- 18, we provide geometric convergence results for DenseNet-40 on different datasets in Fig. 23 in the appendix.
an attempt to improve the convergence speed, we reduce the batch
25

---

Published as a conference paper at ICLR 2024
B1
G1
GI
B2
o
B1
B3
G2
G3
(b)
(c)
Figure simple illustration to explain adding binding examples to batch satisfies the requirements of Cor. 2.1, thus leads to unique OF geometry. and white) classification example with batches. In addition to the data for each batch (included in red), the binding examples 1, added to each batch. (b) Gives the Batch Interaction Graph for the induced subgraph Gi of the class (black) of example and illustrates how batches are connected through examples 1, satisfying on how all three class graphs G1. G2, G3 are connected through the interactions of data points 1, 2, 3, satisfying the second condition of Cor. 2.1.
(a) Gives a 3 class (black, grey
first condition Cor. 2.1. (c) Elaborates
26

---

Published as a conference paper at ICLR 2024
Step
Long-tail
1.4
1.2
R 10
10
100
Sa
0.8
0.6
0.4
0.2
100
200
100
230
300
Epoch
Figure 17: Convergence of Gm to the OF geometry for ResNet- 18 model trained on CIFARIO under Step and Long-tail imbalance, with data augmentation. Results represent the average run results over 5 versions of the experiments with randomly chosen binding examples from each class.
No Batch Shuffling
No Batch Shuftling
Batch Shuffling
Batch-Binding
1.0
Step
Long-tail
250
333
200
300
100
230
301
Epoch
Convergence to the OF geometry for various batching schemes including the analysis- inspired scheme "No Batch Shuffling + Batch-Bindde aatch-Binding". See text for details. Experiments conducted with CIFAR-10 and ResNet-18.
Figure 18:
In another experiment, consider the convergence of ResNet-18 18 embeddings the OF geometry when trained with CIFARIO0 Models are trained for 500 epochs with constant learning rate of 0.1 and batch size of 1024. In Fig. without batch binding, in Fig. 20 ResNet-18 struggles to convergence. With a large number of classes (k 100 for CIFAR1001 it becomes increasingly more likely for the randomly reshuffled batches to not allow sufficient interactions between examples. Particularly for large imbalance ratios (c.g.. R 100, 50), since the number of samples in each minority class could become as low as 5-10, some batches might never encounter examples from minority classes. This is particularly noticeable for STEP-Imbalance and large imbalance ratios, where half the classes are minorities. On the other hand, including the binding examples in each batch (right side) can improve the convergence of the feature geometry to OF.
that ResNet-18 easily converges CIFARIO. However,
E.5.2 CONVERGENCE IN NON-RELU SETTINGS
In addition to observing that batch-binding improves the convergence of SCL with ReLU to OF, we consider the setting without ReLU. In particular in Fig. 21 we compare the convergence of the Gram matrix of mean embeddings to the ETF geometry with and without batch-binding We train ResNet-18 model with CIFAR-10. We observe that with a reduced batch size 128 the convergence be significantly improved with bending-examples.
27

---

Published as a conference paper at ICLR 2024
Batch Shuffling
Batch Shuffling
Batch- Bich-Binding
1.4
CIEARIO
10
0.8
0.6
0.4
1.4
MNIST
0.8
0.6
100 200
300
400
300
100
200
300
100
500
Epoch
Figure 19: Convergence of GM the OF geometry for DenseNet-40 model trained CIFARIO and MNIST under STEP imbalance with and without nutch-binding. While the impact hatch- bindings less noticeable when training on a simple dataset such MNIST, convergence is significantly improved, particularly under imbalance, when training on a more complex dataset such as CIFARIO.
Batch Shuffling
Batch Shuffling
+ Batch- match-Binding
1.4
10
Step
0.8
0.6
1.4
Long-fai
0.8
0.6
100
200
300
200
500
100
200
300
400
500
Epoch
Figure 20: Convergence embeddings geometry as measured by AGM :=
Gm
GM Th || for ResNet-18 18 trained on imbalanced CIFAR100 with SCL. and different batching schemes, Adding binding examples helps with the convergence to the OF geometry, especially under STEP imbalance with larger imbalance ratios.
I
E.6 ON THE CONVERGENCE OF BATCH-SHUFFLINNG TO OF
can be readily observed that fixed mini batch partition without shuffling (No batch-shuffthe shuffling in the experiments above) does satisfy the conditions of Cor. 2.1. Consequently, OF is not the unique
mini-batch
28

---

Published as a conference paper at ICLR 2024
Cuaeargaaon of SCL (un RoLU) to ETV
Baidi Birlia
No Botca Bialit
0.54
0.2
20
Epoda
Figure 21: Geometry convergence of G M M to the ETF geometry with and without binding examples. A batch size of 128 is used, there is no ReLU on the outpul features.
minimizer geometry in such scenarios. This is clearly manifested by the inadequate convergence levels observed in our experiments, as depicted for example in Fig. 18. In contrast, in our experiments, at each epoch (which
when the examples are randomly shuffled prior to partitioning into we call Batch Shuffling), the convergence behavior to OF shows a significant improvement compared to the case of No batch-shuffling. The observed improvement can be attributed to the fact that shuffling enables interactions among examples across epochs. To elaborate on this notion, we present formalization of batch shuffling below.
Let S denote the set of all permutations of [n: = examples. Additionally, let b denote a fixed batch-size. For simplicity, assume n is an integer multiple b. Ah beginning of each training epoch I, we sample permutation 8; = (i1, i2, (x) uniformly from S with replacement. We then define P(s)) {Li,,.... it), fio+1..... i20 h... (in-b+1 .... in)}
where n is the total number of training
as the collection of total z/ mini- maii-batches obtained by partitioning st. These are the mini-batches used at epoch f. Thus, SCL in this epoch can be written as follows:
L. (H)
Σ LH:;), where LB (H:i)) BeP(a.) KB
Σ
log Σ expp(1/ (p. ( h,),
BB
JEB
E E+s
j+i.Uy-M
and recall un. Kiiie B, = c]}. Consider also the loss over all mini-batches obtained by partitioning each permutation of S:
Ls (s)
= Σ E EC sGS BEP(s) 1B
(H:
(10)
Now, since st is uniformly sampled from S. we have the following relation for the expected per-epoch loss to the total loss given in Eqn. (10): E., (H) ( Therefore, training with batch- shuffling can be regarded as a stochastic version of minimizing the total loss in Eq. (10). Regarding the latter, can be checked (see Rem. 3 below) that satisfies the conditions of Cor. 2.1, thereby making OF its unique minimizer geometry. Taken together, these findings suggest that although the per-epoch loss L., obtained through batch-shuffling does not satisfy the conditions of Cor. 2.1 for any specific epoch t. it does satisfy them in expectation. In particular, this can be regarded preliminary justification of the experimental observation of improved convergence with batch- shuf-shuuffiig compared to No batch- shuffling schemes. The latter schemes fail to satisfy Cor. 2.1 under any circumstances, unless in the extreme case of a large batch size where = n and we recover Lim- However, it is important note that the aforementioned argument is a rough approximation, as it is not feasible to average over all |S] = n!! permutations when optimization is typically performed over limited number of epochs. This can explain why, despite exhibiting better convergence compared to the No batch-shuffling schemes, batch-shuffling schemes still demonstrate non-smooth and inconsistent convergence patterns in our experiments. This behavior becomes particularly evident when comparing the convergence levels of batch-shuffling to our batch-binding scheme, which is specifically designed to satisfy the conditions of Cor. 2.1 at each epoch.
Remark 3. To show that Eqn. (10) satisfies Cor. 2.1, we will show that the corresponding batch interaction graph G is a complete graph. This suffices, because the induced subgraphs from Def. 4 for complete graph are connected graphs, and there exist multiple edges between Ge, and Ge for C1 E [k]. To show that G is a complete graph, we argue as follows. Consider the set B of all mini- batches obtained by partitioning all nd elements of S. Fix b 2 2, so that the mini-batcher have least
29

---

Published as a conference paper at ICLR 2024
two examples. We will consider subset B1 B of mani-batchhes and show that the corresponding batch interaction graph for B1 is a complete graph, which would then imply the batch interaction graph of B is also complete graph. Specifically, let B1 - B denote the mini-batches comprising of the first b indices in every permutation 8 E S. In other words, from a given 8= = (21, bn) E S, we let B1 include the mini-batch (i1, (2,....)))). -,). Since S includes all permutations of [,], B1 contains all b-length permutations of the elements of [n], possibly with repetitions. Thus, the batch interaction graph created by the mini-batches in B1 is complete graph. In the definition of the batch interaction graph, since repeated presence of pair of examples does not alter the graph, the batch interaction graph for the same complete graph. Thus, B satisfies Cor. 2.1.
E.7 RELU DOES NOT COMPROMISE ACCURACY
Having discussed the preeserriiirssrvinng value of ReLU activation in SCL training, we now turn towards the impact on generalization accuracy. While the exact relationship between the training feature geometry and generalization is open question, it is conceivable that maintaining symmetry in features is beneficial to generalization (Behnia et al., 2023; Zhu et al., 2022). In fact, our experiments with laas-iimalaance reveal that ReLU improves the test accuracy significantly in a subset of the cases, while not deteriorating in others. There have been suggestions of loss function modifications in such a way that the feature geometry is symmetric in presence of class- imbalance, consequently also achieving improvements in generalization Behnia et al. (2023); Zhu et al. (2022). Since the role of projector block is of interest when training with SCL, we compare the experimental test accuracies both when the projector is retained projst-rojection) as well as discarded (pre-projection) during inference (Khosla et al., 2020; Chen et al., 2020; Zhu et al., 2022).
Pre Projection
Post Projection
MLP
MLP
Step
LongTail
Step
LongTai
Step
LongTail
Step
LongTail
91.88 0.29
92.04 0.10
91. 94 0.04
91.79 0.13
CIFAR-10 10
10
83.70 1.09
83.82 0.70
X3.41 0.64
84.40 0.79
82.35 0.67
82.97 0.83
90.81 1.01
82.99 1.05
100
60.19 a 1.75
68.75 a 0.31
67.12 a 1.04
69.57 1.69
55.31 a 1.98
03.67 0.77
59.83 a 1.75
02.65 a 2.37
CIFAR- (000
72.17 = 0.23
72.32 = 0.60
72.30 0.57
72:25 a 0.35
56.55
57.10
56.27
100
43.49% 0.30
37.19
43,80
39.71
1.64
0.33
37.21
Tiy Inngocee
62.53 1.23
62.36 0,37
62.71 1.10
62:48 0.40
10
50.17 + 1.44
50.8 1.94
49.78 + 0.67
49.81 +: 0.82
48.91 + 1.6
50.23 + 1.83
48.23 0.53
48.76 1.18
100
39.13 a
36.06 i. 0.89 40.02 a 0.62
35.84 1.33
37.41 a: 0.20
34.57 a 0.95
37.19 a 0.68
33.78 1.34
Table 2: Comparison test accuracies for (i) CIFAR- 10 (ii) CIFAR-100 (iii) Tiny ImageNet when ResNet-18 18 is trained with a 2-layer MLP projection head with and without ReLU on the output features. Comparisons in hoth Step and L.T imbalance settings. Additionally comparison of the performance using features before the projection head (pre-projection), and after the head prost-projcction).
Thus, as noted in Table 2 we make a remarkable observation that simply adding ReLU to the network can yield significant improvements in test accuracy in certain cases, while least matching that of the case without ReLU in others. In light of this, our finding of the insensitivity of the embedding geometry in presence of ReLU serves potential explanation. Our analysis therefore paves an important direction for future empirical investigations into the role of the activation functions in the projector network on generalization of SCL..
E.8 RELU HELPS IMPROVE WORST CLASS ACCURACY
As illustrated by Fig. 2 and Fig. 1, the of ReLU results in symmetric OF feature geometry and prevents the collapse of the minority feature vectors. While the direct relationship between NC geometry of train and test features is not yet entirely established, expect that this symmetric
30

---

Published as a conference paper at ICLR 2024
geometry could help improve the worst class test accuracies, especially for cases with large number of classes where the impact of the imbalance on minority classes is more severe. To this end, we compare the average test accuracy of the worst quartile of classes (worst 25 classes) for ResNet-18 model trained on CIFAR-I0 100 (Table 4). Our results indicate that the worst test accuracies improve in the presense of ReLU which could link ReLU's impact achieving the same train margin for all minority classes irrespective of the imbalance ratio.
10
100
MLP RRLUMLP
36,59 a 0.19
23.06 # 0.52
13.29 = 0.83
6.73 a 0.39
4.42 # 0.65
38.52 0.24
27.49 0.37
18.62 0.72
10.41 0.15
6,49 0.35
Table 3: Worst class accuracy for ResNet- 18 model trained using CIFAR-10 10 for different step imbalances.
Stap
Longtail
Stap
Longtail
36.51 0.12
39.70 0.84
38.15 0.52
41.35 0.83
CHEAR 100
10
23.12 a 0.78
27.46 0.18
31,46 a: 0.14
33.76 0.19
20
13.50 +
23.30 1.04
18.66 1.05
25.52
15.39k
5.44
7.15
6.48
8.99
Table 4: Comparison of average test accuracies for the worst 25 classes for ResNet34 trained on CIFAR- 100. For each setup, the test accuracies evaluated over runs. In all cases, can see that the addition ReLU helps improve the worst class accuracies.
E.8.1
IMPACT OF BATCH-BINDING ON GENERALIZATION
Table 5: NCC test accuracy for Mich-Bindiig
Step
Long-tail
Ratio (R)
Reshuffling
Reshuffling Batch- Binddig
Reshuffling
Reshuffling
+ Batch-Biin Batch-Bindiig
10
83.31 0.65%
83.27 0.27%
85.04 = 0.28%
85.03 0.57%
100
64.07 1.79%
64.18 .5%
67.75 + 1.21%
68.08 = 7.72%
To study the effect of binding examples on generalization, we consider the Nearest Class Center (NCC) classifier. For this, each test example is assigned the label of the class center phc learned during training that is closest to it (in Euclidean distance). We consider a simple setup of training ResNet-18 model on CIFAR10 10 under imbalances R 10. 100 both with and without the binding examples. Models are trained with a batch size of 128 for 350 epochs. We ran the experiments with smaller batch size to increase the number of pack-propagation steps as adding data augmentations slows down convergence to OF geometry We start with a learning rate of 0.1, which is reduced by factor of 10 on epochs 250 and 300. Weight decay set data augmentation as it is common practice when considering generalization and test accuracy. In particular, rather than simply adding horizontally flipped images (as described in Sec. 3) we allow for generic augmentations that include simple horizontal or vertical flips and random cropping with probability of 0.5. The NCC test accuracy was measured across versions the experiment with the k binding examples sampled randomly every time and versions without any batch-binding. The results for NCC balanced test accuracies are provided Table 5. While making definitive conclusions regarding the impact of the embeddings geometries and binding examples generalization requires
10-4. addition, we apply
31

---

Published as a conference paper at ICLR 2024
Step
Long-tail
1.0
Zero Eirıs
0.8
10
80
AG
0.4
100
300
400
500
100
300
400
500
Epoch
Figure Convergence of GM the OF geometry for a ViT model trained CIFARIO under Step and Long-tail tail imbalance.
further investigation, this preliminary investigation suggests that batch-binding does not negatively impact NCC test accuracy.
Finally, Fig. 17 shows convergence of embeddings geometry to OF for the same experiments. As expected, the convergence is slightly slower in this case due to the inclusion of data augmentation (random crops and flips).
E.8.2 EXPERIMENTS WITH ADDITIONAL ARCHITECTURES
In order to ensure that our results are not impacted by the specifics of CNN structure of ResNet architecture, we perform series of experiments on DenseNet-40 models and VIT Dosovitskiy et al. (2021). For the VIT models, we used 2 layer vision transformer models with 10 attention heads for the CIFAR-10 10 dataset with exponential and step imbalances. Fig. 22 and Fig. illustrate how the feature geometry converges to OF consistently for both networks, further solidifying our findings and indicating that given a large enough model, assumptions such that UFM, hold for classification tasks regardless of network architecture.
32

---

Published as a conference paper at ICLR 2024
CIFAR1O
MNIST
FMNIST
100
Step
00
0.3
Long-tail
0.5
041
100
230
300
100
500
200
300
600
200
500
$00
503
Epoch
(a) GM Convergence to OF MNIST
CIFARIO
FMNIST
10
Step
10"1
10
Long-tail
10
300
2501
300
300
100
200
(10)
Epoch
(b) Neural Collapse
CIFARIO
MNIST
FMNIST
100
R
3.5
Step
0.6
1.0
0.5
Long-tail
065
0.2
100
2040
301
(10
210
103)
214)
350
00C
100
2011
3001
3.0
000
Epoch
(c) Angles between class-means
Figure 23: Geometric convergence of embeddings (as based on Def. 1, Def. 2, and Def. 3) of DenseNet-40 trained with batch-binding. We measure: (a) AG (b) Bxc. and (c)
Ave (c, defined in text.
33


Reviews:
[{"review_id": "7JPnr6vchT", "reviewer": "Reviewer_PL8u", "summary": "This work studies the training of deep neural networks with a supervised contrastive objective. Empirically, it is shown that if there are ReLU activations after the last layer, then the learned features per class collapse, and that the features of different classes are orthogonal to each other. This is supported by an analytic argument that characterizes the minimum of the loss function as a function of the features (instead of the network weights). This feature geometry is independent of the class ratios. Moreover, a batching strategy is derived that guarantees the emergence of this feature geometry.", "questions": "Are there empirical differences when training models with and without ReLU activations after the last layer. If so, can such differences be explained by the targeted feature geometry? Are such differences more pronounced when using the batching strategy?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "- The finding of this work; that training a deep neural network with a supervised contrastive loss and ReLU activations leads to the same feature geometry irrespective of the class imbalances is very surprising. Prior work showed that this is not the case for the cross entropy loss. Moreover, it is shown that this feature geometry is a *common* minimizer for all batches. This results in a straightforward proof, as there are no 'averaging effects' to account for.\n- The paper is well written\n- Related work (with one exception) is extensively discussed and compared to.", "weaknesses": "- Significance: The work appears to be written on the premise that achieving a symmetric (orthogonal or simplex) feature geometry is beneficial.  For example, a batching strategy is devised that guarantees such a geometry. While symmetry seems reasonable for class-balanced data, for imbalanced data this is not so clear. I did not see an argument supporting this premise, and thus also not for why using the proposed batching strategy  is a good idea. Moreover, there is few empirical support, as the use of ReLUs does not have any effect on classification accuracy, cf. Table 1.\n\n- The analytical argument requires a very strong assumption: unconstrained features. Moreover, only the loss minimizer is derived, but optimization effects for this non-convex problem are not accounted for. On the other hand, this strong assumption allows for results that are architecture independent.\n\n- The empirical part of this work would be stronger, if the theoretical findings are verified on a large scale imbalanced dataset (compared to the mid-scale datasets CIFAR10, CIFAR100 and TinyImagenet). \n\n-  The very related work [Wang & Isola, Understanding Contrastive Representation Learning through Alignment and Uniformity on the Hypersphere, ICML 2020] is not cited. This work studies feature geometries when minimizing the (unsupervised) contrastive loss and should be compared to.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZtWg9l9MXB", "reviewer": "Reviewer_BEvk", "summary": "## Background summary \nThis paper focus on the issue of neural collapse (NC), when hidden features of sample class collapse into one another, and orthogonal framing (OF), when the mean vectors of hidden representation of each class, are orthogonal to one another. For the theoretical framework, the paper chooses the unconstrained features model (UFM), where the loss is minmized over the classification layer and last hidden layers, unconstrained from the other model parameters. \n- *UFM with cross entropy loss:* we have \n$\\min_{w_c,h_i\\in R^d}\\mathcal{L}_{CE}(\\\\{w_c\\\\}, \\\\{h_i\\\\})$ \n- *UFM with supervised contrastive loss (SCL)* for full batch (for mini-batch, it is over $i\\in [n]$):  \n$\\min_{h_i\\in R^d} \\sum_{i\\in B} \\frac{1}{n_{B,y_i} - 1} \\sum_{j\\in B, y_i=y_j, i\\neq j} \\log\\left(1 + \\sum_{\\ell\\neq i, j} \\exp(h_i^\\top h_\\ell - h_i^\\top h_j)\\right)$\n- *Symmetrical solutions under balanced classes.* Intuitively, when classes are balanced (same number of samples per class) both cross entropy & contrastive loss, repel the samples from seperate classes, while absorb the class within one class. This has the effect of the finding the geometry that is symemtric w.r.t. to the class identity, resulting in a equiangular tight frame (ETF). This attraction/repulsion can be seen in the logits $\\sum_{\\ell\\neq i} \\exp(h_i^\\top h_\\ell)$\n- *Non-symmetrical solutions when classes are not balanced*  When the classes are not balanced, the symmetry between class will be broken. For example, if one class completely dominates the training set, both cross-entropy and contrastive loss will lead to solutions, where the majority class is almost anti-parallel to other class means. \n\n## Main contribution: symmetric solutions by passing last hidden layer through ReLU\nThe main contribution of the paper is showing that if the spectral loss is optimised over the space of all-positive hidden layers, namely by passing it through a ReLU activation, denoted as UFM_+, the UFM leads to symmetrical solutions, i.e., NC & OF properties are held again. \n \n- Thm1 proves that for the full batch contrastive loss, if we additionally pass the hidden layers through a ReLU, ie, optimize  last hidden layers over elementwise positive vectors $h_i\\ge 0$, global minima is achieved when for all $i,j \\in [n]$ we have $h_i^\\top h_j = 1(i=j)$. \n- Thm2. proves a similar property for mini-batch contrastive loss, where the collapse and orthogonality property hold within each batch (for samples within each batch $B\\in\\mathcal{B}$ and pair within batch $i,j\\in B$, we have $ h_i^\\top h_j = 1(i = j)$. \n- Cor 2.1: shows a batching propery that if held, the neural collapse and orthgonal frame property hold across the full batch, which can be explained as 3 conditions: \n   1) every batch has a sample from each class \n   2) The batch connectivity graph is connected (nodes $u$ are samples $u\\in [n]$, edges represent $u$ and $v$ appeared in the same batch $u,v\\in B$)\n 3) The sub-graphs corresponding to each class are also connected\n- The paper validates the theories on various ResNet architectures. The experiments show two variations of unbalanced datasts (step and long-tailed) which seem to show that the proposed approach (passing last layer throgh ReLU) does indeed lead to symmetric solutions.\nThe experiments also include a result (Table 1) that shows the proposed approach leads to a higher test accuracy.", "questions": "While the authors present the ideas for symmetrising non-balanced the contrastive loss, to me it seems like the cross-entropy loss will benefit from the same approach. For example, consider the CE loss for $c$-class classification for the full batch is given by:\n$$\\mathcal{L} =\\sum_{i=1}^n\\log\\left( \\frac{\\exp(h_i^\\top w_{y_i} ) }{\\sum_k^c \\exp(h_i^\\top w_k)}\\right)$$\nIt is rather trivial to see that if we pass the last hidden layer through ReLU, since that elementwise condition $h_i\\ge 0$ ensures that $w_k$'s are also positive, and the resulting inner products $h_i^\\top w_i$ will be always non-negative, and the global optimum is achieved when neural collapse and orthognal frames are achieved. For example, if the hidden dimension is larger than number of classes $h,w_k\\in R^d, d\\ge c,$ then we can set $w_1,\\dots, w_c$ to the standard basis $w_k := e_k$ and collapse hidden layers onto the corresponding vectors $h_i := e_{y_i}.$ Therefore, my questions are:\n- Do authors agree with my reasoning? Please correct me if my reasoning above is flawed. \n- If yes, why have authors left out cross entropy in their analysis, given that it's leading to the same result as the contrastive loss results?\n- Is there any pros/cons if we compare symmetrized CE loss, and contrastive loss, both by passing hidden layer through ReLU?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "1. The presentation. The presentation of results, both in terms of writing, mathematical notation and formula, and the figures, is exquisite. The authors take many positive steps in ensuring that the concepts are clear and simple to grasp for the reader. \n2. The selection of the research problem and the ideas presented by the author are (to the best of my knowledge) original. The simplicity of the ideas as well as the presentation of the results make the results intriguing to read.", "weaknesses": "Main issue\n- While the stated objective of the paper, i,e, embeddings that satisfy NC & OF properties, is clear, it not at all clear  that it is a justified goal. For example, consider the simple case of binary classification. The optimal solution without ReLU will map two classes to some anti-parallel vectors $u$ and $v = -u$. This means that the hidden representations for one class will have the maximum degree of separation $\\| u - v\\| = 2$. Now, if we pass them through ReLU, forcing them to be element-wise positive, the solutions will be instead some $u$ and $v$ such that $u^\\top v = 0$. Thus, the their mean-class distance will be $\\| u - v\\| = \\sqrt{2}.$ Without any formalities, what doesn't seem to be a benefit in terms of distance, doesn't seem to be a good approach for generalization either. Even if we make the two classes here very unbalanced, still the anti-parallel geometry seems to be better than the orthogonal approach.  Since the authors do not make any formal/theoretical comments on this topic, the only relevant evidence they provide in this direction is Table 1 test accuracy results. Because of the aforementioned reasons, I consider the results of table 1 to be particularly surprising. Thus, if the the main evidence to support utility of the approach is empirical, I would like much more substantive experiments. Some elaborate explanations or at least speculative comments by the authors will also go a long way to explain the benefits of the orthogonal mean-classes. \n- My second main issue is about the novelty and significance of the contributions. The theory, while stated clean and clearly, is rather thin. The main theorems are not very surprising, and are rather straightforward. As elaborated in the questions section, If I'm not mistaken, similar results can be replicated for the cross-entropy loss in a similar straightforward fashion. On the empirical front, while the results are definitely interesting (notably Table 1), they are not comprehensive to solidify that point. For example, if this approach has any benefits for improving test accuracy, there should be much broader experiments, covering different model architectures, datasets, and most importantly, competing methods, that deal with unbalanced classes. I am not advocating here for state-of-the-art performance. However, it is important to see the empirical results in the context of prior/similar work, so that the reader may assess the empirical benefits for themselves. That said, I am open to reconsidering my views (both on theory and empirical contributions) upon hearing authors responses. \n\nMinor issues: \n- I think there is a slightly overuse of abbreviations in this paper which make it slightly hard to read. For example, I had to go back several times to read what ETF is, or other notations. If possible, please keep the shorthands to a minimum. \n- There is a small discrepancy between eq(1) and eq(2) in the presence of (1+ ... ) within the logarithm and the summation is over $\\ell\\neq i, j$ while another is over $\\ell \\neq i$. Perhaps authors can comment on this? \n- In Figures 2 & 5 axes are not clearly defined, eg. caption can say x-axis ... y-axis ...", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VmYxQe05vC", "reviewer": "Reviewer_q7qd", "summary": "This work proposes a variants of SCL loss to restore the symmetrical geometry of the class-mean learned embeddings in the presence of unbalance class samples, just adding a ReLU activation after the projection layer. It theoretically demonstrates that due to the existence of ReLU, the last-layer feature embeddings of samples with the same labels will be finally aligned, while that of samples with different labels will be orthogonal, when we employ the full-batch training.  While implement mini-batch training, it also showcases that mini-batch selection strategy will significantly influence the learned embedding geometry.", "questions": "The minimum of the loss is not  zero and varies with the number of samples in different classes.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The proposed approach stands out for its simplicity and effectiveness. It offers valuable insights into the subject matter, shedding light on the restoration of symmetrical geometry in the presence of unbalanced class samples.\n\n - The theoretical analysis is well-founded, enhancing the credibility of the approach. Experimental results, while limited to relatively simple datasets like MNIST, CIFAR10/CIFAR100, and TinyImageNet, do support the method's efficacy.", "weaknesses": "The main concerns still lie in the experimental parts.\n\n- The experiments primarily utilize CNN-based architectures. It would be beneficial to explore the applicability of this approach to train other architectures, such as Transformers, to gauge its versatility across neural network types.\n\n-  While the proposed approach demonstrates superiority on simpler datasets, questions remain regarding its performance on more complex and larger datasets. The potential impact of the non-negative constraints imposed by the ReLU activation on representation ability in more intricate tasks needs further investigation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mTiEmUecpk", "reviewer": "Reviewer_ScWC", "summary": "This paper applies RELU activation in the last layer of models in supervised contrastive learning (SCL). With that, SCL can learn representations that converge to the OF geometry irrespective of the level of class imbalance. The author theoretically and empirically verify the effectiveness of this method. Besides, this paper finds that the batch selection is important in representation geometry and they design a batch selection strategy (batch-binding).", "questions": "See my comments above.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The theoretical analysis and empirical results cooperate well. Figure 1, 2, 3 empirically demonstrate the advantages of the additional RELU function.\n2. The motivation is clear and the paper is written well.\n3.  The proposed method is simple and effective.", "weaknesses": "1. Based on the empirical results in the paper (Table 1), the improvements in the test accuracy on CIFAR-100 are not significant, even when the imbalance ratio is large. It would be better to show the advantages and disadvantages of learned representations in different downstream tasks (e.g., Does SCL+ReLU performs better in fine-tuning tasks?).\n2. It would be better to conduct some additional experiments on large scale datasets (e.g., ImageNet100).\n3. It seems that batch size is an important argument when applying sample selections according to the theoretical analysis. It would be better to add the ablation study on that.\n4. Based on the analysis, it seems that various activation functions can achieve the similar effect as ReLU. Is it possible to add a discussion about that?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)