Simplifying DINO via Coding Rate Regularization

TL;DR

SimDINO and SimDINOv2 are highly simplified variants of the DINO and DINOv2 which are more robust to different design choices, such as network architecture and hyperparameters, and they learn even higher-quality representations, measured by performance on downstream tasks, offering a Pareto improvement over the corresponding DINO and DINOv2 models.

Abstract

DINO and DINOv2 are two model families being widely used to learn representations from unlabeled imagery data at large scales. Their learned representations often enable state-of-the-art performance for downstream tasks, such as image classification and segmentation. However, they employ many empirically motivated design choices and their training pipelines are highly complex and unstable -- many hyperparameters need to be carefully tuned to ensure that the representations do not collapse -- which poses considerable difficulty to improving them or adapting them to new domains. In this work, we posit that we can remove most such-motivated idiosyncrasies in the pre-training pipelines, and only need to add an explicit coding rate term in the loss function to avoid collapse of the representations. As a result, we obtain highly simplified variants of the DINO and DINOv2 which we call SimDINO and SimDINOv2, respectively. Remarkably, these simplified models are more robust to different design choices, such as network architecture and hyperparameters, and they learn even higher-quality representations, measured by performance on downstream tasks, offering a Pareto improvement over the corresponding DINO and DINOv2 models. This work highlights the potential of using simplifying design principles to improve the empirical practice of deep learning.

Figures (5)

• Figure 1: The DINO DINO and DINO DINO v2 pipelines are substantially simplified to the respective SimDINO SimDINO and SimDINO SimDINO v2 pipelines. (a) In the DINO DINO pipeline, an input image is turned into patches. Then a global view $v_{g}$ and a local view $v_{c}$ are randomly sampled. The global view is pushed through the teacher encoder, while the other view is through the student encoder. (b) The SimDINO SimDINO pipeline removes the need for expensive post-processing operations present in DINO DINO, such as a dimension-increasing linear layer and a high-dimensional softmax. (c) The DINO DINOv2 pipeline adds masking (here masked patches are denoted by $\times$) and an additional loss on image patch features to the DINO DINO pipeline. (d) The SimDINO SimDINOv2 training operates directly on the learned representations, simplifying the pipeline.
• Figure 2: Evolution of $k$-NN accuracy of ViT-B trained for 100 epochs using DINO and SimDINO paradigms on ImageNet-1K. We omit earlier epochs of similar metrics for better visual clarity.
• Figure 3: Visualization of MaskCut segmentation results from DINO ViT-B/16 (row 1), SimDINO ViT-B/16 (row 2) and SimDINO ViT-L/16 (row 3) on selected images.
• Figure 4: $k$-NN accuracy on ImageNet-1K of ViT-B trained on COCO train2017 using DINO and SimDINO paradigms.
• Figure 5: Visualization of average self-attention maps obtained from both DINO DINO(v2) and SimDINO SimDINO(v2) algorithms.

Theorems & Definitions (4)

• Theorem 3.1: Scale of $\nabla R_{\varepsilon}$
• proof
• Remark 3.2
• Remark 3.3