Follow our official account to discover the beauty of CV technology

This article mainly introduces our work accepted at the ICLR-2025 conference: Accessing Vision Foundation Models via ImageNet-1K.

• Paper link: https://arxiv.org/abs/2407.10366
• Code link: https://github.com/BeSpontaneous/Proteus-pytorch

Existing vision foundation models such as CLIP[1], DINOv2[2], and SynCLR[3] are typically trained on large datasets (CLIP-400M, DINOv2-142M, SynCLR-600M), which not only require significant training resources but also involve datasets that are not publicly available, making it very challenging to train vision foundation models.

Our work aims to solve this problem; we can achieve performance comparable to pre-trained models provided by others by training on just 1.2M images from ImageNet-1K.

The paper, code, and all pre-trained models are now open source; we welcome everyone to communicate and try them out.

TL;DR:

• Motivation: Vision foundation models are known for their strong generalization capability, which is attributed to their large training data. However, these models require substantial training resources, and the training data is often not publicly available, such as in the cases of CLIP and DINOv2.
• Solution: We propose a very simple and general solution called Proteus, which can distill the foundation model into a smaller equivalent model on ImageNet-1K without accessing the original training data.
• Advantages: (1) Low training cost (similar to the DeiT distillation conducted on ImageNet-1K); (2) Strong performance (similar to foundation models trained with large amounts of data); (3) Excellent generalization ability (validated on DINOv2, CLIP, SynCLR).

1. Introduction

By extensively pre-training on diverse and large datasets, vision foundation models [1, 2, 3, 4] have made significant progress in the field of computer vision, aiming to learn comprehensive and versatile visual features that generalize well to various downstream tasks such as classification and segmentation. Therefore, vision foundation models are becoming fundamental components in computer vision research.

Despite these models having released their weights for public use, training foundation models remains difficult for most researchers due to two main factors: (1) The training data for these foundation models is rarely made public. Although there have been attempts to reproduce CLIP [1] using alternative datasets [4], the private nature of data sources has limited exploration into reproducing the training of foundation models like DINOv2 [2] and SynCLR [3]. (2) Even when training data is accessible, training on these large datasets requires significant computational resources, which are difficult for most researchers to obtain. ImageNet-1K [5] has long been a cornerstone of progress in supervised learning, but in the era of foundation models, its relatively ‘small’ size is now less frequently used as a training set. In this work, we attempt to address the following question: Can we reproduce the success of vision foundation models on a smaller dataset (like ImageNet-1K) without sacrificing generalization ability?

Intuitively, utilizing the pre-trained weights of these foundation models is crucial for achieving this task, for example, structured pruning [6, 7]. However, structured pruning requires fine manual design and does not easily generalize to arbitrary architectures, making it difficult to meet the diverse needs of real-world applications. To pursue a more general design, we boldly choose knowledge distillation as the method to achieve this goal, transferring the rich knowledge from the foundation model to a randomly initialized student network. However, unlike methods like TinyCLIP [8], we choose to train on a smaller dataset—ImageNet-1K—rather than the giant datasets of the original foundation model.

Knowledge transfer on ImageNet-1K still presents two key challenges: (1) The exact distribution of those private datasets (e.g., WIT400M [1], LVD-142M [2]) is unknown, and there is likely a distribution shift between ImageNet-1K and these large-scale datasets. This poses a significant challenge to the generalization ability of the target model, as the network tends to memorize training images in a fixed pattern, leading to Dataset Bias [9, 10]. (2) Most vision foundation models [1, 2, 3, 4] are trained through self-supervised learning objectives, which require large amounts of data to be effective. Therefore, directly adopting their optimization strategies in our context may not yield the best results. To address the aforementioned challenges, we propose a simple, general distillation framework called Proteus, which effectively transfers its rich knowledge to the target network by simulating the behavior of vision foundation models.

2. Method

In this section, we will introduce Proteus, a simple and general framework for training vision foundation models on ‘limited’ data (i.e., ImageNet-1K). We first discuss the efforts made to mitigate Dataset Bias so that Proteus can effectively transfer its general representations by mimicking the behavior of pre-trained foundation models. Then, we propose a Proxy Task containing multi-level learning objectives to ensure the model’s applicability across various tasks.

2.1 Proxy Dataset

In the conventional knowledge distillation setup, KL divergence Loss is typically introduced to compute the similarity between the student network’s predictions and the teacher network’s predictions, supplemented by Cross-Entropy Loss to calculate the match between the model’s predicted probability distribution and the dataset’s One-hot Label, guiding the model’s optimization. Empirically, this design performs well in supervised learning scenarios, as it shows excellent performance on ImageNet-1K.

However, we believe this setup hinders knowledge transfer in the following two aspects: (1) Cross-Entropy Loss utilizes information from One-hot Labels, which may lead to Dataset Bias, as the model tends to memorize training images and categories. This memorization makes it difficult for the model to generalize to unseen categories in downstream evaluations. (2) The generation of Class logits implicitly introduces Dataset Bias, as intermediate features are projected onto a predefined dimension, such as the 1000 dimensions of ImageNet-1K, which may be discarded in downstream evaluations. Based on these considerations, we perform knowledge distillation before the Projection head (fully connected layer), utilizing intermediate features for knowledge transfer [11].

2.2 Proxy Task

Foundation models like DINOv2 [2] aim to learn general visual features, excelling not only in high-level classification tasks but also in dense prediction tasks like semantic segmentation. To maximize knowledge transfer capabilities and ensure applicability across various tasks, we conduct distillation on three different levels of training objectives (i.e., token-level, patch-level, and feature-level), transferring rich knowledge by simulating the teacher model’s behavior.

Token-level Objective: To learn discriminative features for high-level understanding, we minimize the L2 distance to align the classification tokens between the teacher model and the student model.

Feature-level Objective: Although the token-level learning objective can obtain discriminative visual features as a Proxy Task, it does not guarantee good performance in dense prediction tasks such as semantic segmentation or depth estimation. To address this issue, we perform feature-level knowledge transfer by minimizing the L2 distance between the features of the teacher model and the student model in a similar manner.

Patch-level Objective: To further mine the hidden knowledge within the foundation model, we draw inspiration from Masked Image Modeling [12, 13, 14] to construct a patch-level learning objective. Given an image, we generate an additional view where some patches are randomly masked, and then send it to the student network to generate intermediate features, restoring the masked areas by minimizing the L2 distance between the patches of the teacher model and the student model.

3. Empirical Validation

We pre-trained on the training set of ImageNet-1K, which contains approximately 1.2 million images distributed across 1000 categories. By default, Proteus is distilled from a foundation model with the same patch size. Following the settings of DINOv2 and SynCLR, we evaluate our method on classification tasks (ImageNet-1K and 12 fine-grained classification datasets) as well as dense prediction tasks (semantic segmentation and depth estimation).

3.1 Accessing DINOv2

DINOv2 is trained on a private large-scale dataset LVD-142M, and we utilize the pre-trained DINOv2 as the teacher model to train a randomly initialized network on ImageNet-1K.

3.1.1 Target Model: ViT-S

Proteus-S significantly outperforms other baseline methods across various tasks while being only slightly behind the Oracle method DINOv2-S with far fewer training data.

3.1.2 Target Model: ViT-B and ViT-L

As we scale up the model, the performance gap between Proteus and the Oracle method DINOv2 narrows. Proteus-L nearly matches the performance of DINOv2-L across various tasks.

3.1.3 Comparison with Distillation in Supervised Learning

Proteus outperforms traditional supervised training in several aspects at similar costs, providing a novel training scheme enhanced by foundation models.

3.2 Accessing SynCLR and CLIP

We test the generalization ability of Proteus by using other foundation models SynCLR and CLIP as teacher networks. SynCLR is trained through contrastive learning on a private dataset of 600 million synthetic images, while CLIP is obtained through contrastive learning on the private dataset WIT-400M with images and corresponding text descriptions.

3.3 Ablation on Proxy Dataset

3.3.1 Dataset Diversity

By increasing the diversity of the Proxy Dataset, the generalization ability of Proteus can be improved. Even in extreme cases where only a single image is used as the Proxy Dataset, Proteus still demonstrates strong robustness.

3.3.2 Scaling Behavior

When we subsample a portion of data from each category or subsample a portion of categories from the total of 1000 categories, Proteus still exhibits strong robustness. This indicates that it is possible to train foundation models even on smaller data scales.

4. Application

First, Proteus comprehensively surpasses supervised learning methods across all metrics, demonstrating its potential to replace traditional supervised learning training paradigms. Additionally, our work supports research on model compression, allowing for the compression of foundation models at a lower cost. Furthermore, we showcase the possibility of training foundation models on datasets smaller than ImageNet-1K, which may also be a direction for future exploration.

Moreover, although our work primarily focuses on pure vision foundation models with image modalities, we hope our work inspires exploration of this idea in large language models (LLMs) and large multimodal models (LMMs) to promote research in the era of foundation models.

Reference

[1] Radford A, Kim J W, Hallacy C, et al. Learning transferable visual models from natural language supervision[C]//International conference on machine learning. PMLR, 2021: 8748-8763.[2] Oquab M, Darcet T, Moutakanni T, et al. Dinov2: Learning robust visual features without supervision[J]. arXiv preprint arXiv:2304.07193, 2023.[3] Tian Y, Fan L, Chen K, et al. Learning vision from models rivals learning vision from data[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024: 15887-15898.[4] Cherti M, Beaumont R, Wightman R, et al. Reproducible scaling laws for contrastive language-image learning[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023: 2818-2829.[5] Deng J, Dong W, Socher R, et al. Imagenet: A large-scale hierarchical image database[C]//2009 IEEE conference on computer vision and pattern recognition. Ieee, 2009: 248-255.[6] Ma X, Fang G, Wang X. Llm-pruner: On the structural pruning of large language models[J]. Advances in neural information processing systems, 2023, 36: 21702-21720.[7] Xia M, Gao T, Zeng Z, et al. Sheared llama: Accelerating language model pre-training via structured pruning[J]. arXiv preprint arXiv:2310.06694, 2023.[8] Wu K, Peng H, Zhou Z, et al. Tinyclip: Clip distillation via affinity mimicking and weight inheritance[C]//Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023: 21970-21980.[9] Torralba A, Efros A A. Unbiased look at dataset bias[C]//CVPR 2011. IEEE, 2011: 1521-1528.[10] Liu Z, He K. A Decade’s Battle on Dataset Bias: Are We There Yet?[J]. arXiv preprint arXiv:2403.08632, 2024.[11] Romero A, Ballas N, Kahou S E, et al. Fitnets: Hints for thin deep nets[J]. arXiv preprint arXiv:1412.6550, 2014.[12] Bao H, Dong L, Piao S, et al. Beit: Bert pre-training of image transformers[J]. arXiv preprint arXiv:2106.08254, 2021.[13] He K, Chen X, Xie S, et al. Masked autoencoders are scalable vision learners[C]//Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022: 16000-16009.[14] Zhou J, Wei C, Wang H, et al. ibot: Image bert pre-training with online tokenizer[J]. arXiv preprint arXiv:2111.07832, 2021.

For the latest AI progress reports, please contact: [email protected]

END

Welcome to join the “Computer Vision“ group👇 Please note:CV