{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": "# Lesson 48: DINOv2 and Self-Distillation\n\nLesson 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."
},
{
"cell_type": "code",
"id": "806c39e0",
"source": "import numpy as np\nimport cv2\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F",
"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. 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\nclass Encoder(nn.Module):\n def __init__(self, out_dim=16):\n super().__init__()\n self.conv = nn.Sequential(\n nn.Conv2d(1, 16, 5, padding=2), nn.ReLU(), nn.MaxPool2d(2),\n nn.Conv2d(16, 32, 5, padding=2), nn.ReLU(), nn.AdaptiveMaxPool2d(1),\n )\n self.proj = nn.Linear(32, out_dim)\n\n def forward(self, x):\n return self.proj(self.conv(x).flatten(1))\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=400, 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_, _, std_centered = train_dino(seed=0, use_centering=True)\n_, _, 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 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.",
"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 47 contrastive pretraining reached ~73% under the same probe setup)')",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"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 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.\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 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.\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 35 (LeNet to AlexNet) and Lesson 47 (SimCLR at toy scale vs. at 1000+ GPU scale), told once more.\n\n### Exercise\n\n1. 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?\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
}