Emerging Properties in Self-Supervised Vision Transformers — part 2/6
Cleaned-text code points [14662, 29730). Bibliography omitted here; source.md includes it.
Several self-supervised methods differ by the operation used to avoid collapse, either through contrastive loss [ 27 ] , clustering constraints [ 8 , 3 ] , predictor [ 12 ] or batch normalizations [ 12 , 58 ] . While our framework can be stabilized with multiple normalizations [ 3 ] , it can also work with only a centering and sharpening of the momentum teacher outputs to avoid model collapse. As shown experimentally in Section 5.3 , centering prevents one dimension to dominate but encourages collapse to the uniform distribution, while the sharpening has the opposite effect. Applying both operations balances their effects which is sufficient to avoid collapse in presence of a momentum teacher. Choosing this method to avoid collapse trades stability for less dependence over the batch: the centering operation only depends on first-order batch statistics and can be interpreted as adding a bias term c c to the teacher: g t ( x ) ← g t ( x ) + c g_{t}(x)\leftarrow g_{t}(x)+c . The center c c is updated with an exponential moving average, which allows the approach to work well across different batch sizes as shown in Section 5.5 : c ← m c + ( 1 − m ) 1 B ∑ i = 1 B g θ t ( x i ) , c\leftarrow mc+(1-m)\frac{1}{B}\sum_{i=1}^{B}g_{\theta_{t}}(x_{i}), (4) where m > 0 m>0 is a rate parameter and B B is the batch size. Output sharpening is obtained by using a low value for the temperature τ t \tau_{t} in the teacher softmax normalization.
### 3.2 Implementation and evaluation protocols
In this section, we provide the implementation details to train with DINO and present the evaluation protocols used in our experiments.
##### Vision Transformer.
We briefly describe the mechanism of the Vision Transformer (ViT) [ 10 , 70 ] and refer to Vaswani et al . [ 70 ] for details about Transformers and to Dosovitskiy et al . [ 10 ] for its adaptation to images. We follow the implementation used in DeiT [ 25 ] . We summarize the configuration of the different networks used in this paper in Table 1 . The ViT architecture takes as input a grid of non-overlapping contiguous image patches of resolution N × N N\times N . In this paper we typically use N = 16 N=16 (“/16”) or N = 8 N=8 (“/8”). The patches are then passed through a linear layer to form a set of embeddings. We add an extra learnable token to the sequence [ 9 , 10 ] . The role of this token is to aggregate information from the entire sequence and we attach the projection head h h at its output. We refer to this token as the class token [CLS] for consistency with previous works [ 9 , 10 , 25 ] , even though it is not attached to any label nor supervision in our case. The set of patch tokens and [CLS] token are fed to a standard Transformer network with a “pre-norm” layer normalization [ 11 , 39 ] . The Transformer is a sequence of self-attention and feed-forward layers, paralleled with skip connections. The self-attention layers update the token representations by looking at the other token representations with an attention mechanism [ 4 ] .
##### Implementation details.
We pretrain the models on the ImageNet dataset [ 21 ] without labels. We train with the adamw optimizer [ 44 ] and a batch size of 1024 1024 , distributed over 16 16 GPUs when using ViT-S/16. The learning rate is linearly ramped up during the first 10 10 epochs to its base value determined with the following linear scaling rule [ 29 ] : l r = 0.0005 ∗ batchsize / 256 lr=0.0005*\text{batchsize}/256 . After this warmup, we decay the learning rate with a cosine schedule [ 43 ] . The weight decay also follows a cosine schedule from 0.04 0.04 to 0.4 0.4 . The temperature τ s \tau_{s} is set to 0.1 0.1 while we use a linear warm-up for τ t \tau_{t} from 0.04 0.04 to 0.07 0.07 during the first 30 30 epochs. We follow the data augmentations of BYOL [ 12 ] (color jittering, Gaussian blur and solarization) and multi-crop [ 3 ] with a bicubic interpolation to adapt the position embeddings to the scales [ 10 , 25 ] . The code and models to reproduce our results is publicly available.
##### Evaluation protocols.
Standard protocols for self-supervised learning are to either learn a linear classifier on frozen features [ 82 , 13 ] or to finetune the features on downstream tasks. For linear evaluations, we apply random resize crops and horizontal flips augmentation during training, and report accuracy on a central crop. For finetuning evaluations, we initialize networks with the pretrained weights and adapt them during training. However, both evaluations are sensitive to hyperparameters, and we observe a large variance in accuracy between runs when varying the learning rate for example. We thus also evaluate the quality of features with a simple weighted nearest neighbor classifier ( k k -NN) as in [ 27 ] . We freeze the pretrain model to compute and store the features of the training data of the downstream task. The nearest neighbor classifier then matches the feature of an image to the k k nearest stored features that votes for the label. We sweep over different number of nearest neighbors and find that 20 20 NN is consistently working the best for most of our runs. This evaluation protocol does not require any other hyperparameter tuning, nor data augmentation and can be run with only one pass over the downstream dataset, greatly simplifying the feature evaluation.
## 4 Main Results
We first validate the DINO framework used in this study with the standard self-supervised benchmark on ImageNet. We then study the properties of the resulting features for retrieval, object discovery and transfer-learning.
### 4.1 Comparing with SSL frameworks on ImageNet
We consider two different settings: comparison with the same architecture and across architectures.
##### Comparing with the same architecture.
In top panel of Table 2 , we compare DINO with other self-supervised methods with the same architecture, either a ResNet-50 [ 34 ] or a ViT-small (which follows the design of DeiT-S [ 25 ] ). The choice of ViT-S is motivated by its similarity with ResNet-50 along several axes: number of parameters (21M vs 23M), throughput (1237/sec VS 1007 im/sec) and supervised performance on ImageNet with the training procedure of [ 25 ] (79.3% VS 79.8%). We explore variants of ViT-S in Appendix D . First, we observe that DINO performs on par with the state of the art on ResNet-50, validating that DINO works in the standard setting. When we switch to a ViT architecture, DINO outperforms BYOL, MoCov2 and SwAV by +3.5% with linear classification and by +7.9% with k k -NN evaluation. More surprisingly, the performance with a simple k k -NN classifier is almost on par with a linear classifier (74.5% versus 77.0%). This property emerges only when using DINO with ViT architectures, and does not appear with other existing self-supervised methods nor with a ResNet-50.
##### Comparing across architectures.
On the bottom panel of Table 2 , we compare the best performance obtained across architectures. The interest of this setting is not to compare methods directly, but to evaluate the limits of a ViT trained with DINO when moving to larger architectures. While training a larger ViT with DINO improves the performance, reducing the size of the patches (“/8” variants) has a bigger impact on the performance. While reducing the patch size do not add parameters, it still leads to a significant reduction of running time, and larger memory usage. Nonetheless, a base ViT with 8 × 8 8\times 8 patches trained with DINO achieves 80.1% top-1 in linear classification and 77.4% with a k k -NN classifier with 10 × 10\times less parameters and 1.4 × 1.4\times faster run time than previous state of the art [ 5 ] .
### 4.2 Properties of ViT trained with SSL
We evaluate properties of the DINO features in terms of nearest neighbor search, retaining information about object location and transferability to downstream tasks.
#### 4.2.1 Nearest neighbor retrieval with DINO ViT
The results on ImageNet classification have exposed the potential of our features for tasks relying on nearest neighbor retrieval. In this set of experiments, we further consolidate this finding on landmark retrieval and copy detection tasks.
##### Image Retrieval.
We consider the revisited [ 53 ] Oxford and Paris image retrieval datasets [ 50 ] . They contain 3 different splits of gradual difficulty with query/database pairs. We report the Mean Average Precision (mAP) for the Medium (M) and Hard (H) splits. In Table 3 , we compare the performance of different off-the-shelf features obtained with either supervised or DINO training. We freeze the features and directly apply k k -NN for retrieval. We observe that DINO features outperform those trained on ImageNet with labels.
An advantage of SSL approaches is that they can be trained on any dataset, without requiring any form of annotations. We train DINO on the 1.2M clean set from Google Landmarks v2 (GLDv2) [ 72 ] , a dataset of landmarks designed for retrieval purposes. DINO ViT features trained on GLDv2 are remarkably good, outperforming previously published methods based on off-the-shelf descriptors [ 68 , 57 ] .
##### Copy detection.
We also evaluate the performance of ViTs trained with DINO on a copy detection task. We report the mean average precision on the “strong” subset of the INRIA Copydays dataset [ 21 ] . The task is to recognize images that have been distorted by blur, insertions, print and scan, etc. Following prior work [ 5 ] , we add 10k distractor images randomly sampled from the YFCC100M dataset [ 66 ] . We perform copy detection directly with cosine similarity on the features obtained from our pretrained network. The features are obtained as the concatenation of the output [CLS] token and of the GeM pooled [ 54 ] output patch tokens. This results in a 1536d descriptor for ViT-B. Following [ 5 ] , we apply whitening on the features. We learn this transformation on an extra 20K random images from YFCC100M, distincts from the distractors. Table 4 shows that ViT trained with DINO is very competitive on copy detection.
#### 4.2.2 Discovering the semantic layout of scenes
As shown qualitatively in Figure 1 , our self-attention maps contain information about the segmentation of an image. In this study, we measure this property on a standard benchmark as well as by directly probing the quality of masks generated from these attention maps.
##### Video instance segmentation.
In Tab. 5 , we evaluate the output patch tokens on the DAVIS-2017 video instance segmentation benchmark [ 52 ] . We follow the experimental protocol in Jabri et al . [ 14 ] and segment scenes with a nearest-neighbor between consecutive frames; we thus do not train any model on top of the features, nor finetune any weights for the task. We observe in Tab. 5 that even though our training objective nor our architecture are designed for dense tasks, the performance is competitive on this benchmark. Since the network is not finetuned, the output of the model must have retained some spatial information. Finally, for this dense recognition task, the variants with small patches (“/8”) perform much better (+ 9.1 % 9.1\% ( 𝒥 CLOSE (\mathcal{J} & OPEN ℱ ) m \mathcal{F})_{m} for ViT-B).
##### Probing the self-attention map.
In Fig. 3 , we show that different heads can attend to different semantic regions of an image, even when they are occluded (the bushes on the third row) or small (the flag on the second row). Visualizations are obtained with 480 480 p images, resulting in sequences of 3601 tokens for ViT-S/8. In Fig. 4 , we show that a supervised ViT does not attend well to objects in presence of clutter both qualitatively and quantitatively. We report the Jaccard similarity between the ground truth and segmentation masks obtained by thresholding the self-attention map to keep 60% of the mass. Note that the self-attention maps are smooth and not optimized to produce a mask. Nonetheless, we see a clear difference between the supervised or DINO models with a significant gap in terms of Jaccard similarities. Note that self-supervised convnets also contain information about segmentations but it requires dedicated methods to extract it from their weights [ 31 ] .
#### 4.2.3 Transfer learning on downstream tasks
In Tab. 6 , we evaluate the quality of the features pretrained with DINO on different downstream tasks. We compare with features from the same architectures trained with supervision on ImageNet. We follow the protocol used in Touvron et al . [ 25 ] and finetune the features on each downstream task. We observe that for ViT architectures, self-supervised pretraining transfers better than features trained with supervision, which is consistent with observations made on convolutional networks [ 3 , 13 , 62 ] . Finally, self-supervised pretraining greatly improves results on ImageNet (+1-2%).
## 5 Ablation Study of DINO
In this section, we empirically study DINO applied to ViT. The model considered for this entire study is ViT-S. We also refer the reader to Appendix for additional studies.
### 5.1 Importance of the Different Components
We show the impact of adding different components from self-supervised learning on ViT trained with our framework.
In Table 7 , we report different model variants as we add or remove components. First, we observe that in the absence of momentum, our framework does not work (row 2 ) and more advanced operations, SK for example, are required to avoid collapse (row 9 ). However, with momentum, using SK has little impact (row 3 ). In addtition, comparing rows 3 and 9 highlights the importance of the momentum encoder for performance. Second, in rows 4 and 5 , we observe that multi-crop training and the cross-entropy loss in DINO are important components to obtain good features. We also observe that adding a predictor to the student network has little impact (row 6 ) while it is critical in BYOL to prevent collapse [ 7 , 12 ] . For completeness, we propose in Appendix B an extended version of this ablation study.
##### Importance of the patch size.
In Fig. 5 , we compare the k k -NN classification performance of ViT-S models trained with different patch sizes, 16 × 16 16\times 16 , 8 × 8 8\times 8 and 5 × 5 5\times 5 . We also compare to ViT-B with 16 × 16 16\times 16 and 8 × 8 8\times 8 patches. All the models are trained for 300 epochs. We observe that the performance greatly improves as we decrease the size of the patch. It is interesting to see that performance can be greatly improved without adding additional parameters. However, the performance gain from using smaller patches comes at the expense of throughput: when using 5 × \times 5 patches, the throughput falls to 44 im/s, vs 180 im/s for 8 × \times 8 patches.
### 5.2 Impact of the choice of Teacher Network
In this ablation, we experiment with different teacher network to understand its role in DINO. We compare models trained for 300 300 epochs using the k k -NN protocol.
##### Building different teachers from the student.