Contrastive Learning: SimCLR, MoCo, and Theory

Module: Machine Learning | Difficulty: Advanced

InfoNCE Loss

Mutual Information Lower Bound

SimCLR Pipeline

1. Sample batch of images
2. Apply augmentations (random crop, color jitter, blur)
3. Encode with backbone
4. Project with MLP head
5. Minimize InfoNCE loss

MoCo (Momentum Contrast)

Use momentum encoder for negative keys:

| Method | Batch Size | Memory | Performance | |--------|-----------|--------|-------------| | SimCLR | 4096 | High | 76.5% | | MoCo-v2 | 256 | Low | 76.8% | | BYOL | 4096 | High | 79.6% | | VICReg | 2048 | Medium | 78.2% |

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimCLR(nn.Module):
    def __init__(self, backbone, projection_dim=128):
        super().__init__()
        self.backbone = backbone
        feat_dim = backbone.fc.in_features
        backbone.fc = nn.Identity()
        self.projector = nn.Sequential(
            nn.Linear(feat_dim, feat_dim), nn.ReLU(),
            nn.Linear(feat_dim, projection_dim))
    def forward(self, x1, x2):
        h1 = self.backbone(x1); h2 = self.backbone(x2)
        z1 = self.projector(h1); z2 = self.projector(h2)
        return z1, z2
    def info_nce_loss(self, z1, z2, temperature=0.5):
        z = F.normalize(torch.cat([z1, z2]), dim=1)
        sim = z @ z.T / temperature
        n = len(z1)
        labels = torch.cat([torch.arange(n, 2*n), torch.arange(n)]).to(sim.device)
        return F.cross_entropy(sim, labels)

Research Insight: Contrastive learning's success depends on the quality of augmentations. The key insight is that good augmentations define what is considered "similar" — they encode invariances of the task.