Self-supervised Learning is More Robust to Dataset Imbalance

Hong Liu; Jeff Z. HaoChen; Adrien Gaidon; Tengyu Ma

Source-linked AI summary

Self-supervised Learning is More Robust to Dataset Imbalance

Hong Liu, Jeff Z. HaoChen, Adrien Gaidon, Tengyu Ma

arXiv:2110.05025v2Oct 11, 2021cs.LGcs.CVstat.ML

TL;DR

This paper studies how self-supervised representations behave under class-imbalanced pre-training, a setting common in large-scale unlabeled data but relatively unexplored. Through systematic experiments, theoretical analysis, and a re-weighted regularization method, it finds SSL more robust than supervised pre-training and further improves SSL on imbalanced data.

Problem
SSL is often trained on balanced datasets, while large-scale unlabeled datasets commonly have long-tailed distributions and the effect of imbalance on SSL remains underexplored.
Method
The paper systematically compares SSL and supervised representations under imbalance, analyzes the phenomenon empirically and theoretically, and proposes a re-weighted regularization method.
Results
SSL representations have smaller balanced-versus-imbalanced performance gaps than supervised representations across dataset sizes, imbalance ratios, and in-domain and out-of-domain evaluations.
Takeaways & Limitations
The findings support using SSL for pre-training or evaluating the impact of imbalanced pre-training on downstream tasks, especially for vision datasets.
Takeaways & Limitations
The experiments mainly focus on vision datasets, while the formulation assumes class-conditional distributions remain unchanged between balanced and imbalanced datasets.

Abstract

from arXiv · show

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.

1 Introduction

The paper studies how self-supervised representations behave under long-tailed unlabeled data and finds them more robust to imbalance than supervised representations. It explains this robustness through richer features from frequent data and proposes re-weighted regularization to further improve SSL.

Empirical findings: SSL remains more robust to class imbalance than supervised representations across dataset sizes, imbalance ratios, and both ID and OOD evaluations.The balance–imbalance performance gap is consistently smaller for SSL than for supervised learning.
Explanation: SSL may learn label-irrelevant-but-transferable features from frequent classes that help classify rare classes and support downstream OOD tasks.The proposed explanation contrasts SSL with supervised learning, which has no incentive to learn label-irrelevant features from frequent examples.
Explanation: Semi-synthetic experiments and theoretical analysis on a simplified setting validate the proposed explanation for SSL’s robustness.The semi-synthetic data exposes label-relevant and label-irrelevant-but-transferable features by design, while the toy example permits a rigorous comparison.
Improvement method: The proposed method closes the small gap between SSL performance on balanced and imbalanced datasets with the same number of examples.This improvement is motivated by the generalization gap between empirical and population pre-training losses on rare data.
Improvement method: A density-aware re-weighted regularization method consistently improves SSL representation quality on imbalanced datasets.The method estimates example density with kernel density estimation and applies stronger sharpness-based regularization to estimated rare examples.

2 Exploring the Effect of Class Imbalance on SSL

The study evaluates representation quality under controlled class imbalance across datasets, pre-training sizes, imbalance ratios, learning paradigms, and ID or OOD protocols. SSL consistently shows a smaller balance–imbalance gap than supervised learning, while supervised learning can lead on sufficiently large ID evaluations and SSL on OOD evaluations.

Problem formulation: The imbalance ratio r is the rarest-class probability divided by the most frequent-class probability, with balanced data defined by r = 1.The class-conditional input distributions are assumed unchanged across balanced and imbalanced datasets.
Evaluation: The evaluation drops the supervised classifier head and compares the resulting feature extractors using balanced ID data and downstream OOD tasks.ID uses a linear classifier trained on balanced data; OOD fine-tunes the feature extractor and head on target distributions.
Experimental design: The experiments vary pre-training size, imbalance ratio, evaluation domain, and SSL algorithm to assess representation quality under class imbalance.The study uses MoCo v2 and SimSiam, with linear probing for ID evaluation and fine-tuning for OOD evaluation.
Results: Across all reported configurations, the SSL balance–imbalance accuracy gap is smaller than the supervised gap for the same number of pre-training examples.This pattern holds for both ID and OOD evaluations and is also shown for CIFAR-10 and ImageNet.
Interpretation: Comparing methods at the same data count favors supervised learning because SSL can be applied more easily to larger unlabeled datasets.SSL does not require labels, avoiding the label-collection requirement for scaling pre-training data.
Results: Supervised representations perform better with reasonably large n in ID evaluation, whereas SSL representations perform better in OOD evaluation.This ID–OOD performance difference is separate from SSL’s greater robustness to imbalance.

3 Analysis

The analysis explains SSL’s robustness to imbalance through richer features learned from frequent data, including label-irrelevant but transferable directions that benefit rare classes. Toy-theoretical and semi-synthetic results support this distinction between SSL and supervised learning.

Motivation: SSL learns richer features from frequent classes that can transfer to rare classes, whereas supervised learning mainly learns features needed for frequent-class discrimination.This mechanism is proposed to explain SSL’s stronger robustness under imbalance.
3.1 Rigorous Analysis on A Toy Setting: In the toy distribution, e1 separates frequent classes while e2 is transferable to the rare class despite not distinguishing the frequent classes.Both directions occur in frequent-class inputs, but only e2 helps classify the rare class.
3.1 Rigorous Analysis on A Toy Setting: For n1, n2 = Θ(poly(d)) and n3 ≤ d^1/5, the theorem states that supervised features have small e2 correlation while SSL recovers e2 with high probability.The guarantees hold for any feature dimension m ≥ 3 with probability at least 1 − O(e^−d^1/10).
3.2 Illustrative Semi-synthetic Experiments: The semi-synthetic experiment gives frequent examples label-relevant left halves plus label-irrelevant right halves, while rare examples use label-relevant right halves.SimCLR activates both sides of frequent examples and the right side of rare examples; supervised learning mainly activates left halves.
3.2 Illustrative Semi-synthetic Experiments: 70.1% vs 44.3%: self-supervised representations outperform supervised representations on rare-class accuracy in the semi-synthetic experiment.The rare-class linear head is trained with 25,000 examples from five rare classes.

4 Improving SSL on Imbalanced Datasets with Regularization

The paper improves imbalanced SSL with reweighted sharpness-aware regularization. The method estimates example density without labels, regularizes estimated rare examples more strongly, and improves representation quality across in-domain and downstream evaluations.

Motivation: Rare examples have a larger pre-training generalization gap than frequent examples: 0.081 vs 0.035 on step-imbalance CIFAR-10.This motivates applying stronger regularization to rare examples.
Reweighted SAM: rwSAM reweights SAM’s inner maximization so the optimization landscape is flatter for rare examples, effectively regularizing them more strongly.The method uses data-dependent weights rather than labels.
Assigning Weight with Kernel Density Estimation: Kernel density estimation assigns weights inversely correlated with estimated representation density, approximating class frequency without access to labels.The density estimate is computed on the learned representations, with bandwidth and exponent selected by cross-validation.
Results: On long-tailed CIFAR-10, rwSAM improves SimSiam performance by over 1% under both step and exponential imbalance.It even surpasses balanced SimSiam trained with the same number of examples.
Results: rwSAM closes the rare-example generalization gap from 0.081 to 0.066 and consistently improves transferability on ImageNet-LT downstream tasks.The ImageNet-LT evaluations use fine-tuning and compare against MoCo v2 and SimSiam.

5 Related Work

Prior work addresses supervised learning under imbalance through resampling, reweighting, specialized losses, training pipelines, and meta-learning. This paper instead focuses on how imbalance affects self-supervised representations.

Classical supervised methods for imbalance include resampling, reweighting, ensemble modifications, and specialized training strategies.
Re-weighted regularization regularizes rare classes more heavily and is typically more effective than reweighting or resampling the losses.
Modern approaches use hard-example-focused losses, specialized pipelines, and meta-learning to address imbalanced recognition.
Existing studies also examine supervised representations and self-training or self-supervised pretraining for imbalanced recognition classifiers.
This paper studies the effect of class imbalance on self-supervised representations, rather than supervised imbalanced recognition classifiers.
Earlier SSL research used context prediction, puzzles, and rotation prediction, while recent methods approach supervised performance on ImageNet and downstream tasks.

6 Conclusion

The paper studies robustness to imbalanced training in self-supervised representations and finds them more robust than supervised representations. It also identifies future directions beyond vision datasets.

The paper presents the first study of robustness to imbalanced training in self-supervised representations.
Self-supervised representations are more robust to class imbalance than supervised representations.
The authors encourage practitioners to use SSL for pre-training or evaluate the impact of imbalanced pre-training on downstream tasks.
The experiments mainly focus on vision datasets, leaving the effect of imbalance on NLP datasets as a future direction.
Future work may analyze imbalance in broader wild settings such as domain shift and inform future unsupervised learning methods.

A.1 Implementation Details

The appendix describes imbalanced CIFAR-10 and ImageNet-LT pre-training, downstream evaluation protocols, and the figures and tables reporting results.

Generating Pre-training Datasets: CIFAR-10 imbalance uses an exponential class distribution with rarest-class ratios r ∈ {0.1, 0.01}, yielding 20,431 or 12,406 total samples.
Training Procedure: Supervised ImageNet-LT pre-training follows standard protocols, while SSL uses a 0.025 initial learning rate, batch size 256, and 300 epochs on standard ImageNet-LT.
Evaluation: ID evaluation trains a linear classifier on balanced CIFAR-10 or ImageNet, while OOD evaluation fine-tunes on four target datasets or probes STL-10.
Additional Results: Figure 6 reports OOD SimSiam results on ImageNet and compares relative gaps with balanced datasets.
Additional Results: Table 2 provides the numerical values corresponding to Figures 2 and 6.

A.2 Additional Results

Additional experiments show that SimSiam is also less sensitive to class imbalance than supervised representations, using ImageNet OOD evaluation and semi-synthetic data analysis.

SimSiam representations are less sensitive to class imbalance than supervised representations on ImageNet OOD evaluation.
Figure 6 reports SimSiam OOD results on ImageNet together with relative performance gaps from balanced datasets.
The semi-synthetic dataset assigns label-relevant and label-irrelevant-but-transferable features to different image halves by design.
The imbalanced semi-synthetic dataset contains 5,000 examples per frequent class and 10 examples per rare class.

C.1 Implementation Details

The section describes implementation choices for semi-synthetic visualization, optimization, and hyperparameter selection. It also specifies the reweighted update used in the method.

Semi-synthetic visualization: Semi-synthetic experiments visualize original CIFAR-10 images and compare Grad-CAM features learned by SimCLR and supervised learning.SimCLR uses both label-relevant and label-irrelevant visual sides, whereas supervised learning emphasizes label-relevant features from frequent data.
Optimization: The model is updated using the loss without reweighting, with φ = φ −η∇φbL(φ + ϵ(φ)).The detailed procedure is referenced in Algorithm 1.
Hyperparameters: Hyperparameters ρ and α are selected by cross-validation, with dataset-specific values for ImageNet-LT, iNaturalist, and CIFAR-10-LT.ImageNet-LT and iNaturalist use ρ = 2 and α = 0.5; CIFAR-10-LT uses ρ = 5 and α = 1.2.

C.2 Additional Results

Additional experiments evaluate SSL representations under imbalanced recognition and detection settings, while introducing the reweighted SAM procedure. The results show smaller imbalance effects for SSL and further gains from rwSAM.

Recognition evaluation: Fine-tuning on imbalanced ImageNet followed by class-aware classifier resampling provides an additional comparison protocol for imbalanced recognition.This protocol compares SSL representations with supervised imbalanced recognition methods.
Recognition evaluation: CRT outperforms direct supervised training on ImageNet-LT, indicating that classifier retraining improves supervised imbalanced recognition.The passage attributes the difference to the composition of the supervised head and learned features.
Reweighted SSL: 55.5 is reported for MoCo V2+rwSAM with a ResNet-50, compared with 55.0 for MoCo V2.These entries provide the reported comparison in the additional results.
Detection evaluation: MoCo V2 representations can match more complicated methods designed for supervised imbalanced recognition, and rwSAM further improves MoCo V2.The comparison uses a simple pre-training and fine-tuning pipeline.
Reweighted SSL: Reweighted SAM computes weights before a second optimization stage that uses the reweighted loss to calculate ϵ(φ).Algorithm 1 defines the input dataset, learned representations, weight computation, and reweighted SAM stage.
Detection evaluation: MoCo V2 has a much smaller imbalance-versus-balance pre-training gap than supervised learning across Pascal VOC detection settings.This pattern holds across the evaluated numbers of examples.

D Proof of Theorem 3.1

The proof analyzes supervised and self-supervised representations in a simplified Gaussian setting. It shows that supervised minimum-norm solutions have limited correlation with a designated feature direction, whereas SSL’s covariance-based objective preserves that direction.

Probabilistic lemmas: The proof assumes independent Gaussian variables and establishes high-probability concentration and weak-correlation properties for the toy dataset.The argument uses Gaussian tail bounds, independence, union bounds, and concentration inequalities.
Supervised learning: A constructed supervised classifier achieves margin at least 1 −O(d^-1 10), with the third classifier vector satisfying ∥w∗ 3∥2 2 ≤O(d^-1 5 ).This provides a feasible classifier with controlled margin and norm for the minimum-norm analysis.
Supervised learning: The minimum-norm supervised classifier has total squared norm at most 2 + O(d^-1 10 ), and its representation has small correlation with e2.The proof transfers the classifier norm bound to the learned representation through the factorization analysis.
Self-supervised learning: The empirical covariance matrix has a large e2-direction quadratic form, e2⊤Me2 ≥Ω(d 2 5 ), while all unit directions orthogonal to e2 have quadratic form O(1).This spectral separation is established using the data matrix properties and concentration bounds.
Self-supervised learning: The self-supervised objective shares a minimizer with matrix factorization, so its representation spans the top m eigenvectors of the empirical covariance matrix.The proof invokes the Eckart–Young–Mirsky theorem and decomposes e2 in the covariance eigenbasis.
Self-supervised learning: The covariance spectral gap forces the self-supervised representation to retain a strong component in the e2 direction.The proof first establishes ζ2 1 ≥1 2 and then bounds the residual component using the separated quadratic forms.