Self-supervised Learning is More Robust to Dataset Imbalance

1 INTRODUCTION . 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 . 2 EXPLORING THE EFFECT OF CLASS IMBALANCE ON SSL . Dataset class imbalance can pose challenge to self-supervised learning in the wild . Without access to labels , we can not 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 . 2.1 PROBLEM FORMULATION . 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 ) . 2.2 EXPERIMENTAL SETUP . 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 .

Summary: 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).