The Details Matter: Preventing Class Collapse in Supervised Contrastive Learning ^†

Abstract

Supervised contrastive learning optimizes a loss that pushes together embeddings of points from the same class while pulling apart embeddings of points from different classes. Class collapse—when every point from the same class has the same embedding—minimizes this loss but loses critical information that is not encoded in the class labels. For instance, the “cat” label does not capture unlabeled categories such as breeds, poses, or backgrounds (which we call “strata”). As a result, class collapse produces embeddings that are less useful for downstream applications such as transfer learning and achieves suboptimal generalization error when there are strata. We explore a simple modification to supervised contrastive loss that aims to prevent class collapse by uniformly pulling apart individual points from the same class. We seek to understand the effects of this loss by examining how it embeds strata of different sizes, finding that it clusters larger strata more tightly than smaller strata. As a result, our loss function produces embeddings that better distinguish strata in embedding space, which produces lift on three downstream applications: 4.4 points on coarse-to-fine transfer learning, 2.5 points on worst-group robustness, and 1.0 points on minimal coreset construction. Our loss also produces more accurate models, with up to 4.0 points of lift across 9 tasks.

1. Introduction

Supervised contrastive learning has emerged as a promising method for training deep models, with strong empirical results over traditional supervised learning [1]. Recent theoretical work has shown that under certain assumptions, class collapse—when the representation of every point from a class collapses to the same embedding on the hypersphere, as in Figure 1—minimizes the supervised contrastive loss $L_{SC}$ [2]. Furthermore, modern deep networks, which can memorize arbitrary labels [3], are powerful enough to produce class collapse.

Although class collapse minimizes $L_{SC}$ and produces accurate models, it loses information that is not explicitly encoded in the class labels. For example, consider images with the label “cat.” As shown in Figure 1, some cats may be sleeping, some may be jumping, and some may be swatting at a bug. We call each of these semantically-unique categories of data—some of which are rarer than others, and none of which are explicitly labeled—a stratum. Distinguishing strata is important; it empirically can improve model performance [4] and fine-grained robustness [5]. It is also critical in high-stakes applications such as medical imaging [6]. However, $L_{SC}$ maps the sleeping, jumping, and swatting cats all to a single “cat” embedding, losing strata information. As a result, these embeddings are less useful for common downstream applications in the modern machine learning landscape, such as transfer learning.

In this paper, we explore a simple modification to $L_{SC}$ that prevents class collapse. We study how this modification affects embedding quality by considering how strata are represented in embedding space. We evaluate our loss both in terms of embedding quality, which we evaluate through three downstream applications, and end model quality.

In Section 3, we present our modification to $L_{SC}$, which prevents class collapse by changing how embeddings are pushed and pulled apart. $L_{SC}$ pushes together embeddings of points from the same class and pulls apart embeddings of points from different classes. In contrast, our modified loss $L_{spread}$ includes an additional class-conditional InfoNCE loss term that uniformly pulls apart individual points from within the same class. This term on its own encourages points from the same class to be maximally spread apart in embedding space, which discourages class collapse (see Figure 1 middle). Even though $L_{spread}$ does not use strata labels, we observe that it still produces embeddings that qualitatively appear to retain more strata information than those produced by $L_{SC}$ (see Figure 2).

In Section 4, motivated by these empirical observations, we study how well $L_{spread}$ preserves distinctions between strata in the representation space. Previous theoretical tools that study the optimal embedding distribution fail to characterize the geometry of strata. Instead, we propose a simple thought experiment considering the embeddings that the supervised contrastive loss generates when it is trained on a partial sample of the dataset. This setup enables us to distinguish strata based on their sizes by considering how likely it is for them to be represented in the sample (larger strata are more likely to appear in a small sample). In particular, we find that points from rarer and more distinct strata are clustered less tightly than points from common strata, and we show that this clustering property can improve embedding quality and generalization error.

In Section 5, we empirically validate several downstream implications of these insights. First, we demonstrate that $L_{spread}$ produces embeddings that retain more information about strata, resulting in lift on three downstream applications that require strata recovery:

• We evaluate how well $L_{spread}$’s embeddings encode fine-grained subclasses with coarse-to-fine transfer learning. $L_{spread}$ achieves up to 4.4 points of lift across four datasets.
• We evaluate how well embeddings produced by $L_{spread}$ can recover strata in an unsupervised setting by evaluating robustness against worst-group accuracy and noisy labels. We use our insights about how $L_{spread}$ embeds strata of different sizes to improve worst-group robustness by up to 2.5 points and to recover 75% performance when 20% of the labels are noisy.
• We evaluate how well we can differentiate rare strata from common strata by constructing limited subsets of the training data that can achieve the highest performance under a fixed training strategy (the coreset problem). We construct coresets by subsampling points from common strata. Our coresets outperform prior work by 1.0 points when coreset size is 30% of the training set.

Next, we find that $L_{spread}$ produces higher-quality models, outperforming $L_{SC}$ by up to 4.0 points across 9 tasks. Finally, we discuss related work in Section 6 and conclude in Section 7.

2. Background

We present our generative model for strata (Section 2.1). Then, we discuss supervised contrastive learning—in particular the SupCon loss $L_{SC}$ from [1] and its optimal embedding distribution [2]—and the end model for classification (Section 2.2).

2.1. Data Setup

We have a labeled input dataset $\mathcal{D}={\left\{(x_i,y_i)\right\}}_{i=1}^N$, where $(x,y)\sim \mathcal{P}$ for $x\in \mathcal{X}$ and $y\in \mathcal{Y}=\{1,\cdots ,K\}$. For a particular data point x, we denote its label as $h\left(x\right)\in \mathcal{Y}$ with distribution $p\left(y\right|x)$. We assume that data is class-balanced such that $p(y=i)=\frac{1}{K}$ for all $i\in \mathcal{Y}$. The goal is to learn a model $\left(y\right|x)$ on $\mathcal{D}$ to classify points.

Data points also belong to categories beyond their labels, called strata. Following [5], we denote a stratum as a latent variable z, which can take on values in $\mathcal{Z}=\{1,\cdots ,C\}$. $\mathcal{Z}$ can be partitioned into disjoint subsets $S_1,\cdots ,S_K$ such that if $z\in S_k$, then its corresponding y label is equal to k. Let $S\left(c\right)$ denote the deterministic label corresponding to stratum c. We model the data generating process as follows. First, the latent stratum is sampled from distribution $p\left(z\right)$. Then, the data point x is sampled from the distribution ${\mathcal{P}}_z=p(·|z)$, and its corresponding label is $y=S\left(z\right)$ (see Figure 2 of [5]). We assume that each class has m strata, and that there exist at least two strata, $z_1,z_2$, where $S\left(z_1\right)\ne S\left(z_2\right)$ and $\mathrm{supp}\left(z_1\right)\cap \mathrm{supp}\left(z_2\right)\ne \varnothing$.

2.2. Supervised Contrastive Loss

Supervised contrastive loss pushes together pairs of points from the same class (called positives) and pulls apart pairs of points from different classes (called negatives) to train an encoder $f:\mathcal{X}\to {\mathbb{R}}^d$. Following previous works, we make three assumptions on the encoder: (1) we restrict the encoder output space to be ${\mathbb{S}}^{d-1}$, the unit hypersphere; (2) we assume $K\le d+1$, which allows Graf et al. [2] to recover optimal embedding geometry; and (3) we assume the encoder f is “infinitely powerful”, meaning that any distribution on ${\mathbb{S}}^{d-1}$ is realizable by $f\left(x\right)$.

2.2.1. SupCon and Collapsed Embeddings

We focus on the SupCon loss $L_{SC}$ from [1]. Denote $\sigma (x,x^{\prime })=f{\left(x\right)}^{\top }f\left(x^{\prime }\right)/\tau$, where $\tau$ is a temperature hyperparameter. Let $\mathcal{B}$ be the set of batches of labeled data on $\mathcal{D}$ and $P(i,B)=\{p\in B\backslash i:h(p)=h(i\left)\right\}$ be the points in B with the same label as $x_i$. For an anchor $x_i$, the SupCon loss is ${\widehat{L}}_{SC}(f,x_i,B)=\frac{-1}{\left|P\right(i,B\left)\right|}{\sum }_{p\in P(i,B)}log\frac{exp\left(\sigma (x_i,x_p)\right)}{{\sum }_{a\in B\backslash i}exp\left(\sigma (x_i,x_a)\right)}$, where $P(i,B)$ forms positive pairs and $B\backslash i$ forms negative pairs.

The optimal embedding distribution that minimizes $L_{SC}$ has one embedding per class, with the per-class embeddings collectively forming a regular simplex inscribed in the hypersphere Graf et al. [2]. Formally, if $h\left(x\right)=i$, then $f\left(x\right)=v_i$ for all $x\in \mathcal{B}$. ${\left\{v_i\right\}}_{i=1}^K$ makes up the regular simplex, defined by: (a) ${\sum }_{i=1}^K{v}_i=0$; (b) $\parallel v_i{\parallel }_2=1$; and (c) $\exists c_K\in \mathbb{R}$ s.t. $v_i^{\top }{v}_j=c_K$ for $i\ne j$. We describe this property as class collapse and define the distribution of $f\left(x\right)$ that satisfies these conditions as collapsed embeddings.

2.2.2. End Model

After the supervised contrastive loss is used to train an encoder, a linear classifier $W\in {\mathbb{R}}^{K\times d}$ is trained on top of the representations $f\left(x\right)$ by minimizing cross-entropy loss over softmax scores. We assume that $\parallel W_y{\parallel }_2\le 1$ for each $y\in \mathcal{Y}$. The end model’s empirical loss can be defined as $\widehat{\mathcal{L}}(W,\mathcal{D})={\sum }_{x_i\in \mathcal{D}}-log\frac{exp\left(f{\left(x_i\right)}^{\top }{W}_{h\left(x_i\right)}\right)}{{\sum }_{j=1}^{K}exp\left(f{\left(x_i\right)}^{\top }{W}_j\right)}$. The model uses softmax scores constructed with $f\left(x\right)$ and W to generate predictions $\widehat{p}\left(y\right|x)$, which we also write as $\widehat{p}\left(y\right|f\left(x\right))$. Finally, the generalization error of the model on $\mathcal{P}$ is the expected cross-entropy between $\widehat{L}\left(y\right|x)$ and $p\left(y\right|x)$, namely $\mathcal{L}(x,y,f)={\mathbb{E}}_{x,y\sim \mathcal{P}}\left[-log\left(y\right|f\left(x\right))\right]$.

3. Method

We now highlight some theoretical problems with class collapse under our generative model of strata (Section 3.1). We then propose and qualitatively analyze the loss function $L_{spread}$ (Section 3.2).

3.1. Theoretical Motivation

We show that the conditions under which collapsed embeddings minimize generalization error on coarse-to-fine transfer and the original task do not hold when distinct strata exist.

Consider the downstream coarse-to-fine transfer task $(x,z)$ of using embeddings $f\left(x\right)$ learned on $(x,y)$ to classify points by fine-grained strata. Formally, coarse-to-fine transfer involves learning an end model with weight matrix $W\in {\mathbb{R}}^{C\times d}$ and fixed $f\left(x\right)$ (as described in Section 2.2) on points $(x,z)$, where we assume the data are class-balanced across z.

Observation 1.

Class collapse minimizes $\mathcal{L}(x,z,f)$ if for all x, (1) $p(y=h(x\left)\right|x)=1$, meaning that each x is deterministically assigned to one class, and (2) $p\left(z\right|x)=\frac{1}{m}$ where $z\in S_{h\left(x\right)}$. The second condition implies that $p\left(x\right|z)=p(x\left|y\right)$ for all $z\in S_y$, meaning that there is no distinction among strata from the same class. This contradicts our data model described in Section 2.1.

Similarly, we characterize when collapsed embeddings are optimal for the original task $(x,y)$.

Observation 2.

Class collapse minimizes $\mathcal{L}(x,y,f)$ if, for all x, $p(y=h(x\left)\right|x)=1$. This contradicts our data model.

Proofs are in Appendix D.1. We also analyze transferability of f on arbitrary new distributions $(x^{\prime },y^{\prime })$ information-theoretically in Appendix C.1, finding that a one-to-one encoder obeys the Infomax principle [7] better than collapsed embeddings on $(x^{\prime },y^{\prime })$. These observations suggest that a distribution over the embeddings that preserves strata distinctions and does not collapse classes is more desirable.

3.2. Modified Contrastive Loss $L_{spread}$

We introduce the loss $L_{spread}$, a weighted sum of two contrastive losses $L_{attract}$ and $L_{repel}$. $L_{attract}$ is a supervised contrastive loss, while $L_{repel}$ encourages intra-class separation. For $\alpha \in [0,1]$,

$$\begin{array}{cc}\hfill L_{spread}& =\alpha L_{attract}+(1-\alpha )L_{repel}.\hfill \end{array}$$

(1)

For a given anchor $x_i$, define $x_i^{aug}$ as an augmentation of the same point as x. Define the set of negative examples for i to be $N(i,B)=\{a\in B\backslash i:h(a)\ne h(i\left)\right\}$. Then,

$$\begin{array}{cc}\hfill {\widehat{L}}_{attract}(f,x_i,B)& =\frac{-1}{\left|P\right(i,B\left)\right|}\times \sum _{p\in P(i,B)}log\frac{exp\left(\sigma (x_i,x_p)\right)}{exp\left(\sigma (x_i,x_p)\right)+{\sum }_{a\in N(i,B)}exp\left(\sigma (x_i,x_a)\right)}\hfill \end{array}$$

(2)

$$\begin{array}{cc}\hfill {\widehat{L}}_{repel}(f,x_i,B)& =-log\frac{exp\left(\sigma (x_i,x_i^{aug})\right)}{{\sum }_{p\in P(i,B)}exp\left(\sigma (x_i,x_p)\right)}.\hfill \end{array}$$

(3)

${\widehat{L}}_{attract}$ is a variant of the SupCon loss, which encourages class separation in embedding space as suggested by Graf et al. [2]. ${\widehat{L}}_{repel}$ is a class-conditional InfoNCE loss, where the positive distribution consists of augmentations and the negative distribution consists of i.i.d samples from the same class. It encourages points within a class to be spread apart, as suggested by the analysis of the InfoNCE loss by Wang and Isola [8].

Qualitative Evaluation

Figure 2 shows t-SNE plots for embeddings produced with $L_{SC}$ versus $L_{spread}$ on the CIFAR10 test set. $L_{spread}$ produces embeddings that are more spread out than those produced by $L_{SC}$ and avoids class collapse. As a result, images from different strata can be better differentiated in embedding space. For example, we show two dogs, one from a common stratum and one from a rare stratum (rare pose). The two dogs are much more distinguishable by distance in the $L_{spread}$ embedding space, which suggests that it helps preserve distinctions between strata.

4. Geometry of Strata

We first discuss some existing theoretical tools for analyzing contrastive loss geometrically and their shortcomings with respect to understanding how strata are embedded. In Section 4.2, we propose a simple thought experiment about the distances between strata in embedding space when trained under a finite subsample of data to better understand our prior qualitative observations. Then, in Section 4.3, we discuss implications of representations that preserve strata distinctions, showing theoretically how they can yield better generalization error on both coarse-to-fine transfer and the original task and empirically how they allow for new downstream applications.

4.1. Existing Analysis

Previous works have studied the geometry of optimal embeddings under contrastive learning [2,8,9], but their techniques cannot analyze strata because strata information is not directly used in the loss function. These works use the infinite encoder assumption, where any distribution on ${\mathbb{S}}^{d-1}$ is realizable by the encoder f applied to the input data. This allows the minimization of the contrastive loss to be equivalent to an optimization problem over probability measures on the hypersphere. As a result, solving this new problem yields a distribution whose characterization is solely determined by information in the loss function (e.g., labels information [2,9]) and is decoupled from other information about the input data x and hence decoupled from strata.

More precisely, if we denote the measure of $x\in \mathcal{X}$ as ${\mu }_{\mathcal{X}}$, minimizing the contrastive loss over the mapping f is equal (at the population level) to minimizing over the pushforward measure ${\mu }_{\mathcal{X}}\circ f^{-1}:{\mathbb{S}}^{d-1}\to [0,1]$. The infinite encoder assumption allows us to relax the problem and instead consider optimizing over any $\mu \in \mathcal{M}\left({\mathbb{S}}^{d-1}\right)$ in the Borel set of probability measures on the hypersphere. Then, the optimal ${\mu }^{★}$ learned is independent of the distribution of the input data $\mathcal{P}$ beyond what is in the relaxed objective function.

This approach using the infinite encoder assumption does not allow for analysis of strata. Strata are unknown at training time and thus cannot be incorporated explicitly into the loss function. Their geometries will not be reflected in the characterization of the optimal distribution obtained from previous theoretical tools. Therefore, we need additional reasoning for our empirical observations that strata distinctions are preserved in embedding space under $L_{spread}$.

4.2. Subsampling Strata

We propose a simple thought experiment based on subsampling the dataset—randomly sampling a fraction of the training data—to analyze strata. Consider the following: we subsample a fraction $t\in [0,1]$ of a training set of N points from $\mathcal{P}$. We use this subsampled dataset ${\mathcal{D}}_t$ to learn an encoder ${\widehat{f}}_t$, and we study the average distance under ${\widehat{f}}_t$ between two strata z and $z^{\prime }$ as t varies.

The average distance between z and $z^{\prime }$ is $\delta ({\widehat{f}}_t,z,z^{\prime })={\parallel {\mathbb{E}}_{x\sim {\mathcal{P}}_z}\left[{\widehat{f}}_t\left(x\right)\right]-{\mathbb{E}}_{x\sim {\mathcal{P}}_{z^{\prime }}}\left[{\widehat{f}}_t\left(x\right)\right]\parallel }_2$ and depends on whether z and $z^{\prime }$ are both in the subsampled dataset. We study when z and $z^{\prime }$ belong to the same class. We have three cases (with probabilities stated in Appendix C.2) based on strata frequency and t—when both, one, or neither of the strata appears in ${\mathcal{D}}_t$:

• Both strata appear in ${\mathcal{D}}_t$ The encoder ${\widehat{f}}_t$ is trained on both z and $z^{\prime }$. For large N, we can approximate this setting by considering ${\widehat{f}}_t$ trained on infinite data from these strata. Points belonging to these strata will be defined in the optimal embedding distribution on the hypersphere, which can be characterized by prior theoretical approaches [2,8,9]. With $L_{spread}$, $\delta ({\widehat{f}}_t,z,z^{\prime })$ depends on $\alpha$, which controls the extent of spread in the embedding geometry. With $L_{SC}$, points from the two strata would asymptotically map to one location on the hypersphere, and $\delta ({\widehat{f}}_t,z,z^{\prime })$ would converge to 0. This case occurs with probability increasing in $p\left(z\right),p\left(z^{\prime }\right),$ and t.
• One stratum but not the other appears in ${\mathcal{D}}_t$ Without loss of generality, suppose that points from z appear in ${\mathcal{D}}_t$ but no points from $z^{\prime }$ do. To understand $\delta ({\widehat{f}}_t,z,z^{\prime })$, we can consider how the end model $\widehat{p}\left(y\right|{\widehat{f}}_t\left(x\right))$ learned using the “source” distribution containing z performs on the “target” distribution of stratum $z^{\prime }$ since this downstream classifier is a function of distances in embedding space. Borrowing from the literature in domain adaptation, the difficulty of this out-of-distribution problem depends on both the divergence between source z and target $z^{\prime }$ distributions and the capacity of the overall model. The $\mathcal{H}\Delta \mathcal{H}$-divergence from Ben-David et al. [10,11], which is studied in lower bounds in Ben-David and Urner [12], and the discrepancy difference from Mansour et al. [13] capture both concepts. Moreover, the optimal geometries of $L_{spread}$ and $L_{SC}$ induce different end model capacities and prediction distributions, with data being more separable under $L_{SC}$, which can help explain why $L_{spread}$ better preserves strata distances. This case occurs with probability increasing in $p\left(z\right)$ and decreasing in $p\left(z^{\prime }\right)$ and t.
• Neither strata appears in ${\mathcal{D}}_t$ The distance $\delta ({\widehat{f}}_t,z,z^{\prime })$ in this case is at most $2D_{\mathrm{TV}}({\mathcal{P}}_z,{\mathcal{P}}_{z^{\prime }})$ (total variation distance) regardless of how the encoder is trained, although differences in transfer from models learned on $\mathcal{Z}\backslash z,z^{\prime }$ to z versus $z^{\prime }$ can be further analyzed. This case occurs with probability decreasing in $p\left(z\right),p\left(z^{\prime }\right),$ and t.

We make two observations from these cases. First, if z and $z^{\prime }$ are both common strata, then as t increases, the distance between them depends on the optimal asymptotic distribution. Therefore, if we set $\alpha =1$ in $L_{spread}$, these common strata will collapse. Second, if z is a common strata and $z^{\prime }$ is uncommon, the second case occurs frequently over randomly sampled ${\mathcal{D}}_t$, and thus the strata are separated based on the difficulty of the respective out-of-distribution problem. We thus arrive at the following insight from our thought experiment:

Common strata are more tightly clustered together, while rarer and more semantically distinct strata are far away from them.

Figure 3 demonstrates this insight. It shows a t-SNE visualization of embeddings from training on CIFAR100 with coarse superclass labels, and with artifically imbalanced subclasses. We show points from the largest subclasses in dark blue and points from the smallest subclasses in light blue. Points from the largest subclasses (dark blue) cluster tightly, whereas points from small subclasses (light blue) are scattered throughout the embedding space.

4.3. Implications

We discuss theoretical and practical implications of our subsampling argument. First, we show that on both the coarse-to-fine transfer task $(x,z)$ and the original task $(x,y)$, embeddings that preserve strata yield better generalization error. Second, we discuss practical implications arising from our subsampling argument that enable new applications.

4.3.1. Theoretical Implications

Consider ${\widehat{f}}_1$, the encoder trained on $\mathcal{D}$ with all N points using $L_{spread}$, and suppose a mean classifier is used for the end model, e.g., $W_y={\mathbb{E}}_{x|y}\left[{\widehat{f}}_1\left(x\right)\right]$ and $W_z={\mathbb{E}}_{x|z}\left[{\widehat{f}}_1\left(x\right)\right]$. On coarse-to-fine transfer, generalization error depends on how far each stratum center is from the others.

Lemma 1.

There exists ${\lambda }_z>0$ such that the generalization error on the coarse-to-fine transfer task is at most

$$\begin{array}{c}\hfill \mathcal{L}(x,z,{\widehat{f}}_1)\le {\mathbb{E}}_z\left[log\left(\sum _{z^{\prime }\in \mathcal{Z}}exp\left(-{\lambda }_z\left(\frac{1}{2}\delta {({\widehat{f}}_1,z,z^{\prime })}^2-1\right)\right)\right)\right]-1,\end{array}$$

(4)

where $\delta ({\widehat{f}}_1,z,z^{\prime })$ is the average distance between strata z and $z^{\prime }$ defined in Section 4.2.

The larger the distances between strata, the smaller the upper bound on generalization error. We now show that a similar result holds on the original task $(x,y)$, but there is an additional term that penalizes points from the same class being too far apart.

Lemma 2.

There exists ${\lambda }_y>0$ such that the generalization error on the original task is at most

$$\begin{array}{cc}\hfill \mathcal{L}(x,y,{\widehat{f}}_1)\le {\mathbb{E}}_z[& {\mathbb{E}}_{z^{\prime }|S\left(z\right)}\left[\frac{1}{2}\delta {({\widehat{f}}_1,z,z^{\prime })}^2-1\right]\hfill \end{array}$$

(5)

$$\begin{array}{cc}\hfill & +log\left(\sum _{y\in \mathcal{Y}}exp\left({\mathbb{E}}_{z^{\prime }|y}\left[-{\lambda }_y\left(\frac{1}{2}\delta {({\widehat{f}}_1,z,z^{\prime })}^2-1\right)\right]\right)\right)].\hfill \end{array}$$

(6)

This result suggests that maximizing distances between strata of different classes is desirable, but less so for distances between strata of the same class as suggested by the first term in the expression. Both results illustrate that separating strata to some extent in the embedding space results in better bounds on generalization error. In Appendix C.3, we provide proofs of these results and derive values of the generalization error for these two tasks under class collapse for comparison.

4.3.2. Practical Implications

Our discussion in Section 4.2 suggests that training with $L_{spread}$ better distinguishes strata in embedding space. As a result, we can use differences between strata of different sizes for downstream applications. For example, unsupervised clustering can help recover pseudolabels for unlabeled, rare strata. These pseudolabels can be used as inputs to worst-group robustness algorithms, or used to detect noisy labels, which appear to be rare strata during training (see Section 5.3 for examples). We can also train over subsampled datasets to heuristically distinguish points that come from common strata from points that come from rare strata. We can then downsample points from common strata to construct minimal coresets (see Section 5.4 for examples).

5. Experiments

This section evaluates $L_{spread}$ on embedding quality and model quality:

• First, in Section 5.2, we use coarse-to-fine transfer learning to evaluate how well the embeddings maintain strata information. We find that $L_{spread}$ achieves lift across four datasets.
• In Section 5.3, we evaluate how well $L_{spread}$ can detect rare strata in an unsupervised setting. We first use $L_{spread}$ to detect rare strata to improve worst-group robustness by up to 2.5 points. We then use rare strata detection to correct noisy labels, recovering 75% performance under 20% noise.
• In Section 5.4, we evaluate how well $L_{spread}$ can distinguish points from large strata versus points from small strata. We downsample points from large strata to construct minimal coresets on CIFAR10, outperforming prior work by 1.0 points at 30% labeled data.
• Finally, in Section 5.5, we show that training with $L_{spread}$ improves model quality, validating our theoretical claims that preventing class collapse can improve generalization error. We find that $L_{spread}$ improves performance in 7 out of 9 cases.

5.1. Datasets and Models

Tabel

Table 1. Summary of the datasets we use for evaluation.

Dataset	Notes
CIFAR10	Standard computer vision dataset
CIFAR10-Coarse	CIFAR10 with animal/vehicle coarse labels
CIFAR100	Standard computer vision dataset
CIFAR100-Coarse	CIFAR100 with standard coarse labels
CIFAR100-Coarse-U	CIFAR100 with standard coarse labels, but with some fine classes
	sub-sampled
MNIST	Standard computer vision dataset
MNIST-Coarse	MNIST with <5 and ≥5 coarse labels
Waterbirds	Robustness dataset mixing up images of birds and their
	backgrounds [14]
ISIC	Images of skin lesions [15]
CelebA	Images of celebrity faces [16]

lists all the datasets we use in our evaluation. CIFAR10, CIFAR100, and MNIST are the standard computer vision datasets. We also use coarse versions of each, wherein classes are combined to create coarse superclasses (animals/vehicles for CIFAR10, standard superclasses for CIFAR100, and <5, ≥5 for MNIST). In CIFAR100-Coarse-U, some subclasses have been artificially imbalanced. Waterbirds, ISIC and CelebA are image datasets with documented hidden strata [5,14,15,16]. We use a ViT model [17] (4 × 4, 7 layers) for CIFAR and MNIST and a ResNet50 for the rest. For the ViT models, we jointly optimize the contrastive loss with a cross entropy loss head. For the ResNets, we train the contrastive loss on its own and use linear probing on the final layer. More details in Appendix E.

5.2. Coarse-to-Fine Transfer Learning

In this section, we use coarse-to-fine transfer learning to evaluate how well $L_{spread}$ retains strata information in the embedding space. We train on coarse superclass labels, freeze the weights, and then use transfer learning to train a linear layer with subclass labels. We use this supervised strata recovery setting to isolate how well the embeddings can recover strata in the optimal setting. For baselines, we compare against training with $L_{SC}$ and the SimCLR loss $L_{SS}$.

Table 2. Performance of coarse-to-fine transfer on various datasets compared against contrastive baselines. In these tasks, we first train a model on coarse task labels, then freeze the representation and train a model on fine-grained subclass labels. $L_{spread}$ produces embeddings that transfer better across all datasets. Best in bold.

	Coarse-to-Fine Transfer
Dataset	${\mathit{L}}_{\mathbf{SS}}$	${\mathit{L}}_{\mathbf{SC}}$	${\mathit{L}}_{\mathbf{spread}}$
CIFAR10-Coarse	71.7	52.5	76.1
CIFAR100-Coarse	62.0	62.4	63.9
CIFAR100-Coarse-U	61.9	59.5	62.4
MNIST-Coarse	97.1	98.8	99.0

reports the results. We find that $L_{spread}$ produces better embeddings for coarse-to-fine transfer learning than $L_{SC}$ and $L_{SS}$. Lift over $L_{SC}$ varies from 0.2 points on MNIST (16.7% error reduction), to 23.6 points of lift on CIFAR10. $L_{spread}$ also produces better embeddings than $L_{SS}$, since $L_{SS}$ does not encode superclass labels in the embedding space.

5.3. Robustness against Worst-Group Accuracy and Noise

In this section, we use robustness to measure how well $L_{spread}$ can recover strata in an unsupervised setting. We use clustering to detect rare strata as an input to worst-group robustness algorithms, and we use a geometric heuristic over embeddings to correct noisy labels.

To evaluate worst-group accuracy, we follow the experimental setup and datasets from Sohoni et al. [5]. We first train a model with class labels. We then cluster the embeddings to produce pseudolabels for hidden strata, which we use as input for a Group-DRO algorithm to optimize worst-group robustness [14]. We use both $L_{SC}$ and cross entropy loss [5] for training the first stage as baselines.

To evaluate robustness against noise, we introduce noisy labels to the contrastive loss head on CIFAR10. We detect noisy labels with a simple geometric heuristic: points with incorrect labels appear to be small strata, so they should be far away from other points of the same class. We then correct noisy points by assigning the label of the nearest cluster in the batch. More details can be found in Appendix E.

Table 3. Unsupervised strata recovery performance (top, F1), and worst-group performance (AUROC for ISIC, Acc for others) using recovered strata. Best in bold.

	Sub-Group Recovery
Dataset	Sohoni et al. [5]	${\mathit{L}}_{\mathbf{SC}}$	${\mathit{L}}_{\mathbf{spread}}$
Waterbirds	56.3	47.2	59.0
ISIC	74.0	92.5	93.8
CelebA	24.2	19.4	24.8
	Worst-Group Robustness
Waterbirds	88.4	86.5	89.0
ISIC	92.0	93.3	92.6
CelebA	55.0	66.1	67.8

shows the performance of unsupervised strata recovery and downstream worst-group robustness. We can see that $L_{spread}$ outperforms both $L_{SC}$ and Sohoni et al. [5] on strata recovery. This translates to better worst-group robustness on Waterbirds and CelebA.

Figure 4 (left) shows the effect of noisy labels on performance. When noisy labels are uncorrected (purple), performance drops by up to 10 points at 50% noise. Applying our geometric heuristic (red) can recover 4.8 points at 50% noise, even without using $L_{spread}$. However, $L_{spread}$ recovers an additional 0.9 points at 50% noise, and an additional 1.6 points at 20% noise (blue). In total, $L_{spread}$ recovers 75% performance at 20% noise, whereas $L_{SC}$ only recovers 45% performance.

5.4. Minimal Coreset Construction

Now we evaluate how well training on fractional samples of the dataset with $L_{spread}$ can distinguish points from large versus small strata by constructing minimal coresets for CIFAR10. We train a ResNet18 on CIFAR10, following Toneva et al. [18], and compare against baselines from Toneva et al. [18] (Forgetting) and Paul et al. [19] (GradNorm, L2Norm). For our coresets, we train with $L_{spread}$ on subsamples of the dataset and record how often points are correctly classified at the end of each run. We bucket points in the training set by how often the point is correctly classified. We then iteratively remove points from the largest bucket in each class. Our strategy removes easy examples first from the largest coresets, but maintains a set of easy examples in the smallest coresets.

Figure 4 (right) shows the results at various coreset sizes. For large coresets, our algorithm outperforms both methods from Paul et al. [19] and is competitive with Toneva et al. [18]. For small coresets, our method outperforms the baselines, providing up to 5.2 points of lift over Toneva et al. [18] at 30% labeled data. Our analysis helps explain this gap; removing too many easy examples hurts performance, since then the easy examples become rare and hard to classify.

5.5. Model Quality

Finally, we confirm that $L_{spread}$ produces higher-quality models and achieves better sample complexity than both $L_{SC}$ and the SimCLR loss $L_{SS}$ from [20].

Table 4. End model performance training with $L_{spread}$ on various datasets compared against contrastive baselines. All metrics are accuracy except for ISIC (AUROC). $L_{spread}$ produces the best performance in 7 out of 9 cases, and matches the best performance in 1 case. Best in bold.

	End Model Perf.
Dataset	${\mathit{L}}_{\mathbf{SS}}$	${\mathit{L}}_{\mathbf{SC}}$	${\mathit{L}}_{\mathbf{spread}}$
CIFAR10	89.7	90.9	91.5
CIFAR10-Coarse	97.7	96.5	98.1
CIFAR100	68.0	67.5	69.1
CIFAR100-Coarse	76.9	77.2	78.3
CIFAR100-Coarse-U	72.1	71.6	72.4
MNIST	99.1	99.3	99.2
MNIST-Coarse	99.1	99.4	99.4
Waterbirds	77.8	73.9	77.9
ISIC	87.8	88.7	90.0

reports the performance of models across all our datasets. We find that $L_{spread}$ achieves better overall performance compared to models trained with $L_{SC}$ and $L_{SS}$ in 7 out of 9 tasks, and matches performance in 1 task. We find up to 4.0 points of lift over $L_{SC}$ (Waterbirds), and up to 2.2 points of lift (AUROC) over $L_{SS}$ (ISIC). In Appendix F, we additionally evaluate the sample complexity of contrastive losses by training on partial subsamples of CIFAR10. $L_{spread}$ outperforms $L_{SC}$ and $L_{SS}$ throughout.

6. Related Work and Discussion

From work in contrastive learning, we take inspiration from [21], who use a latent classes view to study self-supervised contrastive learning. Similarly, [22] considers how minimizing the InfoNCE loss recovers a latent data generating model. We initially started from a debiasing angle to study the effects of noise in supervised contrastive learning inspired by [23], but moved to our current strata-based view of noise instead. Recent work has also analyzed contrastive learning from the information-theoretic perspective [24,25,26], but does not fully explain practical behavior [27], so we focus on the geometric perspective in this paper because of the downstream applications. On the geometric side, we are inspired by the theoretical tools from [8] and [2], who study representations on the hypersphere along with [9].

Our work builds on the recent wave of empirical interest in contrastive learning [20,28,29,30,31] and supervised contrastive learning [1]. There has also been empirical work analyzing the transfer performance of contrastive representations and the role of intra-class variability in transfer learning. [32] find that combining supervised and self-supervised contrastive loss improves transfer learning performance, and they hypothesize that this is due to both inter-class separation and intra-class variability. [33] find that combining cross entropy and self-supervised contrastive loss improves coarse-to-fine transfer, also motivated by preserving intra-class variability.

We derive $L_{spread}$ from similar motivations to losses proposed in these works, and we futher theoretically study why class collapse can hurt downstream performance. In particular, we study why preserving distinctions of strata in embedding space may be important, with theoretical results corroborating their empirical studies. We further propose a new thought experiment for why a combined loss function may lead to better separation of strata.

Our treatment of strata is strongly inspired by [5,6], who document empirical consequences of hidden strata. We are inspired by empirical work that has demonstrated that detecting subclasses can be important for performance [4,34] and robustness [14,35,36].

Each of our downstream applications is a field in itself, and we take inspiration from recent work from each. Our noise heuristic is similar to the ELR [37] and takes inspiration from a various work using contrastive learning to correct noisy labels and for semi-supervised learning [38,39,40]. Our coreset algorithm is inspired by recent work in coresets for modern deep networks [19,41,42], and takes inspiration from [18] in particular.

7. Conclusions

We propose a new supervised contrastive loss function to prevent class collapse and produce higher-quality embeddings. We discuss how our loss function better maintains strata distinctions in embedding space and explore several downstream applications. Future directions include encoding label hierarchies and other forms of knowledge in contrastive loss functions and extending our work to more modalities, models, and applications. We hope that our work inspires further work in more fine-grained supervised contrastive loss functions and new theoretical approaches for reasoning about generalization and strata.