DINO and self-distillation

pip install "torch==2.5.1"

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

K = 8                      # prototypes: slots the head can point at
TAU_S = 0.1                # student temperature
EMA, CENTER_M = 0.99, 0.9  # teacher momentum, centre momentum


def make_batch(n, g):
    """Three textures, and the labels are only ever used to score the result."""
    y = torch.randint(0, 3, (n,), generator=g)
    phase = torch.randint(0, 4, (n, 2), generator=g)
    rows = torch.arange(16).view(16, 1).expand(16, 16)
    cols = torch.arange(16).view(1, 16).expand(16, 16)
    x = torch.zeros(n, 1, 16, 16)
    for i, k in enumerate(y.tolist()):
        a, b = phase[i].tolist()
        if k == 0:
            x[i, 0] = (((rows + a) // 2) % 2).float()
        elif k == 1:
            x[i, 0] = (((cols + b) // 2) % 2).float()
        else:
            x[i, 0] = ((((rows + a) // 2) + ((cols + b) // 2)) % 2).float()
    return x, y


def augment(x, g):
    n = x.shape[0]
    out = torch.empty_like(x)
    shifts = torch.randint(0, 16, (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.6 + 0.6 * torch.rand(n, 1, 1, 1, generator=g))
    return out + 0.1 * torch.randn(out.shape, generator=g)


class Net(nn.Module):
    """Small convolutional trunk, then cosine similarity against K prototypes."""
    def __init__(self):
        super().__init__()
        self.trunk = nn.Sequential(
            nn.Conv2d(1, 16, 3, 2, 1), nn.ReLU(),
            nn.Conv2d(16, 32, 3, 2, 1), nn.ReLU(),
            nn.AdaptiveAvgPool2d(1), nn.Flatten(),
            nn.Linear(32, 64), nn.GELU(), nn.Linear(64, 32))
        self.prototypes = nn.Parameter(torch.randn(K, 32))

    def forward(self, x):
        z = F.normalize(self.trunk(x), dim=1)
        return z @ F.normalize(self.prototypes, dim=1).t()


def run(use_centering, tau_t, steps=600):
    torch.manual_seed(0)
    g = torch.Generator().manual_seed(0)
    student = Net()
    teacher = copy.deepcopy(student)
    for p in teacher.parameters():
        p.requires_grad_(False)
    centre = torch.zeros(K)
    opt = torch.optim.Adam(student.parameters(), lr=3e-3)

    for _ in range(steps):
        x, _ = make_batch(64, g)
        v1, v2 = augment(x, g), augment(x, g)
        with torch.no_grad():
            t1, t2 = teacher(v1), teacher(v2)
            c = centre if use_centering else torch.zeros(K)
            p1 = F.softmax((t1 - c) / tau_t, dim=1)      # centre, then sharpen
            p2 = F.softmax((t2 - c) / tau_t, dim=1)
        s1, s2 = student(v1), student(v2)
        loss = -0.5 * ((p2 * F.log_softmax(s1 / TAU_S, 1)).sum(1).mean()
                       + (p1 * F.log_softmax(s2 / TAU_S, 1)).sum(1).mean())
        opt.zero_grad()
        loss.backward()
        opt.step()
        with torch.no_grad():
            centre = CENTER_M * centre + (1 - CENTER_M) * torch.cat([t1, t2]).mean(0)
            for tp, sp in zip(teacher.parameters(), student.parameters()):
                tp.mul_(EMA).add_(sp.detach(), alpha=1 - EMA)   # teacher trails the student

    x, y = make_batch(600, torch.Generator().manual_seed(123))
    with torch.no_grad():
        p = F.softmax((teacher(x) - (centre if use_centering else 0)) / tau_t, dim=1)
    avg = p.mean(0)
    spread = -(avg * (avg + 1e-12).log()).sum().item()
    per_image = -(p * (p + 1e-12).log()).sum(1).mean().item()
    used = int((avg > 0.01).sum())
    assign = p.argmax(1)
    purity = sum((y[assign == k].mode().values == y[assign == k]).float().sum().item()
                 for k in range(K) if (assign == k).any()) / len(y)
    return spread, per_image, used, purity


print(f"spread entropy: {math.log(K):.3f} means all {K} prototypes shared equally, "
      f"{math.log(3):.3f} means 3 of them, 0 means total collapse")
print(f"{'setting':34s} {'spread':>7s} {'per-image':>10s} {'used':>5s} {'purity':>7s}")
for centering, tau_t, name in [
        (False, 0.02, "sharpening only, no centring"),
        (True, 0.40, "centring, almost no sharpening"),
        (True, 0.02, "centring + sharpening"),
        (True, 0.005, "centring + far too much sharpening")]:
    spread, per_image, used, purity = run(centering, tau_t)
    print(f"{name:34s} {spread:7.3f} {per_image:10.3f} {used:5d} {purity:7.3f}")

import torch
import torch.nn.functional as F
from transformers import AutoImageProcessor, AutoModel, AutoConfig

torch.manual_seed(0)


def shape_image(kind, dx=0, dy=0, bright=1.0):
    """A 224x224 RGB picture holding one plain shape on a grey ground."""
    img = torch.full((3, 224, 224), 0.45)
    ys = torch.arange(224).view(224, 1) - 112 - dy
    xs = torch.arange(224).view(1, 224) - 112 - dx
    if kind == "disc":
        mask = (ys ** 2 + xs ** 2) < 70 ** 2
        colour = torch.tensor([0.85, 0.15, 0.15])
    elif kind == "square":
        mask = (ys.abs() < 60) & (xs.abs() < 60)
        colour = torch.tensor([0.15, 0.30, 0.85])
    else:  # a cross
        mask = ((ys.abs() < 20) & (xs.abs() < 75)) | ((ys.abs() < 75) & (xs.abs() < 20))
        colour = torch.tensor([0.15, 0.65, 0.25])
    img = torch.where(mask.unsqueeze(0), colour.view(3, 1, 1), img)
    return (img * bright).clamp(0, 1)


kinds = ["disc", "square", "cross"]
views, labels = [], []
for k in kinds:
    views.append(shape_image(k, 0, 0, 1.0))
    labels.append(k + " A")
    views.append(shape_image(k, 45, -35, 0.75))       # moved and dimmed
    labels.append(k + " B")
batch = torch.stack(views)

name = "facebook/dinov2-small"
processor = AutoImageProcessor.from_pretrained(name)
mean = torch.tensor(processor.image_mean).view(1, 3, 1, 1)
std = torch.tensor(processor.image_std).view(1, 3, 1, 1)
pixels = (batch - mean) / std

pretrained = AutoModel.from_pretrained(name).eval()
random_init = AutoModel.from_config(AutoConfig.from_pretrained(name)).eval()

print("model:", name)
print("parameters:", f"{sum(p.numel() for p in pretrained.parameters()):,}")
with torch.no_grad():
    out = pretrained(pixel_values=pixels)
print("last_hidden_state:", tuple(out.last_hidden_state.shape),
      "= 1 CLS token + 256 patch tokens of 384 numbers")


def sim_matrix(model):
    with torch.no_grad():
        cls = model(pixel_values=pixels).last_hidden_state[:, 0]
    z = F.normalize(cls, dim=1)
    return z @ z.t()


for title, model in [("pretrained DINOv2", pretrained), ("same network, random weights", random_init)]:
    s = sim_matrix(model)
    print(f"\ncosine similarity between whole-image embeddings, {title}")
    print("            " + "".join(f"{l:>10s}" for l in labels))
    for i, l in enumerate(labels):
        print(f"{l:>12s}" + "".join(f"{s[i, j].item():10.3f}" for j in range(6)))
    same = torch.tensor([s[0, 1], s[2, 3], s[4, 5]]).mean()
    idx = [(0, 2), (0, 4), (2, 4), (1, 3), (1, 5), (3, 5)]
    diff = torch.tensor([s[i, j] for i, j in idx]).mean()
    print(f"  two views of the same shape : {same.item():.3f}")
    print(f"  two different shapes        : {diff.item():.3f}")
    print(f"  gap                         : {same.item() - diff.item():+.3f}")