Lab 04 — ViT + DINO Self-Distillation

Phase: 02 — Model Architecture Mastery | Difficulty: ⭐⭐⭐⭐⭐ | Time: 4–6 hours

DINO (Caron et al., 2021) trains a ViT with no labels and gets attention maps that segment objects. This lab builds both the ViT and the self-distillation machinery from scratch, and makes you prove you understand why it doesn't collapse.

What you build

• PatchEmbed — the conv-as-patch-projection trick (kernel = stride = patch size)
• ViT — CLS token, learned position embeddings, pre-norm encoder blocks, plus get_last_attention for the famous segmentation-without-labels visualization
• DINOHead — MLP → L2-normalized bottleneck → prototype layer (cosine similarity against K learned prototypes)
• DINOLoss — cross-entropy between sharpened teacher and student distributions with EMA centering — the two mechanisms that jointly prevent collapse
• momentum_update — the EMA teacher (a temporal ensemble of past students)
• dino_train_step — multi-crop: teacher sees global views only, student sees all, loss pairs every teacher view with every different student view

Key concepts

Files

Run

pip install -r requirements.txt
pytest test_lab.py -v               # your lab.py
LAB_MODULE=solution pytest test_lab.py -v   # sanity-check the reference

Suggested TODO order

1. PatchEmbed — get the conv-as-projection reshape right first
2. ViT.forward — CLS prepend, pos embed, blocks, norm, return position 0
3. DINOHead — note the L2 normalize before the prototype layer
4. DINOLoss — the detach and the center update are where the learning is
5. momentum_update + the DINO wrapper (teacher frozen!)
6. dino_train_step — pair views correctly, EMA after optimizer step

Success criteria

• All 13 tests pass
• You can explain, without notes, why sharpening alone or centering alone still collapses
• You can state what the teacher would output if you forgot the detach() and trained it by gradient (hint: student and teacher co-collapse to the trivial solution)
• python solution.py style training shows the loss decreasing on repeated identical crops

Full-scale extension (the original spec)

Train ViT-Tiny on CIFAR-10 for 30 epochs (≈1 h single GPU): real augmentation multi-crop, cosine momentum schedule 0.996→1.0, linear-probe accuracy ≥80% without labels, and attention map visualizations showing object segmentation. Add register tokens (DINOv2) and compare.

Interview Q&A

Q: Why does DINO need both centering and sharpening? Centering subtracts a running mean of teacher logits, preventing any prototype from permanently dominating (one-hot collapse) — but pushes toward uniform. Sharpening (low teacher temperature) pushes toward confident outputs — but alone allows one-hot collapse. They balance in opposite directions; remove either and training degenerates.

Q: Why is the teacher an EMA rather than a copy or a separately trained network? EMA gives a temporal ensemble of student checkpoints: smoother and consistently better than the student during training (the paper measures this), providing targets slightly ahead of the student — a bootstrapping effect similar to BYOL's predictor asymmetry.

Q: Why do ViT attention maps segment objects without segmentation labels? The multi-crop objective forces local-crop features to predict global-view prototypes; the CLS token must aggregate from patches that determine object identity, so its attention concentrates on the foreground object. Supervised ViTs show much weaker localization — the objective, not the architecture, produces it.

Q: What changes in DINOv2? Scale (LVD-142M curated data), iBOT-style masked-patch prediction added to the loss, KoLeo regularization, register tokens to fix attention-map artifacts, and efficient implementations — the centering/sharpening core stays.

References

• Emerging Properties in Self-Supervised Vision Transformers (DINO)
• DINOv2: Learning Robust Visual Features without Supervision
• An Image is Worth 16x16 Words (ViT)
• Vision Transformers Need Registers
• Official DINO code