Source-linked AI summary
With a Little Help from My Friends: Nearest-Neighbor Contrastive Learning of Visual Representations
Debidatta Dwibedi, Yusuf Aytar, Jonathan Tompson, Pierre Sermanet, Andrew Zisserman
TL;DR
Instance-discrimination methods typically use augmented views of the same image as contrastive positives, limiting coverage of viewpoints, deformations, and other within-class variation. NNCLR instead uses nearest neighbors from a support set as positives, improving benchmark performance and reducing reliance on complex augmentations.
Problem
Instance discrimination relies on same-image augmentations, which cannot provide positives for all viewpoints, deformations, or similar instances within a semantic class.
Method
NNCLR uses nearest neighbors in a learned representation space as positives alongside different views of the same sample, without clustering.
Results
NNCLR outperforms state-of-the-art methods on 8 of 12 transfer-learning tasks and achieves 73.3% ImageNet accuracy using only random-crop augmentations.
Takeaways & Limitations
Using nearest neighbors increases positive diversity and reduces reliance on data augmentation techniques in self-supervised representation learning.
Takeaways & Limitations
The approach remains constrained by augmentation coverage in baseline instance-discrimination methods, which cannot cover all variances within a class.
Abstract
from arXiv · showhide
Self-supervised learning algorithms based on instance discrimination train encoders to be invariant to pre-defined transformations of the same instance. While most methods treat different views of the same image as positives for a contrastive loss, we are interested in using positives from other instances in the dataset. Our method, Nearest-Neighbor Contrastive Learning of visual Representations (NNCLR), samples the nearest neighbors from the dataset in the latent space, and treats them as positives. This provides more semantic variations than pre-defined transformations. We find that using the nearest-neighbor as positive in contrastive losses improves performance significantly on ImageNet classification, from 71.7% to 75.6%, outperforming previous state-of-the-art methods. On semi-supervised learning benchmarks we improve performance significantly when only 1% ImageNet labels are available, from 53.8% to 56.5%. On transfer learning benchmarks our method outperforms state-of-the-art methods (including supervised learning with ImageNet) on 8 out of 12 downstream datasets. Furthermore, we demonstrate empirically that our method is less reliant on complex data augmentations. We see a relative reduction of only 2.1% ImageNet Top-1 accuracy when we train using only random crops.
1 Google Research, 2 DeepMind
The passage contains a Google Research and DeepMind contact email address.
- The contact address uses the google.com domain.
- The address lists four named contacts before the domain.
- The passage provides contact information rather than research content.
1. Introduction
NNCLR extends instance discrimination by using nearest neighbors from other examples as positives, aiming to capture semantic variations beyond predefined augmentations. The method improves ImageNet, semi-supervised, and transfer-learning performance while reducing reliance on complex data augmentation.
- Instance discrimination trains models to become invariant to multiple transformations of a single sample.
- Random crops and color changes cannot provide positive pairs for different viewpoints, deformations, or other similar instances.
- NNCLR is intended to learn features invariant to different viewpoints, deformations, and intra-class variations without labels or multiple modalities.
- NNCLR uses nearest neighbors in latent space as positives alongside different views of the same sample, without resorting to clustering.
- 8 out of 12 transfer-learning tasks outperform state-of-the-art methods, including supervised ImageNet features.
- 73.3% ImageNet accuracy with only random-crop augmentations reduces the method’s reliance on complex augmentation strategies.
2. Related Work
Related work situates NNCLR within self-supervised and contrastive representation learning, distinguishing it from approaches based on pretext tasks, clustering, prototype centers, or partial nearest-neighbor objectives.
- Self-supervised learning obtains robust representations from raw data without expensive labels or annotations, often through surrogate pretext tasks.
- Contrastive methods learn latent spaces by drawing positive samples together and pushing negative samples apart, sometimes without explicit negatives.
- A related approach uses nearest neighbors of clustered aggregate embeddings as positives, whereas NNCLR uses a fixed memory independent of dataset size without aggregation or clustering.
- SwAV differs from NNCLR by storing prototype centers learned through Sinkhorn clustering throughout pre-training.
- Concurrent work adds a nearest-neighbor term to BYOL, while NNCLR uses the nearest neighbor in all loss terms.
3. Approach
NNCLR extends contrastive instance discrimination by replacing same-image-only positives with nearest neighbors from a support set in latent space. It retains mini-batch negatives and the InfoNCE framework while targeting greater semantic variation and reduced dependence on augmentations.
- Contrastive instance discrimination: InfoNCE pulls positive embeddings together and separates them from negative embeddings, with SimCLR using two augmented views of the same image as positives.The negative pairs are formed from other embeddings in the mini-batch.
- Contrastive instance discrimination: Pre-defined augmentations cannot link potentially same-class samples, limiting invariance to large intra-class variations.This motivates extending positives beyond a single instance.
- Nearest-Neighbor CLR (NNCLR): NNCLR uses the nearest neighbor of an embedding in support set Q as the positive pair instead of relying only on two augmentations.Negative pairs remain sourced from the mini-batch, and the method uses a variant of the InfoNCE loss.
- Nearest-Neighbor CLR (NNCLR): NNCLR applies l2 normalization before nearest-neighbor search and before computing dot products in the loss.The method minimizes the average loss over mini-batch elements and adds a symmetric term as an implementation detail.
- Nearest-Neighbor CLR (NNCLR): A prediction MLP is applied to the positive embedding, and the resulting representation is used in the NNCLR objective.The paper reports that this adds a small performance boost.
- Nearest-Neighbor CLR (NNCLR): The support set is implemented as a first-in-first-out queue sized to approximate the full dataset distribution in embedding space.It is initialized as a random matrix and updated after each training step.
4. Experiments
NNCLR is evaluated against self-supervised and supervised representations on ImageNet, semi-supervised learning, transfer learning, and controlled ablations. Across these evaluations, nearest-neighbor positives improve generalization, reduce reliance on complex augmentations, and expose trade-offs involving neighbor quality, queue size, and computation.
- ImageNet evaluations: NNCLR achieves the best ImageNet linear-evaluation performance among compared two-view methods and improves over MoCo v2 and SimCLR v2 by more than 3.6%.It also improves over InfoMin Aug. by more than 2% in top-1 classification performance.
- ImageNet evaluations: NNCLR outperforms SwAV by 0.3% in Top-1 accuracy in the eight-view comparison, while multi-crop provides only a 0.2% boost for NNCLR.For SwAV, multi-crop provides a 3.5% performance improvement, but increasing crops also raises memory and compute requirements.
- Transfer learning evaluations: NNCLR outperforms supervised ImageNet features on 11 of 12 transfer datasets and improves over BYOL and SimCLR on 8 of 12.The benchmark spans Food101, CIFAR10, CIFAR100, Birdsnap, SUN397, Cars, Aircraft, VOC2007, DTD, Pets, Caltech-101, and Flowers.
- Ablations: 73.3% ImageNet top-1 accuracy is achieved with NNCLR using only random crops after 1000 pre-training epochs, reducing the gap to the fully augmented setting to 1.6%.At the reported shorter setting, NNCLR reaches 68.2% top-1 with only a 4.7% performance drop, compared with larger relative drops for SimCLR and BYOL.
- Ablations: The top-1 nearest neighbor performs best, while sampling beyond k = 1 or using a soft nearest neighbor slightly worsens performance.The authors interpret nearest-neighbor positives as introducing semantic diversity beyond predefined augmentations, although retrievals can also reflect texture or background similarities.
- Discussion: The queue’s nearest-neighbor class accuracy reaches about 57% late in training, whereas an oracle same-class neighbor reaches 75.8% versus 72.9% for the support-set top-1 neighbor at 300 epochs.The authors suggest background-only random crops may limit nearest-neighbor class accuracy and that improved unsupervised neighbor selection could improve performance.
- Discussion: NNCLR begins outperforming SimCLR after about 120 epochs and remains higher through 1000 pre-training epochs, despite requiring learning from harder positives.The training-loss comparison characterizes NNCLR as a more difficult task than using augmented views of the same image as positives.
5. Conclusion
NNCLR increases positive diversity in contrastive self-supervised learning by using nearest neighbors from a support set, while reducing reliance on data augmentation.
- NNCLR increases the diversity of positives in contrastive self-supervised learning by using nearest neighbors from a support set.
- NNCLR achieves state-of-the-art performance across multiple datasets and reduces reliance on data augmentation techniques.
A. Pseudo-code
The NNCLR pseudo-code augments two views, encodes and predicts them, retrieves top-1 nearest neighbors from a queue, computes contrastive losses, and updates the models and queue.
- Two augmented views are encoded into projections and transformed by a prediction MLP.
- Top-1 nearest neighbors for both projections are retrieved from the normalized queue using similarity maximization.
- Backpropagation updates the encoder and predictor, while the queue is refreshed with the latest projections.
- The loss compares each nearest neighbor with the prediction from the opposite view using temperature-scaled logits.
- A momentum encoder can also be used with NNCLR training.
- Nearest neighbors initially reflect color and texture, but later tend to belong to similar semantic classes as the encoder improves.
C. SimSiam with Nearest-neighbor as Positive
NNSiam replaces SimSiam’s stop-gradient branch with a nearest neighbor from the support set, showing that harder positives can improve non-contrastive learning.
- NNSiam replaces SimSiam’s stop-gradient branch with a nearest neighbor from the support set.
- Nearest-neighbor positives implicitly provide stop-gradient behavior because the neighbors are selected as hard matches.
- The experiment compares NNSiam and SimSiam with SimCLR and NNCLR.
- 1.3% improvement in accuracy is achieved under the ImageNet linear evaluation protocol with a non-contrastive loss.
D. Experiments with Vision Transformers
NNCLR is effective for Vision Transformers, achieving higher ImageNet linear-evaluation accuracy than SimCLR under the reported ViT training setup.
- The ViT experiments train for 1000 epochs using two crops, Adam, warmup, cosine decay, weight decay, and stochastic depth.
- 76.5 ViT-B/16 NNCLR Top-1 ImageNet accuracy exceeds 74.5 ViT-B/16 SimCLR accuracy.
- Table 11 reports Vision Transformer Top-1 ImageNet accuracy under the linear protocol.
E. Self-supervised Learning as a Pre-training Step for Supervised Learning
NNCLR pre-training can improve subsequent supervised ImageNet training, making self-supervised initialization a useful pre-training step.
- 79.1% Top-1 ImageNet accuracy was achieved after NNCLR pre-training followed by supervised training, versus 76.2% without self-supervised initialization.Both self-supervised and supervised training used the ImageNet 2012 training set.
- The experiment pre-trained a model with NNCLR for 1000 epochs, then performed regular supervised training for 100 epochs.
- The setup used 100% of ImageNet labels during supervised fine-tuning, unlike limited-label semi-supervised learning.
- Table 12 reports Top-1 ImageNet accuracy for self-supervised pre-training followed by supervised learning.
F. Transfer Learning
The transfer-learning experiments compare supervised and NNCLR-pretrained visual models, including Vision Transformers, across downstream datasets.
- ViT-B/16 with NNCLR outperformed the same architecture trained only with supervised learning on all transfer-learning datasets.
- ViT-B/16 with NNCLR outperformed NNCLR-trained ResNet50 on 8 of 12 transfer-learning datasets.
- Supervised ViT-B/16 transferred worse than supervised ResNet50 on 10 of 12 datasets.
- NNCLR produced notable gains for ViT-B/16 on Birdsnap, SUN397, Cars, Aircraft, DTD, and Flowers.The reported increases were approximately 10.7%, 5.8%, 8.0%, 13.2%, 5.8%, and 9.3%, respectively.
G.3. Comparison of Attention in ResNets vs ViT
The attention comparison examines ResNet-50 and ViT-B/16 under random, supervised, and NNCLR weights to assess object localization and semantic feature behavior.
- Attention maps were compared for ResNet-50 and ViT-B/16 using randomly initialized, ImageNet-supervised, and ImageNet-NNCLR weights.For ViT, the comparison used final-layer [CLS] self-attention; for ResNet-50, it used average-pooled embeddings as query embeddings.
- The transfer-learning table reports Top-1 accuracy for most datasets, with mean per-class accuracy or 11-point MAP for specified exceptions.
- NNCLR-trained ResNets produced attention maps that delineated salient objects in images.
- Supervised ViTs did not highlight objects in their attention maps, whereas supervised and self-supervised ResNets did.
- Color and texture similarities can let randomly initialized models localize objects, so delineated attention does not necessarily indicate semantically meaningful features.The authors suggest combining self-attention and cross-attention maps for more robust interpretation.
- Cross-attention searches use an object feature from one image to locate the same object across other images, with NNCLR features remaining close.