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

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

# Emerging Properties in Self-Supervised Vision Transformers

###### Abstract

In this paper, we question if self-supervised learning provides new properties to Vision Transformer (ViT) [ 10 ] 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 k -NN classifiers, reaching 78.3% top-1 on ImageNet with a small ViT. Our study also underlines the importance of momentum encoder [ 13 ] , multi-crop training [ 3 ] , 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- di stillation 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.

## 1 Introduction

Transformers [ 70 ] have recently emerged as an alternative to convolutional neural networks (convnets) for visual recognition [ 10 , 25 , 83 ] . Their adoption has been coupled with a training strategy inspired by natural language processing (NLP), that is, pretraining on large quantities of data and finetuning on the target dataset [ 9 , 55 ] . The resulting Vision Transformers (ViT) [ 10 ] are competitive with convnets but, they have not yet delivered clear benefits over them: they are computationally more demanding, require more training data, and their features do not exhibit unique properties.

In this paper, we question whether the muted success of Transformers in vision can be explained by the use of supervision in their pretraining. Our motivation is that one of the main ingredients for the success of Transformers in NLP was the use of self-supervised pretraining, in the form of close procedure in BERT [ 9 ] or language modeling in GPT [ 55 ] . These self-supervised pretraining objectives use the words in a sentence to create pretext tasks that provide a richer learning signal than the supervised objective of predicting a single label per sentence. Similarly, in images, image-level supervision often reduces the rich visual information contained in an image to a single concept selected from a predefined set of a few thousand categories of objects [ 21 ] .

While the self-supervised pretext tasks used in NLP are text specific, many existing self-supervised methods have shown their potential on images with convnets [ 3 , 4 , 12 , 13 ] . They typically share a similar structure but with different components designed to avoid trivial solutions (collapse) or to improve performance [ 7 ] . In this work, inspired from these methods, we study the impact of self-supervised pretraining on ViT features. Of particular interest, we have identified several interesting properties that do not emerge with supervised ViTs, nor with convnets:

• Self-supervised ViT features explicitly contain the scene layout and, in particular, object boundaries, as shown in Figure 1 . This information is directly accessible in the self-attention modules of the last block.

• Self-supervised ViT features perform particularly well with a basic nearest neighbors classifier ( k k -NN) without any finetuning, linear classifier nor data augmentation , achieving 78.3% top-1 accuracy on ImageNet.

The emergence of segmentation masks seems to be a property shared across self-supervised methods. However, the good performance with k k -NN only emerge when combining certain components such as momentum encoder [ 13 ] and multi-crop augmentation [ 3 ] . Another finding from our study is the importance of using smaller patches with ViTs to improve the quality of the resulting features.

Overall, our findings about the importance of these components lead us to design a simple self-supervised approach that can be interpreted as a form of knowledge di stillation [ 35 ] with no labels. The resulting framework, DINO, simplifies self-supervised training by directly predicting the output of a teacher network—built with a momentum encoder—by using a standard cross-entropy loss. Interestingly, our method can work with only a centering and sharpening of the teacher output to avoid collapse, while other popular components such as predictor [ 12 ] , advanced normalization [ 3 ] or contrastive loss [ 13 ] add little benefits in terms of stability or performance. Of particular importance, our framework is flexible and works on both convnets and ViTs without the need to modify the architecture, nor adapt internal normalizations [ 58 ] .

We further validate the synergy between DINO and ViT by outperforming previous self-supervised features on the ImageNet linear classification benchmark with 80.1% top-1 accuracy with a ViT-Base with small patches. We also confirm that DINO works with convnets by matching the state of the art with a ResNet-50 architecture. Finally, we discuss different scenarios to use DINO with ViTs in case of limited computation and memory capacity. In particular, training DINO with ViT takes just two 8-GPU servers over 3 days to achieve 76.1 % 76.1\% on ImageNet linear benchmark, which outperforms self-supervised systems based on convnets of comparable sizes with significantly reduced compute requirements [ 3 , 12 ] .

## 2 Related work

##### Self-supervised learning.

A large body of work on self-supervised learning focuses on discriminative approaches coined instance classification [ 4 , 20 , 13 , 27 ] , which considers each image a different class and trains the model by discriminating them up to data augmentations. However, explicitly learning a classifier to discriminate between all images [ 20 ] does not scale well with the number of images. Wu et al . [ 27 ] propose to use a noise contrastive estimator (NCE) [ 32 ] to compare instances instead of classifying them. A caveat of this approach is that it requires comparing features from a large number of images simultaneously. In practice, this requires large batches [ 4 ] or memory banks [ 13 , 27 ] . Several variants allow automatic grouping of instances in the form of clustering [ 2 , 8 , 2 , 36 , 42 , 74 , 80 , 85 ] .

Recent works have shown that we can learn unsupervised features without discriminating between images. Of particular interest, Grill et al . [ 12 ] propose a metric-learning formulation called BYOL, where features are trained by matching them to representations obtained with a momentum encoder. Methods like BYOL work even without a momentum encoder, at the cost of a drop of performance [ 7 , 12 ] . Several other works echo this direction, showing that one can match more elaborate representations [ 26 , 27 ] , train features matching them to a uniform distribution [ 6 ] or by using whitening [ 23 , 81 ] . Our approach takes its inspiration from BYOL but operates with a different similarity matching loss and uses the exact same architecture for the student and the teacher. That way, our work completes the interpretation initiated in BYOL of self-supervised learning as a form of Mean Teacher self-distillation [ 24 ] with no labels.

##### Self-training and knowledge distillation.

Self-training aims at improving the quality of features by propagating a small initial set of annotations to a large set of unlabeled instances. This propagation can either be done with hard assignments of labels [ 41 , 78 , 79 ] or with a soft assignment [ 76 ] . When using soft labels, the approach is often referred to as knowledge distillation [ 7 , 35 ] and has been primarily designed to train a small network to mimic the output of a larger network to compress models. Xie et al . [ 76 ] have shown that distillation could be used to propagate soft pseudo-labels to unlabelled data in a self-training pipeline, drawing an essential connection between self-training and knowledge distillation. Our work builds on this relation and extends knowledge distillation to the case where no labels are available. Previous works have also combined self-supervised learning and knowledge distillation [ 25 , 63 , 5 , 47 ] , enabling self-supervised model compression and performance gains. However, these works rely on a pre-trained fixed teacher while our teacher is dynamically built during training. This way, knowledge distillation, instead of being used as a post-processing step to self-supervised pre-training, is directly cast as a self-supervised objective. Finally, our work is also related to codistillation [ 1 ] where student and teacher have the same architecture and use distillation during training. However, the teacher in codistillation is also distilling from the student, while it is updated with an average of the student in our work.

## 3 Approach

### 3.1 SSL with Knowledge Distillation

The framework used for this work, DINO, shares the same overall structure as recent self-supervised approaches [ 3 , 7 , 4 , 12 , 13 ] . However, our method shares also similarities with knowledge distillation [ 35 ] and we present it under this angle. We illustrate DINO in Figure 2 and propose a pseudo-code implementation in Algorithm 1 .

Knowledge distillation is a learning paradigm where we train a student network g θ s g_{\theta_{s}} to match the output of a given teacher network g θ t g_{\theta_{t}} , parameterized by θ s \theta_{s} and θ t \theta_{t} respectively. Given an input image x x , both networks output probability distributions over K K dimensions denoted by P s P_{s} and P t P_{t} . The probability P P is obtained by normalizing the output of the network g g with a softmax function. More precisely, P s ​ ( x ) ( i ) = exp ⁡ ( g θ s ​ ( x ) ( i ) / τ s ) ∑ k = 1 K exp ⁡ ( g θ s ​ ( x ) ( k ) / τ s ) , P_{s}(x)^{(i)}=\frac{\exp(g_{\theta_{s}}(x)^{(i)}/\tau_{s})}{\sum_{k=1}^{K}\exp(g_{\theta_{s}}(x)^{(k)}/\tau_{s})}, (1) with τ s > 0 \tau_{s}>0 a temperature parameter that controls the sharpness of the output distribution, and a similar formula holds for P t P_{t} with temperature τ t \tau_{t} . Given a fixed teacher network g θ t g_{\theta_{t}} , we learn to match these distributions by minimizing the cross-entropy loss w.r.t. the parameters of the student network θ s \theta_{s} : min θ s ⁡ H ⁡ ( P t ​ ( x ) , P s ​ ( x ) ) , \min_{\theta_{s}}H(P_{t}(x),P_{s}(x)), (2) where H ⁡ ( a , b ) = − a ​ log ⁡ b H(a,b)=-a\log b .

In the following, we detail how we adapt the problem in Eq. ( 2 ) to self-supervised learning. First, we construct different distorted views, or crops, of an image with multi-crop strategy [ 3 ] . More precisely, from a given image, we generate a set V V of different views. This set contains two global views, x 1 g x^{g}_{1} and x 2 g x^{g}_{2} and several local views of smaller resolution. All crops are passed through the student while only the global views are passed through the teacher, therefore encouraging “local-to-global” correspondences. We minimize the loss: min θ s ∑ x ∈ { x 1 g , x 2 g } ∑ x ′ ∈ V x ′ ≠ x H ( P t ( x ) , P s ( x ′ ) ) . \min_{\theta_{s}}\sum_{x\in\{x^{g}_{1},x^{g}_{2}\}}\quad\sum_{\begin{subarray}{c}x^{\prime}\in V\\ x^{\prime}\neq\,x\end{subarray}}\quad H(P_{t}(x),P_{s}(x^{\prime})). (3)

This loss is general and can be used on any number of views, even only 2 2 . However, we follow the standard setting for multi-crop by using 2 global views at resolution 224 2 224^{2} covering a large (for example greater than 50 % 50\% ) area of the original image, and several local views of resolution 96 2 96^{2} covering only small areas (for example less than 50 % 50\% ) of the original image. We refer to this setting as the basic parametrization of DINO, unless mentioned otherwise.

Both networks share the same architecture g g with different sets of parameters θ s \theta_{s} and θ t \theta_{t} . We learn the parameters θ s \theta_{s} by minimizing Eq. ( 3 ) with stochastic gradient descent.

##### Teacher network.

Unlike knowledge distillation, we do not have a teacher g θ t g_{\theta_{t}} given a priori and hence, we build it from past iterations of the student network. We study different update rules for the teacher in Section 5.2 and show that freezing the teacher network over an epoch works surprisingly well in our framework, while copying the student weight for the teacher fails to converge. Of particular interest, using an exponential moving average (EMA) on the student weights, i.e., a momentum encoder [ 13 ] , is particularly well suited for our framework. The update rule is θ t ← λ ​ θ t + ( 1 − λ ) ​ θ s , \theta_{t}\leftarrow\lambda\theta_{t}+(1-\lambda)\theta_{s}, with λ \lambda following a cosine schedule from 0.996 0.996 to 1 1 during training [ 12 ] . Originally the momentum encoder has been introduced as a substitute for a queue in contrastive learning [ 13 ] . However, in our framework, its role differs since we do not have a queue nor a contrastive loss, and may be closer to the role of the mean teacher used in self-training [ 24 ] . Indeed, we observe that this teacher performs a form of model ensembling similar to Polyak-Ruppert averaging with an exponential decay [ 20 , 59 ] . Using Polyak-Ruppert averaging for model ensembling is a standard practice to improve the performance of a model [ 38 ] . We observe that this teacher has better performance than the student throughout the training, and hence, guides the training of the student by providing target features of higher quality. This dynamic was not observed in previous works [ 12 , 58 ] .

##### Network architecture.

The neural network g g is composed of a backbone f f (ViT [ 10 ] or ResNet [ 34 ] ), and of a projection head h h : g = h ∘ f g=h\circ f . The features used in downstream tasks are the backbone f f output. The projection head consists of a 3-layer multi-layer perceptron (MLP) with hidden dimension 2048 2048 followed by ℓ 2 \ell_{2} normalization and a weight normalized fully connected layer [ 22 ] with K K dimensions, which is similar to the design from SwAV [ 3 ] . We have tested other projection heads and this particular design appears to work best for DINO (Appendix C ). We do not use a predictor [ 12 , 7 ] , resulting in the exact same architecture in both student and teacher networks. Of particular interest, we note that unlike standard convnets, ViT architectures do not use batch normalizations (BN) by default. Therefore, when applying DINO to ViT we do not use any BN also in the projection heads, making the system entirely BN-free .

##### Avoiding collapse.