Improving Contrastive Learning on Imbalanced Seed Data via Open-World Sampling

TL;DR

This paper introduces MAK, a principled data sampling framework for contrastive learning that strategically selects unlabeled external data to improve representation quality and class balance, especially in open-world, imbalanced scenarios.

Contribution

We propose MAK, a novel open-world data sampling method based on three principles—tailness, proximity, and diversity—that enhances contrastive learning with external unlabeled data.

Findings

MAK improves representation quality on ImageNet-100-LT.

MAK achieves better class balancedness in learned features.

MAK enhances performance in both full-shot and few-shot settings.

Abstract

Contrastive learning approaches have achieved great success in learning visual representations with few labels of the target classes. That implies a tantalizing possibility of scaling them up beyond a curated "seed" benchmark, to incorporating more unlabeled images from the internet-scale external sources to enhance its performance. However, in practice, larger amount of unlabeled data will require more computing resources due to the bigger model size and longer training needed. Moreover, open-world unlabeled data usually follows an implicit long-tail class or attribute distribution, many of which also do not belong to the target classes. Blindly leveraging all unlabeled data hence can lead to the data imbalance as well as distraction issues. This motivates us to seek a principled approach to strategically select unlabeled data from an external source, in order to learn generalizable, balanced and diverse representations for relevant classes. In this work, we present an open-world unlabeled data sampling framework called Model-Aware K-center (MAK), which follows three simple principles: (1) tailness, which encourages sampling of examples from tail classes, by sorting the empirical contrastive loss expectation (ECLE) of samples over random data augmentations; (2) proximity, which rejects the out-of-distribution outliers that may distract training; and (3) diversity, which ensures diversity in the set of sampled examples. Empirically, using ImageNet-100-LT (without labels) as the seed dataset and two "noisy" external data sources, we demonstrate that MAK can consistently improve both the overall representation quality and the class balancedness of the learned features, as evaluated via linear classifier evaluation on full-shot and few-shot settings. The code is available at: https://github.com/VITA-Group/MAK