Lesson 48: DINOv2 and Self-Distillation¶

Lesson 47's contrastive loss needs explicit negative pairs (other images in the batch) to avoid the trivial solution of mapping every image to the same point. Self-distillation, the mechanism behind DINO (Caron et al., 2021) and its successor DINOv2 (Oquab et al., 2023), removes negatives entirely: a slowly-updated "teacher" network guides a "student" network, with no labels and no negative pairs at all. Without a careful safeguard, this setup collapses to the trivial solution immediately — this lesson builds the safeguard (centering) from scratch and shows exactly why it's needed.

In [1]:

import numpy as np
import cv2
import torch
import torch.nn as nn
import torch.nn.functional as F

Student, teacher, and the collapse problem¶

Both student and teacher are copies of the same small encoder architecture. The student is trained normally, with gradients. The teacher is never trained directly — after each step, its weights are nudged a small amount toward the student's current weights (an exponential moving average, or EMA). The student is trained to match the teacher's output distribution on a different augmented view of the same image.

The obvious failure mode: if the teacher's output doesn't depend on the input at all (always predicts the same constant vector, regardless of image), the student can trivially match it by doing the same — perfect loss, zero information learned. This is representation collapse, and it's the central problem self-distillation has to solve.

In [2]:

def make_image(shape_type, cx, cy, size=16):
    img = np.zeros((size, size), dtype=np.float32)
    if shape_type == 'plus':
        img[cy-1:cy+2, cx-3:cx+4] = 1.0
        img[cy-3:cy+4, cx-1:cx+2] = 1.0
    else:
        yy, xx = np.mgrid[0:size, 0:size]
        img[((xx-cx)**2 + (yy-cy)**2) <= 9] = 1.0
    return img

def make_dataset(rng_local, n, position_range=(4, 12)):
    imgs, labels = [], []
    for _ in range(n):
        shape_type = rng_local.choice(['plus', 'circle'])
        cx, cy = rng_local.integers(*position_range), rng_local.integers(*position_range)
        imgs.append(make_image(shape_type, cx, cy))
        labels.append(0 if shape_type == 'plus' else 1)
    return np.array(imgs, dtype=np.float32), np.array(labels, dtype=np.int64)

def augment(img, rng_local):
    if rng_local.random() < 0.5:
        img = np.fliplr(img).copy()
    angle = rng_local.uniform(-20, 20)
    M = cv2.getRotationMatrix2D((8, 8), angle, 1.0)
    img = cv2.warpAffine(img, M, (16, 16))
    return np.clip(img + rng_local.normal(0, 0.1, img.shape), 0, 1).astype(np.float32)

rng = np.random.default_rng(5)
X_unlabeled, y_unlabeled = make_dataset(rng, 500)

class Encoder(nn.Module):
    def __init__(self, out_dim=16):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 16, 5, padding=2), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 5, padding=2), nn.ReLU(), nn.AdaptiveMaxPool2d(1),
        )
        self.proj = nn.Linear(32, out_dim)

    def forward(self, x):
        return self.proj(self.conv(x).flatten(1))

def dino_loss(student_out, teacher_out, center, student_temp=0.1, teacher_temp=0.04):
    student_logp = F.log_softmax(student_out / student_temp, dim=-1)
    teacher_p = F.softmax((teacher_out - center) / teacher_temp, dim=-1)  # centering happens here
    return -(teacher_p.detach() * student_logp).sum(dim=-1).mean()

dino_loss has two anti-collapse mechanisms built in:

Centering: subtract a running average (center) of recent teacher outputs before the teacher's softmax. This stops the teacher from drifting toward always predicting whatever single class happens to be easiest — a constant output would get centered to exactly zero, canceling itself out.
Sharpening: the teacher's softmax uses a lower temperature (0.04) than the student's (0.1), making the teacher's target distribution more confident/peaked. A very flat, unconfident teacher target is close to a uniform distribution — not quite collapse, but not a useful learning signal either.

Both tricks together are what let DINO train stably without any negative pairs at all.

In [3]:

def train_dino(seed, epochs=400, lr=0.005, batch_size=64, momentum=0.9, center_momentum=0.5, use_centering=True):
    torch.manual_seed(seed)
    student = Encoder()
    teacher = Encoder()
    teacher.load_state_dict(student.state_dict())
    for p in teacher.parameters():
        p.requires_grad_(False)
    opt = torch.optim.Adam(student.parameters(), lr=lr)
    center = torch.zeros(1, 16)
    local_aug_rng = np.random.default_rng(seed + 100)
    n = len(X_unlabeled)
    for epoch in range(epochs):
        idx = np.random.default_rng(epoch).permutation(n)[:batch_size]
        batch = X_unlabeled[idx]
        view1 = np.stack([augment(im, local_aug_rng) for im in batch])
        view2 = np.stack([augment(im, local_aug_rng) for im in batch])
        v1 = torch.tensor(view1).unsqueeze(1)
        v2 = torch.tensor(view2).unsqueeze(1)

        s1, s2 = student(v1), student(v2)
        with torch.no_grad():
            t1, t2 = teacher(v1), teacher(v2)

        c = center if use_centering else torch.zeros_like(center)
        loss = dino_loss(s1, t2, c) / 2 + dino_loss(s2, t1, c) / 2

        opt.zero_grad()
        loss.backward()
        opt.step()

        with torch.no_grad():
            for ps, pt in zip(student.parameters(), teacher.parameters()):
                pt.data.mul_(momentum).add_(ps.data, alpha=1 - momentum)  # EMA teacher update
            if use_centering:
                batch_center = torch.cat([t1, t2], dim=0).mean(dim=0, keepdim=True)
                center = center_momentum * center + (1 - center_momentum) * batch_center

    with torch.no_grad():
        output_std = torch.cat([t1, t2], dim=0).std(dim=0).mean().item()
    return student, teacher, output_std

_, _, std_centered = train_dino(seed=0, use_centering=True)
_, _, std_no_center = train_dino(seed=0, use_centering=False)

print(f'teacher output std, WITH centering:    {std_centered:.4f}')
print(f'teacher output std, WITHOUT centering: {std_no_center:.4f}  (closer to 0 = more collapsed)')

teacher output std, WITH centering:    0.2595
teacher output std, WITHOUT centering: 0.0773  (closer to 0 = more collapsed)

Centering roughly triples the teacher's output spread compared to training without it — the mechanism visibly does what it's supposed to. Now check whether that translates into a useful representation, the same way Lesson 47 did: freeze the teacher and train a linear probe on a handful of labeled examples.

In [4]:

def make_noisy_dataset(rng_local, n, noise=0.35, position_range=(4, 12)):
    imgs, labels = [], []
    for _ in range(n):
        shape_type = rng_local.choice(['plus', 'circle'])
        cx, cy = rng_local.integers(*position_range), rng_local.integers(*position_range)
        img = make_image(shape_type, cx, cy)
        img = np.clip(img + rng_local.normal(0, noise, img.shape), 0, 1).astype(np.float32)
        imgs.append(img)
        labels.append(0 if shape_type == 'plus' else 1)
    return np.array(imgs, dtype=np.float32), np.array(labels, dtype=np.int64)

X_probe_train, y_probe_train = make_noisy_dataset(rng, 8)
X_probe_test, y_probe_test = make_noisy_dataset(rng, 150)

def linear_probe_acc(enc, seed):
    torch.manual_seed(seed)
    with torch.no_grad():
        feat_train = enc(torch.tensor(X_probe_train).unsqueeze(1))
        feat_test = enc(torch.tensor(X_probe_test).unsqueeze(1))
    probe = nn.Linear(feat_train.shape[1], 2)
    opt = torch.optim.Adam(probe.parameters(), lr=0.05)
    ytr = torch.tensor(y_probe_train)
    for _ in range(300):
        opt.zero_grad()
        loss = F.cross_entropy(probe(feat_train), ytr)
        loss.backward()
        opt.step()
    with torch.no_grad():
        preds = probe(feat_test).argmax(1).numpy()
    return (preds == y_probe_test).mean()

probe_accs = []
for seed in range(3):
    _, teacher, _ = train_dino(seed=seed)
    probe_accs.append(linear_probe_acc(teacher, seed=seed + 50))

print(f'linear probe on DINO-style teacher features: {np.mean(probe_accs):.1%}  (+/- {np.std(probe_accs):.1%})')
print(f'(Lesson 47 contrastive pretraining reached ~73% under the same probe setup)')

linear probe on DINO-style teacher features: 50.7%  (+/- 4.3%)
(Lesson 47 contrastive pretraining reached ~73% under the same probe setup)

The collapse safeguard measurably works — the teacher's outputs stay spread out, not constant. But at this toy scale (a few hundred training epochs, a few hundred unlabeled images), the linear probe lands close to chance, well behind Lesson 47's contrastive result on the identical task. This is an honest result, not a bug: self-distillation is known to be more sensitive to hyperparameters (EMA momentum, temperature schedule, centering rate) and generally needs substantially more training signal to reach a useful representation than a contrastive loss with explicit negatives does. Avoiding collapse is necessary but not sufficient for learning something useful — it just clears the way for enough training to eventually do so.

What DINOv2 adds at real scale¶

DINOv2 (Meta AI, 2023) is this exact mechanism — student/teacher self-distillation with centering, plus a few refinements (multiple small "local crops" alongside full-image "global crops", to make the task harder and richer) — scaled up to a curated 142-million-image unlabeled dataset and a Vision Transformer (Lesson 46) backbone with up to 1.1 billion parameters. At that scale, the representation isn't just "usable with a linear probe" — it exhibits striking emergent properties nobody explicitly trained for: attention maps from a DINOv2 ViT often outline object boundaries and parts without ever seeing a segmentation label (a direct preview of Lesson 51), and k-nearest-neighbor classification directly on frozen DINOv2 features rivals supervised training on several benchmarks, all without fine-tuning a single weight.

The throughline from this lesson to DINOv2 is exactly the gap this notebook exposed: the mechanism (student, teacher, centering) is the same code at any scale; what changes between this toy version and a real foundation model is data volume, model capacity, and training duration — the same story as Lesson 35 (LeNet to AlexNet) and Lesson 47 (SimCLR at toy scale vs. at 1000+ GPU scale), told once more.

Exercise¶

Increase epochs in train_dino from 400 to 1200 (this will take longer to run). Does the linear probe accuracy improve noticeably, stay flat, or become unstable — and how does that compare to what more training epochs did for Lesson 47's contrastive approach?
Set teacher_temp=0.1 (matching the student's temperature exactly, removing the sharpening asymmetry) in dino_loss. Does the collapse comparison (with vs. without centering) still show a clear gap, or does removing sharpening make collapse happen even with centering turned on?
Try momentum=0.5 (a much faster-updating teacher) instead of 0.9. A teacher that updates almost as fast as the student loses its main purpose — providing a stable, slowly-changing target. Does training become less stable, and can you see it in the collapse metric?

← Lesson 47: Self-Supervised and Contrastive Learning

Lesson 49: Vision-Language Models →