PaperlayerView exact revision on arXiv

Source-linked AI summary

Unsupervised Learning of Visual Features by Contrasting Cluster Assignments

Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, Armand Joulin

arXiv:2006.09882v5cs.CV

TL;DR

Contrastive learning relies on computationally costly pairwise comparisons, motivating more scalable unsupervised image representation methods. SwAV clusters images online while enforcing consistency across augmented views and introduces multi-crop, improving ImageNet performance by 4.2% and surpassing supervised pretraining on multiple downstream tasks.

  • Problem

    Contrastive learning requires computationally costly pairwise comparisons, while clustering-based alternatives face scalability limitations on large datasets.

  • Method

    SwAV jointly clusters representations online and enforces consistency across augmented views using swapped predictions, complemented by multi-crop augmentations.

  • Results

    +4.2% top-1 accuracy on ImageNet, with SwAV surpassing supervised ImageNet pretraining across multiple downstream tasks.

  • Takeaways & Limitations

    SwAV provides a scalable self-supervised approach that remains effective with small batches and transfers strongly across downstream vision tasks.

Abstract

from arXiv · show

Unsupervised image representations have significantly reduced the gap with supervised pretraining, notably with the recent achievements of contrastive learning methods. These contrastive methods typically work online and rely on a large number of explicit pairwise feature comparisons, which is computationally challenging. In this paper, we propose an online algorithm, SwAV, that takes advantage of contrastive methods without requiring to compute pairwise comparisons. Specifically, our method simultaneously clusters the data while enforcing consistency between cluster assignments produced for different augmentations (or views) of the same image, instead of comparing features directly as in contrastive learning. Simply put, we use a swapped prediction mechanism where we predict the cluster assignment of a view from the representation of another view. Our method can be trained with large and small batches and can scale to unlimited amounts of data. Compared to previous contrastive methods, our method is more memory efficient since it does not require a large memory bank or a special momentum network. In addition, we also propose a new data augmentation strategy, multi-crop, that uses a mix of views with different resolutions in place of two full-resolution views, without increasing the memory or compute requirements much. We validate our findings by achieving 75.3% top-1 accuracy on ImageNet with ResNet-50, as well as surpassing supervised pretraining on all the considered transfer tasks.

1 Introduction

The paper proposes scalable online clustering and multi-crop augmentation to address computational limits in contrastive learning while increasing the number of image views. Combined, these contributions improve ImageNet performance and outperform supervised pretraining on multiple downstream tasks using frozen features.

  • 1 Introduction: Multi-crop consistently improves ImageNet performance between 2% and 4% across SimCLR, DeepCluster, and SeLa.The strategy is reported to generalize across several self-supervised methods.
  • 1 Introduction: 75.3% top-1 accuracy is achieved on ImageNet with a standard ResNet-50, rising to 78.5% with a wider model.These results use the ImageNet linear evaluation protocol.
  • 1 Introduction: The proposed online clustering loss improves performance by +2% on ImageNet and works with large and small batches without a large memory bank or momentum encoder.It replaces extensive pairwise comparisons with an online clustering-based objective.
  • 1 Introduction: Multi-crop increases the number of image views with no computational or memory overhead.It uses smaller-sized images instead of comparing only one pair of transformations per image.
  • 1 Introduction: Combining both contributions improves self-supervised performance by +4.2% on ImageNet and outperforms supervised ImageNet pretraining on multiple downstream tasks.The comparison uses a linear classifier on top of frozen features, without finetuning.

2 Related Work

Prior work on unsupervised visual representation learning includes instance- and contrastive-learning methods, clustering-based pseudo-labeling, and handcrafted pretext tasks. These approaches provide the context for methods that learn representations without manual labels.

  • Instance and contrastive learning: Instance-level methods treat each image as its own class, while memory banks make this formulation more tractable for large datasets.Dosovitskiy et al. assign one class per image; Wu et al. replace the classifier with a memory bank storing previously computed representations and use noise contrastive estimation.
  • Clustering for deep representation learning: Clustering-based methods use k-means assignments or pseudo-labels to learn visual representations and support supervised-network pretraining.Caron et al. demonstrate scalability to large uncurated datasets, while subsequent work casts pseudo-label assignment more formally.
  • Handcrafted pretext tasks: Handcrafted pretext-task methods manipulate inputs to create supervised signals, including encoding jigsaw puzzles as invariances for contrastive learning.Jigsaw tiles are described as non-overlapping crops with small resolution.

3 Method

SwAV learns visual features online by clustering augmented-image representations into prototype codes and enforcing consistency through swapped prediction rather than direct feature comparisons. It prevents collapse with batchwise equipartition, supports small batches, and improves efficiency through soft codes and multi-crop augmentation.

  • Swapped prediction: SwAV computes prototype codes for two augmentations of the same image and predicts each view’s code from the other view.Features are matched to K trainable prototypes, creating a swapped prediction problem between the two augmented views.
  • Swapped prediction: The method compares views through intermediate codes, unlike contrastive learning methods that compare image features directly.SwAV predicts codes obtained from one augmented view using the other view’s representation.
  • Online code computation: Codes are computed online from batch features with an equipartition constraint that assigns examples across prototypes and prevents trivial collapse.The prototypes are shared across batches, allowing SwAV to cluster multiple instances while enforcing distinct codes within each batch.
  • Online code computation: Soft codes outperform rounded discrete codes in the online setting, so SwAV jointly optimizes the encoder and prototypes using the continuous solution.Rounding is a more aggressive optimization step that can converge rapidly but lead to a worse solution.
  • Multi-crop augmentation: Multi-crop uses two standard-resolution crops plus V low-resolution crops, improving several self-supervised methods without substantially increasing memory or compute.Codes are computed only from full-resolution crops because using all crops increases computation and can degrade transfer performance.

4 Main Results

SwAV learns transferable visual features that outperform prior self-supervised and supervised representations across ImageNet evaluation and downstream tasks. These results persist with wider ResNet-50 variants and small-batch training while using modest feature storage.

  • ImageNet evaluation: +4.2% top-1 accuracy over the state of the art is achieved with frozen ImageNet features, only 1.2% below fully supervised performance.SwAV is trained for 800 epochs with batches of 4096 images.
  • ImageNet evaluation: 0.6% is the remaining gap to supervised training for large ResNet-50 architectures.Performance increases with model width and follows a trend similar to supervised learning.
  • Transfer learning: SwAV features outperform supervised features on Places205, VOC07, and iNaturalist2018 linear-classification transfers.The evaluation transfers unlabeled ImageNet features to several downstream vision tasks.
  • Small-batch training: Small-batch SwAV maintains state-of-the-art performance using 256-image batches, a queue of 3, 840 features, and no additional momentum encoder network.MoCov2 requires 65, 536 stored features plus a momentum encoder for good performance in the same setup.

5 Ablation Study

The ablations show that clustering-based methods can outperform contrastive instance learning, while SwAV matches DeepCluster-v2 with comparable accuracy and greater online scalability. Multi-crop consistently improves performance, and SwAV reaches strong accuracy even after relatively short training.

  • Comparing clustering with contrastive instance learning: SwAV and DeepCluster-v2 outperform SimCLR by 2% without multi-crop and by 3.5% with multi-crop under matched training settings.The comparison uses the same data augmentation, epochs, and batch sizes across methods.
  • Advantage of SwAV compared to DeepCluster-v2: 75.2% top-1 accuracy on ImageNet is achieved by DeepCluster-v2 versus 75.3% for SwAV when both use 800 epochs and 8 crops.SwAV remains preferable for extremely large datasets because it operates online, whereas DeepCluster-v2 does not.
  • Impact of training duration: 72.1% top-1 accuracy is reached by SwAV after 100 epochs and 6h15 of training, although longer training further improves performance.Each ResNet-50 was trained on 64 V100 16GB GPUs with batch size 4096.

6 Discussion · A Implementation Details · A.1 Implementation details of SwAV training

The paper situates SwAV within the progress of self-supervised learning and provides publicly available implementation details, including its training loop, optimization settings, and distributed-training advantages.

  • 6 Discussion: Self-supervised learning has progressed rapidly and can surpass supervised learning on transfer tasks, motivating exploration of architectures without supervision.The paper notes that existing architectures were designed for supervised tasks and may not be optimal for self-supervised learning.
  • A Implementation Details: The implementation section provides SwAV pretraining and transfer-learning details, with publicly available code.Code is available at https://github.com/facebookresearch/swav.
  • A.1 Implementation details of SwAV training: The SwAV loop generates two augmented crops, computes prototype scores, and converts them into temperature-scaled probabilities.The pseudocode uses x_t = t(x), x_s = s(x), prototype scores from matrix multiplication, and Softmax(scores / temp).
  • A.1 Implementation details of SwAV training: The assignment computation uses Sinkhorn normalization with ε-based exponentiation and iterative row- and column-marginal normalization.The pseudocode initializes uniform marginals and repeatedly rescales Q before returning column-normalized assignments.
  • A.1 Implementation details of SwAV training: SwAV solves the swapped-prediction objective by matching assignments across views, then updates the network and prototypes with SGD.The loss is -0.5 * mean(q_t * log(p_s) + q_s * log(p_t)), followed by loss.backward(), update(model.params), and update(C).
  • A.1 Implementation details of SwAV training: The optimization uses 10−6 weight decay, LARS, a learning rate of 4.8, 10-epoch warmup, cosine decay to 0.0048, and prototype freezing during epoch one.Batch normalization is synchronized across GPUs, and the prototypes are normalized during training.
  • A.1 Implementation details of SwAV training: 4096 instances are distributed across 64 V100 16Gb GPUs, while τ = 0.1 and ε = 0.05 are used for all runs.Training uses stochastic gradient descent with each GPU processing 64 instances.
  • A.1 Implementation details of SwAV training: SwAV is better suited to multi-node distribution because it shares only Sinkhorn row and column normalization statistics rather than the full feature matrix.SimCLR and MoCo share feature matrices across GPUs at every batch, which can become a bottleneck.

A.2 Data augmentation used in SwAV · A.3 Implementation details of linear classification on ImageNet with ResNet-50 · A.4 Implementation details of semi-supervised learning (finetuning with 1% or 10% labels)

SwAV uses narrowly scaled random crops to create image views, including two global and additional small-resolution views in multi-crop. Its ImageNet evaluation freezes 2048-D ResNet-50 representations for linear classification and also finetunes SwAV-pretrained models with 1% or 10% labels.

  • A.2 Data augmentation used in SwAV: SwAV creates two views using random-sized, random-aspect-ratio crops with scaling s=(0.14, 1).The crops use PyTorch’s RandomResizedCrop and are resized to 224 × 224 pixels unless otherwise specified.
  • A.2 Data augmentation used in SwAV: Multi-crop transforms each image into V + 2 views comprising two global views and V small-resolution zoomed views.The strategy mixes global and small-resolution views rather than using only two full-resolution views.
  • A.3 Implementation details of linear classification on ImageNet with ResNet-50: 75.3 top-1 accuracy is obtained on ImageNet by training a linear classifier on frozen 2048-D final representations from a SwAV-trained ResNet-50.The classifier is trained for 100 epochs with learning rate 0.3, weight decay 10−6, cosine decay, and batch size 256.
  • A.3 Implementation details of linear classification on ImageNet with ResNet-50: Linear evaluation applies standard random-sized, random-aspect-ratio crops and random horizontal flips.The crops use the default RandomResizedCrop parameters for this evaluation.
  • A.4 Implementation details of semi-supervised learning (finetuning with 1% or 10% labels): Semi-supervised finetuning uses either 1% or 10% of ImageNet labeled images with a SwAV-pretrained ResNet-50.The labeled splits follow the official SimCLR code release.
  • A.4 Implementation details of semi-supervised learning (finetuning with 1% or 10% labels): Finetuning follows PCL hyperparameters for 20 epochs with batch size 256, distinct learning rates, and learning-rate drops by factor 0.2 at epochs 12 and 16.No weight decay is applied during finetuning.

A.5 Implementation details of transfer learning on downstream tasks … B.3 Transferring unsupervised features to downstream tasks

The paper specifies transfer-learning protocols across classification and detection, extends SwAV to small batches and larger architectures, and reports computational and downstream evaluations. It also details augmentation ablations and the multi-crop SimCLR loss.

  • A.5 Implementation details of transfer learning on downstream tasks: Linear evaluation uses 2048-D global-average-pooled ResNet-50 features, with a LIBLINEAR linear SVM on VOC07 and SGD linear models on Places205 and iNat18.Images for VOC07 are resized to 256 pixels on the shorter side before a 224 × 224 center crop.
  • A.5 Implementation details of transfer learning on downstream tasks: SwAV features are transferred to object detection using Faster R-CNN on VOC07+12 and DETR on COCO.The VOC07+12 protocol uses an initial learning rate of 0.1, warmed for 1000 iterations, with other hyperparameters following He et al.
  • A.6 Implementation details of training with small batches of 256 images: The small-batch training setup delays queue usage until 15 epochs and simulates batches of size 4096 by storing the last 15 batches.The stored queue contains 3,840 vectors of dimension 128, with weight decay 10^-6, LARS, and learning rate 0.6.
  • A.7 Implementation details of ablation studies: Ablation studies use SimCLR’s augmentation to isolate SwAV’s contributions, including RandomResizedCrop scale s=(0.08, 1) for the two large-resolution views.This replaces the usual s=(0.14, 1) setting for those views.
  • A.8 SimCLR loss with multi-crop augmentation: With multi-crop SimCLR, each instance contributes M crops, producing an effective batch of N = B × M crops; for 2x160+4x96, M = 6.Positive examples are other crops from the same instance, while negatives are crops from other instances in the batch.
  • B.1 Running times: Compute and GPU-memory requirements are measured for ResNet-50 methods trained on 64 V100 16GB GPUs with batch size 4096, mixed precision, and synchronized batch normalization.SwAV performance is also measured across training lengths in hours.
  • B.2 Larger architectures: SwAV benefits from larger architectures, with R50-w2, R50-w4, and R50-w5 trained for 400 epochs using six crops and architecture-specific GPU, batch-size, and queue settings.R50-w2 uses 128 GPUs and batch size 4096; R50-w4 uses 64 GPUs and batch size 2560; R50-w5 uses 128 GPUs and batch size 1536.
  • B.3 Transferring unsupervised features to downstream tasks: Expanded downstream results evaluate frozen-feature linear classification across datasets and ResNet-50 object detection finetuning with Faster R-CNN and DETR.The comparison includes previously and concurrently published self-supervised methods.

B.4 More detection metrics for object detection

Across three detector families, SwAV features outperform ImageNet-supervised pretraining on all reported object-detection metrics. MoCov2 is strong on detection and surpasses SwAV on some metrics, but is not competitive with supervised features on unfine-tuned classification.

  • Detection results: SwAV outperforms ImageNet-supervised pretraining on all reported object-detection evaluation metrics.The comparisons use the additional metrics reported in Tables 7–9.
  • Detection results: The evaluation finetunes ResNet-50 with Mask R-CNN, Faster R-CNN, and DETR.Results are reported in Tables 7, 8, and 9.
  • Comparison with MoCov2: MoCov2 performs particularly well on object detection and surpasses SwAV on some detection metrics, but is not competitive with supervised features on unfine-tuned classification.The classification comparison is reported in Table 6.

B.5 Low-Shot learning on ImageNet for SwAV pretrained on Instagram data · B.6 Image classification with KNN classifiers on ImageNet

SwAV features pretrained on Instagram data substantially improve low-shot ImageNet fine-tuning over training from scratch. On ImageNet KNN evaluation, SwAV surpasses the current state of the art, with fewer neighbors improving performance.

  • B.5 Low-Shot learning on ImageNet for SwAV pretrained on Instagram data: Instagram-pretrained SwAV features are evaluated for low-shot learning on ImageNet using only a few labels per category.The evaluation fine-tunes the pretrained features.
  • B.5 Low-Shot learning on ImageNet for SwAV pretrained on Instagram data: Pretrained Instagram SwAV features considerably improve performance compared with training from scratch.The comparison concerns low-shot ImageNet fine-tuning.
  • B.5 Low-Shot learning on ImageNet for SwAV pretrained on Instagram data: The B.5 experiment tests whether Instagram-pretrained SwAV can function as a pretraining method for low-shot ImageNet learning.Results are reported in Table 9.
  • B.6 Image classification with KNN classifiers on ImageNet: SwAV features are evaluated with K-nearest neighbor classifiers on ImageNet following previous-work protocols.Features come from network outputs for center crops of training and test images.
  • B.6 Image classification with KNN classifiers on ImageNet: The KNN evaluation reports results using 20 and 200 nearest neighbors.These results are presented in Table 11.
  • B.6 Image classification with KNN classifiers on ImageNet: SwAV outperforms the current state of the art in ImageNet KNN classification.The reported evaluation compares the unsupervised features against prior methods.
  • B.6 Image classification with KNN classifiers on ImageNet: Using fewer nearest neighbors boosts the model’s KNN performance.The paper specifically notes this improvement relative to using more neighbors.

C Ablation Studies on Clustering … D Details on Clustering-Based methods: DeepCluster-v2 and SeLa-v2

The ablations show that SwAV is relatively insensitive to prototype count, benefits from learned and soft assignments, and converges with few Sinkhorn-Knopp iterations. The accompanying clustering-based methods use alternating pseudo-label assignment and classification training phases.

  • C.1 Number of prototypes: C.1 Number of prototypes: Varying prototypes from 3k to 100k changes ImageNet performance by at most 0.3.The models use ResNet-50 trained for 400 epochs with 2 × 160 + 4 × 96 crops and frozen-representation linear evaluation.
  • C.2 Learning the prototypes: C.2 Learning the prototypes: Learning prototypes improves SwAV from 73.1 to 73.9.This indicates that adapting prototypes to the dataset distribution improves performance over fixed random prototypes.
  • C.2 Learning the prototypes: C.2 Learning the prototypes: Fixed random prototypes perform almost as well because SwAV uses them to contrast image views without many negative pairwise comparisons.Unlike offline methods such as DeepCluster, SwAV does not strongly encourage prototypes to be categorical.
  • C.3 Hard versus soft assignments: C.3 Hard versus soft assignments: Hard assignments perform worse than SwAV’s default soft assignments.The comparison uses 400-epoch training with 2 × 160 + 4 × 96 crops and frozen-representation linear evaluation.
  • C.4 Impact of the number of iterations in Sinkhorn algorithm: C.4 Impact of the number of iterations in Sinkhorn algorithm: Three Sinkhorn-Knopp iterations are sufficient for convergence.With fewer iterations the loss fails to converge, while more iterations slightly alter transfer performance.
  • D Details on Clustering-Based methods: DeepCluster-v2 and SeLa-v2: D Details on Clustering-Based methods: DeepCluster-v2 and SeLa-v2: Both methods alternate pseudo-label generation in an assignment phase with classification-loss training in a training phase.These implementation details are presented relative to the corresponding original publications.

D.1 Training phase · D.2 Assignment phase

Training minimizes pseudo-label classification losses while assignments remain fixed, then assignments are recomputed across the full dataset in a separate phase. DeepCluster-v2 and SeLa-v2 use different training and assignment mechanisms, with separated phases enabling stable clustering without observed collapse.

  • D.1 Training phase: Both methods minimize the multinomial logistic loss for classifying pseudo-labels q during training.
  • D.1 Training phase: Pseudo-labels remain fixed during training and are updated for the entire dataset once per epoch during assignment.
  • D.1 Training phase: In DeepCluster-v2, the classification head c and convnet weights are trained to classify images into their corresponding pseudo-labels.
  • D.1 Training phase: SeLa-v2 learns prototypes c with stochastic gradient descent during training and optimizes its stated training objective.
  • D.2 Assignment phase: Assignment requires feature representations z for the entire dataset; updating assignments each epoch accounts for one third of total training time in the original implementation.
  • D.2 Assignment phase: DeepCluster-v2 obtains pseudo-labels q by minimizing a spherical k-means objective over normalized features and prototype columns.
  • D.2 Assignment phase: The separated assignment and training phases made cluster reassignments and balanced batch sampling unnecessary, with no collapsing observed during training.
  • D.2 Assignment phase: SeLa-v2 uses Sinkhorn-Knopp to optimize assignments with the same loss used during training, while both clustering methods use three sets of 3,000 prototypes.
Loading 2006.09882v5…