Improving Contrastive Learning on Imbalanced Data via Open-World Sampling

1 Introduction . Contrastive learning has been successfully applied to learning strong visual representations in an unsupervised manner [ 1–5 ] . The promise of label-free learning makes it appealing to scale up contrastive learning with massive unannotated data , e.g. , from internet-scale sources [ 6 ] . However , training with a larger amount of unlabeled data is not cheap . As shown in [ 1 , 2 ] , contrastive learning needs a bigger model and longer training compared to supervised counterparts . With a larger amount of data , it also requires more compute resources ( for bigger model and longer training ) . With a limited computing budget , it is likely that the out-of-distribution data in the wild would suppress the learning of relevant features . Moreover , different from standard benchmarks that are carefully curated and well balanced across classes ( e.g. , ImageNet ) , the data distribution in the open world are extremely diverse and always exhibits long tails [ 7 , 8 ] . A few very recent works [ 9–11 ] have studied whether contrastive learning can still generalize well in those long-tail scenarios . And they found that while contrastive learning learns more balanced feature space than its supervised counterpart , it still exhibits certain vulnerability 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . to the long-tailed data [ 11 ] . Therefore , blindly gathering unlabeled data in the wild may exacerbate the imbalancedness issue of the training data . Here we consider the problem of sampling open-world unlabeled data for improving the representation learning , not just for the head classes but also for the tailed classes . To make it clear , we describe the problem setting in more detail . Problem Setting : Assume that we start from a relatively small ( “ seed ” ) set of unlabeled training data , where the data distribution could be highly skewed yet unspecified . We aim to retrieve an extra set , with a given sampling budget , of freely available images from some external sources ( e.g. , the Web ) , to enhance self-supervised representation learning for targeted distribution ( of seed set ) . It is worth noting that in our setup , we do not use labels for data sampling or training ( e.g . [ 12 ] ) , but only use labels as a way to evaluate the representation . By training on the retrieved unlabeled examples , our goal is to learn “ stronger and fairer ” visual representations that improve not only the overall quality but also the balancedness across various class attributes . It is worth pointing out some unique challenges in our problem setup . 1 ) We may start from a skewed set , and the actual class imbalancedness is unknown to us due to label unavailability , making the most approaches handling imbalance in the ( semi- ) supervised setting like pseudo-labeling , re-sampling , or loss re-weighting [ 13–18 ] inapplicable . 2 ) Adopting a pre-trained backbone trained on imbalanced seed data with tail classes under-learned may amplify unfairness , as it is harder to retrieve images for poorly learned tail classes . 3 ) Widely existing irrelevant outlier samples in the open world are harder to detect given the lack of label information . In this work , we conduct the first systematic study on leveraging additional unlabeled data from external sources when performing contrastive learning on a long-tail seed training set . We observe that , while randomly sampling more unlabeled data can effectively help overall accuracy , it shows only limited and unstable benefits to the balancedness . Thus we aim to seek a principled open-world unlabeled data sampling strategy , and it follows three principles outlined below . The first principle we follow is tailness . Intuitively , we would like to sample more examples that correspond to tail classes . Since neither labels nor class predictions are are unavailable , we seek another proxy for tailness , using each sample ’ s training loss to identify “ hard samples ” . However , compared to similar ideas adopted in supervised learning [ 19 ] , our proxy signal is weaker ( only contrastive loss , and not directly compared to the class label ) and much noisier due to the strong random augmentations in contrastive learning . In view of that , we propose to instead use an empirical expectation of sample loss over multiple random augmentations as the proxy . Our experiments verify that this new proxy of empirical contrastive loss expectation ( ECLE ) can reliably and stably spotlight more samples from tail classes . We further complement the tailness principle with two other principles : proximity and diversity . Without proximity to examples in the seed set , the large-loss samples could simply be outliers , and training with those may hamper the model generalization on the target distribution . Thus we incorporate a feature distance regularizer between new external samples and seed training samples to reject too “ far-away ” samples from the former . Without diversity in the retrieved images , the benefits of the additional data can be significantly lower . We hence include another diversity-promoting term in sample selection . Eventually , our ideas could be mathematically unified into one framework called Model-Aware K-center ( MAK ) ( as shown in Figure 1 ) , which could be viewed as a principled extension of the core-set active learning [ 20 ] , from supervised to unsupervised representation learning . Our contributions are summarized below . • New problem and insight : leveraging external unlabeled data could help contrastive learning on imbalanced data , but needs strategic sampling to avoid backfiring on balancedness . • New tail sample mining for contrastive learning : an empirical contrastive loss expectation ( ECLE ) is proposed to adaptively mine under-learned samples ( usually on the tail classes ) , while alleviating the substantial impact of random data augmentations . • New principled sampling framework : our unified formulation of Model-Aware K-center ( MAK ) seamlessly integrates ECLE , Out-of-Distribution ( OoD ) rejection and diversity promotion . It also extends the ideas of core-set active learning [ 20 ] to unsupervised representation learning . • Compelling empirical results : When training on ImageNet-100-LT , the MAK can consistently yield accuracy and balancedness improvement over the competitive random sampling method across different sampling datasets and budgets . 2 Our Method . 2.1 Overview . Preliminary : In contrastive learning , the representation is learned via enforcing an anchor sample vi to be similar to another positive sample while being different from negative samples . Among various popular contrastive learning frameworks , SimCLR [ 1 ] is easy-to-implement while yielding near state-of-the-art performance . It utilizes two augmented views of the same data as the positive pairs while all other augmented samples in the same batch are treated as negative samples . Assuming a random data augmentation function A ( · , θ ) where θ is the hyperparameter to be sampled from some empirical distribution Θ each time . With θi,1 , θi,2 ∼ Θ , ∀i . The SimCLR loss associated with the i-th sample in the batch could be represented as : LCL , i = − log sτ ( A ( vi , θi,1 ) , A ( vi , θi,2 ) ) sτ ( A ( vi , θi,1 ) , A ( vi , θi,2 ) ) + ∑ v−i ∈V − sτ ( A ( vi , θi,1 ) , v − i ) ( 1 ) where v−i denotes the negative samples sampled from distribution V − . sτ represents the similarity function with a temperature term of τ defined as sτ ( a , b ) = exp ( a · b/τ ) . The overall term measures the entropy for sampling A ( vi , θi,2 ) among all A ( vi , θi,2 ) and v−i . Hereinafter , we focus on SimCLR as the backbone to develop our framework . However , our idea is rather plug-and-play and can be applied with other contrastive learning frameworks adopting the two-branch design with data augmentation views [ 3 , 4 ] , which we will explore as our future work . Principles : As demonstrated in Figure 1 , the principles behind MAK are to ensure tailness , proximity and diversity . While tailness is for sampling more tail class data , proximity ensures the sampled data is in-distribution . Besides , Diversity further enforce the diversity of the sampled data , preventing sampling redundant similar samples . 2.2 Spotting Hard Samples from Tail Classes : A New Proxy for Tailness . In supervised learning , one can spot “ hard examples '' which yield the largest loss values and assign higher weights for accelerated training [ 19 ] , and for handling data imbalance [ 21 ] . However , such a benefit does not extend to contrastive learning straightforwardly . Importantly , the contrastive loss ( 1 ) largely depends on the random augmentations A ( · , θ ) and thus display high randomness . As observed in experiments , the loss variations caused by selecting two similar/dissimilar augmentation views could outweigh the loss difference caused by sample memorization or learning difficulty . Hence naively picking samples with the largest contrastive loss ( 1 ) will only produce highly noisy selections . To eliminate the randomness caused by augmentations , we turn to the following new proxy value for the i-th sample , that is designed to “ smooth out '' random augmentations by integrating over them : LECL , i = Eθi,1 , θi,2∼Θ ( LCL , i ( θi,1 , θi,2 ; τ , vi , V − ) ) , ( 2 ) In practice , the expectation is approximated by the sample mean , e.g. , drawing { θi,1 , θi,2 } for M times and then averaging corresponding LCL , i values . We denote the new loss ( 2 ) as the empirical contrastive loss expectation ( ECLE ) for the i-th sample . Note a similar “ randomized smoothing '' idea was utilized before to learn robust classifiers [ 22 , 23 ] , though in a completely different context from ours . We then sort and choose those with the largest ECLE values ( 2 ) as hard samples . Section 3.4 experimentally verifies that ECLE eliminates the randomness within reasonable computational budgets , and larger ECLE values display correlation with tailness . 2.3 Proximity and Diversity . Proximity : The new ECLE proxy lays the foundation for sampling hard examples . While they are found to correlate with Tailness , the alignment remains to be weak and noisy in the open-world setting , since the external sources can have shifted or non-overlapped class distribution with the target ones , e.g. , containing “ distractor '' classes . Those outliers behave like zero-shot minorities and are likely to incur large losses , since they are from completely unseen and unwanted classes . Adopting only the ECLE proxy might easily pick those outliers , hurting feature learning and generalization on the underlying distribution . We hence construct a regularization term that promotes proximity via rejecting OoD outliers . Let s1 be the new additional set that we have sampled from the external sources , and s0 be the seed training set . We let ∆ ( xi , xj ) denote the feature distance between two samples . We define D ( s0 , s1 ) as the average feature distance between each sample in s1 and its nearest sample in s0 , formalized as : D ( s0 , s1 ) = 1 |s1| ∑ j∈s1 min i∈s0 ∆ ( xi , xj ) ( 3 ) In practice , to compute ∆ ( xi , xj ) , we use the normalized cosine distance , by passing xi , xj through the currently trained backbone and using their extracted features before feeding projection heads . For further efficiency , instead of calculating D ( s0 , s1 ) over the whole s0 , we pre-compute the set of feature prototypes from s0 using K-means clustering , denoted as s0p , and then compute D ( s 0 p , s 1 ) . Diversity : Oversampling too many external images both would add to the training overhead , and might not necessarily help since we might get redundant samples . Assume that we have a sampling budget : s1 : |s1| ≤ K , the challenge is how to attain the informative samples within the size limit . We introduce the following regularization term : H ( s1 ∪ s0 , Sall ) = max i∈Sall min j∈s1∪s0 ∆ ( xi , xj ) ( 4 ) Here Sall denotes all samples from the seed training set and the external sources ( s0 , s1 ∈ Sall ) . Conceptually , minimizing the term ( 4 ) boils down to choosing |s1| center points on top of the given |s0| points , such that the largest distance between any data point from Sall and its nearest center point from s1 ∪ s0 is minimized . It equals to finding a diverse subset cover for the dataset Sall using a minimax facility location formulation [ 24 ] , also known as the name of K-center [ 20 ] . 2.4 Model-Aware K-Center : A Unified Framework Taking together , we end up with solving the following constrained optimization : max s1 : |s1|≤K { ∑ i∈s1 LECL , i −D ( s0 , s1 ) −H ( s1 ∪ s0 , Sall ) } ( 5 ) The goal of ( 5 ) is to find a sample set s1 from the external source , such that it could simultaneously : ( i ) mine more data for tail classes while overcoming augmentation randomness , by sorting external samples by their ECLE values ( tailness ) ; ( ii ) reject the out-of-distribution outliers that might distract Algorithm 1 : A greedy heuristic to efficiently solve MAK . Require : seed training set s0 , external data Sall\s0 , sampling budget K , candidate set size C , feature distance function ∆ ( cosine distance in practice ) , coefficient α ∈ ( 0 , 1 ) . Train feature extractor f on s0 with self-supervised method ; Calculate ECLE LECL , i and average feature distance D ( s0 , s1 ) with f as in Equation 2 and 3 , respectively ; Summarize LECL , i and D ( s0 , s1 ) with score q = αN ( LECL , i ) − ( 1− α ) N ( D ( s0 , s1 ) ) where N ( v ) = v−mean ( v ) std ( v ) is the normalization function ; Construct set S′ : from Sall\S0 , find all samples whose score q are top C largest among all ; // tailness & proximity Construct set s : Initialize s = s0 ; while |s\s0| ≤ K do // Apply K-center greedy algorithm for diversity u = arg maxi∈S′ minj∈s ∆ ( xi , xj ) ; s = s ∪ { u } ; end return s1 = s\s0 training , by constraining feature distances from the seed set ( proximity ) ; and ( iii ) control the sample volume under K while ensuring sample diversity , by K-center sample selection ( diversity ) . We name our optimization ( 5 ) Model-Aware K-center ( MAK ) , since it inherits the vanilla “ datalevel '' K-center objective ( 4 ) , while injecting “ model-level '' sample selection criteria , including largest ( smoothed ) model losses as well as closer feature distances . Those “ model-aware '' terms are necessary in our unsupervised re-balancing goal and our open-world setting , for automatically re-balancing the additional samples in favor the seed set ’ s minority classes , and for rejecting outliers from open-world external data , respectively . Exactly solving ( 5 ) would be NP-hard [ 25 ] , and we instead refer to an empirical greedy routine with a coordinate-descent heuristic . As indicated in Algorithm 1 , we start by training a feature extractor f on the seed training set s0 , which is then used for calculating ELCE loss LECL , i and average feature distance D ( s0 , s1 ) . Afterwards , we sample a preliminary pool S′ with top C largest tailness and proximity samples ( C > K is candidate set size ) . We employ a score q to combine these two metrics , defined as : q = αN ( LECL , i ) − ( 1− α ) N ( D ( s0 , s1 ) ) ( 6 ) where N ( v ) = v−mean ( v ) std ( v ) is a normalization term and α ∈ ( 0 , 1 ) is a coefficient . From S ′ , we then sample a diverse subset s1 of size K , leveraging the K-center greedy algorithm .

Summary: This paper improves the accuracy and balancedness of contrastive learning by proposing a unified sampling framework called Model-Aware K-center. It follows three principles: tailness, proximity, and diversity. As a new problem setting, the authors reported extensive experiments to show their method's consistent improvement over other sampling methods.