What is self-supervised learning and how do contrastive methods like SimCLR learn representations?

Self-supervised learning (SSL) is a form of unsupervised learning where the model is trained on a pretext task defined entirely from the data itself — no human-provided labels. The learned representations can then be transferred to downstream tasks with few or no labels (linear probe, fine-tuning).

Contrastive methods like SimCLR define a pretext task based on augmentation invariance: for each input, create two random augmented views (crops, colour jitter, flips) and train the model so that representations of the two views of the same image are similar (positive pair), while representations of views from different images are dissimilar (negative pairs). The NT-Xent loss (normalised temperature-scaled cross-entropy) implements this: for a batch of N images (2N views), the model is trained to identify the matching view among 2(N-1) negative candidates.

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as T

# Augmentation pipeline: two random views of the same image
augment = T.Compose([
    T.RandomResizedCrop(224),
    T.RandomHorizontalFlip(),
    T.ColorJitter(0.4, 0.4, 0.4, 0.1),
    T.RandomGrayscale(p=0.2),
    T.GaussianBlur(kernel_size=23),
    T.ToTensor(),
])

class SimCLRLoss(nn.Module):
    def __init__(self, temperature=0.07):
        super().__init__()
        self.tau = temperature

    def forward(self, z1, z2):
        # L2-normalise projections to unit sphere
        z1 = F.normalize(z1, dim=1)
        z2 = F.normalize(z2, dim=1)
        # All 2N representations as rows
        z  = torch.cat([z1, z2], dim=0)   # (2N, d)
        # Pairwise cosine similarities / temperature
        sim_matrix = z @ z.T / self.tau   # (2N, 2N)
        # Mask out self-similarities on diagonal
        n = z1.size(0)
        labels = torch.cat([torch.arange(n, 2*n), torch.arange(n)]).to(z.device)
        # Remove diagonal (self-similarity)
        mask = ~torch.eye(2*n, dtype=bool, device=z.device)
        sim_matrix = sim_matrix[mask].view(2*n, -1)
        return F.cross_entropy(sim_matrix, labels)

# After pretraining: linear evaluation
# Freeze backbone, train linear head on downstream task
backbone = resnet50_pretrained
for p in backbone.parameters(): p.requires_grad = False
linear_head = nn.Linear(2048, num_classes)
optimizer = torch.optim.Adam(linear_head.parameters(), lr=1e-3)