Self-supervised Learning is More Robust to Dataset Imbalance
Abstract
Self-supervised learning (SSL) learns general visual representations without the need of labels.
However, large-scale unlabeled datasets in the wild often have long-tailed label distributions, where we know little about the behavior of SSL.
We investigate SSL under dataset imbalance, and ﬁnd out that existing self-supervised representations are more robust to class imbalance than supervised representations.
The performance gap between balanced and imbalanced pre-training with SSL is much smaller than the gap with supervised learning.
Second, to understand the robustness of SSL, we hypothesize that SSL learns richer features from frequent data: it may learn label-irrelevant-but-transferable features that help classify the rare classes.
In contrast, supervised learning has no incentive to learn features irrelevant to the labels of frequent examples.
We validate the hypothesis with semi-synthetic experiments and theoretical analysis on a simpliﬁed setting.
Introduction
Self-supervised learning (SSL) is an important paradigm of machine learning, because it can leverage the availability of large-scale unlabeled datasets to learn representations for a wide range of downstream tasks and datasets [16, 8, 13, 7, 9].
Current SSL algorithms are mostly trained on curated, balanced datasets but large-scale unlabeled datasets in the wild are inevitably imbalanced with a long-tailed label distribution [27, 24].
Curating a class-balanced unlabeled dataset requires the knowledge of labels which defeats the purpose of leveraging unlabeled data by SSL.
The behavior of SSL algorithms under dataset imbalance remains largely underexplored in the literature but extensive studies do not bode well for supervised learning (SL) with imbalanced datasets. The performance of vanilla supervised methods degrades signiﬁcantly on class-imbalanced datasets [11, 5, 3], posing challenges to practical applications such as instance segmentation [29] and depth estimation [37].
Many recent works address this issue with various regularization and re-weighting/re-sampling techniques [1, 33, 18, 11, 5, 6, 30, 17, 32].
In this work, we systematically investigate the representation quality of SSL algorithms under class imbalance.
Perhaps surprisingly, we ﬁnd out that off-the-shelf SSL representations are already more robust to dataset imbalance than the representations learned by supervised pre-training.
We evaluate the representation quality by linear probe on in-domain (ID) data and ﬁnetuning on out-of-domain (OOD) data.
We compare the robustness of SL and SSL representations by computing the gap between the performance of the representations pre-trained on balanced and imbalanced datasets of the same sizes.
We observe that the balance-imbalance gap for SSL is much smaller than SL, under a variety of conﬁgurations with varying dataset sizes and imbalance ratios and with both ID and OOD evaluations (see Figure 1 and Section 2 for more details).
This robustness holds even with the same number of samples for SL and SSL, although SSL does not require labels and hence can be more easily applied to larger datasets than SL.
Why is SSL more robust to dataset imbalance?
We identify the following underlying cause to answer this fundamental question: SSL learns richer features from the frequent classes than SL does.
These features may help classify the rare classes under ID evaluation and are transferable to the downstream tasks under OOD evaluation.
For simplicity, consider the situation where rare classes [Figure] have so limited data that both SL and SSL models overﬁt to the rare data.
In this case, it is important for the models to learn diverse features from the frequent classes which can help classify the rare classes.
Supervised learning is only incentivized to learn those features relevant to predicting frequent classes and may ignore other features.
In contrast, SSL may learn the structures within the frequent classes better—because it is not supervised or incentivized by any labels, it can learn not only the label-relevant features but also other interesting features capturing the intrinsic properties of the input distribution, which may generalize/transfer better to rare classes and downstream tasks.
We empirically validate this intuition by visualizing the features on a semi-synthetic dataset where the label-relevant features and label-irrelevant-but-transferable features are prominently seen by design (cf. Section 3.1).
In addition, we construct a toy example where we can rigorously prove the difference between self-supervised and supervised features in Section 3.2.
We sum up our contributions as follows. (1) We are the ﬁrst to systematically investigate the robustness of self-supervised representation learning to dataset imbalance. (2) We propose and validate an explanation of this robustness of SSL, empirically and theoretically.
Exploring the Effect of Class Imbalance on SSL 2.1 Problem Formulation
Class-imbalanced pre-training datasets.
x denotes the input and y denotes the label.
Supervised pre-training has access to x and y , whereas self-supervised pre-training only observes x .
Given a pre-training distribution P , let r denote the ratio of class imbalance, i.e. the ratio between the probability of the rarest and the most frequent class: r = min
j ∈ [ C ] P ( y = j ) max j ∈ [ C ] P ( y = j ) ≤ 1 .
We use P r to denote the dataset with ratio r .
We also use P bal for the case where r = 1 , i.e. the dataset is balanced.
We assume that for any class j ∈ [ C ] , the class-conditional distribution P r ( x | y =
j ) is the same across datasets for all r .
The pre-training dataset (cid:98) P rn consists of n i.i.d.
samples from P r .
Pre-trained models.
A feature extractor maps the input to the representations.
A linear head classiﬁer is composed with the feature extractor to produce the prediction.
SSL algorithms learn the representations from unlabeled data.
SL pre-training learns the feature extractor and the linear head from labeled data.
We drop the head and only evaluate the quality of feature extractor.
Following the standard protocol in prior works [16, 8], we measure the quality of representations on both in-domain and out-of-domain datasets with either linear probe or ﬁne-tuning , as detailed below.
In-domain (ID) evaluation tests the performance of representations on the balanced in-domain distribution P bal with linear probe.
Given a feature extractor pre-trained on a pre-training dataset (cid:98) P rn with n data points and imbalance ratio r , we train a linear classiﬁer on a balanced dataset and evaluate the representation quality with the accuracy of the learned linear head on P bal .
We denote the ID accuracy of supervised pre-trained representations by A SLID ( n, r ) .
Note that A SLID ( n, 1) stands for the result with balanced pre-training dataset.
For SSL, we denote the accuracy by A SSLID ( n, r ) .
[Figure] Out-
-domain (OOD evaluation tuning trained on a of size n imbalance ratio r ) and a randomly initialized tune
Summary of varying factors.
We aim to study the effect of class imbalance to feature qualities on a diverse set of conﬁgurations with the following varying factors: (1) the number of examples in pre-training n , (2) the imbalance ratio of the pre-training dataset r , (3) ID or OOD evaluation, and (4) self-supervised learning algorithms: MoCo v2 [16], or SimSiam [9].
Experimental Setup.
We pre-train the representations on ImageNet [28] with a wide range of numbers of examples and ratios of imbalance.
Following Liu et al.
[24], we consider the long-tailed
Pareto label distributions with imbalance ratio in { 1 , 0 . 004 , 0 . 0025 } .
For each imbalance ratio, we further downsample the dataset with a sampling ratio in { 0 . 75 , 0 . 5 , 0 . 25 , 0 . 125 } to form datasets with varying sizes.
For ID evaluation, we use the original ImageNet training set for the training phase of linear probe and use the original validation set for the ﬁnal evaluation. For OOD evaluation, we ﬁne-tune the pre-trained feature extractors on the downstream tasks. For self-supervised pre-training, we consider MoCo
2.2 Results: SSL is More Robust than SL to Dataset Imbalance
In Figure 2, for both ID and OOD evaluations, the gap of SSL representations learned on balanced and imbalanced datasets with the same number of pre-training examples, i.e., A SSL ( n, 1) − A SSL ( n, r ) , is smaller than the gap of SL representations, i.e., A SL ( n, 1) − A SL ( n, r ) .
Furthermore, we compute the relative accuracy gap to balanced dataset ∆ SSL ( n, r ) (cid:44) ( A SSL ( n, 1) − A SSL ( n, r )) /A
SSL ( n, 1) in
Figure 1.
With the same number of pre-training examples, the relative gap of SSL representations between balanced and imbalanced datasets is smaller than that of SL representations,
[Equation]
We also note that comparing the robustness with the same number of data is actually in favor of SL, because SSL is more easily applied to larger datasets without the need of collecting labels.
ID vs. OOD.
As shown in Figure 2, representations from SSL perform better than SL pre-training in ID evaluation with reasonably large n , while SSL pre-training is better in OOD evaluation.
This phenomenon is orthogonal to the observation that SSL is more robust to dataset imbalance, and is consistent with recent works (e.g., Chen et al.
[8], He et al.
[16]) which also observed that SSL performs slightly worse than SL on balanced ID evaluation but better on OOD tasks.
Analysis
Where does the robustness stem from?
In this section, we propose a possible reason and justify it with theoretical and empirical analyses.
SSL learns richer features from frequent data that are transferable to rare data.
The rare classes of the imbalanced dataset can contain only a few examples, making it hard to learn proper features from them.
In this case, one may want to resort to [Figure] the features from the frequent classes.
Due to the supervised nature of classiﬁcation, SL mainly learns the features for classifying the frequent classes and may neglect features which can transfer to the rare classes.
In contrast, in SSL, without the bias from the labels, the models learn richer features—both features for classifying the frequent classes and features transferable to the rare classes.
3.1 Illustrative Semi-synthetic Experiments
Dataset.
We consider an imbalanced pre-training dataset with label-relevant and label-irrelevant-buttransferable features modiﬁed from CIFAR-10 as shown in Figure 3 (Left).
Classes 1-5 (frequent) each contains 5000 examples.
Classes 6-10 (rare) each has 10 examples.
The left half of an image of classes 1-5 is from classes 1-5 of the original CIFAR-10 and corresponds to the label of that example.
The right half is from a random label-irrelevant image of CIFAR-10.
In contrast, the left half of an example from classes 6-10 is blank, whereas the right half is label-relevant and from classes 6-10 of the original CIFAR-10.
In this setting, features from the left halves are correlated to the classiﬁcation of the frequent classes, while features from the right halves are label-irrelevant for the frequent classes, but can help classify the rare classes.
Note that features from the right halves cannot be directly learned from the rare examples with only 10 examples per class.
Pre-training.
We pre-train the representations on the aforementioned dataset.
We use SimCLR with
ResNet-50.
After pre-training, we ﬁx the representations evaluate with linear probe on the 5 rare classes to test if the model learns proper features for the rare classes.
Results.
As a sanity check, we also include a supervised model trained on only the 50 rare examples (Supervised-rare in Figure 3 (Right)).
As expected, the accuracy is 36 .
5% , which is almost the same as random representations, indicating the model cannot learn the features for the rare classes due to lack of data.
We then compare supervised learning with self-supervised learning on the imbalanced dataset.
SSL representations performs much better than supervised representations on the rare classes ( 70 . 1% vs 44 .
In Figure 3 (Middle), SL mostly activate the left halves on frequent and rare examples, indicating it learns features on the left.
In sharp contrast, SSL activates the whole image on the frequent examples and the right halves on the rare, indicating it learns features from both parts.
3.2 Rigorous Analysis on A Toy Setting
To justify the above conjecture, we instantiate supervised and self-supervised learning in a setting where the features helpful to classify the frequent classes and features transferable to the rare classes can be clearly separated. In this case, we prove that self-supervised learning learns better features than supervised learning.
Data distribution.
Let e 1 , e 2 be two orthogonal unit-norm vectors in the d -dimensional Euclidean space.
Consider the following pre-training distribution P of a 3-way classiﬁcation problem, where the class label y ∈ [3] .
The input x is generated as follows. Let τ > 0 and ρ > 0 be hyperparameters of the distribution.
First sample q uniformly from { 0 , 1 } and ξ ∼ N (0 , I ) from Gaussian distribution.
For the ﬁrst class ( y = 1 ), set x = e 1 − qτe 2 + ρξ . For the second class ( y = 2 ), set x = − e 1 − qτe 2 + ρξ .
For the third class ( y = 3 ), set x = e 2 + ρξ .
Classes 1 and 2 are frequent classes, while class 3 is the rare class, i.e., P ( y =3) P ( y =1) , P ( y =3) P ( y =2) = o (1) .
See Figure 8 for an illustration of this data distribution.
In this case, both e 1 and e 2 are features from the frequent classes 1 and 2.
However, only e 1 helps classify the frequent classes and only e 2 can be transferred to the rare classes.
Algorithm formulations.
For supervised learning, we train a two-layer linear network f W 1 ,W 2 ( x ) (cid:44)
W 2 W 1 x with weight matrices W 1 ∈ R m × d and W 2 ∈ R 3 × m for some m ≥ 3 , and then use the ﬁrst layer W SL = W 1 as the feature for downstream tasks.
Given a linearly separable labeled dataset,
The direction e 2 learned from frequent data by SSL can help classify the rare class.
we learn such a network with minimal norm (cid:107) W 1 (cid:107) 2 F + (cid:107) W 2 (cid:107) 2 F subject to the margin constraint f W 1 ,W 2 ( x ) y ≥ f W 1 ,W 2 ( x ) y (cid:48)
+ 1 for all data ( x, y ) in the dataset and y (cid:48)
(cid:54) =
y .
1 For self-supervised learned, similar to SimSiam [8], we construct positive pairs ( x + ξ, x + ξ (cid:48) ) where x is from the empirical dataset, ξ and ξ (cid:48) are independent random perturbations.
We learn a matrix W SSL ∈ R m × d which minimizes − ˆ E
[( W ( x + ξ )) T ( W ( x + ξ (cid:48) ))]
+ 12 (cid:107) W (cid:62) W (cid:107) 2 F , where the expectation ˆ E is over the empirical dataset and the randomness of ξ and ξ (cid:48) .
The regularization term 12 (cid:107) W (cid:62) W (cid:107) 2 F is introduced only to make the learned features more mathematically tractable.
We use W SSL x as the feature of data x in the downstream task.
Main intuitions.
We compare the features learned by SSL and supervised learning on an imbalanced dataset that contains an abundant (poly in d ) number of data from the frequent classes but only a small (sublinear in d ) number of data from the rare class.
The key intuition behind our analysis is that supervised learning learns only the e 1 direction (which helps classify class 1 vs. class 2) and some random direction that overﬁts to the rare class.
In contrast, self-supervised learning learns both e 1 and e 2 directions from the frequent classes.
Since how well the feature helps classify the rare class (in ID evaluation) depends on how much it correlates with the e 2 direction, SSL provably learns features that help classify the rare class, while supervised learning fails.
This intuition is formalized by the following theorem.
Theorem 3.1.
Let n 1 , n 2 , n 3 be the number of data from the three classes respectively.
Let ρ = d −and τ
= d 15 in the data generative model.
For n 1 , n 2 =
Θ( poly ( d )) and n 3 ≤ d 15 , with probability at least 1 − O ( e − d 110 ) , the following statements hold for any feature dimension m ≥ 3 :
• Let W SL =
[ w 1 , w 2 , · · · , w m ] (cid:62) be the feature learned by SL, then |(cid:104) e 2 , w i (cid:105)| ≤ O ( d − 120 ) for i ∈ [ m ] .
• Let W SSL =
[ ˜ w 1 , ˜ w 2 , · · · , ˜ w m ] (cid:62) be the feature learned by SSL, then (cid:107)
Π e 2 (cid:107) 2 ≥ 1 − O
( d − 15 ) , where Π projects e 2 onto the row span of W SSL .
Conclusion
We discover that self-supervised representations are more robust to class imbalance than supervised representations and explore the underlying cause of this phenomenon.
We hope our study can inspire analysis of self-supervised learning in broader environments in the wild such as domain shift.
A Details of SectionA.1 Implementation Details
Generating Pre-training Datasets.
CIFAR-10 [22] contains 10 classes with 5000 examples per class.
We use exponential imbalance, i.e. for class c , the number of examples is 5000 × e β ( c − 1) .
we consider imbalance ratio r ∈ { 0 . 1 , 0 . 01 } , i.e. the number of examples belonging to the rarest class is 500 or 50.
The total n s is therefore 20431 or 12406.
ImageNet-LT is constructed by Liu et al. [24], which follows the Pareto distribution with the power value 6.
The number of examples from the rarest class is 5.
We construct a long tailed ImageNet following the Pareto distribution with more imbalance, where the number of examples from the rarest class is 3.
The total number of examples n s is 115846 and 80218 respectively.
For each ratio of imbalance, we further downsample the dataset with the sampling ratio in { 0 . 75 , 0 . 5 , 0 . 25 , 0 . 125 } to formulate different number of examples.
To compare with the balanced setting fairly, we also sample balanced versions of datasets with the same number of examples.
Note that each variant of the dataset is ﬁxed after construction for all algorithms.
See the visualization of label distributions of dataset variants in Figure 5.
Training Procedure.
We use ResNet-50 as backbones.
For supervised pre-training, we follow the standard protocol of He et al.
[15] and Kang et al.
On the standard ImageNet-LT, we train the models for 90 epochs with step learning rate decay.
For down-sampled variants, the training epochs are selected with cross validation.
Fo self-supervised learning, the initial learning rate on the standard ImageNet-LT is set to 0.025 with batch-size 256.
We train the model for 300 epochs on the standard ImageNet-LT and adopt cosine learning rate decay following [16, 9].
We
Evavluation.
For in-domain evaluation (ID), we ﬁrst train the the representations on the aforementioned dataset variants, and then train the linear head classiﬁer on the full balanced CIFAR10 or
ImageNet.
We set the initial learning rate to 30 when training the linear head with batch-sizeand train for 100 epochs in total.
For in-domain out-of-domain evaluation (OOD) on ImageNet, we ﬁrst train the the representations on the aforementioned dataset variants, and then ﬁne-tune the model to CUB-200 [31], Stanford Cars [21], Oxford Pets [26], and Aircrafts [25].
The number of examples of these target datasets ranges from 2k to 10k, which is a reasonable scale as the number of examples of the pre-training dataset variants ranges from 10k to 110k.
The representation quality is evaluated with the average performance on the four tasks.
We set the initial learning rate to 0.1 in ﬁne-tuning train for 150 epochs in total.
For in-domain out-of-domain evaluation (OOD) on CIFAR-10, we use
STL-10 as the downstream target tasks and perform linear probe.
A.2 Additional Results
To validate the phenomenon observed in Section 2 is consistent for different self-supervised learning algorithms, we provide the OOD evaluation results of SimSiam trained on ImageNet variants in Figure 6.
SimSiam representations are also less sensitive to class imbalance than supervised representations.
[Figure] [Table] We also provide the numbers of Figure 2 and Figure 6 in Table 1.
B Details of Section 3.
We ﬁrst generate the balanced semi-synthetic dataset with 5000 examples per class.
The left halves of images from classes 1-5 correspond to the labels, while the right halves are random.
The left halves of images from class 6-10 are blank, whereas the right halves correspond to the labels.
We then generate the imbalanced dataset, which consists of the 5000 examples per class from classes 1- (frequent classes) [Table] and 10 examples per class from classes 6-10 (rare classes).
We use Grad-CAM implementation based on https://github.com/meliketoy/gradcam.pytorch and SimCLR implementation from https://github.com/leftthomas/SimCLR .
We provide examples and
Grad-CAM of the semi-synthetic datasets in Figure 7.
To avoid confusing the left and right parts, we disable the random horizontal ﬂip in data augmentation.
[Figure]
C Proof of Theorem 3.
Data distribution.
Let e 1 , e 2 be two orthogonal unit-norm vectors in the d -dimensional Euclidean space.
Consider the following pre-training distribution P of a 3-way classiﬁcation problem, where the class label y ∈ [3] .
The input x is generated as follows.
Let τ > 0 and ρ > 0 be hyperparameters of the distribution.
First sample q uniformly from { 0 , 1 } and ξ ∼ N (0 , I ) from Gaussian distribution.
For the ﬁrst class ( y = 1 ), set x = e 1 − qτe 2 + ρξ .
For the second class ( y = 2 ), set x = − e 1 − qτe 2 + ρξ .
For the third class ( y = 3 ), set x = e 2 + ρξ .
Classes 1 and 2 are frequent classes, while class 3 is the rare class, i.e., P ( y =3) P ( y =1) , P ( y =3)
P ( y =2) =
o (1) .
See Figure 8 for an illustration of this data distribution.
In this case, both e 1 and e 2 are features from the frequent classes 1 and 2.
However, only e 1 helps classify the frequent classes and only e 2 can be transferred to the rare classes.
Algorithm formulations.
For supervised learning, we train a two-layer linear network f W 1 ,W 2 ( x ) (cid:44)
W 2 W 1 x with weight matrices W 1 ∈ R m × d and W 2 ∈ R 3 × m for some m ≥ 3 , and then use the ﬁrst layer W 1 x as the feature for downstream tasks.
Given a linearly separable labeled dataset, we learn such a network with minimal norm (cid:107) W 1 (cid:107) 2 F + (cid:107) W 2 (cid:107) 2 F subject to the margin constraint f W 1 ,W 2 ( x ) y ≥ f W 1 ,W 2 ( x ) y (cid:48)
+1 for all data ( x, y ) in the dataset and y (cid:48)
(cid:54) =
y .
2 Notice that when the norm is minimized, the row span of W 1 is exactly the same as the row span of matrix W SL (cid:44) [ w 1 , w 2 , w 3 ] (cid:62) ∈ R {+3 × d which minimizes (cid:107) w 1 (cid:107) 
+}
+ (cid:107) w 2 (cid:107) 22 + (cid:107) w 3 (cid:107) 22 subject to the margin constraint w (cid:62) y x ≥ w (cid:62) y (cid:48) x
+ 1 for all empirical data ( x, y ) and y (cid:48) (cid:54) =
y .
Therefore, feature W 1 x is equivalent to W SL x for downstream linear probing, so we only need to analyze the feature quality of W SL x for supervised learning.
For self-supervised learned, similar to SimSiam [8], we construct positive pairs ( x + ξ, x + ξ (cid:48) )
where x is from the empirical dataset, ξ and ξ (cid:48) are independent random perturbations.
We learn a matrix W SSL ∈ R 2 × d which minimizes − ˆ E
[( W ( x + ξ )) T ( W ( x + ξ (cid:48) ))]
+ 12 (cid:107) W (cid:62) W (cid:107) 2 F , where the expectation ˆ E is over the empirical dataset and the randomness of ξ and ξ (cid:48) .
The regularization term 12 (cid:107) W (cid:62) W (cid:107) 2 F is introduced only to make the learned features more mathematically tractable.
We use
W SSL x as the feature of data x in the downstream task.
We notate data from the ﬁrst class as x (1)
i
[Equation]
where i ∈ [ n 1 ] and q (1) i ∈ { 0 , 1 } .
Similarly, we notate data from the second class as x (2)
i
[Equation]
where i ∈ [ n 2 ] and q (1) i ∈ { 0 , 1 } .
We notate data from the third class as x (3) i =
e
[Equation]
3) i where i ∈ [ n 3 ] .
Notice that all ξ ( k ) i are independently sampled from N (0 , I ) .
We ﬁrst introduce the following lemma, which gives some high probability properties of independent
Gaussian random variables.
Lemma C.1.
Let ξ i ∼ N (0 , I ) for i ∈ [ n ] .
Then, for any n ≤ poly ( d ) , with probability at least 1 − e − d 110 and large enough d , we have:
• |(cid:104) ξ i , e 1 (cid:105)| ≤ d 110 , |(cid:104) ξ i , e 2 (cid:105)| ≤ d 110 and |(cid:107)
ξ i (cid:107) 22 − d | ≤ 4 d 34 for all i ∈ [ n ] .
• |(cid:104) ξ i , ξ j (cid:105)| ≤ 3 d 35 for all i (cid:54) =
j .
Proof of Lemma C.1.
Let ξ, ξ (cid:48) ∼ N (0 , I ) be two independent random variables.
By the tail bound of normal distribution, we have
[Equation]
By the tail bound of χ 2 d distribution, we have
[Equation]
Since the directions of ξ and ξ (cid:48) are independent, we can bound their correlation with the norm of ξ times the projection of ξ (cid:48) onto ξ :
[Equation]
Since every ξ i and ξ j are independent when i (cid:54)
=
j , by the union bound, we know that with probability at least 1 − ( n 2 +2 n )( e − d 152 d 110 +2 e −√ d ) , we have |(cid:104) ξ i , e 1 (cid:105)| ≤ d 110 , |(cid:104) ξ i , e 2 (cid:105)| ≤ d 110 and |(cid:107) ξ i (cid:107) 22 − d | ≤ 4 d 34 for all i ∈ [ n ] , and also |(cid:104) ξ i , ξ j (cid:105)| ≤ 3 d 35 for all i (cid:54) =
j .
Since the error probability is exponential in d , for large enough d , the error probability is smaller than e − d 110 , which ﬁnishes the proof.
Using the above lemma, we can prove the following lemma which constructs a linear classiﬁer of the empirical dataset with relatively large margin and small norm.
Lemma C.2.
In the setting of Theorem 3.1, let w ∗ 1 =
e 1 , w ∗ 2 = − e 1 , w ∗ 3 = 1 ρd (cid:80) n 3 i =1 ξ (3) i .
Apply
Lemma C.1 to the set of all ξ ( k ) i where k ∈ [3] and i ∈ [ n k ] .
When the high probability outcome of
Lemma C.1 happens, the margin of classiﬁer { w ∗ 1 , w ∗ 2 , w ∗ 3 } is at least 1 − O ( d − 110 ) .
Furthermore, we have (cid:107) w ∗ 3 (cid:107) 22 ≤ O ( d − 15 ) .
Proof of Lemma C.2.
When the high probability outcome of Lemma C.1 happens, we give a lower bound on the margin for all data in the dataset.
For data x = x (1) i in class 1, we have
[Equation]
So the margin on data ( x (1) i , 1) is
[Equation]
Similarly, for data x (2) i in class 2, the margin is at least 1 − O
( d − 110 ) .
For data x = x (3) i in class 3, we have
[Equation]
On the other hand,
[Equation]
So the margin is
[Equation]
Finally, noticing that (cid:107) w
[Equation]
≤ [Equation] 110 ﬁnishes the proof.
Now we are ready to prove the supervised learning part of Theorem 3.1 using our construction:
Proof of Theorem 3.1 (supervised learning part).
We frist apply Lemma C.1 to the set of all ξ ( k ) i where k ∈ [3] and i ∈ [ n k ] .
We consider the situation when the high probability outcome of
Lemma C.1 holds (which happens with probability at least 1 − e − d 110 ).
By Lemma C.2, the constructed classiﬁer { w ∗ 1 , w ∗ 2 , w ∗ 3 } has margin α ≥ 1 − O
( d − 110 ) in this case.
As a result, { 1 α w ∗ 1 , 1 α w ∗ 2 , 1 α w ∗ 3 } is a classiﬁer with margin 1 and norm bounded by
[Equation]
Let { w 1 , w 2 , w 3 } be the solution of supervised learning.
Since its norm cannot be larger than the constructed one, we have (cid:107) w[-1 (cid:107) 
-]
[Equation] 22 + (cid:107) w 3 (cid:107) 22 ≤ 2 + O ( d − 110 ) .
By standard concentration inequality, when n 1 ≥ poly ( d ) , with probability at least 1 − e d − 110 , we have
[Equation]
14) where the expectation is over all the data from class 1 that satisﬁes q (1) i = 0 .
By the deﬁnition of { w 1 , w 2 , w 3 } we know ( w 1 − w 3 ) (cid:62) x
(1) i ≥ 1 for all i ∈ [ n 1 ] , hence averaging over all the classdata with q (1) i = 0 and using the above inequality gives us
[Equation]
A similar analysis for class 2 data gives us
[Equation]
Now we prove that w 1 , w 2 , w 3 all have small correlation with e 2 .
Without loss of generality, we assume w (cid:62) 3 e 1 (cid:44) t ≥ 0 .
If t ≥ 12 , we have
[Equation]
17) which contradicts with (cid:107) w 1 (cid:107)
[Equation]
+ O ( d − 1 10 ) .
Therefore, there must be t ≤ 1 2 , hence
[Equation]
As a result,
[Equation]
which ﬁnishes the proof
[Equation] To prove the self-supervised learning part of Theorem 3.1, we ﬁrst introduce the following lemma which gives some helpful properties of the empirical data matrix.
Lemma C.3.
In the setting of Theorem 3.1, let M (cid:44) E
x
[ xx (cid:62) ] where the expectation is over empirical data.
Then, when n 1 , n 2 ≥ poly ( d ) , with probability at least 1 − e − d 110 , we have: (1) e (cid:62) {+
+} Me 2 ≥ Ω( d 25 ) , and (2) u (cid:62)[-Mu ≤ O (1) for all u ∈ R d such that u (cid:62) e 
-]
Mu ≤ O (1) for all u ∈ R d such that u (cid:62) e 2 = 0 and (cid:107) u (cid:107) 2 = 1 .
Proof of Lemma C.3.
Let n = n 1 + n 2 + n 3 .
We abuse notation and let ξ i ( i ∈ [ n ] ) be the set of all ξ ( k ) i that appears in the empirical data.
Let matrix M (cid:48)
= 1 n (cid:80) ni =
1 ξ i ξ (cid:62) i .
By standard concentration inequalities and union bound, for n ≥ poly ( d ) , with probability at least 1 − 12 e − d 110 , we have that | M (cid:48)
i,j | ≤ 1 d for all i (cid:54) =
j and | M (cid:48)
i,i − 1 | ≤ 1 d for all i ∈
[ d ] .
In this case, for any vector u ∈ R d such that (cid:107) u (cid:107) 2 = 1 and u (cid:62) e 2 = 0 , we have
[Equation]
On the other hand, by the deﬁnition of data distirbution and standard concentration inequalities, for n ≥ poly ( d ) , with probability at least 1 − 12 e − d 110 we have that: at least 13 of all data either is classwith q (1) i = 1 or class 2 with q (2) i = 1 , and (cid:107) 1 n (cid:80) ni
=1 ξ i (cid:107) 2 ≤ O ( 1 d ) .
In this case, e
[Equation]
[Equation]
[Equation] [Equation] [Equation] Using the above lemma, we can prove the self-supervised learning part of Theorem 3.1.
Proof of Theorem3.1(self-supervised learning part).
Let M =
E
x
[ xx (cid:62) ] be the empirical data matrix, where the expectation is over the dataset.
Notice that self-supervised learning objective has the same minimizer as the matrix factorization objective (cid:107) M − W (cid:62) W (cid:107) 2 F , by Eckart–Young–Mirsky theorem we know that the span of ˜ w 1 , ˜ w 2 is exactly the span of the top 2 eigenvectors of matrix
M .
Let M = (cid:80) di
=1 λ i v i v (cid:62) i where λ i is the i -th largest eigenvalue of M with the corresponding eigenvector v i .
We decompose e 2 in the eigenvector basis as e 2 = (cid:80) di =1 ζ i v i .
By Lemma C.3, we know that with probability at least 1 − e − d 110 , we have e (cid:62)
Me 2 ≥ Ω( d 25 ) and u (cid:62) Mu (cid:107) u (cid:107) 22 ≤ O (1) for all u orthogonal to e 2 .
To prove the result regarding self-supervised learning, we only need to prove that ζ 21 ≥ 1 − O ( d − 15 ) in this case.
We ﬁrst show that ζ 21 ≥ 1 2 .
For contradiction, ﬁrst assume ζ 21 ≤ 1 2 .
Deﬁne vector
[Equation]
[Equation] [Equation] which satisﬁes u (cid:62) e 2 = 0 and (cid:107) u (cid:107) 22 ≤ 4 .
Notice that
[Equation]
Mu
[Equation]
[Equation]
which contradicts to u (cid:62) Mu (cid:107) u (cid:107) 22 ≤ O (1)
Therefore, we have ζ 21 ≥ 12 .
To prove that ζ 21 is close to 1 , we let scalar t = 1 ζ 21 − 1 and deﬁne vector
[Equation]
which satisﬁes u (cid:62) e 2 = {+
+}
Since ζ 21 ≥ 12 , we have t ≤ 1 and (cid:107) u (cid:107) 22 ≤ 1 .
As a result, we have
[Equation]
On the other hand, we know that u (cid:62) Mu (cid:107) u
(cid:107) 22 ≤ O (1) .
Comparing these two bounds gives us t 2 ≤ O ( d − 15 ) , which means ζ 21 ≥ 1 − O
( d − 15 ) .
D Related Work
Supervised Learning with Dataset Imbalance.
There exists a line of works studying supervised imbalanced classiﬁcation.
Ando & Huang [1], Buda et al.
[3] proposed to re-sample the data to make the frequent and rare classes appear with equal frequency in training.
Re-weighting assigns different weights for head and tail classes and eases the optimization difﬁculty under class imbalance [11, 29, 33].
Byrd & Lipton [4], Xu et al.
[35] studied the effect of importance weighting and found out that importance weighting does not change the solution without regularization.
Cao et al.
[5] studied reweighted regularization based on classiﬁer margin, but these techniques are limited to supervised imbalanced recognition.
Cao et al.
[6] proposed to regularize the local curvature of loss on imbalanced and noisy datasets.
Recent works also designed speciﬁc losses or training pipelines for imbalanced recognition [18, 17, 32, 38].
Kang et al.
[20] found out that the representations of supervised learning perform better than the classiﬁer itself with class imbalance.
Yang & Xu [36] studied the effect of self-training and self-supervised pre-training on supervised imbalanced recognition classiﬁers.
In contrast, the focus of our paper is the effect of class imbalance on self-supervised representations , which is orthogonal to the self-supervised learning literature.
Self-supervised Learning.
Recent works on self-supervised learning successfully learn representations that approach the supervised baseline on ImageNet and various downstream tasks.
Contrastive learning methods attract positive pairs and drive apart negative pairs [16, 8].
Siamese networks predict the output of the other branch, and use stop-gradient to avoid collapsing [13, 9].
Clustering methods learn representations by performing clustering on the representations and improve the representations with cluster index [7].
Cole et al.
[10] investigated the effect of data quantity and task granularity on self-supervised representations.
Goyal et al.
[12] studied self-supervised methods on large scale datasets in the wild, but they do not consider dataset imbalance explicitly.
Several works have also theoretically studied the success of self-supervised learning [2, 23, 14, 34].
Ethics Statement
Our paper is the ﬁrst to study the problem of robustness to imbalanced training of self-supervised representations.
This setting is of critical importance to AI Ethics (especially fairness or bias concerns), as large real-world datasets tend to be imbalanced in practice, for instance including less examples from under-represented minorities.
Furthermore, pre-training is a standard practice in deep learning, especially for quickly adapting models to new domains, which corresponds to our OOD evaluation scenario.
Our experiments and theoretical analysis show that SSL is more robust than supervised pre-training, especially in the OOD scenario.
As supervised learning is still the de facto standard for pre-training, our work should have a wide impact, encouraging practitioners to use SSL for pre-training instead, or at least consider evaluating the impact of imbalanced pre-training on their downstream task.
We also remark that the paper does not imply at all that the algorithms proposed or studied can guarantee any form of fairness, and they in fact should still suffer from biases.
The paper should be considered as a step towards studying the important technical issue of dataset imbalance, which is closely related to the fairness or biases questions.