{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# DINOv2 Tutorial — Notebook 3: Training Procedure\n", "\n", "**Series:** DINOv2 Deep Dive | **Prev:** `02_self_supervised_learning_theory.ipynb` | **Next:** `04_dataset_preparation.ipynb`\n", "\n", "---\n", "\n", "## What You Will Learn\n", "- The complete DINOv2 training pipeline with all three loss components\n", "- DINO loss: CLS token self-distillation\n", "- iBOT loss: masked patch prediction\n", "- KoLeo regularizer: entropy maximization\n", "- Multi-crop augmentation strategy\n", "- Training hyperparameters and schedules\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3.1 DINOv2 Training Overview\n", "\n", "DINOv2 combines **three objectives** trained jointly:\n", "\n", "$$\\mathcal{L}_{\\text{total}} = \\mathcal{L}_{\\text{DINO}} + \\lambda_{\\text{iBOT}} \\cdot \\mathcal{L}_{\\text{iBOT}} + \\lambda_{\\text{KoLeo}} \\cdot \\mathcal{L}_{\\text{KoLeo}}$$\n", "\n", "| Loss | Target | What it learns |\n", "|------|--------|----------------|\n", "| **DINO** | CLS token | Global image representation |\n", "| **iBOT** | Masked patches | Dense/local patch features |\n", "| **KoLeo** | Batch embeddings | Spread representations uniformly |\n", "\n", "### Architecture\n", "```\n", "Image ──→ Multi-Crop Augmentation\n", " │\n", " ┌─────────┴──────────┐\n", " │ Global crops (224) │ Local crops (96)\n", " └────────┬───────────┘ │\n", " │ │\n", " ┌────────▼───────────┐ │\n", " │ Teacher ViT │ │\n", " │ (EMA, no grad) │ │\n", " └────────┬───────────┘ │\n", " │ │\n", " Teacher CLS + Patches ┌─────▼──────┐\n", " │ │ Student ViT │\n", " │ │ (trained) │\n", " │ └─────┬───────┘\n", " │ Student CLS + Patches\n", " │ │\n", " ┌────────▼─────────────────────▼────────┐\n", " │ DINO Loss + iBOT Loss + KoLeo Loss │\n", " └───────────────────────────────────────┘\n", "```" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "import matplotlib.pyplot as plt\n", "import matplotlib.patches as mpatches\n", "from PIL import Image, ImageDraw\n", "import torch\n", "import torch.nn.functional as F" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3.2 Multi-Crop Augmentation Strategy\n", "\n", "DINOv2 uses a **multi-crop** strategy:\n", "- **2 global crops**: 224×224 (large, overlapping views)\n", "- **N local crops**: 96×96 (small, focused views, default N=8)\n", "\n", "**Key rule:**\n", "- Teacher sees **only global crops** (high-quality, large context)\n", "- Student sees **all crops** (global + local)\n", "- Loss: student's output on local crops → match teacher's output on global crops\n", "\n", "This forces the student to learn **global context from local views** — a powerful self-supervised signal." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Visualize multi-crop strategy ─────────────────────────────────────────────\n", "def make_scene(size=448):\n", " img = Image.new('RGB', (size, size), (230, 230, 230))\n", " draw = ImageDraw.Draw(img)\n", " # Sky\n", " draw.rectangle([0, 0, size, size//2], fill=(135, 206, 235))\n", " # Ground\n", " draw.rectangle([0, size//2, size, size], fill=(100, 160, 80))\n", " # House\n", " draw.rectangle([100, 200, 300, 380], fill=(180, 120, 80))\n", " draw.polygon([(80, 200), (320, 200), (200, 100)], fill=(160, 60, 60))\n", " # Door\n", " draw.rectangle([175, 300, 225, 380], fill=(100, 60, 30))\n", " # Windows\n", " draw.rectangle([120, 230, 170, 280], fill=(200, 230, 255))\n", " draw.rectangle([230, 230, 280, 280], fill=(200, 230, 255))\n", " # Tree\n", " draw.rectangle([340, 280, 360, 380], fill=(100, 70, 40))\n", " draw.ellipse([300, 200, 400, 300], fill=(50, 140, 50))\n", " # Sun\n", " draw.ellipse([360, 30, 420, 90], fill=(255, 220, 50))\n", " return img\n", "\n", "scene = make_scene(448)\n", "\n", "# Define crops\n", "global_crop1 = (20, 20, 340, 340) # 320×320 → resize to 224\n", "global_crop2 = (80, 80, 400, 400) # 320×320 → resize to 224\n", "local_crops = [\n", " (100, 100, 196, 196), # house detail\n", " (150, 280, 246, 376), # door detail\n", " (300, 180, 396, 276), # tree detail\n", " (340, 20, 436, 116), # sun detail\n", " ( 20, 200, 116, 296), # left side\n", " (200, 20, 296, 116), # sky center\n", " ( 20, 350, 116, 446), # ground left\n", " (330, 350, 426, 446), # ground right\n", "]\n", "\n", "fig, axes = plt.subplots(2, 6, figsize=(20, 8))\n", "\n", "# Original with crop boxes\n", "ax = axes[0, 0]\n", "ax.imshow(scene)\n", "colors_g = ['#e74c3c', '#3498db']\n", "colors_l = ['#2ecc71','#f39c12','#9b59b6','#1abc9c','#e67e22','#34495e','#e91e63','#00bcd4']\n", "for i, (box, col) in enumerate(zip([global_crop1, global_crop2], colors_g)):\n", " x1,y1,x2,y2 = box\n", " rect = mpatches.Rectangle((x1,y1), x2-x1, y2-y1,\n", " linewidth=3, edgecolor=col, facecolor='none')\n", " ax.add_patch(rect)\n", " ax.text(x1+5, y1+15, f'G{i+1}', color=col, fontsize=12, fontweight='bold')\n", "for i, (box, col) in enumerate(zip(local_crops, colors_l)):\n", " x1,y1,x2,y2 = box\n", " rect = mpatches.Rectangle((x1,y1), x2-x1, y2-y1,\n", " linewidth=2, edgecolor=col, facecolor=col, alpha=0.15)\n", " ax.add_patch(rect)\n", " ax.text(x1+2, y1+12, f'L{i+1}', color=col, fontsize=8, fontweight='bold')\n", "ax.set_title('Original Image\\nwith crop boxes', fontweight='bold')\n", "ax.axis('off')\n", "\n", "# Global crops\n", "for i, (box, col) in enumerate(zip([global_crop1, global_crop2], colors_g)):\n", " crop = scene.crop(box).resize((224, 224), Image.BILINEAR)\n", " axes[0, i+1].imshow(crop)\n", " axes[0, i+1].set_title(f'Global Crop {i+1}\\n224×224 → Teacher+Student',\n", " fontweight='bold', color=col)\n", " axes[0, i+1].axis('off')\n", "\n", "# Local crops\n", "for i, (box, col) in enumerate(zip(local_crops[:3], colors_l[:3])):\n", " crop = scene.crop(box).resize((96, 96), Image.BILINEAR)\n", " axes[0, i+3].imshow(crop)\n", " axes[0, i+3].set_title(f'Local Crop {i+1}\\n96×96 → Student only',\n", " fontweight='bold', color=col)\n", " axes[0, i+3].axis('off')\n", "\n", "for i, (box, col) in enumerate(zip(local_crops[3:7], colors_l[3:7])):\n", " crop = scene.crop(box).resize((96, 96), Image.BILINEAR)\n", " axes[1, i].imshow(crop)\n", " axes[1, i].set_title(f'Local Crop {i+4}\\n96×96 → Student only',\n", " fontweight='bold', color=col)\n", " axes[1, i].axis('off')\n", "\n", "# Summary box\n", "axes[1, 4].axis('off')\n", "axes[1, 4].text(0.5, 0.7, 'Multi-Crop Summary', ha='center', fontsize=12, fontweight='bold',\n", " transform=axes[1,4].transAxes)\n", "summary = [\n", " '• 2 global crops (224×224)',\n", " ' → Teacher + Student',\n", " '• 8 local crops (96×96)',\n", " ' → Student only',\n", " '',\n", " 'Loss: student(local) →',\n", " 'match teacher(global)',\n", "]\n", "for i, line in enumerate(summary):\n", " axes[1, 4].text(0.05, 0.55 - i*0.08, line, ha='left', fontsize=10,\n", " transform=axes[1,4].transAxes)\n", "\n", "axes[1, 5].axis('off')\n", "\n", "plt.suptitle('DINOv2 Multi-Crop Augmentation Strategy', fontsize=15, fontweight='bold')\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3.3 DINO Loss (CLS Token Self-Distillation)\n", "\n", "The DINO loss operates on the **CLS token** — the global image representation.\n", "\n", "### Formula\n", "\n", "For each pair of views $(v_1, v_2)$ where $v_1$ is a global crop (teacher) and $v_2$ is any crop (student):\n", "\n", "$$\\mathcal{L}_{\\text{DINO}} = -\\sum_{v_2 \\neq v_1} \\mathbf{p}_t(v_1) \\cdot \\log \\mathbf{p}_s(v_2)$$\n", "\n", "where:\n", "- $\\mathbf{p}_t(x) = \\text{softmax}\\left(\\frac{g_t(x) - c}{\\tau_t}\\right)$ — teacher probability (centered, sharp)\n", "- $\\mathbf{p}_s(x) = \\text{softmax}\\left(\\frac{g_s(x)}{\\tau_s}\\right)$ — student probability\n", "- $c$ — running center (EMA of teacher outputs)\n", "- $\\tau_t = 0.04$ — teacher temperature (sharp)\n", "- $\\tau_s = 0.1$ — student temperature" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Implement and visualize DINO loss ─────────────────────────────────────────\n", "class DINOLoss:\n", " \"\"\"DINO loss: cross-entropy between teacher and student CLS token outputs.\"\"\"\n", " \n", " def __init__(self, out_dim=65536, tau_s=0.1, tau_t=0.04, center_momentum=0.9):\n", " self.tau_s = tau_s\n", " self.tau_t = tau_t\n", " self.center_momentum = center_momentum\n", " self.center = np.zeros(out_dim)\n", " \n", " def update_center(self, teacher_output):\n", " \"\"\"Update running center with EMA.\"\"\"\n", " batch_center = teacher_output.mean(axis=0)\n", " self.center = self.center_momentum * self.center + (1 - self.center_momentum) * batch_center\n", " \n", " def softmax_temp(self, x, T):\n", " x = x - x.max(axis=-1, keepdims=True) # numerical stability\n", " e = np.exp(x / T)\n", " return e / e.sum(axis=-1, keepdims=True)\n", " \n", " def __call__(self, student_out, teacher_out):\n", " \"\"\"Compute DINO loss.\"\"\"\n", " # Teacher: center + sharpen\n", " teacher_centered = teacher_out - self.center\n", " p_t = self.softmax_temp(teacher_centered, self.tau_t)\n", " \n", " # Student: just temperature scaling\n", " p_s = self.softmax_temp(student_out, self.tau_s)\n", " \n", " # Cross-entropy\n", " loss = -np.sum(p_t * np.log(p_s + 1e-8), axis=-1).mean()\n", " \n", " self.update_center(teacher_out)\n", " return loss, p_t, p_s\n", "\n", "# Simulate training dynamics\n", "np.random.seed(42)\n", "D = 20 # small output dim for visualization\n", "dino_loss = DINOLoss(out_dim=D)\n", "\n", "# Simulate: student gradually learns to match teacher\n", "n_steps = 100\n", "losses = []\n", "teacher_logits = np.random.randn(D) * 2 # fixed teacher output\n", "\n", "student_logits = np.random.randn(D) # random init\n", "for step in range(n_steps):\n", " loss, p_t, p_s = dino_loss(student_logits[None], teacher_logits[None])\n", " losses.append(loss)\n", " # Simulate student learning (gradient step toward teacher)\n", " student_logits = student_logits + 0.05 * (teacher_logits - student_logits) + np.random.randn(D)*0.1\n", "\n", "# Final distributions\n", "_, p_t_final, p_s_final = dino_loss(student_logits[None], teacher_logits[None])\n", "\n", "fig, axes = plt.subplots(1, 3, figsize=(16, 5))\n", "\n", "# Loss curve\n", "axes[0].plot(losses, 'b-', linewidth=2)\n", "axes[0].set_title('DINO Loss During Training\\n(simulated)', fontweight='bold')\n", "axes[0].set_xlabel('Training Step')\n", "axes[0].set_ylabel('Cross-Entropy Loss')\n", "axes[0].grid(True, alpha=0.3)\n", "axes[0].fill_between(range(n_steps), losses, alpha=0.2)\n", "\n", "# Teacher vs student distributions\n", "x = np.arange(D)\n", "axes[1].bar(x - 0.2, p_t_final[0], width=0.35, label=f'Teacher (τ={dino_loss.tau_t})', color='crimson', alpha=0.8)\n", "axes[1].bar(x + 0.2, p_s_final[0], width=0.35, label=f'Student (τ={dino_loss.tau_s})', color='steelblue', alpha=0.8)\n", "axes[1].set_title('Teacher vs Student Distributions\\n(after training)', fontweight='bold')\n", "axes[1].set_xlabel('Prototype Index')\n", "axes[1].set_ylabel('Probability')\n", "axes[1].legend()\n", "axes[1].grid(True, alpha=0.3)\n", "\n", "# Effect of temperature\n", "raw = np.array([2.1, 0.3, -0.5, 1.8, 0.1, -1.2, 0.9, 0.4, -0.3, 1.1])\n", "temps = [0.04, 0.1, 0.5, 1.0, 2.0]\n", "for T in temps:\n", " e = np.exp((raw - raw.max()) / T)\n", " p = e / e.sum()\n", " axes[2].plot(p, marker='o', markersize=4, label=f'T={T}', linewidth=1.5)\n", "axes[2].set_title('Effect of Temperature on\\nSoftmax Distribution', fontweight='bold')\n", "axes[2].set_xlabel('Prototype Index')\n", "axes[2].set_ylabel('Probability')\n", "axes[2].legend(fontsize=9)\n", "axes[2].grid(True, alpha=0.3)\n", "\n", "plt.suptitle('DINO Loss: CLS Token Self-Distillation', fontsize=14, fontweight='bold')\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3.4 iBOT Loss (Masked Patch Prediction)\n", "\n", "**iBOT** (Image BERT pre-training with Online Tokenizer) adds a **masked image modeling** objective.\n", "\n", "### How it works:\n", "1. Randomly mask a subset of patches in the **student** input (replace with a learnable `[MASK]` token)\n", "2. The **teacher** sees the unmasked image\n", "3. Loss: student's prediction for masked patches → match teacher's output for those patches\n", "\n", "$$\\mathcal{L}_{\\text{iBOT}} = -\\sum_{i \\in \\mathcal{M}} \\mathbf{p}_t^{(i)} \\cdot \\log \\mathbf{p}_s^{(i)}$$\n", "\n", "where $\\mathcal{M}$ is the set of masked patch indices.\n", "\n", "**Why iBOT?**\n", "- DINO only trains the CLS token → good global features\n", "- iBOT trains **patch tokens** → good dense/local features\n", "- Combined: excellent for both classification AND dense prediction tasks" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Visualize iBOT masking strategy ───────────────────────────────────────────\n", "np.random.seed(7)\n", "\n", "grid_size = 14 # 14×14 = 196 patches for ViT-B/14 with 224px input\n", "n_patches = grid_size * grid_size\n", "mask_ratio = 0.3 # mask 30% of patches\n", "\n", "# Random mask\n", "mask = np.zeros(n_patches, dtype=bool)\n", "mask_indices = np.random.choice(n_patches, int(n_patches * mask_ratio), replace=False)\n", "mask[mask_indices] = True\n", "mask_grid = mask.reshape(grid_size, grid_size)\n", "\n", "# Block mask (more realistic — DINOv2 uses block masking)\n", "block_mask = np.zeros((grid_size, grid_size), dtype=bool)\n", "# Place random blocks\n", "for _ in range(8):\n", " r, c = np.random.randint(0, grid_size-3), np.random.randint(0, grid_size-3)\n", " h, w = np.random.randint(1, 4), np.random.randint(1, 4)\n", " block_mask[r:r+h, c:c+w] = True\n", "\n", "fig, axes = plt.subplots(1, 4, figsize=(18, 5))\n", "\n", "# Random mask\n", "im0 = axes[0].imshow(mask_grid.astype(float), cmap='RdYlGn_r', vmin=0, vmax=1)\n", "axes[0].set_title(f'Random Mask\\n({mask_ratio*100:.0f}% masked = {mask.sum()} patches)',\n", " fontweight='bold')\n", "axes[0].set_xlabel('Patch column'); axes[0].set_ylabel('Patch row')\n", "plt.colorbar(im0, ax=axes[0], fraction=0.046)\n", "\n", "# Block mask\n", "im1 = axes[1].imshow(block_mask.astype(float), cmap='RdYlGn_r', vmin=0, vmax=1)\n", "axes[1].set_title(f'Block Mask\\n({block_mask.sum()} patches masked)',\n", " fontweight='bold')\n", "axes[1].set_xlabel('Patch column')\n", "plt.colorbar(im1, ax=axes[1], fraction=0.046)\n", "\n", "# Student sees masked image\n", "student_view = np.ones((grid_size, grid_size, 3)) * 0.8 # light gray background\n", "# Assign colors to unmasked patches\n", "np.random.seed(42)\n", "for i in range(grid_size):\n", " for j in range(grid_size):\n", " if not block_mask[i, j]:\n", " # Color based on position (simulating image content)\n", " student_view[i, j] = [i/grid_size, j/grid_size, 0.5]\n", " else:\n", " student_view[i, j] = [0.3, 0.3, 0.3] # masked = gray\n", "axes[2].imshow(student_view)\n", "axes[2].set_title('Student Input\\n(masked patches = gray)', fontweight='bold')\n", "axes[2].set_xlabel('Patch column')\n", "\n", "# Teacher sees full image\n", "teacher_view = np.zeros((grid_size, grid_size, 3))\n", "for i in range(grid_size):\n", " for j in range(grid_size):\n", " teacher_view[i, j] = [i/grid_size, j/grid_size, 0.5]\n", "axes[3].imshow(teacher_view)\n", "# Highlight masked positions\n", "for i in range(grid_size):\n", " for j in range(grid_size):\n", " if block_mask[i, j]:\n", " rect = mpatches.Rectangle((j-0.5, i-0.5), 1, 1,\n", " linewidth=2, edgecolor='red', facecolor='none')\n", " axes[3].add_patch(rect)\n", "axes[3].set_title('Teacher Input (full)\\nRed = positions to predict', fontweight='bold')\n", "axes[3].set_xlabel('Patch column')\n", "\n", "plt.suptitle('iBOT: Masked Patch Prediction Strategy', fontsize=14, fontweight='bold')\n", "plt.tight_layout()\n", "plt.show()\n", "\n", "print(f\"Grid size: {grid_size}×{grid_size} = {n_patches} patches\")\n", "print(f\"Block masked: {block_mask.sum()} patches ({block_mask.sum()/n_patches*100:.1f}%)\")\n", "print(f\"Visible to student: {(~block_mask).sum()} patches\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3.5 KoLeo Regularizer\n", "\n", "**KoLeo** (Kozachenko-Leonenko entropy estimator) is a regularizer that **maximizes the entropy** of the batch of representations.\n", "\n", "### Motivation\n", "Without regularization, representations can cluster in a small region of the embedding space. \n", "KoLeo encourages representations to be **spread uniformly** across the hypersphere.\n", "\n", "### Formula\n", "\n", "$$\\mathcal{L}_{\\text{KoLeo}} = -\\frac{1}{n} \\sum_{i=1}^{n} \\log \\left( \\min_{j \\neq i} \\| \\hat{z}_i - \\hat{z}_j \\|_2 \\right)$$\n", "\n", "where $\\hat{z}_i = z_i / \\|z_i\\|_2$ are L2-normalized embeddings.\n", "\n", "This is the **nearest-neighbor distance** in the batch — maximizing it spreads points apart." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Visualize KoLeo regularizer effect ────────────────────────────────────────\n", "def koleo_loss(z):\n", " \"\"\"KoLeo loss: negative log of nearest-neighbor distances.\"\"\"\n", " # L2 normalize\n", " z_norm = z / (np.linalg.norm(z, axis=1, keepdims=True) + 1e-8)\n", " n = len(z_norm)\n", " # Pairwise distances\n", " dists = np.sqrt(((z_norm[:, None] - z_norm[None, :]) ** 2).sum(axis=-1))\n", " np.fill_diagonal(dists, np.inf)\n", " # Nearest neighbor distance\n", " nn_dists = dists.min(axis=1)\n", " return -np.log(nn_dists + 1e-8).mean(), nn_dists\n", "\n", "np.random.seed(42)\n", "n_points = 50\n", "\n", "# Clustered embeddings (bad — high KoLeo loss)\n", "z_clustered = np.random.randn(n_points, 2) * 0.3\n", "z_clustered[:25] += [1.5, 0]\n", "z_clustered[25:] += [-1.5, 0]\n", "\n", "# Spread embeddings (good — low KoLeo loss)\n", "angles = np.linspace(0, 2*np.pi, n_points, endpoint=False)\n", "z_spread = np.column_stack([np.cos(angles), np.sin(angles)])\n", "z_spread += np.random.randn(n_points, 2) * 0.05\n", "\n", "loss_clustered, nn_clustered = koleo_loss(z_clustered)\n", "loss_spread, nn_spread = koleo_loss(z_spread)\n", "\n", "fig, axes = plt.subplots(2, 3, figsize=(16, 10))\n", "\n", "# Clustered\n", "z_c_norm = z_clustered / (np.linalg.norm(z_clustered, axis=1, keepdims=True) + 1e-8)\n", "axes[0,0].scatter(z_c_norm[:,0], z_c_norm[:,1], c=nn_clustered, cmap='RdYlGn', s=60)\n", "theta = np.linspace(0, 2*np.pi, 100)\n", "axes[0,0].plot(np.cos(theta), np.sin(theta), 'k--', alpha=0.3, linewidth=1)\n", "axes[0,0].set_title(f'Clustered Embeddings\\nKoLeo Loss = {loss_clustered:.3f} (HIGH = BAD)',\n", " fontweight='bold', color='red')\n", "axes[0,0].set_aspect('equal'); axes[0,0].grid(True, alpha=0.3)\n", "axes[0,0].set_xlabel('Dim 1'); axes[0,0].set_ylabel('Dim 2')\n", "\n", "# Spread\n", "z_s_norm = z_spread / (np.linalg.norm(z_spread, axis=1, keepdims=True) + 1e-8)\n", "sc = axes[0,1].scatter(z_s_norm[:,0], z_s_norm[:,1], c=nn_spread, cmap='RdYlGn', s=60)\n", "axes[0,1].plot(np.cos(theta), np.sin(theta), 'k--', alpha=0.3, linewidth=1)\n", "axes[0,1].set_title(f'Spread Embeddings\\nKoLeo Loss = {loss_spread:.3f} (LOW = GOOD)',\n", " fontweight='bold', color='green')\n", "axes[0,1].set_aspect('equal'); axes[0,1].grid(True, alpha=0.3)\n", "axes[0,1].set_xlabel('Dim 1'); axes[0,1].set_ylabel('Dim 2')\n", "plt.colorbar(sc, ax=axes[0,1], label='NN distance')\n", "\n", "# NN distance distributions\n", "axes[0,2].hist(nn_clustered, bins=15, alpha=0.7, color='red', label=f'Clustered (mean={nn_clustered.mean():.3f})')\n", "axes[0,2].hist(nn_spread, bins=15, alpha=0.7, color='green', label=f'Spread (mean={nn_spread.mean():.3f})')\n", "axes[0,2].set_title('Nearest-Neighbor Distance Distribution', fontweight='bold')\n", "axes[0,2].set_xlabel('NN Distance'); axes[0,2].set_ylabel('Count')\n", "axes[0,2].legend(); axes[0,2].grid(True, alpha=0.3)\n", "\n", "# Simulate KoLeo effect during training\n", "np.random.seed(0)\n", "z = np.random.randn(n_points, 2) * 0.2 # start clustered\n", "koleo_losses = []\n", "for step in range(200):\n", " loss, nn_d = koleo_loss(z)\n", " koleo_losses.append(loss)\n", " # Gradient: push points apart (simplified)\n", " z_n = z / (np.linalg.norm(z, axis=1, keepdims=True) + 1e-8)\n", " dists = np.sqrt(((z_n[:, None] - z_n[None, :]) ** 2).sum(axis=-1))\n", " np.fill_diagonal(dists, np.inf)\n", " nn_idx = dists.argmin(axis=1)\n", " grad = np.zeros_like(z)\n", " for i in range(n_points):\n", " diff = z[i] - z[nn_idx[i]]\n", " grad[i] = diff / (np.linalg.norm(diff) + 1e-8)\n", " z = z + 0.01 * grad + np.random.randn(*z.shape) * 0.005\n", "\n", "axes[1,0].plot(koleo_losses, 'purple', linewidth=2)\n", "axes[1,0].set_title('KoLeo Loss During Training\\n(decreasing = spreading out)', fontweight='bold')\n", "axes[1,0].set_xlabel('Step'); axes[1,0].set_ylabel('KoLeo Loss')\n", "axes[1,0].grid(True, alpha=0.3)\n", "\n", "# Final spread\n", "z_final_norm = z / (np.linalg.norm(z, axis=1, keepdims=True) + 1e-8)\n", "_, nn_final = koleo_loss(z)\n", "axes[1,1].scatter(z_final_norm[:,0], z_final_norm[:,1], c=nn_final, cmap='RdYlGn', s=60)\n", "axes[1,1].plot(np.cos(theta), np.sin(theta), 'k--', alpha=0.3, linewidth=1)\n", "axes[1,1].set_title('After KoLeo Regularization\\n(more uniform spread)', fontweight='bold', color='green')\n", "axes[1,1].set_aspect('equal'); axes[1,1].grid(True, alpha=0.3)\n", "\n", "# Summary\n", "axes[1,2].axis('off')\n", "summary_text = [\n", " 'KoLeo Regularizer Summary',\n", " '',\n", " 'Formula:',\n", " 'L = -1/n Σ log(min_j ||z_i - z_j||)',\n", " '',\n", " 'Effect:',\n", " '• Maximizes entropy of batch',\n", " '• Spreads embeddings uniformly',\n", " '• Prevents dimensional collapse',\n", " '• Better use of embedding space',\n", " '',\n", " 'Applied to: CLS token features',\n", " 'Weight: λ_KoLeo = 0.1',\n", "]\n", "for i, line in enumerate(summary_text):\n", " weight = 'bold' if i == 0 else 'normal'\n", " axes[1,2].text(0.05, 0.95 - i*0.07, line, transform=axes[1,2].transAxes,\n", " fontsize=10, fontweight=weight, va='top')\n", "\n", "plt.suptitle('KoLeo Regularizer: Spreading Representations Uniformly', fontsize=14, fontweight='bold')\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3.6 Combined Loss and Training Hyperparameters\n", "\n", "### Loss Weights\n", "$$\\mathcal{L} = \\mathcal{L}_{\\text{DINO}} + 1.0 \\cdot \\mathcal{L}_{\\text{iBOT}} + 0.1 \\cdot \\mathcal{L}_{\\text{KoLeo}}$$\n", "\n", "### Key Hyperparameters (from paper)\n", "\n", "| Parameter | Value | Notes |\n", "|-----------|-------|-------|\n", "| Teacher temperature $\\tau_t$ | 0.04 | Sharp teacher distribution |\n", "| Student temperature $\\tau_s$ | 0.1 | Softer student distribution |\n", "| EMA momentum $\\lambda$ | 0.996 → 0.9999 | Cosine schedule |\n", "| Center momentum | 0.9 | For DINO centering |\n", "| Global crop scale | 0.32 – 1.0 | 32%–100% of image |\n", "| Local crop scale | 0.05 – 0.32 | 5%–32% of image |\n", "| Num local crops | 8 | Per image |\n", "| Mask ratio (iBOT) | 0.1 – 0.5 | Random per sample |\n", "| Batch size | 3072 | Across all GPUs |\n", "| Learning rate | 0.001 | AdamW, cosine decay |\n", "| Weight decay | 0.04 → 0.4 | Cosine schedule |\n", "| Training epochs | 500 (ViT-L/g) | On LVD-142M |\n", "| Mixed precision | fp16 | With loss scaling |" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Visualize training schedules ──────────────────────────────────────────────\n", "steps = np.linspace(0, 1, 500)\n", "\n", "# Learning rate: linear warmup + cosine decay\n", "warmup_frac = 0.1\n", "lr_base = 0.001\n", "lr_min = 1e-6\n", "lr = np.where(\n", " steps < warmup_frac,\n", " lr_base * steps / warmup_frac,\n", " lr_min + 0.5 * (lr_base - lr_min) * (1 + np.cos(np.pi * (steps - warmup_frac) / (1 - warmup_frac)))\n", ")\n", "\n", "# Weight decay: cosine from 0.04 to 0.4\n", "wd = 0.04 + 0.5 * (0.4 - 0.04) * (1 - np.cos(np.pi * steps))\n", "\n", "# EMA momentum: cosine from 0.996 to 0.9999\n", "ema = 0.9999 - 0.5 * (0.9999 - 0.996) * (1 + np.cos(np.pi * steps))\n", "\n", "# Teacher temperature: warmup from 0.04 to 0.07 then fixed\n", "tau_t = np.where(steps < 0.3, 0.04 + (0.07 - 0.04) * steps / 0.3, 0.07)\n", "\n", "fig, axes = plt.subplots(2, 2, figsize=(14, 8))\n", "\n", "axes[0,0].plot(steps * 100, lr * 1000, 'b-', linewidth=2)\n", "axes[0,0].axvline(warmup_frac*100, color='red', linestyle='--', alpha=0.7, label='End of warmup')\n", "axes[0,0].set_title('Learning Rate Schedule\\n(warmup + cosine decay)', fontweight='bold')\n", "axes[0,0].set_xlabel('Training Progress (%)'); axes[0,0].set_ylabel('LR × 10³')\n", "axes[0,0].legend(); axes[0,0].grid(True, alpha=0.3)\n", "\n", "axes[0,1].plot(steps * 100, wd, 'orange', linewidth=2)\n", "axes[0,1].set_title('Weight Decay Schedule\\n(cosine: 0.04 → 0.4)', fontweight='bold')\n", "axes[0,1].set_xlabel('Training Progress (%)'); axes[0,1].set_ylabel('Weight Decay')\n", "axes[0,1].grid(True, alpha=0.3)\n", "\n", "axes[1,0].plot(steps * 100, ema, 'green', linewidth=2)\n", "axes[1,0].set_title('EMA Momentum Schedule\\n(cosine: 0.996 → 0.9999)', fontweight='bold')\n", "axes[1,0].set_xlabel('Training Progress (%)'); axes[1,0].set_ylabel('Momentum λ')\n", "axes[1,0].set_ylim(0.994, 1.0001)\n", "axes[1,0].grid(True, alpha=0.3)\n", "\n", "axes[1,1].plot(steps * 100, tau_t, 'purple', linewidth=2)\n", "axes[1,1].set_title('Teacher Temperature Schedule\\n(warmup: 0.04 → 0.07)', fontweight='bold')\n", "axes[1,1].set_xlabel('Training Progress (%)'); axes[1,1].set_ylabel('Temperature τ_t')\n", "axes[1,1].grid(True, alpha=0.3)\n", "\n", "plt.suptitle('DINOv2 Training Hyperparameter Schedules', fontsize=14, fontweight='bold')\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3.7 Full Training Pipeline Summary" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Full pipeline diagram ──────────────────────────────────────────────────────\n", "fig, ax = plt.subplots(figsize=(16, 10))\n", "ax.set_xlim(0, 16); ax.set_ylim(0, 10)\n", "ax.axis('off')\n", "\n", "def fbox(ax, x, y, w, h, fc, ec, text, fs=9):\n", " ax.add_patch(plt.FancyBboxPatch((x,y), w, h, boxstyle='round,pad=0.15',\n", " facecolor=fc, edgecolor=ec, linewidth=2, zorder=2))\n", " ax.text(x+w/2, y+h/2, text, ha='center', va='center', fontsize=fs,\n", " fontweight='bold', color='white', multialignment='center', zorder=3)\n", "\n", "def farrow(ax, x1, y1, x2, y2, label='', color='#555'):\n", " ax.annotate('', xy=(x2,y2), xytext=(x1,y1),\n", " arrowprops=dict(arrowstyle='->', color=color, lw=1.5), zorder=4)\n", " if label:\n", " ax.text((x1+x2)/2, (y1+y2)/2+0.15, label, ha='center', fontsize=8, color=color)\n", "\n", "# Data\n", "fbox(ax, 0.2, 4.5, 2.0, 1.0, '#555', '#333', 'LVD-142M\\nImages', 10)\n", "\n", "# Augmentation\n", "fbox(ax, 3.0, 6.5, 2.0, 0.8, '#2980b9', '#1a5276', 'Global\\nCrop 1\\n224×224', 8)\n", "fbox(ax, 3.0, 5.5, 2.0, 0.8, '#2980b9', '#1a5276', 'Global\\nCrop 2\\n224×224', 8)\n", "fbox(ax, 3.0, 3.5, 2.0, 1.8, '#e67e22', '#a04000', 'Local Crops\\n96×96\\n(×8)', 8)\n", "\n", "# Teacher\n", "fbox(ax, 6.2, 6.0, 2.2, 1.5, '#27ae60', '#1a7a40', 'Teacher ViT\\n(EMA, no grad)\\nθ_t', 9)\n", "\n", "# Student\n", "fbox(ax, 6.2, 3.5, 2.2, 2.0, '#c0392b', '#922b21', 'Student ViT\\n(trained)\\nθ_s', 9)\n", "\n", "# Heads\n", "fbox(ax, 9.2, 7.0, 1.8, 0.7, '#27ae60', '#1a7a40', 'Teacher\\nDINO Head', 8)\n", "fbox(ax, 9.2, 6.0, 1.8, 0.7, '#27ae60', '#1a7a40', 'Teacher\\niBOT Head', 8)\n", "fbox(ax, 9.2, 4.5, 1.8, 0.7, '#c0392b', '#922b21', 'Student\\nDINO Head', 8)\n", "fbox(ax, 9.2, 3.5, 1.8, 0.7, '#c0392b', '#922b21', 'Student\\niBOT Head', 8)\n", "\n", "# Losses\n", "fbox(ax, 11.8, 6.5, 2.0, 0.8, '#8e44ad', '#6c3483', 'DINO Loss\\n(CLS token)', 8)\n", "fbox(ax, 11.8, 5.3, 2.0, 0.8, '#d35400', '#a04000', 'iBOT Loss\\n(patches)', 8)\n", "fbox(ax, 11.8, 4.1, 2.0, 0.8, '#16a085', '#0e6655', 'KoLeo Loss\\n(entropy)', 8)\n", "\n", "# Total loss\n", "fbox(ax, 14.2, 5.0, 1.6, 1.5, '#2c3e50', '#1a252f', 'Total\\nLoss\\nL_total', 9)\n", "\n", "# EMA update\n", "fbox(ax, 6.2, 1.5, 2.2, 0.8, '#7f8c8d', '#555', 'EMA Update\\nθ_t ← λθ_t+(1-λ)θ_s', 8)\n", "\n", "# Arrows\n", "farrow(ax, 2.2, 5.0, 3.0, 6.9, '', '#2980b9')\n", "farrow(ax, 2.2, 5.0, 3.0, 5.9, '', '#2980b9')\n", "farrow(ax, 2.2, 5.0, 3.0, 4.4, '', '#e67e22')\n", "\n", "farrow(ax, 5.0, 6.9, 6.2, 6.8, 'global only', '#27ae60')\n", "farrow(ax, 5.0, 5.9, 6.2, 6.5, '', '#27ae60')\n", "farrow(ax, 5.0, 6.9, 6.2, 4.5, 'all crops', '#c0392b')\n", "farrow(ax, 5.0, 5.9, 6.2, 4.2, '', '#c0392b')\n", "farrow(ax, 5.0, 4.4, 6.2, 3.9, '', '#c0392b')\n", "\n", "farrow(ax, 8.4, 6.8, 9.2, 7.3, '', '#27ae60')\n", "farrow(ax, 8.4, 6.5, 9.2, 6.3, '', '#27ae60')\n", "farrow(ax, 8.4, 4.5, 9.2, 4.8, '', '#c0392b')\n", "farrow(ax, 8.4, 4.2, 9.2, 3.8, '', '#c0392b')\n", "\n", "farrow(ax, 11.0, 7.3, 11.8, 6.9, '', '#8e44ad')\n", "farrow(ax, 11.0, 4.8, 11.8, 5.7, '', '#d35400')\n", "farrow(ax, 11.0, 4.8, 11.8, 4.5, '', '#16a085')\n", "\n", "farrow(ax, 13.8, 6.9, 14.2, 6.0, '', '#2c3e50')\n", "farrow(ax, 13.8, 5.7, 14.2, 5.5, '', '#2c3e50')\n", "farrow(ax, 13.8, 4.5, 14.2, 5.2, '', '#2c3e50')\n", "\n", "# Backprop\n", "ax.annotate('', xy=(7.3, 3.5), xytext=(15.0, 5.0),\n", " arrowprops=dict(arrowstyle='->', color='red', lw=2,\n", " connectionstyle='arc3,rad=-0.3'))\n", "ax.text(12.5, 2.5, 'Backprop → update θ_s only', ha='center', fontsize=10,\n", " color='red', fontweight='bold')\n", "\n", "# EMA\n", "farrow(ax, 7.3, 3.5, 7.3, 2.3, '', 'purple')\n", "farrow(ax, 8.4, 2.3, 7.3, 6.0, '', 'purple')\n", "ax.text(9.5, 1.8, 'EMA: θ_t ← λθ_t + (1-λ)θ_s', ha='center', fontsize=9,\n", " color='purple', style='italic')\n", "\n", "ax.set_title('DINOv2 Complete Training Pipeline', fontsize=15, fontweight='bold')\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3.8 Summary\n", "\n", "| Component | Role | Key Parameter |\n", "|-----------|------|---------------|\n", "| Multi-crop | Creates diverse views | 2 global (224) + 8 local (96) |\n", "| Teacher ViT | Stable target | EMA momentum 0.996→0.9999 |\n", "| Student ViT | Trained model | Backprop through all losses |\n", "| DINO loss | Global representation | τ_t=0.04, τ_s=0.1 |\n", "| iBOT loss | Dense patch features | Mask ratio 10–50% |\n", "| KoLeo loss | Spread embeddings | λ=0.1 |\n", "| Centering | Prevent collapse | EMA center momentum=0.9 |\n", "\n", "---\n", "**Next:** [`04_dataset_preparation.ipynb`](04_dataset_preparation.ipynb) — How DINOv2's LVD-142M dataset was curated from 1.2B web images." ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.10.0" } }, "nbformat": 4, "nbformat_minor": 5 }