DINOv2: Learning Robust Visual Features without Supervision — part 2/4
Cleaned-text code points [15817, 31679). Bibliography omitted here; source.md includes it.
• Sinkhorn-Knopp centering ( Caron et al., 2020 ) . Ruan et al. (2023) recommend to replace the teacher softmax-centering step of DINO and iBot by the Sinkhorn-Knopp (SK) batch normalization of SwAV ( Caron et al., 2020 ) . We run the Sinkhorn-Knopp algorithm steps for 3 iterations. For the student, we apply the softmax normalization.
• KoLeo regularizer ( Sablayrolles et al., 2019 ) . The KoLeo regularizer derives from the Kozachenko-Leonenko differential entropy estimator (see Beirlant et al. (1997) ; Delattre & Fournier (2017) ) and encourages a uniform span of the features within a batch. Given a set of n n vectors ( x 1 , … , x n ) (x_{1},\dots,x_{n}) , it is defined as ℒ koleo = − 1 n ∑ i = 1 n log ( d n , i ) , {\mathcal{L}}_{\mathrm{koleo}}=-\frac{1}{n}\sum_{i=1}^{n}\log(d_{n,i}), where d n , i = min j ≠ i ‖ x i − x j ‖ d_{n,i}=\min_{j\neq i}\|x_{i}-x_{j}\| is the minimum distance between x i x_{i} and any other point within the batch. We also ℓ 2 \ell_{2} -normalize the features before computing this regularizer.
• Adapting the resolution ( Touvron et al., 2019 ) . Increasing image resolution is key to pixel-level downstream tasks such as segmentation or detection, where small objects disappear at low resolutions. However, training at high resolution is time and memory demanding, and instead, we increase the resolution of images to 518 × 518 518\times 518 during a short period at the end of pretraining. This is also similar to UniViT training from Likhomanenko et al. (2021) and FlexiViT training from Beyer et al. (2023) .
## 5 Efficient implementation
We consider several improvements to train models at a larger scale. We train models on A100 GPUs using PyTorch 2.0. The code and pretrained models are made available under Apache 2.0 license 1 1 1 https://github.com/facebookresearch/dinov2 . The details of our models are in the appendix, Table 17 . With the same hardware, compared to the iBOT implementation, the DINOv2 code runs around 2 × 2\times faster using only 1 / 3 1/3 of the memory.
##### Fast and memory-efficient attention.
We implemented our own version of FlashAttention ( Dao et al., 2022 ) to improve memory usage and speed on the self-attention layers. Our version is on par with or better than the original on all cases considered, while covering more use-cases and hardware. Due to the GPU hardware specifics, the efficiency is best when the embedding dimension per head is a multiple of 64, and the matrix operations are even better when the full embedding dimension is a multiple of 256. As a consequence, our ViT-g architecture slightly differs from the architecture proposed by Zhai et al. (2022) in order to maximize compute efficiency, and we use an embedding dimension of 1536 with 24 heads (64 dim/head), rather than 1408 with 16 heads (88 dim/head). Our experiments did not show significant differences in final accuracy, and our ViT-g backbone counts 1.1B parameters.
##### Sequence packing.
The DINO algorithm requires forwarding both large crops (at resolution 224) and small crops (resolution 98). When split into patches, these two groups are represented by token sequences of different lengths and cannot be forwarded together. In order to accelerate training, we use a trick called "sequence packing," which originates from NLP ( Krell et al., 2022 ) . The idea is simple: we concatenate the sequences we must forward through the transformers into a single long sequence. We pass this sequence through the transformer blocks as usual. However, a block-diagonal mask is applied to the self-attention matrix in attention layers, preventing attention between different sequences. This way, the forward is strictly equivalent to forwarding each sequence separately. This trick gives us significant compute efficiency gains compared to using separate forward and backward passes, as in prior implementations. The lower-level components of our setup are available in the xFormers library 2 2 2 https://github.com/facebookresearch/xformers ( Lefaudeux et al. (2022) ).
##### Efficient stochastic depth.
We implement an improved version of stochastic depth ( Huang et al., 2016 ) that skips the computation of the dropped residuals rather than masking the result. This saves memory and compute in proportion approximately equal to the drop rate, thanks to specific fused kernels. With high drop rates ( d = 40 % d=40\% in this work), this allows a drastic improvement in compute efficiency and memory usage. The implementation consists of randomly shuffling the B B samples over the batch dimension, and slicing the first ( 1 − d ) × B (1-d)\times B samples for the computations in the block.
##### Fully-Sharded Data Parallel (FSDP).
Minimizing our objective with the AdamW optimizer requires 4 model replicas in float32 precision – student, teacher, optimizer first moments, optimizer second moments. This sums to 16 GB 16~\mathrm{GB} of memory for a billion-parameter model such as our ViT-g. In order to reduce this memory footprint per GPU, we split the model replicas across GPUs, i.e., sharding 16 GB 16~\mathrm{GB} across GPUs using the PyTorch implementation of FSDP. Consequently, the model size is not bounded by the memory of a single GPU but by the total sum of GPU memory across compute nodes. The Pytorch implementation of FSDP brings a second advantage, which is to save on the cross-GPU communication costs: the weight shards are stored in float32 precision as required by the optimizer, but broadcasting weights and reducing gradients is done in float16 precision for the backbone (MLP heads gradients are reduced in float32 to avoid training instabilities). This leads to approximately 50% reduction in communication costs compared to the float32 gradient all-reduce operation used in DistributedDataParallel (DDP), which is used in other self-supervised pretraining methods ( Caron et al., 2021 ; Zhou et al., 2022a ) . As a consequence, the training procedure scales more efficiently than DDP with float16 autocast when scaling the number of GPU nodes. Overall, Pytorch-FSDP mixed-precision is superior to DDP with autocast in virtually all cases we encountered.
##### Model distillation.
Most of our technical improvements to the training loop aim at improving the training of large models over large quantities of data. For smaller models, we distill them from our largest model, the ViT-g, instead of training them from scratch. Knowledge distillation ( Hinton et al., 2014 ) aims at reproducing the output of a large model with a smaller model by minimizing some distance between both outputs for a set of given inputs. Since our objective function is a form of distillation from the teacher network to the student network, we leverage the same training loop with a few exceptions: we use a larger model as a frozen teacher, keep a spare EMA of the student that we use as our final model, remove the masking and stochastic depth, and, apply the iBOT loss on the two global crops. In our ablations, we observe that this approach achieves better performance than training from scratch, even for a ViT-L. Our distillation method ends up close to the one described by Duval et al. (2023) , except we do not modify the loss terms for distillation and evaluate the EMA of the student.
## 6 Ablation Studies
We present a set of ablations to empirically validate different components of our pipeline: the technical modifications described in Sec. 4 , the pretraining data and the impact of model distillation. We consider various downstream tasks that are described in Sec. 7 .
### 6.1 Improved Training Recipe
Our approach improves over the iBOT method by combining it with several existing components described in Sec. 4 . To evaluate their importance, we train multiple models where we successively add components to a baseline iBOT model. We report the Top-1 accuracy on the validation set of ImageNet-1k with a k-NN and a linear probe in Table 1 . Generally, we observe that each component improves the performance on either k-NN or linear probing and even both in most cases. Only LayerScale and Stochastic Depth incur a performance drop in linear probing but significantly improve the training stability in our experience.
### 6.2 Pretraining Data Source
The quality of features is directly related to the quality of the pretraining data. In this experiment, we probe the impact of LVD-142M compared to ImageNet-22k, a commonly used pretraining dataset, or using directly raw and uncurated data. For the uncurated dataset, we randomly sample 142 142 million images from the same data source as LVD-142M. We train a ViT-g/14 on each dataset for the same number of iterations. We also include a variant of ImageNet-22k obtained by removing the synsets of ImageNet-1k (INet-22k ∖ \setminus INet-1k) for completeness. We report the comparisons in Table 2 .
The most salient observation is that training on a curated set of images works better on most benchmarks than training on uncurated data. This confirms the benefit of curating data, even in the case of self-supervised pretraining. When compared with models trained on ImageNet-22k, training on LVD-142M is also superior on all the benchmarks but ImageNet-1k. This confirms that training on a more diverse set of images improves the quality of the features in domains that are not covered by ImageNet-22k. We also see that training on our curated data increases the performances on domains that are not used for the curation process (INaturalist 2018, 2021 and Places205), proving that scale and diversity can benefit unseen domains. Overall, the conclusion of this ablation is that our dataset provides a good balance of different types of images that leads to the best performance overall.
### 6.3 Model Size and Data
We quantify the importance of scaling data with the model size in Fig. 4 . As the size of models grow, training on LVD-142M becomes more beneficial than training on ImageNet-22k. For instance, a ViT-g trained on LVD-142M matches the performance on ImageNet-1k of a model trained on ImageNet-22k while significantly outperforming it on the other benchmarks.
### 6.4 Loss Components
We validated the proposed technical improvements in Sec. 6.1 by adding them incrementally. This section analyzes the performance hit observed if we ablate specific loss terms, starting from our best-performing model. We ablate the importance of the KoLeo loss and the impact of the masked image modeling term. For both, we report performance on ImageNet-1k using a linear classifier, ADE-20k segmentation using a linear classifier, and nearest-neighbor image retrieval on Oxford-M. Table 3(a) shows the impact of using the KoLeo loss. We see that the instance retrieval performance improves by more than 8 % 8\% , confirming that this term helps spread features in the output space. At the same time, the other metrics do not suffer from this regularization. In Table 3(b) , we show the impact of using the masked image modeling term from iBOT. This term is critical for dense prediction tasks, leading to almost 3 % 3\% performance improvement.
### 6.5 Impact of Knowledge Distillation
For small architectures, we distill larger models instead of training them from scratch. We use the distillation procedure described in Sec. 5 . We evaluate the effectiveness of this approach by comparing a ViT-L/14 trained from scratch with one distilled from a ViT-g/14 over 12 benchmarks in Fig. 5 . We also report the performance of the ViT-g/14 used for distillation as a topline. The distilled model outperforms the one trained from scratch on all 12 benchmarks, validating our pretraining approach for small models.
### 6.6 Impact of Resolution
We measure the impact of changing the resolution during the pretraining on the performance of image and patch-level features. We consider models trained from scratch using a fixed resolution of either 224 × 224 224\times 224 or 416 × 416 416\times 416 , and a model trained from scratch at 224 × 224 224\times 224 , then resumed for 10k more iterations at 416 × 416 416\times 416 . High-resolution training is compute-intensive, so we conduct this ablation on a small setup: a ViT-L/16 trained on ImageNet1k. In Fig. 6 , we report the performance of a linear probe on ImageNet-1k and ADE-20k, evaluated at various resolutions. The model trained on high-resolution images performs best across resolutions, but this comes at a high cost: training at 416 416 is approximately 3 × 3\times{} more compute-intensive than training at 224 224 . On the other hand, training at high resolution for only 10k iterations at the end of the training is almost as good and only requiring a fraction of the compute. As a consequence, we include this step at the end of the training rather than training at a high resolution from scratch.
## 7 Results
In this section, we present the empirical evaluation of our models on many image understanding tasks. We evaluate both global and local image representations, on category and instance-level recognition, semantic segmentation, monocular depth prediction, and action recognition. We detail the list of benchmarks in Appendix C . The goal of this evaluation is twofold. First, we show that our self-supervised features outperform the current state of the art by a very large margin. Second, we show that they match, or surpass the performance of weakly-supervised ones on a substantial number of tasks.
Baselines. In our comparisons, we use two kinds of models as baselines. We compare to the best performing self-supervised models that are openly available. First, we run our evaluations for MAE ( He et al., 2022 ) , DINO ( Caron et al., 2021 ) , SEERv2 ( Goyal et al., 2022a ) , MSN ( Assran et al., 2022 ) , EsViT ( Li et al., 2022a ) , Mugs ( Zhou et al., 2022b ) and iBOT ( Zhou et al., 2022a ) . When several architectural variants were proposed for a given method, we report results for the one that leads to best top-1 accuracy on ImageNet-1k. Second, we report performance of open-source weakly-supervised models such as CLIP ( Radford et al., 2021 ) , OpenCLIP ( Ilharco et al., 2021 ; Cherti et al., 2023 ) , and SWAG ( Singh et al., 2022 ) . When evaluating models on ImageNet-1k, we report the performance for each of the aforementioned methods. For all other evaluations, we report the four best-performing models amongst SSL ones. Also, for reference, we report the best performing OpenCLIP-G for weakly-supervised ones.
### 7.1 ImageNet Classification
As a first evaluation, we probe the quality of the holistic image representation produced by the model on the ImageNet-1k classification dataset. We evaluate the quality of features by training a simple classifier over a frozen backbone, and do not perform finetuning of the backbone weights. Following previous work, we use a linear model for simplicity, ensuring a reproducible evaluation, despite the fact that classes may not be linearly separable. Because most SSL methods were developped using ImageNet-1k validation performance as a debugging signal, we also report the top-1 accuracy on ImageNet-ReaL and ImageNet-V2. In order to report this additional validation performance, for all models, we run the evaluation with our code. We compare our frozen features to the best publicly available SSL features in Table 4 , regardless of architecture or pretraining data. We see the components proposed in this work lead to a very significant improvement ( + 4.2 % +4.2\% ) over the previous state of the art (iBOT ViT-L/16 trained on ImageNet-22k) on linear evaluation. At the same time, we also see that the performance increase on the alternative test sets is larger for our method, indicating stronger generalization. We describe details of our linear evaluation in Appendix B.3 .
##### How far are we from weakly-supervised models?