Contrastive learning for images

pip install "torch==2.5.1"

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


def nt_xent(z, temperature=0.5):
    """z holds 2N rows: rows 0..N-1 are view A, rows N..2N-1 are view B of the same images."""
    z = F.normalize(z, dim=1)                 # every embedding onto the unit circle
    n = z.shape[0] // 2
    sim = z @ z.t() / temperature             # cosine similarity of every pair
    sim.fill_diagonal_(-1e9)                  # an embedding must not match itself
    target = torch.cat([torch.arange(n, 2 * n), torch.arange(0, n)])
    return F.cross_entropy(sim, target)


# three images, two views each, laid out by hand in two dimensions
good = torch.tensor([
    [1.0,  0.1], [-0.1, 1.0], [-1.0, -0.1],    # view A of images 1, 2, 3
    [1.0, -0.1], [ 0.1, 1.0], [-1.0,  0.1],    # view B of the same three
])
scrambled = good[[3, 4, 5, 1, 2, 0]]           # same points, wrong pairings

print(f"views agree     : loss {nt_xent(good).item():.4f}")
print(f"views scrambled : loss {nt_xent(scrambled).item():.4f}")
print(f"pure chance     : loss {math.log(2 * 3 - 1):.4f}   (= ln(2N-1) with N=3)")

print("\ntemperature changes how harshly near-misses are punished")
for t in (0.05, 0.1, 0.5, 1.0, 10.0):
    print(f"  temperature {t:5.2f} -> loss {nt_xent(good, t).item():.4f}")

print("\nchance loss grows with batch size, so the task gets harder as N grows")
for n in (8, 64, 256, 4096):
    print(f"  N = {n:5d} -> ln(2N-1) = {math.log(2 * n - 1):.3f}")

print("\ncollapse scores well on positives and terribly on this loss")
collapsed = torch.ones(6, 2)
print(f"  every embedding identical -> loss {nt_xent(collapsed).item():.4f}")
print(f"  cosine of the positive pair -> {F.cosine_similarity(collapsed[:3], collapsed[3:]).mean():.4f}")

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


def make_batch(n, g):
    """Three shapes on a 16x16 canvas: a bar, a column and a diagonal."""
    y = torch.randint(0, 3, (n,), generator=g)
    x = torch.zeros(n, 1, 16, 16)
    for i, k in enumerate(y.tolist()):
        r, c = torch.randint(2, 10, (2,), generator=g).tolist()
        if k == 0:
            x[i, 0, r, c:c + 6] = 1.0
        elif k == 1:
            x[i, 0, r:r + 6, c] = 1.0
        else:
            for d in range(6):
                x[i, 0, r + d, c + d] = 1.0
    return x, y


def augment(x, g):
    """Two random views of the same picture: shift, brightness, noise."""
    n = x.shape[0]
    out = torch.empty_like(x)
    shifts = torch.randint(-3, 4, (n, 2), generator=g)
    for i in range(n):
        out[i] = torch.roll(x[i], shifts=tuple(shifts[i].tolist()), dims=(1, 2))
    out = out * (0.5 + torch.rand(n, 1, 1, 1, generator=g))
    return out + 0.15 * torch.randn(out.shape, generator=g)


class Encoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.body = nn.Sequential(
            nn.Conv2d(1, 16, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(16, 32, 3, stride=2, padding=1), nn.ReLU(),
            nn.AdaptiveAvgPool2d(1), nn.Flatten())
        self.head = nn.Sequential(nn.Linear(32, 32), nn.ReLU(), nn.Linear(32, 16))

    def forward(self, x):
        h = self.body(x)
        return h, self.head(h)         # features, and the projection the loss uses


def nt_xent(z, temperature=0.5):
    z = F.normalize(z, dim=1)
    n = z.shape[0] // 2
    sim = z @ z.t() / temperature
    sim.fill_diagonal_(-1e9)
    target = torch.cat([torch.arange(n, 2 * n), torch.arange(0, n)])
    return F.cross_entropy(sim, target)


def positive_only(z):
    """Pull the two views together and do nothing else. Watch what happens."""
    n = z.shape[0] // 2
    return -F.cosine_similarity(z[:n], z[n:]).mean()


def diagnose(enc):
    x, _ = make_batch(300, torch.Generator().manual_seed(99))
    with torch.no_grad():
        feats, proj = enc(x)
    p = F.normalize(proj, dim=1)
    off = p @ p.t()
    mean_cos = ((off.sum() - off.diag().sum()) / (300 * 299)).item()
    return feats.std(dim=0).mean().item(), mean_cos


for name, loss_fn in [("positive pairs only", positive_only), ("NT-Xent with negatives", nt_xent)]:
    torch.manual_seed(0)
    g = torch.Generator().manual_seed(0)
    enc = Encoder()
    s0, c0 = diagnose(enc)
    opt = torch.optim.Adam(enc.parameters(), lr=3e-3)
    for step in range(300):
        x, _ = make_batch(64, g)
        views = torch.cat([augment(x, g), augment(x, g)])
        _, z = enc(views)
        loss = loss_fn(z)
        opt.zero_grad()
        loss.backward()
        opt.step()
    s1, c1 = diagnose(enc)
    print(f"{name:24s} before: feature spread {s0:.4f}, cosine between different images {c0:+.4f}")
    print(f"{'':24s} after : feature spread {s1:.4f}, cosine between different images {c1:+.4f}"
          f"   final loss {loss.item():.3f}")