Contrastive Learning, SimCLR, and MoCo for Visual Representations

Module: Computer Vision | Difficulty: Premium

InfoNCE Loss (SimCLR)

Momentum Encoder (MoCo)

where typically.

BYOL Loss (No Negatives)

Evaluation: Linear Probe

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(inplace=True),
            nn.Linear(feat_dim, projection_dim),
        )
        self.temperature = 0.07

    def nt_xent_loss(self, z1, z2):
        B = z1.shape[0]
        z = torch.cat([z1, z2], dim=0)
        z = F.normalize(z, dim=1)
        sim = z @ z.T / self.temperature
        sim.fill_diagonal_(-1e9)
        labels = torch.cat([
            torch.arange(B, 2*B), torch.arange(0, B)]).to(z.device)
        return F.cross_entropy(sim, labels)

    def forward(self, x1, x2):
        z1 = self.projector(self.backbone(x1))
        z2 = self.projector(self.backbone(x2))
        return self.nt_xent_loss(z1, z2)

class DINOHead(nn.Module):
    def __init__(self, in_dim, out_dim=65536,
                 hidden_dim=2048, bottleneck_dim=256):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(in_dim, hidden_dim), nn.GELU(),
            nn.Linear(hidden_dim, hidden_dim), nn.GELU(),
            nn.Linear(hidden_dim, bottleneck_dim),
        )
        self.last_layer = nn.utils.weight_norm(
            nn.Linear(bottleneck_dim, out_dim, bias=False))
        self.last_layer.weight_g.data.fill_(1)

    def forward(self, x):
        x = self.mlp(x)
        x = F.normalize(x, dim=1)
        return self.last_layer(x)

Research Insight: Self-supervised vision has reached a tipping point where pretrained representations rival or exceed supervised ones. DINOv2 demonstrated that with careful curation of training data (142M images via automatic data mining), self-supervised ViT features produce high-quality dense features rivaling specialized segmentation models. The key open question is whether these methods can scale to the same data regimes as language models (trillions of tokens vs. hundreds of millions of images), which would require efficient web-scale visual data curation.