SimCLR is one of the most influential concrete contrastive learning frameworks โ a deliberately simple recipe (as its name suggests: "Simple framework for Contrastive Learning of visual Representations") that achieved strong results with a clean, minimal design.
The SimCLR Pipeline
- Take an image \(\mathbf{x}\) and apply two independent random augmentations, producing two "views" \(\tilde{\mathbf{x}}_i\) and \(\tilde{\mathbf{x}}_j\) โ a positive pair.
- Pass both views through a shared encoder network (typically a CNN, like a ResNet) to get representations \(\mathbf{h}_i, \mathbf{h}_j\).
- Pass those representations through a small "projection head" (an additional MLP) to get \(\mathbf{z}_i, \mathbf{z}_j\), the vectors actually used in the InfoNCE loss.
- Train the whole thing end-to-end with the InfoNCE loss from Contrastive Learning, treating every other image in the (large) batch as negatives.
A Subtle but Important Detail: The Projection Head
SimCLR found that computing the contrastive loss on \(\mathbf{z}\) (after the projection head) rather than directly on \(\mathbf{h}\) (the encoder's raw output) produces meaningfully better downstream representations โ even though \(\mathbf{h}\), not \(\mathbf{z}\), is what actually gets reused for downstream tasks afterward. The intuition: the contrastive task itself discards some information that's actually useful downstream, and the projection head absorbs that specific loss, protecting \(\mathbf{h}\) from being distorted by it.
Diagram
Two augmented views pass through a shared encoder and projection head; the loss is computed on the projected z, but h is what actually gets reused downstream.
Code
import torch.nn as nn
import torchvision.transforms as T
augmentation = T.Compose([
T.RandomResizedCrop(224), T.RandomHorizontalFlip(), T.ColorJitter(0.4, 0.4, 0.4, 0.1), T.RandomGrayscale(p=0.2)
])
class SimCLRModel(nn.Module):
def __init__(self, encoder, feature_dim, proj_dim=128):
super().__init__()
self.encoder = encoder # this is what gets REUSED downstream
self.projection_head = nn.Sequential(
nn.Linear(feature_dim, feature_dim), nn.ReLU(), nn.Linear(feature_dim, proj_dim)
) # this gets DISCARDED after pretraining
def forward(self, x):
h = self.encoder(x)
z = self.projection_head(h)
return h, z # h for downstream use, z for the contrastive loss
Why SimCLR Needs Large Batches
Since every other example in the batch serves as a negative, SimCLR's effective negative pool size is directly tied to batch size โ small batches provide a weak, less-informative contrastive signal. This is exactly the practical limitation (requiring very large batches, and correspondingly large GPU memory) that MoCo, the next note, was designed to overcome.
Common Mistakes
- Discarding the projection head's benefit by using \(\mathbf{h}\) directly in the loss โ SimCLR's empirical finding specifically favors computing the loss on the projected \(\mathbf{z}\), then discarding the projection head and keeping \(\mathbf{h}\) for downstream use.
- Using too small a batch size and expecting SimCLR-level performance โ the framework's design assumes a large negative pool, which fundamentally requires large batches.
Interview Relevance
Q: "Why does SimCLR use a separate projection head, computing the contrastive loss on its output rather than the encoder's raw representation directly?" The contrastive task's specific objective can distort or discard information that's actually useful for downstream tasks. By computing the loss on a separate projected representation \(\mathbf{z}\) and discarding the projection head afterward, the encoder's own output \(\mathbf{h}\) is protected from this distortion, empirically producing better downstream transfer performance than training the loss directly on \(\mathbf{h}\).
Practice Question
Why does SimCLR's reliance on other batch members as negatives directly tie its effectiveness to batch size?