Self-Damaging Contrastive Learning

Contrastive learning (Chen et al., 2020a; He et al., 2020; Grill et al., 2020; Jiang et al., 2020; You et al., 2020) recently prevails for deep neural networks (DNNs) to learn powerful visual representations from unlabeled data. The state-of-the-art contrastive learning frameworks consistently benefit from using bigger models and training on more task-agnostic unlabeled data (Chen et al., 2020b). The predominant promise implied by those successes is to leverage contrastive learning techniques to pre-train strong and transferable representations from internet-scale sources of unlabeled data. However, going from the controlled benchmark data to uncontrolled real-world data will run into several gaps. For example, most natural image and language data exhibit a Zipf long-tail distribution where various feature attributes have very different occurrence frequencies (Zhu et al., 2014; Feldman, 2020). Broadly speaking, such imbalance is not only limited to the standard single-label classification with majority versus minority class (Liu et al., 2019), but also can extend to multi-label problems along many attribute dimensions (Sarafianos et al., 2018). That naturally questions whether contrastive learning can still generalize well in those long-tail scenarios.

We are not the first to ask this important question. Earlier works  (Yang & Xu, 2020; Kang et al., 2021) pointed out that when the data is imbalanced by class, contrastive learning can learn more balanced feature space than its supervised counterpart. Despite those preliminary successes, we find that the state-of-the-art contrastive learning methods remain certain vulnerability to the long-tailed data (even indeed improving over vanilla supervised learning), after digging into more experiments and imbalance settings (see Sec 4). Such vulnerability is reflected on the linear separability of pre-trained features (the instance-rich classes has much more separable features than instance-scarce classes), and affects downstream tuning or transfer performance. To conquer this challenge further, the main hurdle lies in the absence of class information; therefore, existing approaches for supervised learning, such as re-sampling the data distribution (Shen et al., 2016; Mahajan et al., 2018) or re-balancing the loss for each class (Khan et al., 2017; Cui et al., 2019; Cao et al., 2019), cannot be straightforwardly made to work here.

Our overall goal is to find a bold push to extend the loss re-balancing and cost-sensitive learning ideas (Khan et al., 2017; Cui et al., 2019; Cao et al., 2019) into an unsupervised setting. The initial hypothesis arises from the recent observations that DNNs tend to prioritize learning simple patterns (Zhang et al., 2016; Arpit et al., 2017; Liu et al., 2020; Yao et al., 2020; Han et al., 2020; Xia et al., 2021). More precisely, the DNN optimization is content-aware, taking advantage of patterns shared by more training examples, and therefore inclined towards memorizing the majority samples. Since long-tail samples are underrepresented in the training set, they will tend to be poorly memorized, or more “easily forgotten” by the model - a characteristic that one can potentially leverage to spot long-tail samples from unlabeled data in a model-aware yet class-agnostic way.

However, it is in general tedious, if ever feasible, to measure how well each individual training sample is memorized in a given DNN (Carlini et al., 2019). One blessing comes from the recent empirical finding (Hooker et al., 2020) in the context of image classification. The authors observed that, network pruning, which usually removes the smallest-magnitude weights in a trained DNN, does not affect all learned classes or samples equally. Rather, it tends to disproportionally hamper the DNN memorization and generalization on the long-tailed and most difficult images from the training set. In other words, long-tail images are not “memorized well” and may be easily “forgotten” by pruning the model, making network pruning a practical tool to spot the samples not yet well learned or represented by the DNN.

Inspired by the aforementioned, we present a principled framework called Self-Damaging Contrastive Learning (SDCLR), to automatically balance the representation learning without knowing the classes. The workflow of SDCLR is illustrated in Fig. 1. In addition to creating strong contrastive views by input data augmentation, SDCLR introduces another new level of contrasting via“model augmentation, by perturbing the target model’s structure and/or current weights. In particular, the key innovation in SDCLR is to create a dynamic self-competitor model by pruning the target model online, and contrast the pruned model’s features with the target model’s. Based on the observation (Hooker et al., 2020) that pruning impairs model ability to predict accurately on rare and atypical instances, those samples in practice will also have the largest prediction differences before then pruned and non-pruned models. That effectively boosts their weights in the contrastive loss and leads to implicit loss re-balancing. Moreover, since the self-competitor is always obtained from the updated target model, the two models will co-evolve, which allows the target model to spot diverse memorization failures at different training stages and to progressively learn more balanced representations. Below we outline our main contributions:

• •

  Seeing that unsupervised contrastive learning is not immune to the imbalance data distribution, we design a Self-Damaging Contrastive Learning (SDCLR) framework to address this new challenge.

• •

  SDCLR innovates to leverage the latest advances in understanding DNN memorization. By creating and updating a self-competitor online by pruning the target model during training, SDCLR provides an adaptive online mining process to always focus on the most easily forgotten (long tailed) samples throughout training.

• •

  Extensive experiments across multiple datasets and imbalance settings show that SDCLR can significantly improve not only the balancedness of the learned representation.

Contrastive Learning. Contrastive learning learns visual representation via enforcing similarity of the positive pairs $(v_{i},v_{i}^{+})$ and enlarging the distance of negative pairs $(v_{i},v_{i}^{-})$. Formally, the loss is defined as

where $s\left(v_{i},v_{i}^{+},\tau\right)$ indicates the similarity between positive pairs while $s\left(v_{i},v_{i}^{-},\tau\right)$ is the similarity between negative pairs. $\tau$ represents the temperature hyper-parameter. The negative samples $v_{i}^{-}$ are sampled from negative distribution $V^{-}$. The similarity metric is typically defined as

SimCLR (Chen et al., 2020a) is one of the state-of-the-art contrastive learning frameworks. For an input image, SimCLR would augment it twice with two different augmentations, and then process them with two branches that share the same architecture and weights. Two different versions of the same image are set as positive pairs, and the negative image is sampled from the rest images in the same batch.

Pruning Identified Exemplars. (Hooker et al., 2020) systematically investigates the model output changes introduced by pruning and finds that certain examples are particularly sensitive to sparsity. These images most impacted after pruning are termed as Pruning Identified Exemplars (PIEs), representing the difficult-to-memorize samples in training. Moreover, the authors also demonstrate that PIEs often show up at the long-tail of a distribution.

We extend (Hooker et al., 2020)’s PIE hypothesis from supervised classification to the unsupervised setting for the first time. Moreover, instead of pruning a trained model and expose its PIEs once, we are now integrating pruning into the training process as an online step. With PIEs dynamically generated by pruning a target model under training, we expect them to expose different long-tail examples during training, as the model continues to be trained. Our experiments show that PIEs answer well to those new challenges.

Observation: Contrastive learning is NOT immune to imbalance. Long-tail distribution fails many supervised approaches build on balanced benchmarks (Kang et al., 2019). Even contrastive learning does not rely on class labels, it still learns the transformation invariances in a data-driven manner, and will be affected by dataset bias (Purushwalkam & Gupta, 2020). Particularly for long-tail data, one would naturally hypothesize that the instance-rich head classes may dominate the invariance learning procedure and leaves the tail classes under-learned.

The concurrent work (Kang et al., 2021) signaled that using the contrastive loss can obtain a balanced representation space that has similar separability (and downstream classification performance) for all the classes, backed by experiments on ImageNet-LT (Liu et al., 2019) and iNaturalist (Van Horn et al., 2018). We independently reproduced and validated their experimental findings. However, we have to point out that it was pre-mature to conclude “contrastive learning is immune to imbalance”.

To see that, we present additional experiments in Section 4.3. While that conclusion might hold for a moderate level of imbalance as presented in current benchmarks, we have constructed a few heavily imbalanced data settings, in which cases contrastive learning will become unable to produce balanced features. In those case, the linear separability of learned representation can differ a lot between head and tail classes. We suggest that our observations complement those in (Yang & Xu, 2020; Kang et al., 2021), that while (vanilla) contrastive learning can to some extent alleviate the imbalance issue in representation learning, it does not possess full immunity and calls for further boosts.

Our SDCLR Framework. Figure 1 overviews the high-level workflow of the proposed SDCLR framework. By default, SDCLR is built on top of the simCLR pipeline (Chen et al., 2020a), and follows its most important components such as data augmentations and non-linear projection head. The main difference between simCLR and SDCLR lies in that, simCLR feeds the two augmented images into the same target network backbone (via weight sharing); while SDCLR creates a “self-competitor” by pruning the target model online, and lets the two different branches take the two augmented images to contrast their features.

Specifically, at each iteration we will have a dense branch $N_{1}$, and a sparse branch $N^{p}_{2}$ by pruning $N_{1}$, using the simplest magnitude-based pruning as described in (Han et al., 2015), following the practice of (Hooker et al., 2020). Ideally, the pruning mask of $N^{p}_{2}$ could be updated per iteration after the model weights are updated. In practice, since the backbone is a large DNN and its weights will not change much for a single iteration or two, we set the pruning mask to be lazy-updated at the beginning of every epoch, to save computational overheads; all iterations in the same epoch then adopt the same mask¹¹1We also tried to update pruning masks more frequently, and did not find observable performance boosts.. Since the self-competitor is always obtained and updated from the latest target model, the two branches will co-evolve during training.

We sample and apply two different augmentation chains to the input image $I$, creating two different versions [$\hat{I}_{1}$, $\hat{I}_{2}$]. They are encoded by [$N_{1}$, $N^{p}_{2}$], and their output features [$f_{1}$, $f^{p}_{2}$] are fed into the nonlinear projection heads to enforce similarity be under the NT-Xent loss (Chen et al., 2020a). Ideally, if the sample is well-memorized by $N_{1}$, pruning $N_{1}$ will not “forget” it – thus little extra perturbation will be caused and the contrasting is roughly the same as in the original simCLR. Otherwise, for rare and atypical instances, SDCLR will amplify the prediction differences between then pruned and non-pruned models – hence those samples’ weights be will implicitly increased in the overall loss.

When updating the two branches, note that [$N_{1}$, $N^{p}_{2}$] will share the same weights in the non-pruned part, and $N_{1}$ will independently update the remaining part (corresponding to weights pruned to zero in $N^{p}_{2}$). Yet, we empirically discover that it helps to let either branch have its independent batch normalization layers, as the features in dense and sparse may show different statistics (Yu et al., 2018).

SDCLR can work with more contrastive learning frameworks. We focus on implementing SDCLR on top of simCLR for proving the concept. However, our idea is rather plug-and-play and can be applied with almost every other contrastive learning framework adopting the the two-branch design (He et al., 2020; LeCun et al., 1990; Grill et al., 2020). We will explore combining SDCLR idea with them as our immediate future work.

Pruning is NOT for model efficiency in SDCLR. To avoid possible confusion, we stress that we are NOT using pruning for any model efficiency purpose. In our framework, pruning would be better described as “selective brain damage”. It is mainly used for effectively spotting samples not yet well memorized and learned by the current model. However, as will be shown in Section 4.9, the pruned branch can have a “side bonus”, that sparsity itself can be an effective regularizer that improves few-shot tuning.

SDCLR benefits beyond standard class imbalance. We also want to draw awareness that SDCLR can be extended seamlessly beyond the standard single-class label imbalance case. Since SDCLR relies on no label information at all, it is readily applicable to handling various more complicated forms of imbalance in real data, such as the multi-label attribute imbalance (Sarafianos et al., 2018; Yun et al., 2021).

Moreover, even in artificially class-balanced datasets such as ImageNet, there hide more inherent forms of “imbalance”, such as the class-level difficulty variations or instance-level feature distributions (Bilal et al., 2017; Beyer et al., 2020). Our future work would explore SDCLR in those more subtle imbalanced learning scenarios in the real world.

Our experiments are based on three popular imbalanced datasets at varying scales: long-tail CIFAR-10, long-tail CIFAR-100 and ImageNet-LT. Besides, to further stretch out contrastive learning’s imbalance handling ability, we also consider a more realistic and more challenging benchmark long-tail ImageNet-100 as well as another long tail ImageNet with a different exponential sampling rule. The long-tail ImageNet-100 contains less classes, which decreases the number of classes that looks similar and thus can be more vulnerable to imbalance.

Long-tail CIFAR10/CIFAR100: The original CIFAR-10/ CIFAR-100 datasets consist of 60000 32$\times$32 images in 10/100 classes. Long tail CIFAR-10/ CIFAR-100 (CIFAR10-LT / CIFAR100-LT) were first introduced in (Cui et al., 2019) by sampling long tail subsets from the original datasets. Its imbalance factor is defined as the class size of the largest class divided by the smallest class. We by default consider a challenging setting with the imbalance factor set as 100. To alleviate randomness, all experiments are conducted with five different long tail sub-samplings.

ImageNet-LT: ImageNet-LT is a widely used benchmark introduced in (Liu et al., 2019). The sample number of each class is determined by a Pareto distribution with the power value $\alpha=6$. The resultant dataset contains 115.8K images, with the sample number per class ranging from 1280 to 5.

ImageNet-LT-exp: Another long tail distribution of ImageNet we considered is given by an exponential function (Cui et al., 2019), where the imbalanced factor set as 256 to ensure the the minor class scale is the same as ImageNet-LT. The resultant dataset contains 229.7K images in total. ²²2Refer to our code for details of ImageNet-LT-exp and ImageNet-100-LT.

Long tail ImageNet-100: In many fields such as medical, material, and geography, constructing an ImageNet scale dataset is expensive and even impossible. Therefore, it is also worth considering a dataset with a small scale and large resolution. We thus sample a new long tail dataset called ImageNet-100-LT from ImageNet-100 (Tian et al., 2019). The sample number of each class is determined by a down-sampled (from 1000 classes to 100 classes) Pareto distribution used for ImageNet-LT. The dataset contains 12.21K images, with the sample number per class ranging from 1280 to 52.

To evaluate the influence brought by long tail distribution, for each long tail subset, we would sample a balanced subset from the corresponding full dataset with the same total size as the long tail one to disentangle the influences of long tail and sample size.

For all pre-training, we follow the SimCLR recipe (Chen et al., 2020a) including its augmentations, projection head structures. The default pruning ratio is 90% for CIFAR and 30% for ImageNet. We adopt Resnet-18 (He et al., 2016) for small datasets (CIFAR10/CIFAR100), and Resnet-50 for larger datasets (ImageNet-LT/ImageNet-100-LT), respectively. More details on hyperparameters can be found in the supplementary.

The balancedness of a feature space can be reflected by the linear separability w.r.t. all classes. To measure the linear separability, we identically follow (Kang et al., 2021) to employ a three-step protocol: i) learn the visual representation $f_{v}$ on the training dataset with $\mathcal{L}_{CL}$. ii) training a linear classifier layer $L$ on the top of $f_{v}$ with a labeled balanced dataset (by default, the full dataset where the imbalanced subset is sampled from). iii) evaluating the accuracy of the linear classifier $L$ on the testing set. Hereinafter, we define such accuracy measure as the linear separability performance.

To better understand the influence of the balancedness for down-stream tasks, we consider the important practical application of few-shot learning (Chen et al., 2020b). The only difference between measuring few-shot learning performance and measuring linear separability accuracy lies in step ii): we use only 1% samples of the full dataset from which the pre-training imbalanced dataset is sampled. Hereinafter, we define the accuracy measure with this protocol as the few-shot performance.

We further divide each dataset to three disjoint groups in terms of the size of classes: {Many, Medium, Few}. In subsets of CIFAR10/CIFAR100, Many and Few each include the largest and smallest $\frac{1}{3}$ classes, respectively. For instance in CIFAR-100: the classes with [500-106, 105-20, 19-5] samples belong to [Many (34 classes), Medium (33 classes), Few (33 classes)] categories, respectively. In subsets of ImageNet, we follow OLTR (Liu et al., 2019) to define Many as classes each with over training 100 samples, Medium as classes each with 20-100 training samples and Few as classes under 20 training samples. We report the average accuracy for each specified group, and also use the standard deviation (Std) among the three groups’ accuracies as another balancedness measure.

We now investigate if contrastive learning is vulnerable to the long-tail distribution. In this section, we use $D_{i}$ to represent the long tail split of a dataset while $D_{b}$ denotes its balanced counterpart.

As shown in Table. 1, for both CIFAR10 and CIFAR100, models pre-trained on $D_{i}$ show larger imbalancedness than that on $D_{b}$. For instance, in CIFAR100, while models pre-trained on $D_{b}$ show almost the same accuracy for three groups, the accuracy gradually drops from many to few when pre-training subset switches from $D_{b}$ to $D_{i}$. This indicates that the balancedness of contrastive learning is still fragile when trained over the long tail distributions.

We next explore if the imbalanced representation would influence the downstream few-shot learning applications. As shown in Table. 2, in CIFAR10, the few shot performance of Many drops by 1.07% when switching from $D_{b}$ to $D_{i}$ while that of Medium and Few decrease by 5.30% and 6.12%. In CIFAR100, when pre-training with $D_{b}$, the few-shot performance on three groups are similar, and it would become imbalanced when the pre-training dataset switches from $D_{b}$ to $D_{i}$. In a word, the balancedness of few-shot performance is consistent with the representation balancedness. Moreover, the bias would become even more serious: The gap between Many and Few enlarge from 6.81% to 21.86% on CIFAR10 and from 6.65% to 14.73% on CIFAR100.

We further study if the imbalance can also influence large scale dataset like ImageNet in Table. 3. For ImageNet-LT and Imagenet-LT-exp, while the imbalancedness of linear separability performance shows weak, that problem becomes much more significant for few-shot performance. Especially, for Imagenet-LT-exp, the few-shot performance of Many is 7.96% higher than that of Few. The intuition behind this is that the large volume of the balanced fine-tuning dataset could mitigate the influence of imbalancedness from the pre-trained model. When the scale decreases to 100 classes (ImageNet-100), the imbalancedness consistently exists and it can be reflected via both linear separability performance and few-shot performance.

We compare the proposed SDCLR with SimCLR (Chen et al., 2020a) on the datasets that are most easily to be impacted by long tail distribution: CIFAR10-LT, CIFAR100-LT, and ImageNet-100-LT. As shown in Table. 4, the proposed SDCLR leads to a significant linear separability performance improvement of $\textbf{6.45}\%$ in CIFAR10-LT and $\textbf{8.37}\%$ in CIFAR100-LT. Meanwhile, SDCLR also improve the balancedness by reducing the Std by $0.07\%$ in CIFAR10 and $0.23\%$ in CIFAR100. In Imagenet-100-LT, SDCLR achieve an improvement on linear separability performance of $1.02\%$ while reducing the Std by $0.29$.

On few-shot settings, as shown in Table. 5, the proposed SDCLR consistently improves the few-shot performance by [$3.39\%$, $2.31\%$, $0.22\%$] while decreasing the Std by [$2.81$, $2.29$, $0.45$] in [CIFAR10, CIFAR100, Imagenet-100-LT], respectively.

SDCLR is a pre-training approach that is fully compatible with almost any existing long-tail algorithm. To show that, on CIFAR-100-LT with the imbalance factor of 100, we use SDCLR as pre-training, to fine-tune a SOTA long-tail algorithm RIDE (Wang et al., 2020b) on its top. With SDCLR pre-training, the overall accuracy can reach 50.56%, super-passing the original RIDE result by 1.46%. Using SimCLR pre-training for RIDE only spots 50.01% accuracy.

Even balanced in sample numbers per class, existing datasets can still suffer from more hidden forms of “imbalance”, such as sampling bias, and different classes’ difficulty/ambiguity levels, e.g., see (Bilal et al., 2017; Beyer et al., 2020). To evaluate whether the proposed SDCLR can address such imbalancedness, we further run the proposed framework on balanced datasets: The full dataset of CIFAR10 and CIFAR100. We compare SDCLR with SimCLR following standard linear evaluation protocol (Chen et al., 2020a) (On the same dataset, it first pre-trains the backbone, and then finetunes one linear layer on the top of the output features).

The results are shown in Table 6. Note here the $Std$ denotes the standard deviation of classes as we are studying the imbalance caused by the varying difficulty of classes. The proposed SDCLR can boost the linear evaluation accuracy by [$0.39\%$, $3.48\%$] while reducing the Std by [$1.0$, $0.16$] in [CIFAR10, CIFAR100], respectively, proving that the proposed method can also help to improve the balancedness even in the balanced datasets.

We then measure the distribution of PIEs mined by the proposed SDCLR. Specifically, when pre-training on long tail splits of CIFAR100, we sample top 1% testing data that is most easily influenced by pruning and then evaluate the percentage of many, medium and minor in it under different training epochs. The difficulty of forgetting a sample is defined by the features’ cosine similarity before and after pruning. Figure 2 shows the minor and medium are much more likely to be impacted comparing to many. In particular, while the group distributions of the found PIEs show some variations along with training epochs, in general, we find samples from the minor group to gradually increase, while the many group samples continue to stay low percentage especially when it is close to convergence.

Random dropout baseline: To verify whether pruning is necessary, we compare with using random dropout (Srivastava et al., 2014) to generate the sparse branch. Under dropout ratio of 0.9, [linear separability, few-shot accuracy] are [21.99$\pm$0.35%, 15.48$\pm$0.42%], which are much worse than both SimCLR and SDCLR reported in Tab. 4 and 5. In fact, the dropout baseline is often hard to converge.

Focal loss baseline: We also compare with the popular focal loss (Lin et al., 2017) for conducting this suggested sanity check. With the best grid searched gamma of 2.0 ( grid is [0.5, 1.0, 2.0, 3.0]), it decreases the [accuracy,std] of linear separability from [47.33$\pm$0.33%, 2.70$\pm$1.25%] to [46.48$\pm$0.51%, 2.99$\pm$1.01%], respectively. Further analysis shows the contrastive loss scale is not tightly connected with the major or minor class membership as we hypothesized. A possible reason is that the randomness of SimCLR augmentations also notably affects the loss scale.

Extending to Moco pre-training: We try MocoV2 (He et al., 2020; Chen et al., 2020c) on CIFAR100-LT. The [accuracy,std] of linear separability is [48.23$\pm$0.20%, 3.50$\pm$0.98%] and [accuracy,std] of few-shot performance is [24.68$\pm$0.36%, 6.67$\pm$1.45%], respectively, which is worse than SDCLR in Tab 4 and 5.

We study the linear separability performance under different pruning ratios in one imbalance subset of CIFAR100. As shown in Figure. 3, the overall accuracies consistently increase with the pruning ratio until it exceeds 90%, which will lead to a quick drop. That shows a trade-off for the sparse branch between being stronger (i.e., needing larger capacity) and being effective in spotting more difficult examples (i.e., needing being sparse).

We also explore the linear separability and few-shot performance of the sparse branch in Figure 4. In the linear separability case (a), the sparse branch quickly lags behind the dense branch when the sparsity goes above 70$\%$, due to limited capacity. Interestingly, even a “weak” sparse branch can still assist the learning of its dense branch. The few shot performance also shows the similar trend.

The Sparse Branch Architecture The visualization of pruned ratio for each layer is illustrated in Figure 6. Overall, we find the sparse branch’s deeper layers to be more heavily pruned. This is aligned with the intuition that higher-level features are more class-specific.

Visualization for SDCLR We visualize the features of SDCLR and SimCLR on minor classes with Grad-CAM (Selvaraju et al., 2017) in the Figure 5. SDCLR shows to better localize class-discriminative regions for tail samples.

The datasets we employed are CIFAR10/100, and ImageNet. Their downloading links can be found in Table. 7.

For CIFAR10, CIFAR100, ImageNet, and ImageNet-100, the testing dataset is set as its official validation datasets. We also randomly select [10000, 20000, 2000] samples from the official training datasets of [CIFAR10/CIFAR100, ImageNet, ImageNet-100] as validation datasets, respectively.

We identically follow SimCLR (Chen et al., 2020a) for pre-training settings except the epochs number. On the full dataset of CIFAR10/CIFAR100, we pre-train for 1000 epochs. In contrast, on sub-sampled CIFAR10/CIFAR100, we would enlarge the pre-training epochs number to 2000 given the dataset size is small. Moreover, the pre-training epochs of ImageNet-LT-exp/ImageNet-100-LT is set as 500.

We employ SGD with momentum 0.9 as the optimizer for all fine-tuning. We follow (Chen et al., 2020c) employing learning rate of 30 and remove the weight decay for all fine-tuning. When fine-tuning for linear separability performance, we train for 30 epochs and decrease the learning rate by 10 times at epochs 10 and 20 as we find more epochs could lead to over-fitting. However, when fine-tuning for few-shot performance, we would train for 100 epochs and decrease the learning rate at epoch 40 and 60, given the training set is far smaller.