{
"cells": [
{
"cell_type": "markdown",
"id": "ad7d154b",
"metadata": {},
"source": "# Lesson 51: DINO and Self-Distillation\n\nLesson 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."
},
{
"cell_type": "code",
"id": "806c39e0",
"source": [
"import numpy as np\n",
"import cv2\n",
"import torch\n",
"import torch.nn as nn\n",
"import torch.nn.functional as F\n",
"import matplotlib.pyplot as plt"
],
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "4623af21",
"source": "## Student, teacher, and the collapse problem\n\nBoth student and teacher are copies of the same small encoder architecture — a ViT, exactly Lesson {{lesson:self_supervised_learning}}'s architecture (Lesson {{lesson:vision_transformers}}'s pattern: patchify, linearly embed, a `[CLS]` token gathers information via self-attention). 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.\n\nThe 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.",
"metadata": {}
},
{
"cell_type": "code",
"id": "dd49b70d",
"source": "def make_image(shape_type, cx, cy, size=16):\n img = np.zeros((size, size), dtype=np.float32)\n if shape_type == 'plus':\n img[cy-1:cy+2, cx-3:cx+4] = 1.0\n img[cy-3:cy+4, cx-1:cx+2] = 1.0\n else:\n yy, xx = np.mgrid[0:size, 0:size]\n img[((xx-cx)**2 + (yy-cy)**2) <= 9] = 1.0\n return img\n\ndef make_dataset(rng_local, n, position_range=(4, 12)):\n imgs, labels = [], []\n for _ in range(n):\n shape_type = rng_local.choice(['plus', 'circle'])\n cx, cy = rng_local.integers(*position_range), rng_local.integers(*position_range)\n imgs.append(make_image(shape_type, cx, cy))\n labels.append(0 if shape_type == 'plus' else 1)\n return np.array(imgs, dtype=np.float32), np.array(labels, dtype=np.int64)\n\ndef augment(img, rng_local):\n if rng_local.random() < 0.5:\n img = np.fliplr(img).copy()\n angle = rng_local.uniform(-20, 20)\n M = cv2.getRotationMatrix2D((8, 8), angle, 1.0)\n img = cv2.warpAffine(img, M, (16, 16))\n return np.clip(img + rng_local.normal(0, 0.1, img.shape), 0, 1).astype(np.float32)\n\nrng = np.random.default_rng(5)\nX_unlabeled, y_unlabeled = make_dataset(rng, 500)\n\ndef patchify(img, patch_size):\n B, C, H, W = img.shape\n P = patch_size\n patches = img.unfold(2, P, P).unfold(3, P, P) # (B, C, H/P, W/P, P, P)\n patches = patches.contiguous().view(B, C, -1, P, P).permute(0, 2, 1, 3, 4)\n return patches.reshape(B, -1, C * P * P) # (B, num_patches, C*P*P)\n\ndef positional_encoding(T, D):\n pos = torch.arange(T).unsqueeze(1).float()\n i = torch.arange(D).unsqueeze(0).float()\n angle_rates = 1.0 / (10000 ** (2 * (i // 2) / D))\n angles = pos * angle_rates\n pe = torch.zeros(T, D)\n pe[:, 0::2] = torch.sin(angles[:, 0::2])\n pe[:, 1::2] = torch.cos(angles[:, 1::2])\n return pe\n\nclass Encoder(nn.Module):\n def __init__(self, out_dim=16, img_size=16, patch_size=4, embed_dim=32, n_heads=4, n_layers=2):\n super().__init__()\n self.patch_size = patch_size\n n_patches = (img_size // patch_size) ** 2\n patch_dim = patch_size * patch_size\n self.embed = nn.Linear(patch_dim, embed_dim)\n self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim) * 0.02)\n self.register_buffer('pos_enc', positional_encoding(n_patches + 1, embed_dim))\n layer = nn.TransformerEncoderLayer(embed_dim, n_heads, dim_feedforward=embed_dim * 2,\n batch_first=True, dropout=0.0)\n self.encoder = nn.TransformerEncoder(layer, num_layers=n_layers)\n self.proj = nn.Linear(embed_dim, out_dim)\n\n def forward(self, x):\n p = patchify(x, self.patch_size)\n tok = self.embed(p)\n cls = self.cls_token.expand(x.shape[0], -1, -1)\n tok = torch.cat([cls, tok], dim=1) + self.pos_enc\n out = self.encoder(tok)\n return self.proj(out[:, 0]) # project the [CLS] token's final state to out_dim\n\ndef dino_loss(student_out, teacher_out, center, student_temp=0.1, teacher_temp=0.04):\n student_logp = F.log_softmax(student_out / student_temp, dim=-1)\n teacher_p = F.softmax((teacher_out - center) / teacher_temp, dim=-1) # centering happens here\n return -(teacher_p.detach() * student_logp).sum(dim=-1).mean()",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "0d866242",
"source": "`dino_loss` has two anti-collapse mechanisms built in:\n- **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.\n- **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.\n\nBoth tricks together are what let DINO train stably without any negative pairs at all.",
"metadata": {}
},
{
"cell_type": "code",
"id": "3b887ca5",
"source": "def train_dino(seed, epochs=200, lr=0.005, batch_size=64, momentum=0.9, center_momentum=0.5, use_centering=True):\n torch.manual_seed(seed)\n student = Encoder()\n teacher = Encoder()\n teacher.load_state_dict(student.state_dict())\n for p in teacher.parameters():\n p.requires_grad_(False)\n opt = torch.optim.Adam(student.parameters(), lr=lr)\n center = torch.zeros(1, 16)\n local_aug_rng = np.random.default_rng(seed + 100)\n n = len(X_unlabeled)\n for epoch in range(epochs):\n idx = np.random.default_rng(epoch).permutation(n)[:batch_size]\n batch = X_unlabeled[idx]\n view1 = np.stack([augment(im, local_aug_rng) for im in batch])\n view2 = np.stack([augment(im, local_aug_rng) for im in batch])\n v1 = torch.tensor(view1).unsqueeze(1)\n v2 = torch.tensor(view2).unsqueeze(1)\n\n s1, s2 = student(v1), student(v2)\n with torch.no_grad():\n t1, t2 = teacher(v1), teacher(v2)\n\n c = center if use_centering else torch.zeros_like(center)\n loss = dino_loss(s1, t2, c) / 2 + dino_loss(s2, t1, c) / 2\n\n opt.zero_grad()\n loss.backward()\n opt.step()\n\n with torch.no_grad():\n for ps, pt in zip(student.parameters(), teacher.parameters()):\n pt.data.mul_(momentum).add_(ps.data, alpha=1 - momentum) # EMA teacher update\n if use_centering:\n batch_center = torch.cat([t1, t2], dim=0).mean(dim=0, keepdim=True)\n center = center_momentum * center + (1 - center_momentum) * batch_center\n\n with torch.no_grad():\n output_std = torch.cat([t1, t2], dim=0).std(dim=0).mean().item()\n return student, teacher, output_std\n\n_, teacher_centered, std_centered = train_dino(seed=0, use_centering=True)\n_, teacher_no_center, std_no_center = train_dino(seed=0, use_centering=False)\n\nprint(f'teacher output std, WITH centering: {std_centered:.4f}')\nprint(f'teacher output std, WITHOUT centering: {std_no_center:.4f} (closer to 0 = more collapsed)')",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "8fb753a1",
"source": "Centering keeps the teacher's output spread several times larger than 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 {{lesson:self_supervised_learning}} did: freeze the teacher and train a linear probe on a handful of labeled examples.",
"metadata": {}
},
{
"cell_type": "code",
"id": "c8c2149d",
"source": "def make_noisy_dataset(rng_local, n, noise=0.35, position_range=(4, 12)):\n imgs, labels = [], []\n for _ in range(n):\n shape_type = rng_local.choice(['plus', 'circle'])\n cx, cy = rng_local.integers(*position_range), rng_local.integers(*position_range)\n img = make_image(shape_type, cx, cy)\n img = np.clip(img + rng_local.normal(0, noise, img.shape), 0, 1).astype(np.float32)\n imgs.append(img)\n labels.append(0 if shape_type == 'plus' else 1)\n return np.array(imgs, dtype=np.float32), np.array(labels, dtype=np.int64)\n\nX_probe_train, y_probe_train = make_noisy_dataset(rng, 8)\nX_probe_test, y_probe_test = make_noisy_dataset(rng, 150)\n\ndef linear_probe_acc(enc, seed):\n torch.manual_seed(seed)\n with torch.no_grad():\n feat_train = enc(torch.tensor(X_probe_train).unsqueeze(1))\n feat_test = enc(torch.tensor(X_probe_test).unsqueeze(1))\n probe = nn.Linear(feat_train.shape[1], 2)\n opt = torch.optim.Adam(probe.parameters(), lr=0.05)\n ytr = torch.tensor(y_probe_train)\n for _ in range(300):\n opt.zero_grad()\n loss = F.cross_entropy(probe(feat_train), ytr)\n loss.backward()\n opt.step()\n with torch.no_grad():\n preds = probe(feat_test).argmax(1).numpy()\n return (preds == y_probe_test).mean()\n\nprobe_accs = []\nfor seed in range(3):\n _, teacher, _ = train_dino(seed=seed)\n probe_accs.append(linear_probe_acc(teacher, seed=seed + 50))\n\nprint(f'linear probe on DINO-style teacher features: {np.mean(probe_accs):.1%} (+/- {np.std(probe_accs):.1%})')\nprint(f'(Lesson {{lesson:self_supervised_learning}} contrastive pretraining reached ~56% under the same probe setup)')",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "tsne_intro49",
"metadata": {},
"source": [
"## Visualizing the representation, directly\n",
"\n",
"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 {{lesson:visualizing_cnns}}) 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.\n"
]
},
{
"cell_type": "code",
"id": "tsne_algo49",
"metadata": {},
"execution_count": null,
"outputs": [],
"source": [
"def tsne(X, n_iter=500, perplexity=15.0, lr=100.0, seed=0):\n",
" rng = np.random.default_rng(seed)\n",
" n = X.shape[0]\n",
" sq_dists = ((X[:, None, :] - X[None, :, :]) ** 2).sum(-1)\n",
"\n",
" # high-dimensional affinities P: for each point, binary-search a Gaussian bandwidth so its\n",
" # neighbor distribution has the target perplexity (an implicit \"how many neighbors matter\" knob)\n",
" target_entropy = np.log(perplexity)\n",
" P = np.zeros((n, n))\n",
" for i in range(n):\n",
" others = np.arange(n) != i\n",
" d_i = sq_dists[i, others]\n",
" lo, hi = 1e-4, 1e4\n",
" for _ in range(50):\n",
" beta = (lo + hi) / 2\n",
" p = np.exp(-d_i * beta)\n",
" p /= p.sum() + 1e-12\n",
" entropy = -np.sum(p * np.log(p + 1e-12))\n",
" if entropy > target_entropy:\n",
" lo = beta # entropy too high (too spread out) -> need a larger beta to sharpen it\n",
" else:\n",
" hi = beta\n",
" P[i, others] = p\n",
" P = (P + P.T) / (2 * n) # symmetrize into one joint distribution over pairs\n",
" P = np.maximum(P, 1e-12) * 4.0 # early exaggeration: temporarily inflate P so true neighbors clump together faster\n",
"\n",
" # low-dimensional embedding Y, fit by gradient descent on KL(P || Q)\n",
" Y = rng.normal(0, 1e-2, (n, 2))\n",
" velocity = np.zeros_like(Y)\n",
" for it in range(n_iter):\n",
" if it == 100:\n",
" P /= 4.0 # turn off early exaggeration once clusters have separated\n",
" d = ((Y[:, None, :] - Y[None, :, :]) ** 2).sum(-1)\n",
" 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\"\n",
" np.fill_diagonal(num, 0)\n",
" Q = np.maximum(num / num.sum(), 1e-12)\n",
" coeff = (P - Q) * num\n",
" grad = 4 * (coeff[:, :, None] * (Y[:, None, :] - Y[None, :, :])).sum(1)\n",
" momentum = 0.5 if it < 100 else 0.8\n",
" velocity = momentum * velocity - lr * grad\n",
" Y = Y + velocity\n",
" return Y"
]
},
{
"cell_type": "code",
"id": "tsne_usage49",
"metadata": {},
"execution_count": null,
"outputs": [],
"source": [
"torch.manual_seed(99) # a third baseline: never trained at all\n",
"untrained_encoder = Encoder()\n",
"\n",
"Xv = torch.tensor(X_probe_test).unsqueeze(1)\n",
"\n",
"def teacher_features(encoder):\n",
" with torch.no_grad():\n",
" return encoder(Xv).numpy()\n",
"\n",
"Y_untrained = tsne(teacher_features(untrained_encoder), seed=0)\n",
"Y_no_center = tsne(teacher_features(teacher_no_center), seed=0)\n",
"Y_centered = tsne(teacher_features(teacher_centered), seed=0)\n",
"\n",
"fig, axes = plt.subplots(1, 3, figsize=(12, 4))\n",
"titles = ['untrained', 'no centering (collapsed)', 'with centering']\n",
"for ax, Y, title in zip(axes, [Y_untrained, Y_no_center, Y_centered], titles):\n",
" ax.scatter(*Y[y_probe_test == 0].T, c='tab:blue', s=14, label='plus')\n",
" ax.scatter(*Y[y_probe_test == 1].T, c='tab:orange', s=14, label='circle')\n",
" ax.set_title(title, fontsize=9)\n",
" ax.legend(fontsize=7)\n",
"fig.suptitle('t-SNE of teacher embeddings on the noisy probe test set, colored by true shape', y=1.02)\n",
"plt.show()"
]
},
{
"cell_type": "markdown",
"id": "tsne_disc49",
"metadata": {},
"source": "None of the three panels shows the clean two-blob separation Lesson {{lesson:visualizing_cnns}} 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 57% for the untrained baseline, 59% for the collapsed no-centering teacher, and 49% for the properly-trained centered teacher — all close to the 50% chance floor, and the centered teacher is, if anything, slightly *worse* than the two baselines it's supposed to beat.\n\nThis 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 200 training epochs on 500 unlabeled toy images with a tiny ViT encoder — an architecture Lesson {{lesson:vision_transformers}} already showed needs far more data than a CNN to learn useful structure at all — there simply isn't enough training signal for centering's benefit to show up as *useful, noise-robust structure* in the embedding space. The result is an even starker version of the same lesson the linear-probe accuracy above already pointed to: avoiding collapse is necessary, but nowhere near sufficient, for learning something useful. 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.\n"
},
{
"cell_type": "markdown",
"id": "7d7aceef",
"source": "The collapse safeguard measurably works — the teacher's outputs stay spread out, not constant. But at this toy scale (a couple hundred training epochs, a few hundred unlabeled images), the linear probe lands at chance, behind Lesson {{lesson:self_supervised_learning}}'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.\n\n## What DINOv2 adds at real scale\n\n**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 {{lesson:vision_transformers}}) 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 {{lesson:segment_anything}}), and k-nearest-neighbor classification directly on frozen DINOv2 features rivals supervised training on several benchmarks, all without fine-tuning a single weight.\n\nThe 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 {{lesson:classic_architectures}} (LeNet to AlexNet) and Lesson {{lesson:self_supervised_learning}} (SimCLR at toy scale vs. at 1000+ GPU scale), told once more.\n\n### Exercises\n\n1. Increase `epochs` in `train_dino` from 200 to 800 (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 {{lesson:self_supervised_learning}}'s contrastive approach?\n2. 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?\n3. 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?",
"metadata": {}
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python",
"version": "3.10.0"
}
},
"nbformat": 4,
"nbformat_minor": 5
}