Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update NNCLR model example in docs #1790

Merged
merged 4 commits into from
Jan 31, 2025
Merged
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 27 additions & 1 deletion docs/source/examples/nnclr.rst
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,33 @@
NNCLR
=====

Example implementation of the NNCLR architecture.
NNCLR is a self-supervised framework for visual representation learning that builds upon contrastive methods. It shares similarities with SimCLR, such as using two augmented views of the same image, projection and prediction heads, and a contrastive loss. However, it introduces key modifications:

1. Nearest Neighbor Replacement: Instead of directly comparing two augmented views of the same sample, NNCLR replaces each sample with its nearest neighbor in a support set (or memory bank). This increases semantic variation in the learned representations.
2. Symmetric Loss: The contrastive loss is made symmetric to improve training stability.
3. Architectural Adjustments: NNCLR employs different sizes for projection and prediction head layers compared to SimCLR.

These improvements result in significantly better performance across multiple self-supervised learning benchmarks. Compared to SimCLR and other self-supervised methods, NNCLR achieves:
- Higher ImageNet linear evaluation accuracy.
- Improved semi-supervised learning results.
- Superior performance on transfer learning tasks, outperforming BYOL, SimCLR, and even supervised ImageNet pretraining in 8 out of 12 benchmarked cases.
guarin marked this conversation as resolved.
Show resolved Hide resolved

Key Components
--------------

- **Data Augmentations**: NNCLR applies the same transformations as SimCLR, including random cropping, resizing, color jittering, and Gaussian blur, to create diverse views of the same image.
- **Backbone**: A convolutional neural network (typically ResNet) encodes augmented images into feature representations.
- **Projection Head**: A multilayer perceptron (MLP) maps features into a contrastive space, improving representation learning.
- **Memory Bank**: NNCLR maintains a first-in, first-out (FIFO) memory bank, storing past feature representations. Older features are gradually discarded, ensuring a large and diverse set approximating the full dataset.
- **Nearest Neighbor Sampling**: Each feature representation is replaced by its nearest neighbor from the memory bank, introducing additional semantic variation beyond standard augmentations.
- **Contrastive Loss**: NNCLR employs normalized temperature-scaled cross-entropy loss (NT-Xent), encouraging alignment between positive pairs and separation from negative pairs.

Good to Know
----------------

- **Optimized for CNNs**: NNCLR is specifically designed for convolutional neural networks (CNNs), particularly ResNet. It is not recommended for transformer-based architectures.
- **Augmentation Robustness**: Compared to SimCLR, NNCLR is less dependent on strong augmentations since nearest neighbor sampling introduces natural semantic variation. However, performance still benefits from well-chosen augmentations and larger batch sizes.


Reference:
`With a Little Help from My Friends: Nearest-Neighbor Contrastive Learning of Visual Representations, 2021 <https://arxiv.org/abs/2104.14548>`_
Expand Down