Emerging Properties in Self-Supervised Vision Transformers — part 3/6

Cleaned-text code points [29730, 45498). Bibliography omitted here; source.md includes it.

In Fig. 6 (right), we compare different strategies to build the teacher from previous instances of the student besides the momentum teacher. First we consider using the student network from a previous epoch as a teacher. This strategy has been used in a memory bank [ 27 ] or as a form of clustering hard-distillation [ 8 , 2 , 14 ] . Second, we consider using the student network from the previous iteration, as well as a copy of the student for the teacher. In our setting, using a teacher based on a recent version of the student does not converge. This setting requires more normalizations to work. Interestingly, we observe that using a teacher from the previous epoch does not collapse, providing performance in the k k -NN evaluation competitive with existing frameworks such as MoCo-v2 or BYOL. While using a momentum encoder clearly provides superior performance to this naive teacher, this finding suggests that there is a space to investigate alternatives for the teacher.

##### Analyzing the training dynamic.

To further understand the reasons why a momentum teacher works well in our framework, we study its dynamic during the training of a ViT in the left panel of Fig. 6 . A key observation is that this teacher constantly outperforms the student during the training, and we observe the same behavior when training with a ResNet-50 (Appendix D ). This behavior has not been observed by other frameworks also using momentum [ 13 , 12 ] , nor when the teacher is built from the previous epoch. We propose to interpret the momentum teacher in DINO as a form of Polyak-Ruppert averaging [ 20 , 59 ] with an exponentially decay. Polyak-Ruppert averaging is often used to simulate model ensembling to improve the performance of a network at the end of the training [ 38 ] . Our method can be interpreted as applying Polyak-Ruppert averaging during the training to constantly build a model ensembling that has superior performances. This model ensembling then guides the training of the student network [ 24 ] .

### 5.3 Avoiding collapse

We study the complementarity role of centering and target sharpening to avoid collapse. There are two forms of collapse: regardless of the input, the model output is uniform along all the dimensions or dominated by one dimension. The centering avoids the collapse induced by a dominant dimension, but encourages an uniform output. Sharpening induces the opposite effect. We show this complementarity by decomposing the cross-entropy H H into an entropy h h and the Kullback-Leibler divergence (“KL”) D K ​ L D_{KL} : H ⁡ ( P t , P s ) = h ⁡ ( P t ) + D K ​ L ​ ( P t | P s ) . H(P_{t},P_{s})=h(P_{t})+D_{KL}(P_{t}|P_{s}). (5) A KL equal to zero indicates a constant output, and hence a collapse. In Fig. 7 , we plot the entropy and KL during training with and without centering and sharpening. If one operation is missing, the KL converges to zero, indicating a collapse. However, the entropy h h converges to different values: 0 0 with no centering and − log ⁡ ( 1 / K ) -\log(1/K) with no sharpening, indicating that both operations induce different form of collapse. Applying both operations balances these effects (see study of the sharpening parameter τ t \tau_{t} in Appendix D ).

### 5.4 Compute requirements

In Tab. 8 , we detail the time and GPU memory requirements when running ViT-S/16 DINO models on two 8 8 -GPU machines. We report results with several variants of multi-crop training, each having a different level of compute requirement. We observe in Tab. 8 that using multi-crop improves the accuracy / running-time tradeoff for DINO runs. For example, the performance is 72.5 % 72.5\% after 46 46 hours of training without multi-crop (i.e. × 224 2 2\!\times\!224^{2} ) while DINO in × 224 2 + × 96 2 2\!\times\!224^{2}+10\!\times\!96^{2} crop setting reaches 74.6 % 74.6\% in 24 24 hours only. This is an improvement of + 2 % +2\% while requiring × 2\!\times less time, though the memory usage is higher ( 15.4 ​ G 15.4G versus 9.3 ​ G 9.3G ). We observe that the performance boost brought with multi-crop cannot be caught up by more training in the × 224 2 2\!\times\!224^{2} setting, which shows the value of the “local-to-global” augmentation. Finally, the gain from adding more views diminishes (+.2% form × 6\!\times to × 10\!\times 96 2 96^{2} crops) for longer trainings.

Overall, training DINO with Vision Transformers achieves 76.1 76.1 top-1 accuracy using two 8-GPU servers for 3 days. This result outperforms state-of-the-art self-supervised systems based on convolutional networks of comparable sizes with a significant reduction of computational requirements [ 12 , 3 ] . Our code is available to train self-supervised ViT on a limited number of GPUs.

### 5.5 Training with small batches

In Tab. 9 , we study the impact of the batch size on the features obtained with DINO. We also study the impact of the smooth parameter m m used in the centering update rule of Eq. 4 in Appendix D . We scale the learning rate linearly with the batch size [ 29 ] : l ​ r = 0.0005 ∗ batchsize / 256 lr=0.0005*\text{batchsize}/256 . Tab. 9 confirms that we can train models to high performance with small batches. Results with the smaller batch sizes ( b ​ s = 128 bs=128 ) are slightly below our default training setup of b ​ s = 1024 bs=1024 , and would certainly require to re-tune hyperparameters like the momentum rates for example. Note that the experiment with batch size of 128 128 runs on only 1 1 GPU. We have explored training a model with a batch size of 8 8 , reaching 35.2 % 35.2\% after 50 50 epochs, showing the potential for training large models that barely fit an image per GPU.

## 6 Conclusion

In this work, we have shown the potential of self-supervised pretraining a standard ViT model, achieving performance that are comparable with the best convnets specifically designed for this setting. We have also seen emerged two properties that can be leveraged in future applications: the quality of the features in k k -NN classification has a potential for image retrieval where ViT are already showing promising results [ 22 ] . The presence of information about the scene layout in the features can also benefit weakly supervised image segmentation. However, the main result of this paper is that we have evidences that self-supervised learning could be the key to developing a BERT-like model based on ViT. In the future, we plan to explore if pretraining a large ViT model with DINO on random uncurated images could push the limits of visual features [ 28 ] .

##### Acknowledgement.

We thank Mahmoud Assran, Matthijs Douze, Allan Jabri, Jure Zbontar, Alaaeldin El-Nouby, Y-Lan Boureau, Kaiming He, Thomas Lucas as well as the Thoth and FAIR teams for their help, support and discussions around this project. Julien Mairal was funded by the ERC grant number 714381 (SOLARIS project) and by ANR 3IA MIAI@Grenoble Alpes (ANR-19-P3IA-0003).

## Appendix

### A Additional Results

##### k k -NN classification.

In Tab. 10 , we evaluate the frozen representations given by ResNet-50 or ViT-small pre-trained with DINO with two evaluation protocols: linear or k k -NN. For both evaluations, we extract representations from a pre-trained network without using any data augmentation. Then, we perform classification either with weighted k k -NN or with a linear regression learned with cyanure library [ 16 ] . In Tab. 10 we see that ViT-S accuracies are better than accuracies obtained with RN50 both with a linear or a k k -NN classifier. However, the performance gap when using the k k -NN evaluation is much more significant than when considering linear evaluation. For example on ImageNet 1%, ViT-S outperforms ResNet-50 by a large margin of + 14.1 % +14.1\% with k k -NN evaluation. This suggests that transformers architectures trained with DINO might offer more model flexibility that benefits the k k -NN evaluation. K K -NN classifiers have the great advantage of being fast and light to deploy, without requiring any domain adaptation. Overall, ViT trained with DINO provides features that combine particularly well with k k -NN classifiers.

##### Self-supervised ImageNet pretraining of ViT.

In this experiment, we study the impact of pretraining a supervised ViT model with our method. In Tab. 11 , we compare the performance of supervised ViT models that are initialized with different pretraining or guided during training with an additional pretrained convnet. The first set of models are pretrained with and without supervision on the large curated dataset composed of 300M images. The second set of models are trained with hard knowledge distillation from a pretrained supervised RegNetY [ 56 ] . The last set of models do not use any additional data nor models, and are initialized either randomly or after a pretraining with DINO on ImageNet. Compare to random initialization, pretraining with DINO leads to a performance gain of +1%. This is not caused by a longer training since pretraining with supervision instead of DINO does not improve performance. Using self-supervised pretraining reduces the gap with models pretrained on extra data or distilled from a convnet.

##### Low-shot learning on ImageNet.

We evaluate the features obtained with DINO applied on ViT-S on low-shot learning. In Tab. 12 , we report the validation accuracy of a logistic regression trained on frozen features ( frozen ) with 1% and 10% labels. The logistic regression is trained with the cyanure library [ 16 ] . When comparing models with a similar number of parameters and image/sec, we observe that our features are on par with state-of-the-art semi-supervised models. Interestingly, this performance is obtained by training a multi-class logistic regression on frozen features, without data augmentation nor finetuning .

### B Methodology Comparison

We compare the performance of different self-supervised frameworks, MoCo-v2 [ 6 ] , SwAV [ 3 ] and BYOL [ 12 ] when using convnet or ViT. In Tab. 13 , we see that when trained with ResNet-50 (convnet), DINO performs on par with SwAV and BYOL. However, DINO unravels its potential with ViT, outperforming MoCo-v2, SwAV and BYOL by large margins (+4.3% with linear and +6.2% with k-NN evaluations). In the rest of this section, we perform ablations to better understand the performance of DINO applied to ViT. In particular, we provide a detailed comparison with methods that either use a momentum encoder, namely MoCo-v2 and BYOL, and methods that use multi-crop, namely SwAV.

##### Relation to MoCo-v2 and BYOL.

In Tab. 14 , we present the impact of ablating components that differ between DINO, MoCo-v2 and BYOL: the choice of loss, the predictor in the student head, the centering operation, the batch normalization in the projection heads, and finally, the multi-crop augmentation. The loss in DINO is a cross-entropy on sharpened softmax outputs ( CE ) while MoCo-v2 uses the InfoNCE contrastive loss ( INCE ) and BYOL a mean squared error on l2-normalized outputs ( MSE ). No sharpening is applied with the MSE criterion. Though, DINO surprisingly still works when changing the loss function to MSE , but this significantly alters the performance (see rows ( 1 , 2 ) and ( 4 , 9 )). We also observe that adding a predictor has little impact ( 1 , 3 ). However, in the case of BYOL, the predictor is critical to prevent collapse ( 7 , 8 ) which is consistent with previous studies [ 7 , 12 ] . Interestingly, we observe that the teacher output centering avoids collapse without predictor nor batch normalizations in BYOL ( 7 , 9 ), though with a significant performance drop which can likely be explained by the fact that our centering operator is designed to work in combination with sharpening. Finally, we observe that multi-crop works particularly well with DINO and MoCo-v2, removing it hurts performance by 2 − 4 % 2-4\% ( 1 versus 4 and, 5 versus 6 ). Adding multi-crop to BYOL does not work out-of-the-box ( 7 , 10 ) as detailed in Appendix E and further adaptation may be required.

##### Relation to SwAV.

In Tab. 15 , we evaluate the differences between DINO and SwAV: the presence of the momentum encoder and the operation on top of the teacher output. In absence of the momentum, a copy of the student with a stop-gradient is used. We consider three operations on the teacher output: Centering , Sinkhorn-Knopp or a Softmax along the batch axis. The Softmax is similar to a single Sinkhorn-Knopp iteration as detailed in the next paragraph. First, these ablations show that using a momentum encoder significantly improves the performance for ViT ( 3 versus 6 , and 2 versus 5 ). Second, the momentum encoder also avoids collapse when using only centering (row 1 ). In the absence of momentum, centering the outputs does not work ( 4 ) and more advanced operations are required ( 5 , 6 ). Overall, these ablations highlight the importance of the momentum encoder, not only for performance but also to stabilize training, removing the need for normalization beyond centering.

##### Details on the Softmax(batch) variant.

The iterative Sinkhorn-Knopp algorithm [ 8 ] used in SwAV [ 3 ] is implemented simply with the following PyTorch style code. ⬇

When performing a single Sinkhorn iteration ( num_iters=1 ) the implementation can be highly simplified into only two lines of code, which is our softmax(batch) variant:

We have seen in Tab. 15 that this highly simplified variant of SwAV works competitively with SwAV. Intuitively, the softmax operation on the batch axis allows to select for each dimension (or “cluster”) its best matches in the batch.

##### Validating our implementation.

We observe in Tab. 13 that our reproduction of BYOL, MoCo-v2, SwAV matches or outperforms the corresponding published numbers with ResNet-50. Indeed, we obtain 72.7 % 72.7\% for BYOL while [ 12 ] report 72.5 % 72.5\% in this 300 300 -epochs setting. We obtain 71.1 % 71.1\% for MoCo after 300 300 epochs of training while [ 6 ] report 71.1 % 71.1\% after 800 800 epochs of training. Our improvement compared to the implementation of [ 6 ] can be explained by the use of a larger projection head (3-layer, use of batch-normalizations and projection dimension of 256 256 ).

##### Relation to other works.

DINO is also related to UIC [ 14 ] that use outputs from the previous epoch as hard pseudo-labels for “unsupervised classification”. However, we use centering to prevent collapse while UIC resorts to balance sampling techniques as in [ 8 ] . Our work can be interpreted as a soft UIC variant with momentum teacher.

The concurrent work CsMI [ 77 ] also exhibits strong performance with simple k-NN classifiers on ImageNet, even with convnets. As DINO, CsMI combines a momentum network and multi-crop training, which we have seen are both crucial for good k-NN performance in our experiments with ViTs. We believe studying this work would help us identifying more precisely the components important for good k k -NN performance and leave this investigation for future work.

### C Projection Head

Similarly to other self-supervised frameworks, using a projection head [ 4 ] improves greatly the accuracy of our method. The projection head starts with a n n -layer multi-layer perceptron (MLP). The hidden layers are 2048d and are with gaussian error linear units (GELU) activations. The last layer of the MLP is without GELU. Then we apply a ℓ 2 \ell_{2} normalization and a weight normalized fully connected layer [ 7 , 22 ] with K K dimensions. This design is inspired from the projection head with a “prototype layer” used in SwAV [ 3 ] . We do not apply batch normalizations.

##### BN-free system.

Unlike standard convnets, ViT architectures do not use batch normalizations (BN) by default.