[Paper Review] Supervised Contrastive Learning

Prannay Khosla, Piotr Teterwak+7 more|arXiv (Cornell University)|Apr 23, 2020

Domain Adaptation and Few-Shot LearningComputer Science62 references117 citations

TL;DR

The paper introduces SupCon, a supervised extension of contrastive learning that uses all samples from the same class as positives, achieving state-of-the-art results on ImageNet and improved robustness and hyperparameter stability.

ABSTRACT

Contrastive learning applied to self-supervised representation learning has seen a resurgence in recent years, leading to state of the art performance in the unsupervised training of deep image models. Modern batch contrastive approaches subsume or significantly outperform traditional contrastive losses such as triplet, max-margin and the N-pairs loss. In this work, we extend the self-supervised batch contrastive approach to the fully-supervised setting, allowing us to effectively leverage label information. Clusters of points belonging to the same class are pulled together in embedding space, while simultaneously pushing apart clusters of samples from different classes. We analyze two possible versions of the supervised contrastive (SupCon) loss, identifying the best-performing formulation of the loss. On ResNet-200, we achieve top-1 accuracy of 81.4% on the ImageNet dataset, which is 0.8% above the best number reported for this architecture. We show consistent outperformance over cross-entropy on other datasets and two ResNet variants. The loss shows benefits for robustness to natural corruptions and is more stable to hyperparameter settings such as optimizers and data augmentations. Our loss function is simple to implement, and reference TensorFlow code is released at https://t.ly/supcon.

Motivation & Objective

• Leverage label information to pull together same-class embeddings while pushing apart different classes.
• Extend self-supervised contrastive losses to the fully supervised setting with multiple positives per anchor.
• Analyze and identify the best-performing supervised contrastive loss formulation.
• Demonstrate state-of-the-art ImageNet results and improved robustness and stability over cross-entropy.

Proposed method

• Apply data augmentation to create two views of each input.
• Encode views with a shared encoder to obtain normalized representations on a unit sphere.
• Project representations through a projection network and normalize the projection for contrastive learning.
• Define and compare two supervised contrastive loss formulations that aggregate positives from the same class.
• Show that the out-of-batch supervision loss (L_out^sup) outperforms the in-batch version (L_in^sup) due to normalization effects and gradient structure.
• Train in two stages (contrastive pretraining followed by a linear classifier) to evaluate representations.

Experimental results

Research questions

• RQ1 Can supervised contrastive losses outperform cross-entropy on large-scale datasets like ImageNet?
• RQ2 Does incorporating multiple positives per anchor improve clustering of same-class samples in embedding space?
• RQ3 How do different supervised contrastive loss formulations compare in performance and training stability?
• RQ4 Is SupCon robust to augmentations, optimizers, and data quantity, compared to traditional supervised losses?
• RQ5 What is the transferability of SupCon representations to other datasets and tasks?

Key findings

• SupCon achieves 81.4% top-1 accuracy on ImageNet with ResNet-200, 0.8 percentage points above the previous state of the art for that architecture.
• SupCon outperforms cross-entropy and other baselines on CIFAR-10, CIFAR-100, and ImageNet across architectures.
• SupCon improves robustness to natural corruptions on ImageNet-C and demonstrates lower sensitivity to hyperparameters and augmentations.
• Using multiple positives per anchor and many negatives enhances gradient signals, enabling effective learning without explicit hard-negative mining.
• A memory-efficient setup with 8192 memory size and batch size 256 yields 79.1% top-1 on ImageNet with ResNet-50, surpassing the 6144-batch baseline without memory.
• Two-stage training (contrastive pretraining plus linear evaluation) achieves competitive results and can be performed with as few as 10 epochs for the linear stage.