Proceeding Paper

The Details Matter: Preventing Class Collapse in Supervised
Contrastive Learning †Daniel Y. Fu *,‡ , Mayee F. Chen ‡ , Michael Zhang, Kayvon Fatahalian and Christopher Ré
Department of Computer Science, Stanford University, Stanford, CA 94035, USA;mfchen@cs.stanford.edu (M.F.C.); mzhang@cs.stanford.edu (M.Z.); kayvonf@cs.stanford.edu (K.F.);
chrismre@cs.stanford.edu (C.R.)* Correspondence: danfu@cs.stanford.edu† Presented at the AAAI Workshop on Artificial Intelligence with Biased or Scarce Data (AIBSD), Online,
28 February 2022.
‡ These authors contributed equally to this work.

Citation: Fu, D.Y.; Chen, M.F.; Zhang,M.; Fatahalian, K.; Ré, C. The Details

Abstract: Supervised contrastive learning optimizes a loss that pushes together embeddings ofpoints 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 ofthis 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 distinguishstrata in embedding space, which produces lift on three downstream applications: 4.4 points oncoarse-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.

Matter: Preventing Class Collapse inSupervised Contrastive Learning.
Comput. Sci. Math. Forum 2022, 3, 4.

Keywords: contrastive learning; supervised contrastive learning; transfer learning; robustness; noisylabels; coresets

https://doi.org/10.3390/
cmsf2022003004
Academic Editors: Kuan-Chuan Peng
and Ziyan Wu
Published: 15 April 2022
Publisher’s Note: MDPI stays neutral
with regard to jurisdictional claims inpublished maps and institutional affiliations.

Copyright: © 2022 by the authors.
Licensee MDPI, Basel, Switzerland.
This article is an open access article
distributed under the terms andconditions of the Creative Commons
Attribution (CC BY) license (https://creativecommons.org/licenses/by/4.0/).

1. Introduction
Supervised contrastive learning has emerged as a promising method for training deepmodels, 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 LSC [2]. Furthermore, modern
deep networks, which can memorize arbitrary labels [3], are powerful enough to produceclass collapse.
Although class collapse minimizes LSC and produces accurate models, it loses information that is not explicitly encoded in the class labels. For example, consider images withthe 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 ofdata—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, LSC maps the sleeping, jumping, and swatting cats all to a single“cat” embedding, losing strata information. As a result, these embeddings are less usefulComput. Sci. Math. Forum 2022, 3, 4. https://doi.org/10.3390/cmsf2022003004

https://www.mdpi.com/journal/csmf

Comput. Sci. Math. Forum 2022, 3, 4

2 of 27for common downstream applications in the modern machine learning landscape, such astransfer learning.
In this paper, we explore a simple modification to LSC that prevents class collapse.
We study how this modification affects embedding quality by considering how strata arerepresented in embedding space. We evaluate our loss both in terms of embedding quality,
which we evaluate through three downstream applications, and end model quality.

Sleeping

Class
Collapse

Head-OnClasses,
StrataSpreadApart

Higher-Quality
Embeddings

Swatting
Jumping
Supervised Contrastive
Learning
Unlabeled Strata

Transfer Learning

Robustness

New Loss
Function

Higher-Quality
ModelsMinimalCoresets

Applications ofEmbeddings

Figure 1. Classes contain critical information that is not explicitly encoded in the class labels. Supervised contrastive learning (left) loses this information, since it maps unlabeled strata such as
sleeping cats, jumping cats, and swatting cat to a single embedding. We introduce a new loss function
Lspread that prevents class collapse and maintains strata distinctions. Lspread produces higher-quality
embeddings, which we evaluate with three downstream applications.

In Section 3, we present our modification to LSC , which prevents class collapse bychanging how embeddings are pushed and pulled apart. LSC pushes together embeddingsof points from the same class and pulls apart embeddings of points from different classes.
In contrast, our modified loss Lspread 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 Lspread does
not use strata labels, we observe that it still produces embeddings that qualitatively appear
to retain more strata information than those produced by LSC (see Figure 2).
In Section 4, motivated by these empirical observations, we study how well Lspreadpreserves distinctions between strata in the representation space. Previous theoretical toolsthat study the optimal embedding distribution fail to characterize the geometry of strata.
Instead, we propose a simple thought experiment considering the embeddings that thesupervised 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 asmall sample). In particular, we find that points from rarer and more distinct strata areclustered 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 Lspread produces embeddings that retain more information about
strata, resulting in lift on three downstream applications that require strata recovery:•
We evaluate how well Lspread ’s embeddings encode fine-grained subclasses with coarseto-fine transfer learning. Lspread achieves up to 4.4 points of lift across four datasets.
•
We evaluate how well embeddings produced by Lspread can recover strata in an unsupervised setting by evaluating robustness against worst-group accuracy and noisylabels. We use our insights about how Lspread 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 undera fixed training strategy (the coreset problem). We construct coresets by subsamplingComput. Sci. Math. Forum 2022, 3, 4

3 of 27

points from common strata. Our coresets outperform prior work by 1.0 points whencoreset size is 30% of the training set.
Next, we find that Lspread produces higher-quality models, outperforming LSC by up
to 4.0 points across 9 tasks. Finally, we discuss related work in Section 6 and conclude
in Section 7.
0.98 Simw/ CenterAvg Cos.
Sim: 0.88

Plane
Car
Bird
Cat
Deer

0.76 Simw/ Center

Common Stratum

Dog
Frog
Horse
Ship
Truck

0.92 Simw/ Center

LSC (Class Collapse)

0.22 Simw/ Center

Rare Stratum(Rare Pose)

Avg Cos.
Sim: 0.49

Lspread (No Class Collapse)

Figure 2. Lspread produces embeddings that are qualitatively better than those produced by LSC .
We show t-SNE visualizations of embeddings for the CIFAR10 test set and report cosine similarity
metrics (average intracluster cosine similarities, and similarities between individual points and theclass cluster). Lspread produces lower intraclass cosine similarity and embeds images from rare strata
further out over the hypersphere than LSC .

2. BackgroundWe present our generative model for strata (Section 2.1). Then, we discuss supervised
contrastive learning—in particular the SupCon loss LSC 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 D = {( xi , yi )}iN=1 , where ( x, y) ∼ P for x ∈ X andy ∈ Y = {1, . . . , K }. For a particular data point x, we denote its label as h( x ) ∈ Y with
distribution p(y| x ). We assume that data is class-balanced such that p(y = i ) = K1 for alli ∈ Y . The goal is to learn a model p̂(y| x ) on 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 Z = {1, . . . , C }. Z
can be partitioned into disjoint subsets S1 , . . . , SK such that if z ∈ Sk , then its corresponding
y label is equal to k. Let S(c) 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(z). Then, the data point x is sampled from the distribution Pz = p(·|z),and its corresponding label is y = S(z) (see Figure 2 of [5]). We assume that each class
has m strata, and that there exist at least two strata, z1 , z2 , where S(z1 ) 6= S(z2 ) andsupp(z1 ) ∩ supp(z2 ) 6= ∅.
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 : X → Rd . Following previous works, we make three assumptions on theencoder: (1) we restrict the encoder output space to be Sd−1 , the unit hypersphere; (2) we
assume K ≤ 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
Sd−1 is realizable by f ( x ).
2.2.1. SupCon and Collapsed Embeddings
We focus on the SupCon loss LSC from [1]. Denote σ ( x, x 0 ) = f ( x )> f ( x 0 )/τ, whereτ is a temperature hyperparameter. Let B be the set of batches of labeled data on D andComput. Sci. Math. Forum 2022, 3, 4

4 of 27P(i, B) = { p ∈ B\i : h( p) = h(i )} be the points in B with the same label as xi . For an1
anchor xi , the SupCon loss is L̂SC ( f , xi , B) = | P(−i,Blog ∑)| ∑ p∈ P(i,B)

exp(σ ( xi ,x p ))a∈ B\i exp( σ ( xi ,x a )), whereP(i, B) forms positive pairs and B\i forms negative pairs.
The optimal embedding distribution that minimizes LSC has one embedding per class,with the per-class embeddings collectively forming a regular simplex inscribed in thehypersphere Graf et al. [2]. Formally, if h( x ) = i, then f ( x ) = vi for all x ∈ B . {vi }iK=1
makes up the regular simplex, defined by: a) ∑iK=1 vi = 0; b) kvi k2 = 1; and c) ∃cK ∈ R s.t.
vi> v j = cK for i 6= j. We describe this property as class collapse and define the distribution off ( x ) 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 ∈ RK ×d is trained on top of the representations f ( x ) by minimizing cross-entropy loss
over softmax scores. We assume that kWy k2 ≤ 1 for each y ∈ Y . The end model’s empiricalloss can be defined as L̂(W, D) = ∑ xi ∈D − logexp( f ( xi )> Wh( x ) )i∑Kj=1 exp( f ( xi )> Wj ). The model uses softmax

scores constructed with f ( x ) and W to generate predictions p̂(y| x ), which we also write asp̂(y| f ( x )). Finally, the generalization error of the model on P is the expected cross-entropy
between p̂(y| x ) and p(y| x ), namely L( x, y, f ) = Ex,y∼P [− log p̂(y| f ( x ))].
3. Method
We now highlight some theoretical problems with class collapse under our generativemodel of strata (Section 3.1). We then propose and qualitatively analyze the loss function
Lspread (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 ( x )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 ∈ RC×d and fixed f ( x ) (as describedin Section 2.2) on points ( x, z), where we assume the data are class-balanced across z.
Observation 1. Class collapse minimizes L( x, z, f ) if for all x, (1) p(y = h( x )| x ) = 1, meaningthat each x is deterministically assigned to one class, and (2) p(z| x ) = m1 where z ∈ Sh( x) . The
second condition implies that p( x |z) = p( x |y) for all z ∈ Sy , 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 originaltask ( x, y).
Observation 2. Class collapse minimizes L( x, y, f ) if, for all x, p(y = h( x )| x ) = 1. This
contradicts our data model.
Proofs are in Appendix D.1. We also analyze transferability of f on arbitrary new
distributions ( x 0 , y0 ) information-theoretically in Appendix C.1, finding that a one-to-one
encoder obeys the Infomax principle [7] better than collapsed embeddings on ( x 0 , y0 ).
These observations suggest that a distribution over the embeddings that preserves strata
distinctions and does not collapse classes is more desirable.

Comput. Sci. Math. Forum 2022, 3, 4

5 of 27

3.2. Modified Contrastive Loss Lspread
We introduce the loss Lspread , a weighted sum of two contrastive losses L attract and
Lrepel . L attract is a supervised contrastive loss, while Lrepel encourages intra-class separation.
For α ∈ [0, 1],

Lspread = αL attract + (1 − α) Lrepel .

(1)

aug

For a given anchor xi , define xi as an augmentation of the same point as x. Define
the set of negative examples for i to be N (i, B) = { a ∈ B\i : h( a) 6= h(i )}. Then,

L̂ attract ( f , xi , B) =exp(σ ( xi , x p ))−1
× ∑ log| P(i, B)| p∈ P(i,B) exp(σ( xi , x p )) + ∑ a∈ N (i,B) exp(σ( xi , x a ))

(2)

aug

exp(σ( xi , xi )).
L̂repel ( f , xi , B) = − log∑ p∈ P(i,B) exp(σ ( xi , x p ))(3)

L̂ attract is a variant of the SupCon loss, which encourages class separation in embedding
space as suggested by Graf et al. [2]. L̂repel is a class-conditional InfoNCE loss, where thepositive distribution consists of augmentations and the negative distribution consists ofi.i.d samples from the same class. It encourages points within a class to be spread apart, assuggested by the analysis of the InfoNCE loss by Wang and Isola [8].
Qualitative Evaluation
Figure 2 shows t-SNE plots for embeddings produced with LSC versus Lspread on
the CIFAR10 test set. Lspread produces embeddings that are more spread out than those
produced by LSC and avoids class collapse. As a result, images from different strata canbe better differentiated in embedding space. For example, we show two dogs, one from acommon stratum and one from a rare stratum (rare pose). The two dogs are much more
distinguishable by distance in the Lspread 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 ourprior 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 notdirectly used in the loss function. These works use the infinite encoder assumption, where
any distribution on Sd−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 overprobability 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 ∈ X as µX , minimizing the contrastive
loss over the mapping f is equal (at the population level) to minimizing over the pushfor-

Comput. Sci. Math. Forum 2022, 3, 4

6 of 27

ward measure µX ◦ f −1 : Sd−1 → [0, 1]. The infinite encoder assumption allows us to relax
the problem and instead consider optimizing over any µ ∈ M(Sd−1 ) in the Borel set ofprobability measures on the hypersphere. Then, the optimal µ? learned is independent of
the distribution of the input data P beyond what is in the relaxed objective function.
This approach using the infinite encoder assumption does not allow for analysis ofstrata. 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 needadditional reasoning for our empirical observations that strata distinctions are preserved inembedding space under Lspread .
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 ∈ [0, 1] of a training set of N points from P . We use this subsampleddataset Dt to learn an encoder fˆt , and we study the average distance under fˆt between two
strata z and z0 as t varies.
The average distance between z and z0 is δ( fˆt , z, z0 ) = kEx∼Pz [ fˆt ( x )] − Ex∼Pz0 [ fˆt ( x )]k2
and depends on whether z and z0 are both in the subsampled dataset. We study when z andz0 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 Dt :1.
Both strata appear in Dt The encoder fˆt is trained on both z and z0 . For large N, we
can approximate this setting by considering fˆt trained on infinite data from these strata.

2.

3.

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 Lspread , δ( fˆt , z, z0 ) depends on α, which controls the extent of spread in theembedding geometry. With LSC , points from the two strata would asymptoticallymap to one location on the hypersphere, and δ( fˆt , z, z0 ) would converge to 0. Thiscase occurs with probability increasing in p(z), p(z0 ), and t.
One stratum but not the other appears in Dt Without loss of generality, suppose thatpoints from z appear in Dt but no points from z0 do. To understand δ( fˆt , z, z0 ), we
can consider how the end model p̂(y| fˆt ( x )) learned using the “source” distribution
containing z performs on the “target” distribution of stratum z0 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 z0 distributions and the capacityof the overall model. The H ∆H-divergence from Ben-David et al. [10,11], which isstudied 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
Lspread and LSC induce different end model capacities and prediction distributions,with data being more separable under LSC , which can help explain why Lspread better
preserves strata distances. This case occurs with probability increasing in p(z) anddecreasing in p(z0 ) and t.
Neither strata appears in Dt The distance δ( fˆt , z, z0 ) in this case is at most 2DTV (Pz , Pz0 )(total variation distance) regardless of how the encoder is trained, although differencesin transfer from models learned on Z \z, z0 to z versus z0 can be further analyzed. Thiscase occurs with probability decreasing in p(z), p(z0 ), and t.

We make two observations from these cases. First, if z and z0 are both common
strata, then as t increases, the distance between them depends on the optimal asymptotic
distribution. Therefore, if we set α = 1 in Lspread , these common strata will collapse.
Second, if z is a common strata and z0 is uncommon, the second case occurs frequently
over randomly sampled Dt , and thus the strata are separated based on the difficulty of theComput. Sci. Math. Forum 2022, 3, 4

7 of 27

respective out-of-distribution problem. We thus arrive at the following insight from ourthought experiment:
Common strata are more tightly clustered together, while rarer and more semanticallydistinct 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 thesmallest subclasses in light blue. Points from the largest subclasses (dark blue) cluster
tightly, whereas points from small subclasses (light blue) are scattered throughout theembedding space.

Subclass Size
500 Points
250 Points
100 Points
50 Points

Figure 3. Points from large subclasses cluster tightly; points from small subclasses scatter (CIFAR100Coarse, unbalanced subclasses).

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 fˆ1 , the encoder trained on D with all N points using Lspread , and supposehihia mean classifier is used for the end model, e.g., Wy = Ex|y fˆ1 ( x ) and Wz = Ex|z fˆ1 ( x ) .
On coarse-to-fine transfer, generalization error depends on how far each stratum center isfrom the others.
Lemma 1. There exists λz > 0 such that the generalization error on the coarse-to-fine transfer task
is at most




1
0 2
ˆ
ˆ
− 1,
(4)
L( x, z, f 1 ) ≤ Ez log ∑ exp − λz δ( f 1 , z, z ) − 12z0 ∈Z
where δ( fˆ1 , z, z0 ) is the average distance between strata z and z0 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 anadditional term that penalizes points from the same class being too far apart.
Lemma 2. There exists λy > 0 such that the generalization error on the original task is at most1L( x, y, fˆ1 ) ≤ Ez Ez0 |S(z) δ( fˆ1 , z, z0 )2 − 1
2


h
1
i
+ log ∑ exp Ez0 |y − λy δ( fˆ1 , z, z0 )2 − 1
.
2
y∈Y

(5)
(6)

Comput. Sci. Math. Forum 2022, 3, 4

8 of 27

This result suggests that maximizing distances between strata of different classes isdesirable, but less so for distances between strata of the same class as suggested by the firstterm in the expression. Both results illustrate that separating strata to some extent in theembedding 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 twotasks under class collapse for comparison.
4.3.2. Practical Implications
Our discussion in Section 4.2 suggests that training with Lspread 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 recoverpseudolabels for unlabeled, rare strata. These pseudolabels can be used as inputs to worstgroup 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 comefrom 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 Lspread on embedding quality and model quality:
First, in Section 5.2, we use coarse-to-fine transfer learning to evaluate how well theembeddings maintain strata information. We find that Lspread achieves lift across four
datasets.
In Section 5.3, we evaluate how well Lspread can detect rare strata in an unsupervised
setting. We first use Lspread 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 Lspread can distinguish points from large strata
versus points from small strata. We downsample points from large strata to constructminimal coresets on CIFAR10, outperforming prior work by 1.0 points at 30% labeled
data.
Finally, in Section 5.5, we show that training with Lspread improves model quality, validating our theoretical claims that preventing class collapse can improve generalization
error. We find that Lspread improves performance in 7 out of 9 cases.

5.1. Datasets and Models
Tabel 1 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–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.

Comput. Sci. Math. Forum 2022, 3, 4

9 of 27

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

CIFAR10CIFAR10-Coarse
CIFAR100
CIFAR100-Coarse
CIFAR100-Coarse-U

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

5.2. Coarse-to-Fine Transfer Learning
In this section, we use coarse-to-fine transfer learning to evaluate how well Lspreadretains 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 canrecover strata in the optimal setting. For baselines, we compare against training with LSC
and the SimCLR loss LSS .
Table 2 reports the results. We find that Lspread produces better embeddings for coarseto-fine transfer learning than LSC and LSS . Lift over LSC varies from 0.2 points on MNIST(16.7% error reduction), to 23.6 points of lift on CIFAR10. Lspread also produces better
embeddings than LSS , since LSS does not encode superclass labels in the embedding space.
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 representationand train a model on fine-grained subclass labels. Lspread produces embeddings that transfer better
across all datasets. Best in bold.
Coarse-to-Fine Transfer
DatasetLSS

LSC

Lspread

CIFAR10-Coarse
CIFAR100-Coarse
CIFAR100-Coarse-U
MNIST-Coarse

71.7
62.0
61.9
97.1

52.5
62.4
59.5
98.8

76.1
63.9
62.4
99.0

5.3. Robustness Against Worst-Group Accuracy and NoiseIn this section, we use robustness to measure how well Lspread can recover strata
in an unsupervised setting. We use clustering to detect rare strata as an input to worstgroup 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 datasetsfrom 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 LSC 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 withincorrect labels appear to be small strata, so they should be far away from other points ofthe same class. We then correct noisy points by assigning the label of the nearest cluster inthe batch. More details can be found in Appendix E.

Comput. Sci. Math. Forum 2022, 3, 4

10 of 27

Table 3 shows the performance of unsupervised strata recovery and downstream worstgroup robustness. We can see that Lspread outperforms both LSC 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 areuncorrected (purple), performance drops by up to 10 points at 50% noise. Applying ourgeometric heuristic (red) can recover 4.8 points at 50% noise, even without using Lspread .
However, Lspread recovers an additional 0.9 points at 50% noise, and an additional 1.6 points
at 20% noise (blue). In total, Lspread recovers 75% performance at 20% noise, whereas LSC
only recovers 45% performance.
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
DatasetSohoni et al. [5]LSC

Lspread

56.3
74.0
24.2

47.2
92.5
19.4

59.0
93.8
24.8

Waterbirds
ISIC
CelebA

Worst-Group Robustness
88.4
92.0
55.0

Waterbirds
ISIC
CelebA

86.5
93.3
66.1

Coreset Performance, CIFAR10

Noisy Label Performance, CIFAR10

Acc

90

96

86
82
78
0%

89.0
92.6
67.8

92
Correction, Lspread
Correction, LSC
No Correction
20%
Noise Rate

88
40%

100%

Ours
Forgetting
GradNorm
L2NormRandom60%
Coreset Size

20%

Figure 4. (Left) Performance of models under various amounts of label noise for the contrastive
loss head. (Right) Performance of a ResNet18 trained with coresets of various sizes. Our coreset
algorithm is competitive with the state-of-the-art in the large coreset regime (from 40–90% coresets),
but maintains performance for small coresets (smaller than 40%). At the 10% coreset, our algorithm
outperforms [18] by 32 points and matches random sampling.

5.4. Minimal Coreset Construction
Now we evaluate how well training on fractional samples of the dataset with Lspreadcan 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 Lspread 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 thelargest 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

Comput. Sci. Math. Forum 2022, 3, 4

11 of 27

too many easy examples hurts performance, since then the easy examples become rare andhard to classify.
5.5. Model Quality
Finally, we confirm that Lspread produces higher-quality models and achieves bettersample complexity than both LSC and the SimCLR loss LSS from [20]. Table 4 reports theperformance of models across all our datasets. We find that Lspread achieves better overall
performance compared to models trained with LSC and LSS in 7 out of 9 tasks, and matches
performance in 1 task. We find up to 4.0 points of lift over LSC (Waterbirds), and up to
2.2 points of lift (AUROC) over LSS (ISIC). In Appendix F, we additionally evaluate thesample complexity of contrastive losses by training on partial subsamples of CIFAR10.
Lspread outperforms LSC and LSS throughout.
Table 4. End model performance training with Lspread on various datasets compared against contrastive baselines. All metrics are accuracy except for ISIC (AUROC). Lspread produces the best
performance in 7 out of 9 cases, and matches the best performance in 1 case. Best in bold.
End Model Perf.
DatasetLSS

LSC

Lspread

CIFAR10CIFAR10-Coarse
CIFAR100
CIFAR100-Coarse
CIFAR100-Coarse-U
MNIST
MNIST-Coarse
Waterbirds
ISIC

89.7
97.7
68.0
76.9
72.1
99.1
99.1
77.8
87.8

90.9
96.5
67.5
77.2
71.6
99.3
99.4
73.9
88.7

91.5
98.1
69.1
78.3
72.4
99.2
99.4
77.9
90.0

6. Related Work and Discussion
From work in contrastive learning, we take inspiration from [21], who use a latentclasses 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–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 thehypersphere along with [9].
Our work builds on the recent wave of empirical interest in contrastive learning [20,28–31]
and supervised contrastive learning [1]. There has also been empirical work analyzing thetransfer 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 bothinter-class separation and intra-class variability. [33] find that combining cross entropy
and self-supervised contrastive loss improves coarse-to-fine transfer, also motivated bypreserving intra-class variability.
We derive Lspread from similar motivations to losses proposed in these works, andwe 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 separationof strata.

Comput. Sci. Math. Forum 2022, 3, 4

12 of 27

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 inspirationfrom 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 forsemi-supervised learning [38–40]. Our coreset algorithm is inspired by recent work incoresets 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 incontrastive 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.
Author Contributions: Conceptualization, D.Y.F. and M.F.C.; methodology, D.Y.F. and M.F.C.; software, D.Y.F.; validation, D.Y.F. and M.Z.; formal analysis, M.F.C.; investigation, D.Y.F., M.F.C. and
M.Z.; resources, D.Y.F. and M.F.C.; data curation, D.Y.F.; writing—original draft preparation, D.Y.F.,
M.F.C. and M.Z.; writing—review and editing, D.Y.F., M.F.C. and M.Z.; visualization, D.Y.F.; supervision, K.F. and C.R.; project administration, D.Y.F. and M.F.C.; funding acquisition, D.Y.F. and M.F.C.
All authors have read and agreed to the published version of the manuscript.
Funding: This research was funded by NIH under No. U54EB020405 (Mobilize), NSF under Nos.
CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML); ONR under
No. N000141712266 (Unifying Weak Supervision); ONR N00014-20-1-2480: Understanding andApplying Non-Euclidean Geometry in Machine Learning; N000142012275 (NEPTUNE); the Moore
Foundation, NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi,
BASF, Accenture, Ericsson, Qualcomm, Analog Devices, the Okawa Foundation, American Family
Insurance, Google Cloud, Salesforce, Total, the HAI-GCP Cloud Credits for Research program, theStanford Data Science Initiative (SDSI), Department of Defense (DoD) through the National Defense
Science and Engineering Graduate Fellowship (NDSEG) Program, and members of the Stanford
DAWN project: Facebook, Google, and VMWare. The Mobilize Center is a Biomedical Technology
Resource Center, funded by the NIH National Institute of Biomedical Imaging and Bioengineering
through Grant P41EB027060. The U.S. Government is authorized to reproduce and distribute reprints
for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings,
and conclusions or recommendations expressed in this material are those of the authors and do notnecessarily reflect the views, policies, or endorsements, either expressed or implied, of NIH, ONR, or
the U.S. Government.
Institutional Review Board Statement: Not Applicable.
Informed Consent Statement: Not Applicable.
Data Availability Statement: Datasets used in this paper are publicly available and describedin Appendix E.
Acknowledgments: We thank Nimit Sohoni for helping with coreset and robustness experiments,
and we thank Beidi Chen and Tri Dao for their helpful comments.
Conflicts of Interest: The authors declare no conflict of interest. The funders had no role in the designof the study; in the collection, analyses, or interpretation of data; in the writing of the manuscript, orin the decision to publish the results.

Comput. Sci. Math. Forum 2022, 3, 4

13 of 27

We provide a glossary in Appendix A. Then we provide definitions of terms in
Appendix B. We discuss additional theoretical results in Appendix C. We provide proofs
in Appendix D. We discuss additional experimental details in Appendix E. Finally, we
provide additional experimental results in Appendix F.
Appendix A. Glossary
The glossary is given in Table A1 below.
Table A1. Glossary of variables and symbols used in this paper.
SymbolUsed for

LSC
LspreadxyDh( x )

SupCon (see Section 2.2), a supervised contrastive loss introduced by [1].
Our modified loss function defined in Section 3.2.
Input data x ∈ X .
Class label y ∈ Y = {1, . . . , K }.
Dataset of N points {( xi , yi )}iN=1 drawn i.i.d. from P .
The class that x belongs to, i.e., h( x ) is a label drawm from p(y| x ). This label
information is used as input in the supervised contrastive loss.
The end model’s predicted distribution over y given x.
A stratum is a latent variable z ∈ Z = {1, . . . , C } that further categorizes data
beyond labels.
The set of all strata corresponding to label k (deterministic).
The label corresponding to strata c (deterministic).
The distribution of input data belonging to stratum z, i.e., x ∼ p(·|z).
The number of strata per class.
Dimension of the embedding space.
The encoder f : X → Rd maps input data to an embedding space and is learned byminimizing the contrastive loss function.
The unit hypersphere, formally {v ∈ Rd : kvk2 = 1}.
Temperature hyperparameter in contrastive loss function.

p̂(y| x )
z
Sk
S(c)
Pz
m
d
f

Sd −1
τ
σ( x, x 0 )BP(i, B){vi }iK=1
W

L̂(W, D)L( x, y, f )L attract
Lrepelαx aug
N (i, B)tDtfˆt
δ( fˆt , z, z0 )f ( x )> f ( x 0 )Notation for.
τ
Set of batches of labeled data on D .
Points in B with the same label as xi , formally { p ∈ B\i : h( p) = h(i )}.
A regular simplex inscribed in the hypersphere (see Definition A1).
The weight matrix that parametrizes the downstream linear classifier(end model) learned on f ( x ).
The empirical cross entropy loss used to learn W over dataset D (see (A1)).
The generalization error of the end model of predicting output y on x using
encoder f (see (A2) and (A3)).
A variant on SupCon that is used in Lspread that pushes points of a class together(see (2)).
A class-conditional InfoNCE loss that is used in Lspread to pull apart points
within a class (see (3)).
Hyperparameter α ∈ [0, 1] controls how to balance L attract and Lrepel .
An augmentation of data point x.
Points in B with a label different from that of xi , formally { a ∈ B\i : h( a) 6= h(i )}.
Fraction of training data t ∈ [0, 1] that is varied in our thought experiment.
Randomly sampled dataset from P with size equal to t · N fraction of D .
Encoder trained on sampled dataset Dt .
The distance between centers of strata z and z0 under encoder fˆt ,namely δ( fˆt , z, z0 ) = kEx∼Pz [ fˆt ( x )] − Ex∼Pz0 [ fˆt ( x )]k2 .

Appendix B. Definitions
We restate definitions used in our proofs.
Definition A1 (Regular Simplex). The points {vi }iK=1 form a regular simplex inscribed in thehypersphere if1.

∑iK=1 vi = 0

Comput. Sci. Math. Forum 2022, 3, 4

14 of 27

2.
3.

kvi k = 1 for all i∃cK ≤ 1 s.t. vi> v j = cK for i 6= j

Definition A2 (Downstream model). Once an encoder f ( x ) is learned, the downstream model
consists of a linear classifier trained using the cross-entropy loss:L̂(W, D) = ∑ − logxi ∈Dexp( f ( xi )> Wh( xi ) )∑Kj=1 exp( f ( xi )> Wj )

.

(A1)

Define Ŵ := argminkW k2 ≤1 L̂(W, D). Then, the end model’s outputs are the probabilitiesp̂(y| x ) = p̂(y| f ( x )) =exp( f ( x )> Ŵy )∑Kj=1 , exp( f ( x )> Ŵj )

,

(A2)

and the generalization error isL( x, y, f ) = Ex,y [− log p̂(y| f ( x ))].

(A3)

Appendix C. Additional Theoretical Results
Appendix C.1. Transfer Learning on ( x 0 , y0 )We now show an additional transfer learning result on new tasks ( x 0 , y0 ). Formally,recall that we learn the encoder f on ( x, y) ∼ P . We wish to use it on a new task with target
distribution ( x 0 , y0 ) ∼ P 0 . We find that an injective encoder f ( x ) is more appropriate to be
used on new distributions than collapsed embeddings based on the Infomax principle [7].
Observation A1. Define f c (y) as the mapping to collapsed embeddings and f 1−1 ( x ) as an injective
mapping, both learned on P . Construct a new variable ye with joint distribution ( x 0 , ye) ∼ p(y| x ) ·p0 ( x 0 ) and suppose that ye ⊥⊥ y0 | x 0 . Then, by the data processing inequality, it holds that I (ye, y0 ) ≤00
I ( x , y ) where I (·, ·) is the mutual information between two random variables. We apply f c to ye
and f 1−1 to x 0 to get thatI ( f c (ye), y0 ) ≤ I ( f 1−1 ( x 0 ), y0 ).
Therefore, f 1−1 obeys the Infomax principle [7] better on P 0 than f c . Via Fano’s inequality,
this statement implies that the Bayes risk for learning y0 from x 0 is lower using f 1−1 than f c .
Appendix C.2. Probabilities of Strata z, z0 Appearing in Subsampled Dataset
As discussed in Section 4.2, the distance between strata z and z0 in embedding spacedepends on if these strata appear in the subsampled dataset Dt that the encoder was trainedon. We define the exact probabilities of the three cases presented. Let Pr(z, z0 ∈ Dt ) be theprobability that both strata are seen, Pr(z ∈ Dt , z0 ∈/ Dt ) be the probability that only z isseen, and Pr(z, z0 ∈/ Dt be the probability that neither are seen.
First, the probability of neither strata appearing in Dt is easy to compute. In particular,
we have that Pr(z, z0 ∈/ Dt ) = (1 − p(z) − p(z0 ))tN . This quantity decreases in p(z) and0p(z ), confirming that it is less likely for two common strata to not appear in Dt .
Second, the probability of z being in Dt and z0 not being in Dt can be expressed
as Pr(z ∈ Dt |z0 ∈/ Dt ) · Pr(z0 ∈/ Dt ). Pr(z0 ∈/ Dt ) is equal to (1 − p(z0 ))tN , and Pr(z ∈00
Dt |z ∈/ Dt ) = 1 − Pr(z ∈/ Dt |z ∈/ Dt ) = 1 − (1 − p(z|z ∈ Z \z0 ))tN . Finally, notep(z)that p(z|z ∈ Z \z0 ) = 1− p(z0 ) . Putting this together, we get that Pr(z ∈ Dt , z0 ∈/ Dt ) =(1 − p(z0 ))tN − (1 − p(z0 ) − p(z))tN , and we can similarly construct Pr(z0 ∈ Dt , z ∈/ D t ).

Comput. Sci. Math. Forum 2022, 3, 4

15 of 27

This quantity depends on the difference between p(z) and p(z0 ), so this case is common
when one stratum is common and one is rare.
Lastly, the probability of both z and z0 being in Dt is thus Pr(z, z0 ∈ Dt ) = 1 − Pr(z, z0 ∈// Dt ) − Pr(z ∈ Dt , z0 ∈/ Dt ) = 1 + (1 − p(z0 ) − p(z))tN − (1 −Dt ) − Pr(z0 ∈ Dt , z ∈p(z0 ))tN − (1 − p(z))tN . This quantity increases in p(z) and p(z0 ).
Appendix C.3. Performance of Collapsed Embeddings on Coarse-to-Fine Transfer and Original Task
Lemma A1. Denote f c to be the encoder that collapses embeddings such that f c ( x ) = vy for any( x, y) ∼ P . Then, the generalization error on the coarse-to-fine transfer task using f c and a linear
classifier learned using cross entropy loss is at leastL( x, z, f c ) ≥ log(m exp(1) + (C − m) exp(cK ) − 1,
where cK is the dot product of any two different class-collapsed embeddings. The generalization
error on the original task under the same setup is at leastL( x, y, f c ) ≥ log(exp(1) + (K − 1) exp(cK )) − 1.
Proof. We first bound generalization error on the coarse-to-fine transfer task. For collapsed
embeddings, f ( x ) = vi when h( x ) = i, where h( x ) is information available at training
time that follows the distribution p(y| x ). We thus denote the embedding f ( x ) as vh( x) .
Therefore, we write the generalization error with an expectation over h( x ) and factorize
the expectation according to our generative model.
CK

Ex,z,h(x) [− log p̂(z| f ( x ))] = − ∑

∑

Z

C

K

p( x, z, h( x )) log p̂(z|h( x ))dxz=1 h( x )=1

=−∑

∑

Z

C

K

Z

p(z) p( x |z) p(h( x )| x ) log p̂(z|h( x ))dxz=1 h( x )=1=−∑∑p(z) p( x |z) p(h( x )| x ) log

z=1 h( x )=1

C= ∑ p(z)Ex∼Pz
z =1

"

K

exp( f h>( x) Wz )∑iC=1 exp( f h>( x) Wi )#C

∑ p(y|x) − v>y Wz + log ∑ exp(v>y Wi )

y =1

dx



.

i =1

Furthermore, since the W learned over collapsed embeddings satisfies Wz = vy forS(z) = y, we have that log ∑iC=1 exp(v>y Wi ) = m exp(1) + (C − m ) exp( cK ) for any y, andour expected generalization error isC∑ p(z)Ex∼Pz [− p(y = S(z)|x) − p(y 6= S(z)|x)δ + log(m exp(1) + (C − m) exp(cK ))]

z =1

C

= log(m exp(1) + (C − m) exp(cK )) − cK − (1 − cK ) ∑ p(z)Ex∼Pz [ p(y = S(z)| x )].
z =1This tells us that the generalization error is at most log(m exp(1) + (C − m) exp(cK )) −cK and at least log(m exp(1) + (C − m) exp(cK )) − 1.
For the original task, we can apply this same approach to the case where m = 1, C = K
to get that the average generalization error ishiEh(x) L( x, y, fˆ1 ) = log(exp(1) + (K − 1) exp(cK ))C− cK − (1 − cK ) ∑ p(z)Ex∼Pz [ p(y = S(z)| x )].
z =1

Comput. Sci. Math. Forum 2022, 3, 4

16 of 27

This is at least log(exp(1) + (K − 1) exp(cK )) − 1 and at most log(exp(1) + (K −1) exp(cK )) − cK .
Appendix D. Proofs
Appendix D.1. Proofs for Theoretical Motivation
We provide proofs for Section 3.1. First, we characterize the optimal linear classifier(for both the coarse-to-fine transfer task and the original task) learned on the collapsed
embeddings. Note that this result appears similar to Corollary 1 of [2], but their result
minimizes the cross entropy loss over both the encoder and downstream weights (i.e., in aclassical supervised setting where only cross entropy is used in training).
Lemma A2 (Downstream linear classifier for coarse-to-fine task). Suppose the dataset Dz isclass-balanced across z, and the embeddings satisfy f ( x ) = vi if h( x ) = i where {vi }iK=1 form theregular simplex. Then the optimal weight matrix W ? ∈ RC×d that minimizes L̂(W, Dz ) satisfies
Wz? = vy for y = S(z).
Proof. Formally, the convex optimization problem we are solving isKminimize − ∑ ∑ log
y =1 z ∈ Sy

exp(v>y Wz )∑Cj=1 exp(v>y Wj )

s.t. kWz k22 ≤ 1 ∀z ∈ Z

(A4)
(A5)

The Lagrangian of this optimization problem isK∑ ∑y =1 z ∈ Sy−v>
y Wz + m

K

 C

y =1

j =1

∑ log ∑

exp(v>y Wj )C+ ∑ λi (kWi k22 − 1),i =1

and the stationarity condition w.r.t. Wz isK− vS(z) + m ∑vy exp(v>y Wz )

C>
y=1 ∑ j=1 exp( vy Wj )Substituting Wz = vS(z) , we get −vS(z) + m ∑yK=1+ 2λz Wz = 0.

(A6)vy exp(v>y vS(z) )+ 2λz vS(z) = 0. Using∑Cj=1 exp(v>y vS( j) )v
exp(1)+exp(δ) ∑y6=S(z) vy
the fact that vi> v j = δ for all i 6= j, this equals −vS(z) + m · S(zm) exp(1)+(C−m) exp+(δ)

exp(1)−exp(δ)
K
1
2λz vS(z) = 0. Next, recall that ∑i=1 vi = 0. Then, λz = 2 1 − m · m exp(1)+(C−m) exp(δ) ≥ 0,
satisfying the dual constraint. We can further verify complementary slackness and primal
feasibility, since kWz? k22 = 1, to confirm that an optimal weight matrix satisfies Wz? = vy for
y = S ( z ).
Corollary A1. When we apply the above proof to the case when m = 1, we recover that the optimal
weight matrix W ? ∈ RK ×d that minimizes L̂(W, D) for the original task on ( x, y) ∼ P satisfies
Wy? = vy for all y ∈ Y .
We now prove Observation 1 and 2. Then, we present an additional result on transfer
learning on collapsed embeddings to general tasks of the form ( x 0 , y0 ) ∼ P 0 .

Comput. Sci. Math. Forum 2022, 3, 4

17 of 27

Proof of Observation 1. We write out the generalization error for the downstream task,L( x, z, f ) = Ex,z [− log p̂(z| x )] using our conditions that p(y = h( x )| x ) = 1 and p(z| x ) = m1 .

L( x, z, f ) = −ZCp( x ) ∑ p(z| x ) log p̂(z| f ( x ))dx
z =1

=−

Z

C

exp( f ( x )> Wz )z =1

∑iC=1 exp( f ( x )> Wi )p( x ) ∑ p(z| x ) log

K Z

=−∑

y=1 x:h( x )=yp( x ) ·dx

1
exp( f ( x )> Wz )
.
log C
∑
m z∈S
∑i=1 exp( f ( x )> Wi )y

To minimize this, f ( x ) should be the same across all x where h( x ) is the same value,
since p(z| x ) does not change across fixed h( x ) and thus varying f ( x ) will not further
decrease the value of this expression. Therefore, we rewrite f ( x ) as f h( x) . Using the factthat y is class balanced, our loss is now

L( x, y, z) = −

1 K
∑
m y∑
=1 z ∈ S

y=−Z
x:h( x )=yp( x ) log

exp( f h>( x) Wz )∑iC=1 exp( f h>( x) Wz )dxexp( f y> Wz )
1 K
.
log
∑
C y∑
∑iC=1 exp( f y> Wi )=1 z ∈ Sy

We claim that f y = vy and Wz = vy for all S(z) = y minimizes this convex function.
The corresponding Lagrangian isK∑ ∑y =1 z ∈ Sy

− f y> Wz + m

K

 C

y =1

i =1

∑ log ∑

exp( f y> Wi )



K

C

y =1

i =1

+ ∑ νy (k f y k22 − 1) + ∑ λi (kWi k22 − 1).

The stationarity condition with respect to Wz is the same as (A6), and we have alreadydemonstrated that the feasibility constraints and complementary slackness are satisfied on
W. The stationarity condition with respect to f y is

− ∑ Wz + m ·z ∈ Sy∑iC=1 Wi exp( f y> Wi )∑iC=1 exp( f y> Wi )+ 2λy f y = 0.

Substituting in Wi = vS(i) and f y = vy , we get − ∑z∈Sy vy + m ·∑iC=1 vS(i) exp(v>y v S (i ) )+∑iC=1 exp(v>y v S (i ) )mvy exp(1)−mvy exp(δ)
2λy vy = 0. From the regular simplex definition, this is −mvy + m m exp(1)+(C−m) exp(δ) +m(exp(1)−exp(δ))
2λy vy = 0. We thus have that λy = m2 1 − m exp(1)+(C−m) exp(δ) , and the feasibility

constraints are satisfied. Therefore, f y = Wz = vy for y = S(z) minimizes the generalization
error L( x, z, f ) when p(h( x )| x ) = 1 and p(z| x ) = m1 .
RRp(z| x ) = m1 and p(y = h( x )| x ) = 1, so p(z)= x:h( x)=S(z) p(z, x )dx= m1 x:h( x)=S(z) p( x )p(z| x ) p( x )11= Kp( x ) =mK = C . p ( z ) being class balanced means that p ( x | z ) =p(z)p(y| x ) p( x )= p( x |y). Therefore, this condition suggests that there is no distinction amongp(y)=the strata within a class.
Proof of Observation 2. This observation follows directly from Observation 1 by repeating
the proof approach with z = y, m = 1.
Lastly, suppose it is not true that pR(y = h( x )| x ) = 1. Then, the generalization error
on the original task is L( x, y, f ) = − X ∑yK=1 p( x ) p(y| x ) log p̂(y| f ( x )), which is mini-Comput. Sci. Math. Forum 2022, 3, 4

18 of 27

mized when p̂(y| f ( x )) = p(y| x ). Intuitively, a model constructed with label information,p̂(y|h( x )), will not improve over one that uses x itself to approximate p(y| x ).
Appendix D.2. Proofs for Theoretical Implications
We provide proofs for Section 4.3.
Proof of Lemma 1. The generalization error is##""exp( fˆ1 ( x )> Wz )ˆL( x, z, f 1 ) = −Ez Ex∼Pz log C∑i=1 exp( fˆ1 ( x )> Wi )
"
"
##
C
>
>
= Ez Ex∼P − fˆ1 ( x ) Wz + log ∑ exp( fˆ1 ( x ) Wi ) .
z

i =1

Using the definition of the mean classifier,
"

#
C

>
ˆ
ˆ
ˆ
L( x, z, f 1 ) = Ez −1 + Ex∼Pz log ∑ exp f 1 ( x ) Ex∼Pi [ f 1 ( x )]
i =1

"

C



= −1 + Ez Ex∼Pz log ∑ exp fˆ1 ( x )> Ei [ fˆ1 ( x )]



#
.

i =1

Since fˆ1 ( x ) is bounded, there exists a constant λ > 0 such that
h
 C
C
i

Ex∼Pz log ∑ exp fˆ1 ( x )> Ei [ fˆ1 ( x )] ≤ log ∑ exp λEz [ fˆ1 ( x )]> Ei [ fˆ1 ( x )] .
i =1

i =1

We can also rewrite the dot product between mean embeddings per strata in terms ofthe distance between them:
"
#
 C


L( x, z, fˆ1 ) ≤ −1 + Ez log ∑ exp λEz [ fˆ1 ( x )]> Ei [ fˆ1 ( x )]
i =1

"

#

λ
2
= −1 + Ez log ∑ exp − kEz [ fˆ1 ( x )] − Ei [ fˆ1 ( x )]k + λ
.
2
i =1
 C



This directly gives us our desired bound.
Proof of Lemma 2. The generalization error is""##exp( fˆ1 ( x )> WS(z) )ˆL( x, y, f 1 ) = −Ez Ex∼Pz log K∑i=1 exp( fˆ1 ( x )> Wi )""##K
= Ez Ex∼P − fˆ1 ( x )> W
+ log ∑ exp( fˆ1 ( x )> Wi ) .
S(z)

z

i =1

We substitute in the definition of the mean classifier to getˆL( x, y, f 1 ) = Ez − ∑ p(z0 |S(z))Ez [ fˆ1 ( x )]> Ez0 [ fˆ1 ( x )]z 0 ∈ SS ( z )+ Ex∼Pz

K



i =1

z 0 ∈ Silog ∑ exp ∑ p(z |Si ) fˆ1 ( x )> Ez0 [ fˆ1 ( x )]
0



.

Comput. Sci. Math. Forum 2022, 3, 4

19 of 27

We can rewrite the dot product between mean embeddings per strata in terms of the
distance between them:1ˆL( x, y, f 1 ) =Ez ∑ p(z0 |S(z)) ·kEz [ fˆ1 ( x )] − Ez0 [ fˆ1 ( x )]k2 − 1
2
z0 ∈S
S(z)

K



+ Ex∼Pz log ∑ exp



i =1

∑ p(z |Si ) fˆ1 (x)> Ez0 [ fˆ1 (x)]
0



.

z 0 ∈ Si

We can write kEz [ fˆ1 ( x )] − Ez0 [ fˆ1 ( x )]k in the above expression as δ( fˆ1 , z, z0 ), which wehave analyzed:L( x, y, fˆ1 ) =Ez1∑ p(z0 |S(z)) · 2 δ( fˆ1 , z, z0 )2 − 10z ∈ SS ( z )




K
+ Ex∼Pz log ∑ exp ∑ p(z0 |Si ) fˆ1 ( x )> Ez0 [ fˆ1 ( x )].
i =1

z 0 ∈ Si

From our previous proof, there exists λ > 0 such that this is at mostL( x, y, fˆ1 ) ≤Ez1∑ p(z0 |S(z)) · 2 δ( fˆ1 , z, z0 )2 − 10z ∈ SS ( z ) K∑ exp( ∑ p(z |Si )λEz [ fˆ1 (x)]> Ez0 [ fˆ1 (x)])

+ log

i =1



=Ez

0

z 0 ∈ Si∑ p(z0 |S(z)) ·12

z 0 ∈ SS ( z )

 K

+ log

∑ exp

i =1





δ( fˆ1 , z, z0 )2 − 1



 λ

0
2
ˆ
ˆ
p
(
z
|
S
)
−
kE
[
f
(
x
)]
−
E
[
f
(
x
)]k
+
λ
.
0
z 1
i
∑
z 1
2
z0 ∈S
i

We can write each weighted summation over p(z0 |S(z)) and p(z0 |Si ) as an expectation
and use the definition of δ( fˆ1 , z, z0 ) to obtain our desired bound.
Appendix E. Additional Experimental Details
Appendix E.1. Datasets
We first describe all the datasets in more detail:
•
•
•

•
•

CIFAR10, CIFAR100, and MNIST are all the standard computer vision datasets.
CIFAR10-Coarse consists of two superclasses: animals (dog, cat, deer, horse, frog,bird) and vehicles (car, truck, plane, boat).
CIFAR100-Coarse consists of twenty superclasses. We artificially imbalance subclasses
to create CIFAR100-Coarse-U. For each superclass, we select one subclass to keep
all 500 points, select one subclass to subsample to 250 points, select one subclass to
subsample to 100 points, and select the remaining two to subsample to 50 points. We
use the original CIFAR100 class index to select which subclasses to subsample: thesubclass with the lowest original class index keeps all 500 points, the next subclasskeeps 250 points, etc.
MNIST-Coarse consists of two superclasses: <5 and ≥5.
Waterbirds [14] is a robustness dataset designed to evaluate the effects of spurious
correlations on model performance. The waterbirds dataset is constructed by croppingout birds from photos in the Caltech-UCSD Birds dataset [43], and pasting them onbackgrounds from the Places dataset [44]. It consists of two categories: water birds andland birds. The water birds are heavily correlated with water backgrounds and theComput. Sci. Math. Forum 2022, 3, 4

20 of 27••

land birds with land backgrounds, but 5% of the water birds are on land backgrounds,
and 5% of the land birds are on water backgrounds. These form the (imbalanced)
hidden strata.
ISIC is a public skin cancer dataset for classifying skin lesions [15] as malignant or benign. 48% of the benign images contain a colored patch, which form the hidden strata.
CelebA is an image dataset commonly used as a robustness benchmark [14,16]. The
task is blonde/not blonde classification. Only 6% of blonde faces are male, which
creates a rare stratum in the blonde class.

Appendix E.2. Hyperparameters
For all model quality experiments for Lspread , we first fixed τ = 0.5 and swept
α ∈ [0.16, 0.25, 0.33, 0.5, 0.67]. We then took the two best-performing values and sweptτ ∈ [0.1, 0.3, 0.5, 0.7, 0.9]. For LSC and LSS , we swept τ ∈ [0.1, 0.3, 0.5, 0.7, 0.9]. Final
hyperparameter values for (τ, α) for Lspread were (0.9, 0.67) for CIFAR10, (0.5, 0.16) for
CIFAR10-coarse, (0.5, 0.33) for CIFAR100, (0.5, 0.25) for CIFAR100-Coarse, (0.5, 0.25) for
CIFAR100-Coarse-U, (0.5, 0.5) for MNIST, (0.5, 0.5) for MNIST-coarse, (0.5, 0.5) for ISIC,
and (0.5, 0.5) for waterbirds.
For coarse-to-fine transfer learning, we fixed τ = 0.5 for all losses and swept α ∈[0.16, 0.25, 0.33, 0.5, 0.67]. Final hyperparameter values for α were 0.25 for CIFAR10-Coarse,
0.25 for CIFAR100-Coarse, 0.25 for CIFAR100-Coarse-U, and 0.5 for MNIST-Coarse.
Appendix E.3. Applications
We describe additional experimental details for the applications.
Appendix E.3.1. Robustness Against Worst-Group Performance
We follow the evaluation of [5]. First, we train a model on the standard class labels. We
evaluate different loss functions for this step, including Lspread , LSC , and the cross entropy
loss LCE . Then we project embeddings of the training set using a UMAP projection [45], andcluster points to discover unlabeled subgroups. Finally, we use the unlabeled subgroups ina Group-DRO algorithm to optimize worst-group robustness [14].
Appendix E.3.2. Robustness Against Noise
We use the same training setup as we use to evaluate model quality, and introduce
symmetric noise into the labels for the contrastive loss head. We train the cross entropy
head with a fraction of the full training set. In Section 5.3, we report results from training
with 20% labels to cross entropy. We report additional levels in Appendix F.
We detect noisy labels with a simple geometric heuristic: for each point, we computethe cosine similarity between the embedding of the point and the center of all the otherpoints in the batch that have the same class. We compare this similarity value to the averagecosine similarity with points in the batch from every other class, and rank the points by thedifference between these two values. Points with incorrect labels have a small difference
between these two values (they appear to be small strata, so they are far away from pointsof the same class). Given the noise level e as an input, we rank the points by this heuristic
and mark the e fraction of the batch with the smallest scores as noisy. We then correct theirlabels by adopting the label of the closest cluster center.
Appendix E.3.3. Minimal Coreset Construction
We use the publicly-available evaluation framework for coresets from [18] (https://github.com/mtoneva/example_forgetting, accessed on 1 October 2021). We use the official
repository from [19] (https://github.com/mansheej/data_diet, accessed on 1 October 2021)
to recreate their coreset algorithms.

Comput. Sci. Math. Forum 2022, 3, 4

21 of 27

Our coreset algorithm proceeds in two parts. First, we give each point a difficulty
rating based on how likely we are to classify it correctly under partial training. Then we
subsample the easiest points to construct minimal coresets.
First, we mirror the set up from our thought experiment and train with Lspread on
random samples of t% of the CIFAR10 training set, taking three random samples for each oft ∈ [10, 20, 50] (and we train the cross entropy head with 1% labeled data). For each run, werecord which points are classified correctly by the cross entropy head at the end of training,and bucket points the training set by how often the point was correctly classified. To
construct a coreset of size t%, we iteratively remove points from the largest bucket in eachclass. Our strategy removes easy examples first from the largest coresets, but maintains a
set of easy examples in the smallest coresets.
Appendix F. Additional Experimental Results
In this section, we report three sets of additional experimental results: the performanceof using L attract on its own to train models, sample complexity of Lspread compared to LSC ,and additional noisy label results (including a bonus de-noising algorithm).
Appendix F.1. Performance of L attract
In an early iteration of this project, we experienced success with using L attract on its
own to train models, before realizing the benefits of adding in an additional term to preventclass collapse. As an ablation, we report on the performance of using L attract on its own
in Table A2. L attract can outperform LSC , but Lspread outperforms both. We do not report
the results here, but L attract also performs significantly worse than LSC on downstream
applications, since it more direclty encourages class collapse.
Table A2. Performance of Lspread compared to LSC and using L attract on its own. Best in bold.
End Model Perf.
DatasetLSS

LSC

L attract

Lspread

CIFAR10CIFAR100

89.7
68.0

90.9
67.5

91.3
68.9

91.5
69.1

Appendix F.2. Sample Complexity
Figure A1 shows the performance of training ViT models with various amountsof labeled data for Lspread , LSC , and LSS . In these experiments, we train the cross entropy head with 1% labeled data to isolate the effect of training data on the contrastive
losses themselves.
Lspread outperforms LSC and LSS throughout. At 10% labeled data, Lspread outperforms LSS by 13.9 points, and outperforms LSC by 0.5 points. By 100% labeled data (forthe contrastive head), Lspread outperforms LSS by 25.4 points, and outperforms LSC by
10.3 points.

Comput. Sci. Math. Forum 2022, 3, 4

22 of 27

Sample Complexity, CIFAR10
L_spread
L_SC
L_SS

Acc

80
70
60
50

0.2

0.4
0.6
0.8
% labeled data

1.0

Figure A1. Performance of training ViT with Lspread compared to training with LSC and LSS on
CIFAR10 at various amounts of labeled data. Lspread outperforms the baselines at each point. The
cross entropy head here is trained with 1% labeled data to isolate the effect of training data on thecontrastive losses.

Appendix F.3. Noisy Labels
In Section 5.3, we reported results from training the contrastive loss head with noisylabels and the cross entropy loss with clean labels from 20% of the training data.
In this section, we first discuss a de-noising algorithm inspired by [23] that we initially developed to correct for noisy labels, but that we did not observe strong empiricalresults from. We hope that reporting this result inspires future work into improving
contrastive learning.
We then report additional results with larger amounts of training data for the crossentropy head.
Appendix F.3.1. Debiasing Noisy Contrastive Loss
First, we consider the triplet loss and show how to debias it in expectation under noise.
Then we present an extension to supervised contrastive loss.
Noise-Aware Triplet Loss
Consider the triplet loss:
Ltriplet = Ex∼P ,x+ ∼ p+ (·| x), − log
x − ∼ p− (·| x )
exp(σ ( x, x + )).
exp(σ ( x, x + )) + exp(σ( x, x − ))(A7)

Now suppose that we do not have access to true labels but instead have noisy labels
denoted by the weak classifier ye := eh( x ). We adopt a simple model of symmetric noise
where pe = Pr(noisy label is correct).
e + and P
e − as p( x + |e
We use ye to construct Ph( x ) = eh( x + )) and p( x − |eh( x ) 6= eh( x − )).
For simplicity, we start by looking at how the triplet loss in (A7) is impacted when noise istriplete + and Pe−.
not addressed in the binary setting. Define Lnoisy as Ltriplet used with P
noisy

Lemma A3. When class-conditional noise is uncorrected, Ltriplet is equivalent toComput. Sci. Math. Forum 2022, 3, 4

23 of 27"3

3

( pe + (1 − pe) ) Ltriplet + pe(1 − pe)E

− log

x ∼P
x1+ ,x2+ ∼ p+ (·| x )x ∼Px1− ,x2− ∼ p− (·| x )#exp(σ ( x, x1− ))#exp(σ ( x, x1+ )) + exp(σ( x, x2+ ))"+ pe(1 − pe)Eexp(σ ( x, x1+ ))− log

exp(σ ( x, x1− )) + exp(σ( x, x2− ))exp(σ( x, x − ))+ pe(1 − pe)E x∼P− log.
exp(σ ( x, x + )) + exp(σ ( x, x − ))x + ∼ p+ (·| x )x − ∼ p− (·| x )noisy

Proof. We split Ltriplet depending on if the noisy positive and negative pairs are truly
positive and negative.
noisy



Ltriplet = E

x ∼Pxe+ ∼ pe+ (·| x )xe− ∼ pe− (·| x )− log

exp(σ( x, xe+ ))exp(σ( x, xe+ )) + exp(σ( x, xe− ))= p(h( x ) = h( xe+ ), h( x ) 6= h( xe− ))Ex ∼Px + ∼ p+ (·| x )x − ∼ p− (·| x )

− log



exp(σ( x, x + ))exp(σ( x, x + )) + exp(σ( x, x − ))

"
e−

e+

+ p(h( x ) = h( x ), h( x ) = h( x ))E

x ∼P
x1+ ,x2+ ∼ p+ (·| x )

− log

e−

+ p(h( x ) 6= h( x ), h( x ) 6= h( x ))Ex ∼P
x1− ,x2− ∼ p− (·| x )exp(σ( x, x1+ ))#exp(σ( x, x1− ))#exp(σ( x, x1+ )) + exp(σ( x, x2+ ))

"
e+



− log

exp(σ( x, x1− )) + exp(σ( x, x2− ))exp(σ( x, x − )).
+ p(h( x ) 6= h( xe+ ), h( x ) = h( xe− ))E x∼P− log
exp(σ( x, x + )) + exp(σ( x, x − ))x + ∼ p+ (·| x )x − ∼ p− (·| x )Define pe = p(noisy label is correct). Note thatp(h( x ) = h( xe+ ), h( x ) 6= h( xe− )) = pe3 + (1 − pe)3 ,(i.e., all three points are correct or all reversed, such that their relative pairings arecorrect). In addition, the other three probabilities above are all equal to pe(1 − pe).
We now show that there exists a weighted loss function that in expectation equals
Ltriplet .
Lemma A4. DefineeLtriplet = E

x ∼P ,e + (·| x )xe1+ ,ex2+ ∼P− − e−
ex1 ,ex2 ∼P (·| x )− w+ σ( x, xe1+ ) + w− σ( x, xe1− )






+ w1 log exp σ( x, xe1+ ) + exp σ( x, xe1− )


+
+
−
−
− w2 log (exp(σ( x, xe1 )) + exp(σ( x, xe2 ))) · (exp(σ( x, xe1 )) + exp(σ( x, xe2 ))) ,whereComput. Sci. Math. Forum 2022, 3, 4

24 of 27

pe2 + (1 − pe)22ep(1 − pe)w− =(2e
p − 1)2(2ep − 1)2hiThen, E e
Ltriplet = Ltriplet .

w+ =

w1 =

pe2 + (1 − pe)2(2e
p − 1)2w2 =pe(1 − pe).
(2ep − 1)2


Proof. We evaluate E −w1 σ ( x, xe1+ ) + w2 σ ( x, xe1− ) and the other terms separately. Using
the same probabilities as computed in Lemma A3,

E −w1 σ( x, xe1+ ) + w2 σ( x, xe1− ) = −( pe2 + (1 − pe)2 )w1 E σ( x, x1+ )






− 2e
p(1 − pe)w1 E σ ( x, x1− ) + ( pe2 + (1 − pe)2 )w2 E σ ( x, x1− ) + 2ep(1 − pe)w2 E σ ( x, x1+ )= −E σ( x, x1+ ) .
We evaluate the remaining terms:
h




i
E w3 log exp σ( x, xe1+ ) + exp σ( x, xe1− )
=
i
h




( pe2 + (1 − pe)2 )w3 E log exp σ( x, x1+ ) + exp σ( x, x1− )hi+ pe(1 − pe)w3 E log (exp(σ( x, xe1+ )) + exp(σ( x, xe2+ ))) · (exp(σ( x, xe1− )) + exp(σ( x, xe2− ))) .

and
h




i
E w4 log exp σ( x, xe1+ ) + exp σ( x, xe2+ )
h




i
+ E w4 log exp σ( x, xe1− ) + exp σ( x, xe2− )
=
h




i
( pe2 + (1 − pe)2 )w4 E log exp σ( x, x1+ ) + exp σ( x, x2+ )
i


h


+ 4e
p(1 − pe)w4 E log exp σ ( x, x1+ ) + exp σ( x, x1− )
i


h


.
+ ((1 − pe)2 + pe2 )w4 E log exp σ( x, x1− ) + exp σ( x, x2− )Examining the coefficients, we see that( pe2 + (1 − pe)2 )w3 − 4ep(1 − pe)w4 =pe(1 − pe)w3 − ( pe2 + (1 − pe)2 )w4 =( pe2 + (1 − pe)2 )2 4ep2 (1 − pe)2−=1(2e
p − 1)2(2ep − 1)2( pe2 + (1 − pe)2 ) pe(1 − pe)pe(1 − pe)( pe2 + (1 − pe)2 )
−
= 0,
(2e
p − 1)2(2ep − 1)2

h




i
which shows that only the term E log exp σ ( x, x1+ ) + exp σ ( x, x1− )persists. This
completes our proof.
We now show the general case for debiasing L attract , which uses more negative samples.
Proposition A1. Define m = n + 1 (as the “batch size” in the denominator), and
eL attract = E x∼P{ xei+ }im=1m{ xe−j } j =1

− w+ σ( x, xe1+ ) + w− σ( x, xe1− )
 m−k


+
−
e
e
∑ exp σ(x, xi ) + ∑ exp σ(x, x j ) .

m

 k

k =0

i =1

+ ∑ wk log



j =1

(A8)

(A9)

Comput. Sci. Math. Forum 2022, 3, 4

25 of 27

~ = {w0 , . . . wm } ∈ Rm+1 is the solution to thew+ and w− are defined in the same was as before. w
system Pw = e2 where e2 is the standard basis vector in Rm+1 where the 2nd index is 1 and allothers are 0. The i, jth element of P is Pij = peQi,j + (1 − pe)Qm−i,j where

Qi,j =

∑min{ j,m−i} ( j )( m− j )(1 − pe)i− j+2k pem+ j−i−2kk =0k i − j+k∑min{i,m− j} (m− j)( j )(1 − pe) j−i+2k pem− j+i−2kk =0k
j −i + k

j≤i
j>i

h
i
Then, E e
L attract = L attract .
We do not present the proof for Proposition A1, but the steps are very
similartohithe proof for the triplet loss case. We also note that a different form of E e
L attract must
be computed for the multi-class case, which we do not present here (but can be derived
through computation).
Observation A2. Note that the values of Qi,j have high variance in the noise rate as m increases.
Additionally, note that the number of terms in the summation of Qi,j increase combinatorially withm. We found this de-noising algorithm very unstable as a result.
Appendix F.3.2. Additional Noisy Label Results
Now we report the performance of denoising algorithms with additional amountsof labeled data for the cross entropy loss head. We also report the performance of usinge
L attract to debias noisy labels.
Figure A2 shows the results. Our geometric correction together with Lspread works
the most consistently. Using the geometric correction with LSC can be unreliable, since
LSC can learn memorize noisy labels early on in training. The expectation-based debiasingalgorithm e
L attract occasionally shows promise but is unreliable, and is very sensitive tohaving the correct noise rate as an input.

Acc

Acc

Noisy label performance, 20% labels to L_CE
90
88
86
84
82
80
78
0.0
92
91
90
89
88
87
86
85
0.0

Geometry-correction, L_spread
Geometry-correction, L_SC
Expectation-based noisy correction
L_SC, no correction

0.1

0.2
0.3
Noise Rate

0.4

0.5

Noisy label performance, 60% labels to L_CE
Geometry-correction, L_spread
Geometry-correction, L_SC
Expectation-based noisy correction
L_SC, no correction

0.1

0.2
0.3
Noise Rate

0.4

0.5

Figure A2. Performance of models under various amounts of label noise for the contrastive loss head,and various amounts of clean training data for the cross entropy loss.

Comput. Sci. Math. Forum 2022, 3, 4

26 of 27

References
1.
2.
3.
4.
5.
6.

7.
8.
9.
10.
11.
12.

13.
14.
15.

16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.

Khosla, P.; Teterwak, P.; Wang, C.; Sarna, A.; Tian, Y.; Isola, P.; Mschinot, A.; Liu, C.; Krishnan, D. Supervised Contrastive Learning.
Adv. Neural Inf. Process. Syst. 2020, 33, 18661–18673.
Graf, F.; Hofer, C.; Niethammer, M.; Kwitt, R. Dissecting Supervised Constrastive Learning. Proc. Int. Conf. Mach. Learn. PMLR.
2021, 139, 3821–3830.
Zhang, C.; Bengio, S.; Hardt, M.; Recht, B.; Vinyals, O. Understanding deep learning requires rethinking generalization. Commun.
ACM 2016, 64, 107–115. [CrossRef]
Hoffmann, A.; Kwok, R.; Compton, P. Using subclasses to improve classification learning. In European Conference on Machine
Learning; Springer: Berlin/Heidelberg, Germany, 2001; pp. 203–213.
Sohoni, N.; Dunnmon, J.; Angus, G.; Gu, A.; Ré, C. No Subclass Left Behind: Fine-Grained Robustness in Coarse-GrainedClassification Problems. Adv. Neural Inf. Process. Syst. 2020, 33, 19339–19352.
Oakden-Rayner, L.; Dunnmon, J.; Carneiro, G.; Ré, C. Hidden stratification causes clinically meaningful failures in machinelearning for medical imaging. In Proceedings of the Proceedings of the ACM conference on health, inference, and learning,
Toronto, ON, Canada, 2–4 April 2020; pp. 151–159.
Linsker, R. Self-organization in a perceptual network. Computer 1988, 21, 105–117. [CrossRef]
Wang, T.; Isola, P. Understanding contrastive representation learning through alignment and uniformity on the hypersphere.
Proc. Int. Conf. Mach. Learn. PMLR 2020, 119, 9929–9939.
Robinson, J.; Chuang, C.Y.; Sra, S.; Jegelka, S. Contrastive learning with hard negative samples. arXiv 2020, arXiv:2010.04592.
Ben-David, S.; Blitzer, J.; Crammer, K.; Kulesza, A.; Pereira, F.; Vaughan, J.W. A theory of learning from different domains. Mach.
Learn. 2010, 79, 151–175. [CrossRef]
Ben-David, S.; Blitzer, J.; Crammer, K.; Pereira, F. Analysis of representations for domain adaptation. Adv. Neural Inf. Process. Syst.
2007, 19, 137.
Ben-David, S.; Urner, R. On the Hardness of Domain Adaptation and the Utility of Unlabeled Target Samples. In Proceedings ofthe 23rd International Conference, Lyon, France, 29–31 October 2012; Bshouty, N.H., Stoltz, G., Vayatis, N., Zeugmann, T., Eds.;
Springer: Berlin/Heidelberg, Germany, 2012; pp. 139–153.
Mansour, Y.; Mohri, M.; Rostamizadeh, A. Domain adaptation: Learning bounds and algorithms. arXiv 2009, arXiv:0902.3430.
Sagawa, S.; Koh, P.W.; Hashimoto, T.B.; Liang, P. Distributionally Robust Neural Networks for Group Shifts: On the Importanceof Regularization for Worst-Case Generalization. arXiv 2019, arXiv:1911.08731.
Codella, N.; Rotemberg, V.; Tschandl, P.; Celebi, M.E.; Dusza, S.; Gutman, D.; Helba, B.; Kalloo, A.; Liopyris, K.; Marchetti, M.;
et al. Skin lesion analysis toward melanoma detection 2018: A challenge hosted by the international skin imaging collaboration(isic). arXiv 2019, arXiv:1902.03368.
Liu, Z.; Luo, P.; Wang, X.; Tang, X. Deep Learning Face Attributes in the Wild. In Proceedings of the Proceedings of International
Conference on Computer Vision (ICCV), Santiago, Chile, 7 December 2015.
Dosovitskiy, A.; Beyer, L.; Kolesnikov, A.; Weissenborn, D.; Zhai, X.; Unterthiner, T.; Dehghani, M.; Minderer, M.; Heigold, G.;
Gelly, S.; et al. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. arXiv 2020, arXiv:2010.11929.
Toneva, M.; Sordoni, A.; des Combes, R.T.; Trischler, A.; Bengio, Y.; Gordon, G.J. An Empirical Study of Example Forgetting
during Deep Neural Network Learning. arXiv 2018, arXiv:1812.05159.
Paul, M.; Ganguli, S.; Dziugaite, G.K. Deep Learning on a Data Diet: Finding Important Examples Early in Training. arXiv 2021,
arXiv:2107.07075.
Chen, T.; Kornblith, S.; Norouzi, M.; Hinton, G. A simple framework for contrastive learning of visual representations. Proc. Int.
Conf. Mach. Learn. PMLR 2020, 119, 1597–1607
Arora, S.; Khandeparkar, H.; Khodak, M.; Plevrakis, O.; Saunshi, N. A theoretical analysis of contrastive unsupervised
representation learning. arXiv 2019, arXiv:1902.09229.
Zimmermann, R.S.; Sharma, Y.; Schneider, S.; Bethge, M.; Brendel, W. Contrastive Learning Inverts the Data Generating Process.
arXiv 2021, arXiv:2012.08850.
Chuang, C.Y.; Robinson, J.; Torralba, A.; Jegelka, S. Debiased Contrastive Learning. Adv. Neural Inf. Process. Syst. 2020, 33,
8765–8775.
Oord, A.v.d.; Li, Y.; Vinyals, O. Representation learning with contrastive predictive coding. arXiv 2018, arXiv:1807.03748.
Tian, Y.; Sun, C.; Poole, B.; Krishnan, D.; Schmid, C.; Isola, P. What makes for good views for contrastive learning? arXiv 2020,
arXiv:2005.10243.
Tsai, Y.H.H.; Wu, Y.; Salakhutdinov, R.; Morency, L.P. Self-supervised Learning from a Multi-view Perspective. arXiv 2020,
arXiv:2006.05576.
Tschannen, M.; Djolonga, J.; Rubenstein, P.K.; Gelly, S.; Lucic, M. On Mutual Information Maximization for Representation
Learning. arXiv 2019, arXiv:1907.13625.
He, K.; Fan, H.; Wu, Y.; Xie, S.; Girshick, R. Momentum Contrast for Unsupervised Visual Representation Learning. arXiv 2019,arXiv:1911.05722.
Chen, X.; Fan, H.; Girshick, R.; He, K. Improved Baselines with Momentum Contrastive Learning. arXiv 2020, arXiv:2003.04297.
Goyal, P.; Caron, M.; Lefaudeux, B.; Xu, M.; Wang, P.; Pai, V.; Singh, M.; Liptchinsky, V.; Misra, I.; Joulin, A.; et al. Self-supervised
Pretraining of Visual Features in the Wild. arXiv 2021, arXiv:2103.01988.

Comput. Sci. Math. Forum 2022, 3, 4

31.
32.
33.

34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.

27 of 27

Caron, M.; Misra, I.; Mairal, J.; Goyal, P.; Bojanowski, P.; Joulin, A. Unsupervised Learning of Visual Features by Contrasting
Cluster Assignments. Adv. Neural Inf. Process. Syst. 2020, 33, 9912–9924.
Islam, A.; Chen, C.F.; Panda, R.; Karlinsky, L.; Radke, R.; Feris, R. A Broad Study on the Transferability of Visual Representations
with Contrastive Learning. arXiv 2021, arXiv:2103.13517.
Bukchin, G.; Schwartz, E.; Saenko, K.; Shahar, O.; Feris, R.; Giryes, R.; Karlinsky, L. Fine-grained Angular Contrastive Learningwith Coarse Labels. In Proceedings of the 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR),
Virtual, 19 June 2021; [CrossRef]d’Eon, G.; d’Eon, J.; Wright, J.R.; Leyton-Brown, K. The Spotlight: A General Method for Discovering Systematic Errors in Deep
Learning Models. arXiv 2021, arXiv:2107.00758.
Duchi, J.; Hashimoto, T.; Namkoong, H. Distributionally robust losses for latent covariate mixtures. arXiv 2020, arXiv:2007.13982.
Goel, K.; Gu, A.; Li, Y.; Re, C. Model Patching: Closing the Subgroup Performance Gap with Data Augmentation. arXiv 2020,
arXiv:2008.06775.
Liu, S.; Niles-Weed, J.; Razavian, N.; Fernandez-Granda, C. Early-Learning Regularization Prevents Memorization of Noisy
Labels. Adv. Neural Inf. Process. Syst. 2020, 33, 20331–20342.
Li, J.; Xiong, C.; Hoi, S.C. Semi-supervised Learning with Contrastive Graph Regularization. In Proceedings of the IEEE/CVF
International Conference on Computer Vision (ICCV), Virtual, 11 October 2021.
Ciortan, M.; Dupuis, R.; Peel, T. A Framework using Contrastive Learning for Classification with Noisy Labels. arXiv 2021,
arXiv:2104.09563.
Li, J.; Socher, R.; Hoi, S.C. DivideMix: Learning with Noisy Labels as Semi-supervised Learning. arXiv 2020, arXiv:2002.07394.
Ju, J.; Jung, H.; Oh, Y.; Kim, J. Extending Contrastive Learning to Unsupervised Coreset Selection. arXiv 2021, arXiv:2103.03574.
Sener, O.; Savarese, S. Active Learning for Convolutional Neural Networks: A Core-Set Approach. arXiv 2017, arXiv:1708.00489.
Welinder, P.; Branson, S.; Mita, T.; Wah, C.; Schroff, F.; Belongie, S.; Perona, P. Caltech-UCSD Birds 200. In Technical Report
CNS-TR-2010-001; California Institute of Technology: Pasadena, CA, USA, 2010.
Zhou, B.; Lapedriza, A.; Xiao, J.; Torralba, A.; Oliva, A. Learning Deep Features for Scene Recognition using Places Database.
Adv. Neural Inf. Process. Syst. 2014, 27, 487–495.
McInnes, L.; Healy, J.; Saul, N.; Großberger, L. UMAP: Uniform Manifold Approximation and Projection. J. Open Source Softw.
2018, 3. [CrossRef]