Emerging Properties in Self-Supervised Vision Transformers

Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, Armand Joulin

Seminal April 2021

TL;DR DINO trains a Vision Transformer without labels by having a student network match a momentum teacher's output on different crops of the same image. The [CLS] token's attention maps segment objects on their own, and the features reach 80.1% ImageNet top-1 with a linear classifier and 78.3% with plain k-NN.

What they do

The authors ask whether self-supervised pretraining gives Vision Transformers (ViTs) properties that supervised ViTs and convnets lack. With DINO, self-distillation with no labels, they find two: the [CLS] token’s last-layer attention outlines objects, and frozen features classify nearly as well with k-NN as with a linear layer.

How it works

Student and teacher. Both networks have the same architecture (a ViT or ResNet backbone plus an MLP projection head) but different weights. Each outputs a K-dimensional softmax, and the student is trained with cross-entropy to match the teacher’s output on a different view of the image.
Multi-crop. Two global 224x224 crops and several local 96x96 crops all go through the student. Only the global crops go through the teacher, which encourages local-to-global correspondence.
Momentum teacher. The teacher’s weights are an exponential moving average of the student’s, and no gradient flows into it.
Avoiding collapse. Only centering (subtracting a running mean of the teacher’s outputs) and sharpening (a low teacher temperature). There is no predictor, contrastive loss or batch norm.

Results

ImageNet with ViT-S: 77.0% linear and 74.5% k-NN, versus 73.5% and 66.6% for the best of BYOL, MoCo-v2 and SwAV on the same backbone. ViT-B/8 reaches 80.1% linear. On ResNet-50, DINO’s 75.3% ties SwAV for the best result among the self-supervised methods compared.
Attention masks scored against PASCAL VOC12 ground truth: Jaccard 45.9 for DINO ViT-S/16, versus 27.3 for a supervised ViT and 22.0 for a random one.
With frozen features and no finetuning: DAVIS-2017 video segmentation (J&F)m of 71.4 with ViT-B/8, and 51.5 mAP on revisited Oxford (Medium split) when pretrained on Google Landmarks v2.
Ablations on ViT-S/16: without the momentum teacher, k-NN accuracy falls to 0.1% (collapse). Replacing cross-entropy with MSE drops it from 72.8% to 52.6%.

Limitations

Small patches give much of the gain but cost speed: throughput falls from 180 im/s with 8x8 patches to 44 im/s with 5x5. On DAVIS, the 16x16-patch models (61.8, 62.3) score below the self-supervised STC baseline (67.6). Multi-crop raises peak memory per GPU from 9.3 GB to 15.4 GB. Teachers copied from the current or previous-iteration student do not converge. Pretraining used curated ImageNet; uncurated images are left to future work.

Figures

From Figure 1: Self-attention of the [CLS] token on the last-layer heads of a ViT with 8x8 patches trained with no su... — From Figure 1: Self-attention of the [CLS] token on the last-layer heads of a ViT with 8x8 patches trained with no supervision. The maps outline the objects, giving unsupervised object segmentations.

From Figure 4: Mask from thresholding the self-attention map to keep 60% of its mass, best head of a ViT-S/8 trained ... — From Figure 4: Mask from thresholding the self-attention map to keep 60% of its mass, best head of a ViT-S/8 trained with supervision.

From Figure 4: The same mask from the best head of a ViT-S/8 trained with DINO.

From Figure 8: Last-block self-attention of a ViT-S/8 trained with DINO, queried at a set of reference points. The ne... — From Figure 8: Last-block self-attention of a ViT-S/8 trained with DINO, queried at a set of reference points. The network separates the objects without any supervision.

Abstract

In this paper, we question if self-supervised learning provides new properties to Vision Transformer (ViT) that stand out compared to convolutional networks (convnets). Beyond the fact that adapting self-supervised methods to this architecture works particularly well, we make the following observations: first, self-supervised ViT features contain explicit information about the semantic segmentation of an image, which does not emerge as clearly with supervised ViTs, nor with convnets. Second, these features are also excellent k-NN classifiers, reaching 78.3% top-1 on ImageNet with a small ViT. Our study also underlines the importance of momentum encoder, multi-crop training, and the use of small patches with ViTs. We implement our findings into a simple self-supervised method, called DINO, which we interpret as a form of self-distillation with no labels. We show the synergy between DINO and ViTs by achieving 80.1% top-1 on ImageNet in linear evaluation with ViT-Base.

Summaries are written automatically from each paper's text. Figures and videos come from the paper and its project page; all credit goes to the authors.

Matt Strong