Lesson 50: DINOv2 and Self-Distillation¶

Lesson 49'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
import matplotlib.pyplot as plt

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

_, teacher_centered, std_centered = train_dino(seed=0, use_centering=True)
_, teacher_no_center, 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 49 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 49 contrastive pretraining reached ~73% under the same probe setup)')

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

Visualizing the representation, directly¶

The output-std numbers above measure collapse on the clean, easy synthetic images used for training. But the linear probe just above was evaluated on noisier, harder images, and landed close to chance. t-SNE (van der Maaten & Hinton, 2008★, introduced in Lesson 39) can show directly whether that weak probe result is a property of the representation itself, not just the small 8-example probe set: project each noisy test image's teacher embedding down to 2D for three encoders — an untrained baseline, the collapsed (no-centering) teacher, and the properly-trained (centered) teacher — and look at whether plus and circle images end up in separate regions. True labels (y_probe_test) are used only to color the plot, never to train anything.

In [5]:

def tsne(X, n_iter=500, perplexity=15.0, lr=100.0, seed=0):
    rng = np.random.default_rng(seed)
    n = X.shape[0]
    sq_dists = ((X[:, None, :] - X[None, :, :]) ** 2).sum(-1)

    # high-dimensional affinities P: for each point, binary-search a Gaussian bandwidth so its
    # neighbor distribution has the target perplexity (an implicit "how many neighbors matter" knob)
    target_entropy = np.log(perplexity)
    P = np.zeros((n, n))
    for i in range(n):
        others = np.arange(n) != i
        d_i = sq_dists[i, others]
        lo, hi = 1e-4, 1e4
        for _ in range(50):
            beta = (lo + hi) / 2
            p = np.exp(-d_i * beta)
            p /= p.sum() + 1e-12
            entropy = -np.sum(p * np.log(p + 1e-12))
            if entropy > target_entropy:
                lo = beta  # entropy too high (too spread out) -> need a larger beta to sharpen it
            else:
                hi = beta
        P[i, others] = p
    P = (P + P.T) / (2 * n)          # symmetrize into one joint distribution over pairs
    P = np.maximum(P, 1e-12) * 4.0   # early exaggeration: temporarily inflate P so true neighbors clump together faster

    # low-dimensional embedding Y, fit by gradient descent on KL(P || Q)
    Y = rng.normal(0, 1e-2, (n, 2))
    velocity = np.zeros_like(Y)
    for it in range(n_iter):
        if it == 100:
            P /= 4.0  # turn off early exaggeration once clusters have separated
        d = ((Y[:, None, :] - Y[None, :, :]) ** 2).sum(-1)
        num = 1.0 / (1.0 + d)  # Student-t kernel -- heavy tails let moderately-distant points repel more strongly than a Gaussian would, avoiding the "crowding problem"
        np.fill_diagonal(num, 0)
        Q = np.maximum(num / num.sum(), 1e-12)
        coeff = (P - Q) * num
        grad = 4 * (coeff[:, :, None] * (Y[:, None, :] - Y[None, :, :])).sum(1)
        momentum = 0.5 if it < 100 else 0.8
        velocity = momentum * velocity - lr * grad
        Y = Y + velocity
    return Y

In [6]:

torch.manual_seed(99)  # a third baseline: never trained at all
untrained_encoder = Encoder()

Xv = torch.tensor(X_probe_test).unsqueeze(1)

def teacher_features(encoder):
    with torch.no_grad():
        return encoder(Xv).numpy()

Y_untrained = tsne(teacher_features(untrained_encoder), seed=0)
Y_no_center = tsne(teacher_features(teacher_no_center), seed=0)
Y_centered = tsne(teacher_features(teacher_centered), seed=0)

fig, axes = plt.subplots(1, 3, figsize=(12, 4))
titles = ['untrained', 'no centering (collapsed)', 'with centering']
for ax, Y, title in zip(axes, [Y_untrained, Y_no_center, Y_centered], titles):
    ax.scatter(*Y[y_probe_test == 0].T, c='tab:blue', s=14, label='plus')
    ax.scatter(*Y[y_probe_test == 1].T, c='tab:orange', s=14, label='circle')
    ax.set_title(title, fontsize=9)
    ax.legend(fontsize=7)
fig.suptitle('t-SNE of teacher embeddings on the noisy probe test set, colored by true shape', y=1.02)
plt.show()

No description has been provided for this image

All three panels look similarly unconvincing — none shows the clean two-blob separation Lesson 39 found for a supervised CNN's features. A quick 1-nearest-neighbor check in the raw 16-d feature space (does each point's nearest neighbor share its true label?) confirms what the plots suggest: about 59% for the centered teacher, 63% for the collapsed no-centering teacher, and 62% for the untrained baseline — all close to the 50% chance floor, with no meaningful gap between "collapsed" and "properly trained."

This is worth sitting with rather than explaining away: the output-std metric earlier in this lesson is real and correctly shows centering prevents collapse on the training distribution, but that's a necessary condition, not a sufficient one, and this visualization is direct evidence of the gap between them. At only 400 training epochs on 500 unlabeled toy images with a tiny encoder, there simply isn't enough training signal for centering's benefit to show up as useful, noise-robust structure in the embedding space — exactly the same limitation the linear-probe accuracy above already pointed to, now visible directly in the feature space itself rather than through a downstream accuracy number. It's a concrete illustration of why DINOv2 needs the scale described below: collapse-avoidance is table stakes, and the actual representation quality only emerges with far more data, capacity, and training than a toy notebook can offer.

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 49'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 47) 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 54), 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 37 (LeNet to AlexNet) and Lesson 49 (SimCLR at toy scale vs. at 1000+ GPU scale), told once more.

Exercises¶

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 49'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 49: Self-Supervised and Contrastive Learning

Lesson 51: Masked Autoencoders →