Year: 2021

Abstract:
Self-supervised learning (SSL) is a scalable way to learn general visual representations since it learns without labels. However, large-scale unlabeled datasets in the wild often have long-tailed label distributions, where we know little about the behavior of SSL. In this work, we systematically investigate self-supervised learning under dataset imbalance. First, we find via extensive experiments that off-the-shelf self-supervised representations are already more robust to class imbalance than supervised representations. The performance gap between balanced and imbalanced pre-training with SSL is significantly smaller than the gap with supervised learning, across sample sizes, for both in-domain and, especially, out-of-domain evaluation. Second, towards understanding the robustness of SSL, we hypothesize that SSL learns richer features from frequent data: it may learn label-irrelevantbut-transferable features that help classify the rare classes and downstream tasks. In contrast, supervised learning has no incentive to learn features irrelevant to the labels from frequent examples. We validate this hypothesis with semi-synthetic experiments as well as rigorous mathematical analyses on a simplified setting. Third, inspired by the theoretical insights, we devise a re-weighted regularization technique that consistently improves the SSL representation quality on imbalanced datasets with several evaluation criteria, closing the small gap between balanced and imbalanced datasets with the same number of examples.

heading: 1 INTRODUCTION
text: Self-supervised learning (SSL) is an important paradigm for 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 (He et al., 2020; Chen et al., 2020; Grill et al., 2020; Caron et al., 2020; Chen & He, 2021). 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 (Reed, 2001). Curating a class-balanced unlabeled dataset requires the knowledge of labels, which defeats the purpose of leveraging unlabeled data by SSL.
The robustness of SSL algorithms to dataset imbalance remains largely underexplored in the literature, but extensive studies on supervised learning (SL) with imbalanced datasets do not bode well. The performance of vanilla supervised methods degrades significantly on imbalanced datasets (Cui et al., 2019; Cao et al., 2019; Buda et al., 2018), posing challenges to practical applications such as instance segmentation (Tang et al., 2020) and depth estimation (Yang et al., 2021). Many recent works address this issue with various regularization and re-weighting/re-sampling techniques (Ando & Huang, 2017; Wang et al., 2017; Jamal et al., 2020; Cui et al., 2019; Cao et al., 2019; 2021; Wang et al., 2020; Tian et al., 2020; Hong et al., 2021).
In this work, we systematically investigate the representation quality of SSL algorithms under class imbalance. Perhaps surprisingly, we find 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 finetuning 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 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 configurations 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 robust to dataset imbalance? We hypothesize 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 have so limited data that either SL or SSL models overfit 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.2). In addition, we construct a toy example where we can rigorously prove the difference between self-supervised and supervised features in Section 3.1.
Finally, given our theoretical insights, we take a step towards further improving SSL algorithms, closing the small gap between balanced and imbalanced datasets. We identify the generalization gap between the empirical and population pre-training losses on rare data as the key to improvements.
To this end, we design a simple algorithm that first roughly estimates the density of examples with kernel density estimation and then applies a larger sharpness-based regularization (Foret et al., 2020) to the estimated rare examples. Our algorithm consistently improves the representation quality under several evaluation protocols.
We sum up our contributions as follows. (1) We are the first 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. (3) We propose a principled method to improve SSL under unknown dataset imbalance.

heading: 2 EXPLORING THE EFFECT OF CLASS IMBALANCE ON SSL
text: Dataset class imbalance can pose challenge to self-supervised learning in the wild. Without access to labels, we cannot know in advance whether a large-scale unlabeled dataset is imbalanced. Hence, we need to study how SSL will behave under dataset imbalance to deploy SSL in the wild safely. In this section, we systematically investigate the effect of class imbalance on the self-supervised representations with experiments.

heading: 2.1 PROBLEM FORMULATION
text: Class-imbalanced pre-training datasets. We assume the datapoints / inputs are in Rd and come from C underlying classes. Let x denote the input and y denote the corresponding label. Supervised pre-training algorithms have access to the inputs and corresponding labels, whereas self-supervised pre-training only observes the inputs. Given a pre-training distribution P over over Rd × [C], let r denote the ratio of class imbalance. That is, r is the ratio between the probability of the rarest class and the most frequent class: r = minj∈[C] P(y=j)maxj∈[C] P(y=j) ≤ 1. We will construct distributions with varying imbalance ratios and use Pr to denote the dataset with ratio r. We also use Pbal for the case where r = 1, i.e. the dataset is balanced. Large scale data in the wild often follow heavily long-tailed label distributions, i.e., r is small. Throughout this paper we assume that for any class j ∈ [C], the class-conditional distribution Pr(x|y = j) is the same across balanced and imbalanced datasets for all r. The pre-training dataset P̂rn consists of n i.i.d. samples from Pr. Pre-trained models. A feature extractor is a function fφ : Rd → Rm parameterized by neural network parameters φ, which maps inputs to representations. A linear head is a linear function gθ : Rm → RC , which can be composed with fφ to produce the label. SSL algorithms learn φ from unlabeled data. Supervised 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 φ.1
Following the standard evaluation protocol in prior works (He et al., 2020; Chen et al., 2020), we measure the quality of learned representations on both in-domain and out-of-domain datasets with either linear probe or fine-tuning, as detailed below.
In-domain (ID) evaluation tests the performance of representations on the balanced in-domain distribution Pbal with linear probe. Given a feature extractor fφ pre-trained on a pre-training dataset P̂rn with n data points and imbalance ratio r, we train a C-way linear classifier θ on a balanced dataset2 sampled i.i.d. from Pbal. We evaluate the representation quality with the top-1 accuracy of the learned linear head on Pbal. We denote the ID accuracy of supervised pre-trained representations by ASLID (n, r). Note that A SL ID (n, 1) stands for the result with balanced pre-training dataset. For SSL representations, we denote the accuracy by ASSLID (n, r).
Out-of-domain (OOD) evaluation tests the performance of representations by fine-tuning the feature extractor and the head on a (or multiple) downstream target distribution Pt. Starting from a feature extractor fφ (pre-trained on a dataset of size n and imbalance ratio r) and a randomly initialized classifier θ, we fine-tune φ and θ on the target dataset P̂t, and evaluate the representation quality by the expected top-1 accuracy on Pt. We use ASLOOD(n, r) and ASSLOOD(n, r) to denote the resulting accuracies of supervised and self-supervised representations, respectively.
Summary of varying factors. We aim to study the effect of class imbalance to feature qualities on a diverse set of configurations 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 (He et al., 2020), or SimSiam (Chen & He, 2021).

heading: 2.2 EXPERIMENTAL SETUP
text: Datasets. We pre-train the representations on variants of ImageNet (Russakovsky et al., 2015) or CIFAR-10 (Krizhevsky & Hinton, 2009) with a wide range of numbers of examples and ratios of imbalance. Following Liu et al. (2019), we consider exponential and Pareto distributions, which closely simulate the natural long-tailed distributions. We consider imbalance ratio in {1, 0.004, 0.0025} for ImageNet and {1, 0.1, 0.01} for CIFAR-10. 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. Note that we fix the variant of the dataset when comparing different algorithms. For ID evaluation, we use the original CIFAR-10 or ImageNet training set for the training phase of linear probe and use
1It is well-known that the composition of the head and features learned from supervised learning is more sensitive to imbalanced dataset than the quality of feature extractor φ (Cao et al., 2019; Kang et al., 2020). Please also see Table 4 in Appendix C for a comparison.
2We essentially use the largest balanced labeled ID dataset for this evaluation, which oftentimes means the entire curated training dataset, such as CIFAR-10 with 50,000 examples and ImageNet with 1,281,167 examples.
the original validation set for the final evaluation. For OOD evaluation of representations learned on CIFAR-10, we use STL-10 (Coates et al., 2011) as the target /downstream dataset. For OOD evaluation of representations learned on ImageNet, we fine-tune the pre-trained feature extractors on CUB-200 (Wah et al., 2011), Stanford Cars (Krause et al., 2013), Oxford Pets (Parkhi et al., 2012), and Aircrafts (Maji et al., 2013), and measure the representation quality with average accuracy on the downstream tasks.
Models. We use ResNet-18 on CIFAR-10 and ResNet-50 on ImageNet as backbones. For supervised pre-training, we follow the standard protocol of He et al. (2016) and Kang et al. (2020). For selfsupervised pre-training, we consider MoCo v2 (He et al., 2020) and SimSiam (Chen & He, 2021). We run each evaluation experiment with 3 seeds and report the average and standard deviation in the figures. Further implementation details are deferred to Section A.1.

heading: 2.3 RESULTS: SELF-SUPERVISED LEARNING IS MORE ROBUST THAN SUPERVISED LEARNING TO DATASET IMBALANCE
text: In Figure 2, we plot the results of ID and OOD evaluations, respectively. 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., ASSL(n, 1)−ASSL(n, r), is smaller than the gap of supervised representations, i.e., ASL(n, 1)−ASL(n, r), consistently in all configurations. Furthermore, we compute the relative accuracy gap to balanced dataset ∆SSL(n, r) , (ASSL(n, 1)−ASSL(n, r))/ASSL(n, 1) in Figure 1. We observe that 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 across the board,
∆SSL(n, r) , ASSL(n, 1)−ASSL(n, r) ASSL(n, 1) ∆SL(n, r) , A SL(n, 1)−ASL(n, r) ASL(n, 1) . (1)
4The maximum n is smaller for extreme imbalance. The standard deviation comes only from the randomness of evaluation. We do not include the stddev for ImageNet ID due to limitation of computation resources.
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, we observe that representations from supervised pre-training perform better than self-supervised pre-training in ID evaluation with reasonably large n, while self-supervised pre-training is better in OOD evaluation. This phenomenon is orthogonal to our observation that SSL is more robust to dataset imbalance, and is consistent with recent works (e.g., Chen et al. (2020); He et al. (2020)) which also observed that SSL performs slightly worse than supervised learning on balanced ID evaluation but better on OOD tasks.

heading: 3 ANALYSIS
text: We have found out with extensive experiments that self-supervised representations are more robust to class imbalance than supervised representations. A natural and fundamental question arises: 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 for the rare classes. In this case, one may want to resort to the features learned from the frequent classes for help. However, due to the supervised nature of classification tasks, the supervised model mainly learns the features that help classify the frequent classes and may neglect other features which can transfer to the rare classes and potentially the downstream tasks. Partly because of this, Jamal et al. (2020) explicitly encourage the model to learn features transferable from the frequent to the rare classes with meta-learning. In contrast, in self-supervised learning, without the bias or incentive from the labels, the models can learn richer features that capture the intrinsic structures of the inputs—both features useful for classifying the frequent classes and features transferable to the rare classes simultaneously. In the following subsections, we study this reason in a setting where the two sets of features can be clearly separated. In Section 3.1, we proved SSL learns the with Gaussian input. In Section 3.2, we validate the hypothesis on a variant of CIFAR-10.

heading: 3.1 RIGOROUS ANALYSIS ON A TOY SETTING
text: 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 separated. In this case, we prove that self-supervised learning learns better features than supervised learning.
Data distribution. Let e1, e2 be two orthogonal unit-norm vectors in the d-dimensional Euclidean space. Consider the following pre-training distribution P of a 3-way classification 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 first class (y = 1), set x = e1 − qτe2 + ρξ. For the second class (y = 2), set x = −e1 − qτe2 + ρξ. For the third class (y = 3), set x = e2 + ρξ. 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 3 for an illustration of this data distribution. In this case, both e1 and e2 are features from the frequent classes 1 and 2. However, only e1 helps classify the frequent classes and only e2 can be transferred to the rare classes.
Algorithm formulations. For supervised learning, we train a two-layer linear network fW1,W2(x) , W2W1x with weight matrices W1 ∈ Rm×d and W2 ∈ R3×m for some m ≥ 3, and then use the first layer W1x as the feature for downstream tasks. Given a linearly separable labeled dataset, we learn such a network with minimal norm ‖W1‖2F +‖W2‖2F subject to the margin constraint fW1,W2(x)y ≥ fW1,W2(x)y′+1 for all data (x, y) in the dataset and y
′ 6= y.5 Notice that when the norm is minimized, the row span of W1 is exactly the same as the row span of matrix WSL , [w1, w2, w3]> ∈ R3×d which minimizes ‖w1‖22 + ‖w2‖22 + ‖w3‖22 subject to the margin constraint w>y x ≥ w>y′x+ 1 for all empirical data (x, y) and y′ 6= y. Therefore, featureW1x is equivalent toWSLx for downstream linear probing, so we only need to analyze the feature quality of WSLx for supervised learning. For selfsupervised learned, similar to SimSiam (Chen et al., 2020), we construct positive pairs (x+ ξ, x+ ξ′) where x is from the empirical dataset, ξ and ξ′ are independent random perturbations. We learn a matrix WSSL ∈ R2×d which minimizes −Ê[(W (x+ ξ))T (W (x+ ξ′))] + 12‖W
>W‖2F , where the expectation Ê is over the empirical dataset and the randomness of ξ and ξ′. The regularization term 1 2‖W
>W‖2F is introduced only to make the learned features more mathematically tractable. We use WSSLx 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 sampled from the data distribution P 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 the e1 direction (which helps classify class 1 vs. class 2) and some random direction that overfits to the rare class. In contrast, self-supervised learning learns both e1 and e2 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 e2 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 n1, n2, n3 be the number of data from the three classes respectively. Let ρ = d− 1 5 and τ = d 1 5 in the data generative model. For n1, n2 = Θ(poly(d)) and n3 ≤ d 1 5 , with probability at least 1−O(e−d 1 10 ), the following statements hold:
• Let WSL = [w1, w2, w3]> be the feature learned by SL, then |〈e2, wi〉| ≤ O(d− 1 20 ) for i ∈ [3]. • Let WSSL = [w̃1, w̃2]> be the feature learned by SSL, then ‖Πe2‖2 ≥ 1 − O(d− 1 5 ), where Π
projects e2 onto the span of w̃1 and w̃2.
Supervised learning results in features WSL whose rows have small correlation with the transferable feature e2, indicating that supervised learning only learns features for classifying the frequent classes and ignore the potential transferable features. In contrast, self-supervised learning recovers e2 well, even though e2 is not relevant to classifying the frequent classes.

heading: 3.2 ILLUSTRATIVE SEMI-SYNTHETIC EXPERIMENTS
text: In the previous subsection, we have shown that self-supervised learning provably learns labelirrelevant-but-transferable features from the frequent classes which can help classify the rare class in the toy case, while supervised learning mainly focuses on the label-relevant features. However, in real-world datasets, it is intractable to distinguish the two groups of features. To amplify this effect in a real-world dataset and highlight the insight of the theoretical analysis, we design a semi-synthetic experiment on SimCLR (Chen et al., 2020) to validate our conclusion.
5Previous work shows that deep linear networks trained with gradient descent using logistic loss converge to this min norm solution in direction (Ji & Telgarsky, 2018).
Dataset. In the theoretical analysis above, the frequent classes contain both features related to the classification of frequent classes and features transferable to the the rare classes. Similarly, we consider an imbalanced pre-training dataset with two groups of features modified from CIFAR-10 as shown in Figure 4 (Left). We construct classes 1-5 as the frequent classes, where each class contains 5000 examples. Classes 6-10 are the rare classes, where each class has 10 examples. In this case, the ratio of imbalance r = 0.002. Each image from classes 1-5 consists of a left half and a right half. The left half of an example 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 image of CIFAR-10, which is label-irrelevant. 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 of the images are correlated to the classification 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 classes since they have only 10 examples per class. This is consistent with the setting of Theorem 3.1.
Pre-training. We pre-train the representations on the semi-synthetic imbalanced dataset. For supervised learning, we use ResNet-50 on this 10-way classification task. For self-supervised learning, we use SimCLR with ResNet-50. To avoid confusing the left and right parts, we disable the random horizontal flip in the data augmentation. After pre-training, we fix the representations and train a linear classifier on top of the representations with balanced data from the 5 rare classes (25000 examples in total) to test if the model learns proper features for the rare classes during pre-training. In Figure 4 (Right), we test the classifier on the rare classes. In Figure 4 (Middle), we further visualize the Grad-CAM (Selvaraju et al., 2017) of the representations on the held-out set6.
Results. As a sanity check, we first pre-train a supervised model with only the 50 rare examples and train the linear head classifier with 25000 examples from the rare classes (5-way classification) to see if the model can learn proper features for the rare classes with only rare examples (Supervised-rare in Figure 4 (Right)). As expected, the accuracy is 36.5%, which is almost the same as randomly initialized representations with trained head classifier, indicating that the model cannot learn the features for the rare classes with only rare examples due to the limited number of examples. We then compare supervised learning with self-supervised learning on the whole semi-synthetic dataset. In Figure 4 (Right), self-supervised representations performs much better than supervised representations on the rare classes (70.1% vs 44.3%). We further visualize the activation maps of representations with Grad-CAM. Supervised learning mostly activate the left halves of the examples for both frequent and rare classes, indicating that it mainly learn features on the left. In sharp contrast, self-supervised
6CIFAR images are of low resolution. For visualization, we use high resolution version of the CIFAR-10 images in Figure 4 (Middle). We also provide the visualization on original CIFAR-10 images in Figure 7.
learning activates the whole image on the frequent examples and the right part on the rare examples, indicating that it learns features from both parts.

heading: 4 IMPROVING SSL ON IMBALANCED DATASETS WITH REGULARIZATION
text: In this section, we aim to further improve the performance of SSL to close the gap between imbalanced and balanced datasets. Many prior works on imbalanced supervised learning regularize the rare classes more strongly, motivated by the observation that the rare classes suffer from more overfitting (Cao et al., 2019; 2021). Inspired by these works, we compute the generalization gaps (i.e., the differences between empirical and validation pre-training losses) on frequent and rare classes for the stepimbalance CIFAR-10 datasets (where 5 classes are frequent class with 5000 examples per class and the rest are rare with 50 examples per class). Indeed, as shown in Table 1, we still observe a similar phenomenon—the frequent classes have much smaller pre-training generalization gap than the rare classes (0.035 vs. 0.081), which indicates the necessity of more regularization on the rare classes.
We need a data-dependent regularizer that can have different effects on rare and frequent examples. Thus, weight decay or dropout (Srivastava et al., 2014) are not suitable. The prior work of Cao et al. (2019) regularizes the rare classes more strongly with larger margin, but it does not apply to SSL where no labels are available. Inspired by Cao et al. (2021), we adapt sharpness-aware minimization (SAM) (Foret et al., 2021) to imbalanced SSL.
Reweighted SAM (rwSAM). SAM improves model generalization by penalizing loss sharpness. Suppose the training loss of the representation fφ is L̂(φ), i.e. L̂(φ) = 1n ∑n j=1 `(xj , φ). SAM seeks parameters where the loss is uniformly low in the neighboring area,
min φ L̂(φ+ (φ)), where (φ) = arg max ‖ ‖<ρ >∇φL̂(φ). (2)
To take the weight of different examples into account, we add reweighting to the inner maximization step of SAM. Intuitively, we wish the optimization landscape to be flatter for rare examples, which is in effect regularizing the model more on rare examples. Concretely, consider the reweighted training loss associated with weight vector w ∈ Rn, L̂w(φ) = 1n ∑n j=1 wj`(xj , φ). The reweighted SAM
objective re-weights the regularization-related terms (e.g., w) but not the training loss L̂:
min φ L̂(φ+ w(φ)), where w(φ) = arg max ‖ ‖<ρ >∇φL̂w(φ). (3)
Assigning Weight with Kernel Density Estimation. The weight wj of an example xj should be inversely correlated with the frequency of the corresponding class yj . However, we have no access to the labels. In order to approximate the frequency of examples, we use kernel density estimation on top of the representations fφ. Concretely, denote by K(·, h) the Gaussian density with bandwidth h. We assign wi to be inversely correlated with the estimated density, i.e., wi =( 1 n ∑n j=1K(fφ(xi) − fφ(xj), h) )−α where h and α > 0 are hyperparameters selected by cross validation.

heading: 4.1 EXPERIMENTS
text: We test the proposed rwSAM on CIFAR-10 with step or exponential imbalance and ImageNetLT (Liu et al., 2019). After self-supervised pre-training on the long-tailed dataset, we evaluate the representations by (1) linear probing on the balanced in-domain dataset and (2) fine-tuning on downstream target datasets. For (1) and (2), we compare with SSL, SSL+SAM (w/o reweighting), and SSL balanced, which learns the representations on the balanced dataset with the same number of examples. Implementation details and additional results are deferred to Section C.
Results. Table 1 summarizes results on long tailed CIFAR-10. With both step and exponential imbalance, rwSAM improves the performance of SimSiam over 1%, and even surpasses the performance of SimSiam on balanced CIFAR-10 with the same number of examples. Gap Freq. refers to the gap between training and test SimSiam loss on the 5 frequent classes, while Gap Rare refers to that on the 5 rare classes. Note that compared to SimSiam, rwSAM closes the generalization gap of pre-training loss on rare examples from 0.081 to 0.066, which verifies the effect of re-weighted regularization. In Table 2, we provide the result of fine-tuning on downstream tasks with representations pre-trained on ImageNet-LT. The proposed method improves the transferability of representations to downstream tasks consistently.

heading: 5 RELATED WORK
text: Supervised Learning with Dataset Imbalance. There exists a line of works studying supervised imbalanced classification. Ando & Huang (2017); Buda et al. (2018) 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 difficulty under class imbalance (Cui et al., 2019; Tang et al., 2020; Wang et al., 2017). Byrd & Lipton (2019); Xu et al. (2021) studied the effect of importance weighting and found out that importance weighting does not change the solution without regularization. Cao et al. (2019) studied reweighted regularization based on classifier margin, but these techniques are limited to supervised imbalanced recognition. Cao et al. (2021) proposed to regularize the local curvature of loss on imbalanced and noisy datasets. Recent works also designed specific losses or training pipelines for imbalanced recognition (Jamal et al., 2020; Hong et al., 2021; Wang et al., 2021; Zhang et al., 2021). Kang et al. (2020) found out that the representations of supervised learning perform better than the classifier itself with class imbalance. Yang & Xu (2020) studied the effect of self-training and self-supervised pre-training on supervised imbalanced recognition classifiers. 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 (He et al., 2020; Chen et al., 2020). Siamese networks predict the output of the other branch, and use stop-gradient to avoid collapsing (Grill et al., 2020; Chen & He, 2021). Clustering methods learn representations by performing clustering on the representations and improve the representations with cluster index (Caron et al., 2020). Cole et al. (2021) investigated the effect of data quantity and task granularity on self-supervised representations. Goyal et al. (2021) 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 (Arora et al., 2019; Lee et al., 2020; HaoChen et al., 2021; Wei et al., 2021).

heading: 6 CONCLUSION
text: In this paper, we study self-supervised learning under class imbalance. We discover that selfsupervised 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 selfsupervised learning in broader environments in the wild such as domain shift, and provide insights for the design of future unsupervised learning methods.
REPRODUCIBILITY STATEMENT
To ensure reproducibility, we describe the implementation details of the algorithms and the construction of the datasets in Section A.1 and Section C. The code of the experiments in Section 2, Section 3.2 and Section 4 is provided in the supplementary material. We describe the setting and data assumptions of the toy case in Section 3.1 and provide the proof in Section D.
ETHICS STATEMENT
Our paper is the first 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.

heading: A DETAILS OF SECTION 2
text: A.1 IMPLEMENTATION DETAILS
Generating Pre-training Datasets. CIFAR-10 (Krizhevsky & Hinton, 2009) 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 ns is therefore 20431 or 12406. ImageNet-LT is constructed by Liu et al. (2019), 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 ns 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 fixed after construction for all algorithms. See the visualization of label distributions of dataset variants in Figure 5.
Training Procedure. For supervised pre-training, we follow the standard protocol of He et al. (2016) and Kang et al. (2020). 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 (He et al., 2020; Chen & He, 2021). We train the models for more epochs on the down sampled variants to ensure the same number of total iterations. The code on CIFAR-10 LT is adapted from https://github.com/Reza-Safdari/SimSiam-91. 9-top1-acc-on-CIFAR10.
Evavluation. For in-domain evaluation (ID), we first train the the representations on the aforementioned dataset variants, and then train the linear head classifier on the full balanced CIFAR10 or ImageNet. We set the initial learning rate to 30 when training the linear head with batch-size 4096 and train for 100 epochs in total. For in-domain out-of-domain evaluation (OOD) on ImageNet, we first train the the representations on the aforementioned dataset variants, and then fine-tune the model to CUB-200 (Wah et al., 2011), Stanford Cars (Krause et al., 2013), Oxford Pets (Parkhi et al., 2012), and Aircrafts (Maji et al., 2013). 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 fine-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.
We also provide the numbers of Figure 2 and Figure 6 in Table 3.

heading: B DETAILS OF SECTION 3.2
text: We first 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-5 (frequent classes), 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.
C DETAILS OF SECTION 4
C.1 IMPLEMENTATION DETAILS
We use the same implementation as Section 2 for supervised and self-supervised learning baselines. We implement sharpness-aware minimization following (Foret et al., 2021). In each step of update, we first compute the reweighted loss L̂w(φ) = 1ns ∑ns j=1 wj`(x s j , φ) and compute its gradient w.r.t. φ,
∇φL̂w(φ). Then we can compute φ as φ = ρsgn(∇φL̂w(φ))|∇φL̂w(φ)|q−1/ ( ‖∇φL̂w(φ)‖qq )1/p ,
where 1p + 1 q = 1. Finally, we update the model on the loss L̂(φ) by φ = φ− η∇φL̂(φ+ (φ)). We select the hyperparameters ρ and α with cross validation. On ImageNet-LT and iNaturalist, ρ = 2 and α = 0.5. On CIFAR-10-LT, ρ = 5 and α = 1.2.
C.2 ADDITIONAL RESULTS
We further introduce another evaluation protocol of the representations: following the protocol of Kang et al. (2020); Yang & Xu (2020), fine-tuning on imbalanced dataset with supervision, and then re-training the linear classifier with class-aware resampling, to compare with supervised imbalanced recognition methods. In MoCo V2 pre-training, we use the standard data augmentation following He et al. (2020). In fine-tuning, we use RandAugment (Cubuk et al., 2020). For (3), we further compare with CRT (Kang et al., 2020), LADE (Hong et al., 2021), and RIDE (Wang et al., 2021), which are
strong methods tailored to supervised imbalanced recognition. Results are provided in Table 4. Supervised here refers to training the feature extractor and linear classifier with supervision on the imbalanced dataset directly. CRT first trains the feature extractor with supervision, and then re-trains the classifier with resampled loss. Note that CRT is performing better than Supervised, indicating that the composition of the head and features learned from supervised learning is more sensitive to imbalanced dataset than the quality of feature extractor itself.
Even with a simple pre-training and fine-tuning pipeline, MoCo V2 representations can be comparable with much more complicated state-of-the-arts tailored to supervised imbalanced recognition, further corroborating the power of SSL under class imbalance. With rwSAM, we can further improve the result of MoCo V2.

heading: D PROOF OF THEOREM 3.1
text: We notate data from the first class as x(1)i = e1 − q (1) i τe2 + ρξ (1) i where i ∈ [n1] and q (1) i ∈ {0, 1}. Similarly, we notate data from the second class as x(2)i = −e1 − q (2) i τe2 + ρξ (2) i where i ∈ [n2] and q (1) i ∈ {0, 1}. We notate data from the third class as x (3) i = e2 + ρξ (3) i where i ∈ [n3]. Notice that all ξ(k)i are independently sampled from N (0, I). We first introduce the following lemma, which gives some high probability properties of independent Gaussian random variables. Lemma D.1. Let ξi ∼ N (0, I) for i ∈ [n]. Then, for any n ≤ poly(d), with probability at least 1− e−d 1 10 and large enough d, we have:
• |〈ξi, e1〉| ≤ d 1 10 , |〈ξi, e2〉| ≤ d 1 10 and |‖ξi‖22 − d| ≤ 4d 3 4 for all i ∈ [n].
• |〈ξi, ξj〉| ≤ 3d 3 5 for all i 6= j.
Proof of Lemma D.1. Let ξ, ξ′ ∼ N (0, I) be two independent random variables. By the tail bound of normal distribution, we have
Pr ( |〈ξ, e1〉| ≥ d 1 10 ) ≤ d− 110 · e− d 1 5 2 . (4)
By the tail bound of χ2d distribution, we have Pr ( |‖ξ‖22 − d| ≥ 4d 3 4 ) ≤ 2e− √ d. (5)
Since the directions of ξ and ξ′ are independent, we can bound their correlation with the norm of ξ times the projection of ξ′ onto ξ:
Pr ( |〈ξ, ξ′〉| ≥ 3d 35 ) ≤ Pr ( ‖ξ‖2 ≥ √ d+ 2d 3 8 ) + Pr ( |〈ξ′, ξ ‖ξ‖ 〉| ≥ d 110 ) ≤ e − d 1 5 2 d 1 10 + 2e− √ d.
(6)
Since every ξi and ξj are independent when i 6= j, by the union bound, we know that with probability
at least 1− (n2 + 2n)( e − d
1 5 2
d 1 10
+ 2e− √ d), we have |〈ξi, e1〉| ≤ d 1 10 , |〈ξi, e2〉| ≤ d 1 10 and |‖ξi‖22−d| ≤
4d 3 4 for all i ∈ [n], and also |〈ξi, ξj〉| ≤ 3d 3 5 for all i 6= j. Since the error probability is exponential in d, for large enough d, the error probability is smaller than e−d 1 10 , which finishes the proof.
Using the above lemma, we can prove the following lemma which constructs a linear classifier of the empirical dataset with relatively large margin and small norm.
Lemma D.2. In the setting of Theorem 3.1, let w∗1 = e1, w∗2 = −e1, w∗3 = 1ρd ∑n3 i=1 ξ (3) i . Apply Lemma D.1 to the set of all ξ(k)i where k ∈ [3] and i ∈ [nk]. When the high probability outcome of Lemma D.1 happens, the margin of classifier {w∗1 , w∗2 , w∗3} is at least 1−O(d− 1 10 ). Furthermore, we have ‖w∗3‖22 ≤ O(d− 1 5 ).
Proof of Lemma D.2. When the high probability outcome of Lemma D.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
w∗1 >x = 1 + 〈ξ(1)i , e1〉ρ ≥ 1− ρd 1 10 , (7)
w∗2 >x = −1 + 〈ξ(1)i , e1〉ρ ≤ −1 + ρd 1 10 , (8)
w∗3 >x =
1
ρd
( e1 − q(1)i τe2 + ρξ (1) i )> n3∑ j=1 ξ (3) j  ≤ n3(τ + 1) ρd d 1 10 + 3n3 d d 3 5 . (9)
So the margin on data (x(1)i , 1) is
w∗1 >x− w∗3 >x ≥ 1− ρd 110 − n3(τ + 1) ρd d 1 10 − 3n3 d d 3 5 ≥ 1−O(d− 110 ). (10)
Similarly, for data x(2)i in class 2, the margin is at least 1−O(d− 1 10 ).
For data x = x(3)i in class 3, we have
w∗3 >x =
1
ρd  n3∑ j=1 ξ (3) j > (e2 + ρξ(3)i ) ≥ 1d‖ξ(3)i ‖22 − 3n3d d 35 − n3d 1 10 ρd ≥ 1−O(d− 15 ). (11)
On the other hand,
w∗1 >x = 〈ρξ(3)i , e1〉 ≤ ρd 1 10 , (12)
w∗2 >x = 〈ρξ(3)i ,−e1〉 ≤ ρd 1 10 . (13)
So the margin is
w∗3 >x−max{w∗1 >x,w∗2 >x} ≥ w∗3 >x− ρd 110 ≥ 1−O(d− 110 ). (14)
Finally, noticing that ‖w∗3‖2 ≤ 2n3 √ d ρd ≤ 2d − 110 finishes 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 D.1 to the set of all ξ(k)i where k ∈ [3] and i ∈ [nk]. We consider the situation when the high probability outcome of Lemma D.1 holds (which happens with probability at least 1−e−d 1 10 ). By Lemma D.2, the constructed classifier {w∗1 , w∗2 , w∗3} has margin α ≥ 1−O(d− 1 10 ) in this case. As a result, { 1αw ∗ 1 , 1 αw ∗ 2 , 1 αw ∗ 3} is a classifier with margin 1 and norm bounded by
‖ 1 α w∗1‖22 + ‖ 1 α w∗2‖22 + ‖ 1 α w∗3‖22 = 2 + ‖w∗3‖22 α2 ≤ 2 +O(d− 110 ). (15)
Let {w1, w2, w3} be the solution of supervised learning. Since its norm cannot be larger than the constructed one, we have ‖w1‖22 + ‖w2‖22 + ‖w3‖22 ≤ 2 + O(d− 1 10 ). By standard concentration inequality, when n1 ≥ poly(d), with probability at least 1− ed − 1
10 , we have∣∣∣Ei∈[n1],q(1)i =0[x(1)i ]− e1∣∣∣ ≤ d− 110 , (16) where the expectation is over all the data from class 1 that satisfies q(1)i = 0. By the definition of {w1, w2, w3} we know (w1 − w3)>x(1)i ≥ 1 for all i ∈ [n1], hence averaging over all the class 1 data with q(1)i = 0 and using the above inequality gives us
(w1 − w3)>e1 ≥ 1− ‖w1 − w3‖2 · d− 1 10 ≥ 1−O(d− 110 ). (17)
A similar analysis for class 2 data gives us
(w2 − w3)>(−e1) ≥ 1−O(d− 1 10 ). (18)
Now we prove that w1, w2, w3 all have small correlation with e2. Without loss of generality, we assume w>3 e1 , t ≥ 0. If t ≥ 12 , we have
〈w1, e1〉2 + 〈w2, e1〉2 + 〈w3, e1〉2 ≥ ( t+ 1−O(d− 110 ) )2 > 2.25−O(d− 110 ), (19)
which contradicts with ‖w1‖22 + ‖w2‖22 + ‖w3‖22 ≤ 2 +O(d− 1 10 ). Therefore, there must be t ≤ 12 , hence
〈w1, e1〉2 + 〈w2, e1〉2 + 〈w3, e1〉2 (20) ≥ ( 1 + t−O(d− 110 ) )2 + ( 1− t−O(d− 110 ) )2 + t2 (21)
≥ 2 + 3t2 −O(d− 110 ) (22)
≥ 2−O(d− 110 ). (23) As a result,
〈w1, e2〉2 + 〈w2, e2〉2 + 〈w3, e2〉2 (24) ≤ ‖w1‖22 + ‖w2‖22 + ‖w3‖22 − 〈w1, e1〉2 − 〈w2, e1〉2 − 〈w3, e1〉2 (25)
≤ ( 2 +O(d− 1 10 ) ) − ( 2−O(d− 110 ) )
(26)
≤ O(d− 110 ), (27) which finishes the proof.
To prove the self-supervised learning part of Theorem 3.1, we first introduce the following lemma which gives some helpful properties of the empirical data matrix.
Lemma D.3. In the setting of Theorem 3.1, let M , Ex[xx>] where the expectation is over empirical data. Then, when n1, n2 ≥ poly(d), with probability at least 1 − e−d 1 10 , we have: (1) e>2 Me2 ≥ Ω(d 2 5 ), and (2) u>Mu ≤ O(1) for all u ∈ Rd such that u>e2 = 0 and ‖u‖2 = 1.
Proof of Lemma D.3. Let n = n1 + n2 + n3. We abuse notation and let ξi (i ∈ [n]) be the set of all ξ (k) i that appears in the empirical data. Let matrix M ′ = 1n ∑n i=1 ξiξ > i . By standard concentration inequalities and union bound, for n ≥ poly(d), with probability at least 1− 12e −d 1 10 , we have that |M ′i,j | ≤ 1d for all i 6= j and |M ′ i,i − 1| ≤ 1d for all i ∈ [d]. In this case, for any vector u ∈ R
d such that ‖u‖2 = 1 and u>e2 = 0, we have
u>Mu ≤ 2‖u‖22 + 2u>(ρ2M ′)u ≤ 2 + 2ρ2 + 2ρ2‖M ′ − I‖F ≤ O(1). (28)
On the other hand, by the definition of data distirbution and standard concentration inequalities, for
n ≥ poly(d), with probability at least 1− 12e −d 1 10 we have that: at least 13 of all data either is class 1 with q(1)i = 1 or class 2 with q (2) i = 1, and ‖ 1n ∑n i=1 ξi‖2 ≤ O( 1 d ). In this case,
e>2 Me2 = Ex[(e>2 x)2] ≥ (Ex[e>2 x])2 ≥
( 1
3 τ − e>2
( 1
n n∑ i=1 ξi
))2 ≥ Ω(τ2) = Ω(d 25 ). (29)
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 = Ex[xx>] 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 ‖M −W>W‖2F , 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 = ∑d i=1 λiviv > i where λi is the i-th largest eigenvalue of M with the corresponding
eigenvector vi. We decompose e2 in the eigenvector basis as e2 = ∑d i=1 ζivi.
By Lemma D.3, we know that with probability at least 1− e−d 1 10 , we have e>2 Me2 ≥ Ω(d 2 5 ) and u>Mu ‖u‖22
≤ O(1) for all u orthogonal to e2. To prove the result regarding self-supervised learning, we only need to prove that ζ21 ≥ 1−O(d− 1 5 ) in this case.
We first show that ζ21 ≥ 12 . For contradiction, first assume ζ 2 1 ≤ 12 . Define vector
u = ζ1v1 − ∑d i=2 ζivi
(1− ζ21 ) . (30)
which satisfies u>e2 = 0 and ‖u‖22 ≤ 4. Notice that
u>Mu
‖u‖22 ≥ e
> 2 Me2
4 ≥ Ω(d 25 ), (31)
which contradicts to u >Mu ‖u‖22 ≤ O(1). Therefore, we have ζ21 ≥ 12 .
To prove that ζ21 is close to 1, we let scalar t = 1 ζ21 − 1 and define vector
u = −tζ1v1 + d∑ i=2 ζivi, (32)
which satisfies u>e2 = 0. Since ζ21 ≥ 12 , we have t ≤ 1 and ‖u‖ 2 2 ≤ 1. As a result, we have
u>Mu
‖u‖22 ≥ t2λ1ζ21 + d∑ i=2 λiζ 2 i ≥ t2e>2 Me2 ≥ Ω(t2d 2 5 ). (33)
On the other hand, we know that u >Mu ‖u‖22 ≤ O(1). Comparing these two bounds gives us t2 ≤ O(d− 15 ), which means ζ21 ≥ 1−O(d− 1 5 ).

Review summaries:
The authors explore the robustness of feature learning via self-supervised learning (SSL) and supervised learning (SL) with imbalance datasets.  Generally, SL can learn better features than SSL, and features are better from balanced than from imbalance datasets.  However, with imbalance data, SSL is more robust than SL in terms of the performance difference of features learned from balance and imbalance datasets.  That is, the performance does not drop as much for SSL with imbalance data from balanced data.   The robustness is observed from both in-domain (ID)  and out-of-domain (OOD) tasks (different downstream tasks).  They hypothesized that "SSL learns richer features from frequent data that are transferable to rare data."  Using a synthetic dataset with 2 frequent classes and 1 rare class, they observe that SL learns one feature to distinguish the two frequent classes and one feature to overfit the rare class.  However, SSL learns two features that can classify the 3 classes well. That is, SSL can learn label-irrelevant-but-transferable features from the frequent classes which can help classify the rare class.   To illustrate the observation from a synthetic dataset also exists in the real world, they construct a semi-synthetic dataset.  From CIFAR 10, they generate 5 fequent classes and 5 rare classes.  For the frequent classes, the right half of the image is replaced by a random half image, ie not relevant to the label.  For the rare classes, the left half is replaced by blank.  Using activation maps, they observe that SL learn features from only the left half (relevant to the label) but not from the right half (irrelevant to the label).   However, SSL learns features from both halves.  To further improve robustness, they adapt SAM to the imbalance scenario: rwSAM. SAM improves model generalization by penalizing loss sharpness. For SAM, the loss is uniformly low in the neighboring area.   To have a flatter region for rare examples, rwSAM reweights rare examples in the inner maximization step of SAM.  Since labels are not available, they estimate kernel density on feature vectors and the weight is inversely proportional to the density.  On two datasets, empirical results indicate rwSAM improves performance over SAM.  

The paper studies the differences between representations pretrained with Self-Supervised Learning (SSL) and the standard Supervised Learning (SL) frameworks in the In-Domain (ID) and Out-Of-Domain (OOD) settings of the image classification task with an imbalanced dataset.   The paper presents different results that are relevant to the computer vision and machine learning community:  1) In the ID setting, representations pretrained with SL outperform those pretrained with SSL. However, in the OOD setting, representations pretrained with SSL outperform those pretrained with SL and are even robust to the imbalance factor. 2) An analysis in Section 3.1 on a toy dataset and semi-synthetic data (by cropping and fusing half images during pretraining) in Section 3.2 validate that SSL methods learn representations that do not overfit to labels and show improved accuracy at test time. 3) The paper also proposes a regularization framework (called rwSAM) that promotes flatter landscape for rare examples to improve performance at test time (see Table 1 and 2).

This paper studies and compares the performance of self-supervised representation learning methods against supervised representation learning when there are class imbalance. It shows empirically that self-supervised methods are more robust to class imbalance, i.e. the performance gap between models trained on balanced and imbalanced datasets is smaller. The paper also investigates the quality of the features learnt by SSL vs SL, and shows both theoretically in a limited setting and empirically on synthetic datasets that features learnt by SSL are diverse and capture characteristics that might be useful for rare classes. Finally, they provide a regularisation method for SSL to encourage model to learn better features for rare classes without having access to labels.