Emerging Properties in Self-Supervised Vision Transformers

In plain words

Vision Transformers (ViTs) read an image as a grid of small patches, and they are usually trained with human labels. This paper asks whether a ViT trained without any labels learns something special, and it finds that it does. Its features separate the object in a picture from the background, and a very simple nearest-neighbour lookup on those features classifies images well.

The method is called DINO. It has two copies of the same network. The student sees several crops of one picture and learns to predict what the teacher says about it. The teacher is not trained directly; it is a slowly moving average of the student's past weights. The teacher only looks at the two large crops, so the student has to match a whole-image view from a small piece. Two simple tricks keep the training from collapsing to a constant answer: the teacher's outputs are re-centred and sharpened.

The point of the work is that these ideas were built for convolutional networks and needed extra machinery there. Paired with ViTs, the recipe is simpler and the features are strong for lookup and retrieval.

TL;DR

DINO is a label-free self-distillation method. A student ViT matches a teacher built as an exponential moving average of the student, using cross-entropy on softmax outputs across several crops. Teacher centering and sharpening stop collapse, so the approach needs no contrastive negatives, no predictor and no batch normalisation in the ViT case. The authors report strong k-NN classification and object-segmentation-like attention maps from the frozen features.

Problem

Existing self-supervised methods were built and tuned mainly for convnets, and several of them rely on contrastive negatives, predictor heads or batch normalisation §3.1. The paper asks whether adapting such methods to ViTs works and whether the resulting ViT features differ from convnet features §1. The paper asks whether self-supervised ViTs show properties that convnets do not, and whether a simple method can exploit them §1. Evaluation is done with linear probes, fine-tuning and a frozen weighted k-NN classifier §3.2.

Method

Framing. DINO is a self-distillation scheme. A student network g_{\theta_s} matches the output distribution of a teacher network g_{\theta_t}. Both produce K-dimensional softmax distributions with temperatures \tau_s and \tau_t §3.1 eq. 1. Training minimises cross-entropy H(a, b) = -a \log b with respect to the student only §3.1 eq. 2.

Views. From each image the method draws two global crops, which cover a large part of the image and are resized to 224^2, and several local crops at 96^2 that cover a small part of the image §3.1 §3.2. All crops go through the student; only the two global crops go through the teacher. The loss compares each teacher output with every student output on a different view, which encourages local-to-global correspondence §3.1 eq. 3 alg. 1.

Teacher. There is no fixed teacher. The teacher is an exponential moving average of the student weights, with a momentum \lambda that follows a cosine schedule from 0.996 to 1 during training §3.1. The authors observe that this teacher outperforms the student at every point in training and interpret it as a form of Polyak-Ruppert averaging §5.2 app. D. Copying the student into the teacher fails to converge, and a teacher that is frozen for one epoch works surprisingly well §3.1 §5.2.

Avoiding collapse. The paper uses no contrastive term, no predictor and no normalisation in the ViT case. Collapse is prevented by two operations on the teacher output §3.1:

• Centering: subtract a running mean c of the teacher outputs. The mean is updated with an exponential moving average using rate m §3.1 eq. 4. Centering alone pushes the output towards a uniform distribution.
• Sharpening: use a low teacher temperature \tau_t, which pushes the output towards one dimension. Sharpening alone pushes towards one dominant dimension §3.1 §5.3.

Applying both balances the two failure modes. The paper reports the entropy of the teacher target and the KL between teacher and student as diagnostics for this balance §5.3 eq. 5.

Architecture. The network is a backbone f (a ViT, or a ResNet for the convnet comparison) followed by a projection head h, so g = h \circ f. Downstream features are taken from f §3.1. The head is a 3-layer MLP with hidden width 2048 and GELU activations, a final bottleneck of width 256, an \ell_2 normalisation, and a weight-normalised linear layer to K outputs §3.1 app. C. For ViT the head has no batch normalisation at all, so the whole system is BN-free §3.1 app. C. The authors use K = 65536 by default app. C. Student and teacher share the same architecture §3.1.

Training step. The pseudocode in alg. 1 is the single-global-view version without multi-crop; the full loss in eq. 3 adds the local views. The pseudocode below is the multi-crop version in our own notation, with the details taken from the official code main_dino.py main_dino.py.

# One DINO training step with multi-crop  [alg. 1] [eq. 3]
# views: list of crops; views[0:2] are global (224^2), the rest local (96^2)
# Output: loss scalar; student, teacher and center are updated in place

def dino_step(views, student, teacher, center, tau_s, tau_t, lam, m):
    # Teacher sees only the two global views, without gradients  [§3.1]
    # [code: main_dino.py#L318]
    teacher_logits = [teacher(v) for v in views[:2]]  # (B, K) each
    student_logits = [student(v) for v in views]      # (B, K) each

    # Teacher: subtract the centre, then sharpen with tau_t  [§3.1] [eq. 4]
    # [code: main_dino.py#L387-L390]
    teacher_probs = [
        softmax((t - center) / tau_t, dim=-1).detach()
        for t in teacher_logits
    ]                                                 # (B, K) each

    # Cross-entropy for every pair of different views  [eq. 3]
    # [code: main_dino.py#L392-L402]
    total = 0.0
    pairs = 0
    for i, teacher_p in enumerate(teacher_probs):
        for j, student_z in enumerate(student_logits):
            if i == j:
                continue
            # Student log-probabilities at temperature tau_s  [eq. 1]
            log_q = log_softmax(student_z / tau_s, dim=-1)
            total += mean(-sum(teacher_p * log_q, dim=-1))
            pairs += 1
    loss = total / pairs  # paper's eq. 3 sums; code divides by pair count

    # Update the centre with the mean of all teacher logits  [eq. 4]
    # [code: main_dino.py#L407-L416]
    batch_mean = mean(concat(teacher_logits), dim=0)  # (K,)
    center = m * center + (1 - m) * batch_mean

    # Teacher = exponential moving average of student weights, in place  [§3.1]
    # [code: main_dino.py#L347-L350]
    for p_t, p_s in zip(teacher.params, student.params):
        set_value(p_t, lam * p_t + (1 - lam) * p_s)

    return loss, center

Schedules. Learning rate, weight decay and teacher momentum follow cosine schedules, and the student's learning rate has a linear warm-up §3.2 main_dino.py. The teacher temperature is warmed up linearly over the first epochs §3.2 app. D. The exact values are in the Hyperparameters section.

Hyperparameters

Values not stated in the paper are marked not stated. Values that exist only in the code are marked code only. Where the code's default differs from the paper, see "Paper vs code".

Results

Author-reported. Setup: ImageNet-1k pretraining without labels, with evaluation by linear probe, weighted k-NN, retrieval, DAVIS segmentation, transfer fine-tuning and copy detection, as stated per table §4 §5. Numbers are top-1 unless noted.

ImageNet, same architecture tab. 2. Columns: Arch., parameters (M), images per second on a V100 at batch 128, linear top-1, k-NN top-1.

Starred baselines were run by the authors; the k-NN column was run for models with released weights tab. 2. The appendix text says SwAV, MoCo-v2 and DINO were trained for 800 epochs in this table app. D. The source also contains a commented-out sentence saying the other DINO models used 300 epochs; it is not active text, so we did not use it.

Image retrieval (mAP, revisited Oxford and Paris, medium and hard) tab. 3.

Copy detection (mAP, Copydays "strong") tab. 4.

DAVIS 2017 video segmentation, frozen features, nearest-neighbour propagation tab. 5. J&F, J mean, F mean.

Attention segmentation, PASCAL VOC12 Jaccard, 60% attention mass fig. 4. Table in the figure: ViT-S/16 random 22.0, supervised 27.3, DINO 45.9; ViT-S/8 random 21.8, supervised 23.7, DINO 44.7.

Transfer by fine-tuning, top-1 tab. 6.

Ablation: components for ViT-S/16, 300 epochs tab. 7. k-NN and linear top-1. Row 1 is DINO. "Mom." momentum teacher, "SK" Sinkhorn-Knopp, "MC" multi-crop, "Pred." predictor.

Patch size, k-NN (figure) fig. 5. Text reports: at 5×5 patches the throughput falls to 44 im/s, against 180 im/s at 8×8 §5.1.

Teacher type, k-NN, ViT-S, 300 epochs fig. 6.

Compute, ViT-S/16 on two 8-GPU machines tab. 8. Top-1 is linear evaluation; time is total wall clock; mem. is peak memory per GPU.

Batch size, k-NN, 100 epochs, no multi-crop tab. 9.

Appendix results.

• Frozen features, k-NN vs logistic regression, ViT-S/16 vs ResNet-50 app. A: ImageNet 100%: linear 75.7 vs 72.1, k-NN 74.5 vs 67.5; ImageNet 1%: linear 64.5 vs 55.1, k-NN 61.3 vs 47.2; VOC07 linear 89.2 vs 88.9, k-NN 88.0 vs 84.9. Average gain 5.6 points for k-NN app. A.
• Low-shot, frozen features, logistic regression on ImageNet app. A: DINO-frozen ViT-S/16 64.5 (1%) and 72.2 (10%). Other rows use fine-tuning or semi-supervision.
• Method comparison at 300 epochs app. B: ResNet-50 linear / k-NN: DINO 74.5 / 65.6; ViT-S DINO 76.1 / 72.8; MoCo-v2 ViT-S 71.6 / 62.0; BYOL ViT-S 71.4 / 66.6; SwAV ViT-S 71.8 / 64.7.
• Projection head depth, with and without ℓ2 bottleneck, 100 epochs, K = 4096 app. C: with bottleneck 62.2 / 68.0 / 69.3 for 2, 3, 4 linear layers; without bottleneck 61.6 / 62.9 for 1 and 2 layers, and 0.1 for 3 and 4 layers. Default is 4 linear layers.
• Output dimension K app. C: 1024 → 67.8; 4096 → 69.3; 16384 → 69.2; 65536 → 69.7; 262144 → 69.1 (k-NN top-1).
• Centre momentum m app. D: 0 → 69.1; 0.9 → 69.7; 0.99 → 69.4; 0.999 → 0.1 (k-NN top-1; epoch count not given in this table).
• Teacher temperature \tau_t app. D: 0 → 43.9; 0.02 → 66.7; 0.04 → 69.6; 0.06 → 68.7; 0.08 → 0.1; 0.04 → 0.07 warm-up → 69.7 (k-NN top-1; epoch count not given in this table).
• Multi-crop scale range s app. E: s = 0.08 → 65.6; 0.16 → 68.0; 0.24 → 69.7; 0.32 → 69.8; 0.48 → 69.5 (k-NN top-1; epoch count not given).
• Multi-crop in other frameworks, ViT-S/16, 300 epochs app. E: DINO without multi-crop 67.9 k-NN / 72.5 linear, with multi-crop 72.7 / 75.9. BYOL drops from 66.6 / 71.4 to 59.8 / 64.8 with multi-crop.
• Linear evaluation on ViT-S: concatenating the last 4 [CLS] tokens reaches 77.0; a single final layer gives 76.1 app. F. ViT-B: [CLS] only 78.0, with average-pooled patch tokens 78.2 app. F.
• Longer training, ViT-S k-NN app. D: 100 epochs 70.9; 300 epochs 72.8; 800 epochs 74.5.

Paper vs code

Code: official repository facebookresearch/dino at commit 7c446df5b9f45747937fb0d72314eb9f7b66930a, licensed Apache-2.0. The paper links this repository in its opening footnote [abstract].

• Teacher temperature schedule. The paper uses a warm-up from 0.04 to 0.07 over 30 epochs §3.2 app. D. The code's defaults are start 0.04, final 0.04 and warm-up epochs 0 main_dino.py. The help text for the warm-up length says "Default: 30" while the actual default is 0 main_dino.py. Use the 30-epoch schedule from the paper for faithful reproduction; the code default is effectively constant 0.04.
• Number of local crops. The paper's standard multi-crop is 2 × 224² plus 6 × 96² app. E. The code default is 8 local crops main_dino.py. The paper's experiments use 6 local views, so set local_crops_number explicitly.
• Global and local scale. The appendix sets global crops to (s, 1) and local crops to (0.05, s), and finds s near 0.3 best (the 0.32 column is highest in its table) app. E. The code defaults are global (0.4, 1) and local (0.05, 0.4) main_dino.py. The main text describes global views as covering more than 50% of the image and local views as covering less than 50% §3.2, which does not map directly onto the scale parameters.
• Loss normalisation. Equation 3 sums the cross-entropy over all pairs of different views eq. 3. The code averages over the pairs main_dino.py. The code's loss is the paper's sum divided by 2(|V| − 1). Because the code clips gradients at norm 3.0 main_dino.py, the scale does matter for the clipping step.
• Weight-normalised output layer. The paper describes a weight-normalised output layer §3.1 app. C. The code freezes the norm parameter at 1 by default (norm_last_layer=True) vision_transformer.py main_dino.py. The code's own help text says this is set to False for ViT-S in their experiments. The paper does not state whether the norm is trainable.
• Augmentation details. The paper follows BYOL's augmentations, namely colour jitter, Gaussian blur and solarisation, without probabilities §3.2. The code sets probabilities as listed in the Hyperparameters table main_dino.py. Solarisation is applied only to the second global view in the code.
• Training details absent from the paper. Stochastic depth 0.1, gradient clipping at 3.0, freezing the output layer for one epoch, and fp16 are set in the code only main_dino.py main_dino.py. The paper does not mention them.
• Teacher without stochastic depth. The student is built with drop-path 0.1, while the teacher is built without it, so its drop-path is the constructor default of 0 main_dino.py vision_transformer.py.
• Evaluation k. The paper uses k = 20 §3.2 app. F. The code evaluates k in {10, 20, 100, 200} and reports all of them eval_knn.py.

Pitfalls

• Headline numbers belong to different models. The abstract's 78.3% k-NN is ViT-S/8 (not ViT-S/16, which gets 74.5%) and the 80.1% linear is ViT-B/8 (ViT-B/16 gets 78.2%) tab. 2.
• Collapse depends on both operations. Without centering, the entropy goes to 0; without sharpening, the entropy goes to -\log(1/K) §5.3. The paper reports that a fixed teacher temperature above 0.06 drives the loss to ln K unless the temperature is warmed up from a lower value app. D. The table shows k-NN 0.1 at 0.08 app. D.
• Momentum is required. Using the student itself, or the student from the previous iteration, as the teacher does not converge §5.2. The centre momentum also matters: m = 0.999 collapses app. D.
• Small batches need a different momentum. Batch 128 reaches 57.9 k-NN, against 59.9 at 1024 tab. 9. The paper says hyperparameters such as the momentum would need re-tuning §5.5. The code help suggests a higher base momentum (0.9995) at batch 256 main_dino.py. Treat this as an unvalidated suggestion from the code.
• Patch size costs throughput. /8 patches give better features but run at 180 im/s against 1007 for /16 for ViT-S, and 5 × 5 patches run at 44 im/s §5.1. Memory also rises with more multi-crop views (9.3 G to 15.4 G per GPU) tab. 8. Half precision is discouraged for patch sizes below 16 in the code main_dino.py main_dino.py.
• Multi-crop is not free for every method. The same multi-crop recipe that helps DINO hurts BYOL in the authors' experiments app. E. Do not assume it transfers.
• Long training. The active appendix text says DINO in the main table was trained for 800 epochs, the same as SwAV and MoCo-v2 app. D tab. 2. A commented-out sentence in the source says other DINO models used 300 epochs; it is not active, so we did not use it. Compare rows with care.
• Compute reported in the text. The text says the ViT training reached 76.1 top-1 on two 8-GPU servers in 3 days §5.4. The table gives the same 76.1 at 300 epochs with 72.6 hours for the largest multi-crop setting tab. 8. Use the table for run time.
• Frozen-feature evaluation is sensitive. The paper notes that linear probes vary a lot with learning rate, so it uses k-NN as a hyperparameter-free check §3.2. Linear results need a learning-rate sweep app. F.
• Attention "segmentation" is qualitative plus one benchmark. The masks come from thresholding attention to 60% (or 80% in the appendix) of its mass, and the Jaccard numbers depend on that threshold fig. 4 app. D.
• Hardware-specific timings. Throughput numbers are on a V100 at batch 128 tab. 1 tab. 2. Compare them only with the same setup.
• Unstated augmentation probabilities. The paper does not give probabilities for colour jitter, blur or solarisation. Use the code values main_dino.py if you need to match the released models.

Lineage

• Builds on: knowledge distillation, where a student matches a fixed teacher §3.1; the momentum encoder of MoCo, which DINO uses as an EMA teacher with a different role, because there is no queue or contrastive loss §3.1; the mean teacher and Polyak-Ruppert averaging §3.1 §5.2; multi-crop and the projection-head design with a prototype layer from SwAV §3.1 app. C; the BYOL augmentations and cosine schedules §3.2; and the ViT architecture with the DeiT training recipe (arXiv:2010.11929) §3.2.
• Compared against: MoCo-v2, BYOL, SwAV, SimCLR, InfoMin, BarlowTwins, OBoW, DCv2 and supervised training on ImageNet tab. 2. The paper also compares with Multigrain tab. 4 and the video methods CT, MAST, STC and STM tab. 5.
• Related work noted by the paper: CsMI, which also gives strong k-NN results with a momentum network and multi-crop app. B; UIC, which uses the previous epoch as hard pseudo-labels app. B.