[Paper Review] Self-supervised Learning is More Robust to Dataset Imbalance

Hong Liu, Jeff Z. HaoChen+2 more|arXiv (Cornell University)|Oct 11, 2021

Domain Adaptation and Few-Shot LearningComputer Science68 references63 citations

TL;DR

The paper shows that self-supervised learning (SSL) representations are more robust to class imbalance than supervised representations, provides theoretical and empirical explanations, and introduces a reweighting-based regularization to further improve SSL on imbalanced data.

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 out 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-irrelevant-but-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 and theoretical 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.

Motivation & Objective

• Investigate how class imbalance affects self-supervised representations when pre-training.
• Compare robustness of SSL vs. supervised pre-training using in-domain and out-of-domain evaluations.
• Develop theoretical and empirical explanations for SSL robustness under imbalance.
• Propose a reweighted regularization technique to improve SSL performance on imbalanced data.

Proposed method

• Systematically evaluate SSL (MoCo v2 and SimSiam) and supervised pre-training on CIFAR-10 and ImageNet with varying imbalance ratios and sample sizes.
• Assess representation quality via linear probing on balanced in-domain data and fine-tuning on downstream out-of-domain data.
• Provide a toy theoretical setting to contrast features learned by SSL and supervised learning under imbalance.
• Conduct semi-synthetic experiments to visualize transferable vs. label-relevant features learned by SSL vs. SL.
• Introduce reweighted sharpness-aware minimization (rwSAM) with kernel-density-based instance weighting to improve SSL on imbalanced data.

Experimental results

Research questions

• RQ1 How does dataset imbalance affect the quality of SSL versus supervised representations in ID and OOD evaluations?
• RQ2 Why does SSL tend to learn transferable features from frequent classes that help with rare classes?
• RQ3 Can reweighted regularization improve SSL performance on imbalanced datasets, and how does it affect generalization on rare examples?
• RQ4 Do theoretical toy settings and semi-synthetic experiments support the claim that SSL captures label-irrelevant-but-transferable features?

Key findings

• SSL representations are more robust to class imbalance than supervised representations across multiple configurations.
• The robustness gap between balanced and imbalanced pre-training is smaller for SSL than for SL in both ID and OOD evaluations.
• SSL tends to learn richer, transferable features from frequent classes that can help with rare classes and downstream tasks.
• A reweighted sharpness-aware minimization (rwSAM) method consistently improves SSL representation quality on imbalanced datasets and closes part of the gap to balanced data.