ICLR 2025｜如何在ImageNet-1K上训练视觉基础模型？

前言仅在ImageNet-1K的1.2M图像上训练就可以在多项任务上达到跟别人所提供的预训练模型相当的性能。

欢迎关注公众号CV技术指南，专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读、CV招聘信息。

本文转载自极市平台

仅用于学术分享，若侵权请联系删除

本文主要介绍我们最近被ICLR2025接收的工作：Accessing Vision Foundation Models via ImageNet-1K

现有视觉基础模型例如CLIP[1], DINOv2[2], SynCLR[3]通常是在巨额数据量（CLIP-400M, DINOv2-142M, SynCLR-600M）下训练得到的，这不仅对训练资源有着非常大的需求，同时这些数据集也处于未公开的状态，让训练视觉基础模型非常困难。我们本次工作就是为了解决这个问题，我们仅在ImageNet-1K的1.2M图像上训练就可以在多项任务上达到跟别人所提供的预训练模型相当的性能。论文，代码和所有预训练模型现已开源，欢迎大家交流和试用。

TL, DR:

动机：视觉基础模型因其强大的泛化能力而著称，这得益于其庞大的训练数据。然而，这些模型需要巨大的训练资源，而且训练数据通常是未公开的，例如CLIP和DINOv2。

解决：我们提出了一个非常简单且通用的解决方案，名为Proteus，可以在不访问原始训练数据的情况下，将基础模型在ImageNet-1K上蒸馏成较小的等效模型。

优点： (1) 低训练成本（类似于在ImageNet-1K上进行的DeiT蒸馏）；(2) 强大的性能（类似于使用大量数据训练的基础模型）；(3) 优秀的泛化能力（在DINOv2、CLIP、SynCLR上验证）。

1. Introduction

通过在多样且庞大的数据集上进行广泛的预训练，视觉基础模型[1，2，3，4]在计算机视觉领域取得了显著进展，旨在学习全面且多功能的视觉特征，能够很好地泛化到各种下游任务，如分类、分割等。因此，视觉基础模型正成为计算机视觉研究中的基础组件。

尽管这些模型已经发布了其权重供公众使用，但由于两个主要因素，训练基础模型对于大多数研究人员来说仍然难以实现：（1）这些基础模型的训练数据很少公开。尽管已经有尝试使用替代数据集[4]来重现CLIP [1]，但由于数据源私密，重现DINOv2 [2]和SynCLR [3]等基础模型的训练仍然鲜有探索。（2）即使训练数据可以获取，使用这些庞大的数据集进行训练需要大量的计算资源，这对于大多数研究人员来说是难以获取的。ImageNet-1K [5]，长期以来一直是监督学习领域进步的基石，但在基础模型时代，由于其相对较“小”的规模，现在较少被用作训练集。在这项工作中，我们试图解决以下问题：是否可以在不牺牲泛化能力的情况下，在更小的数据集（如ImageNet-1K）上重现视觉基础模型的成功？

直观上，利用这些基础模型的预训练权重对于完成这一任务至关重要, 例如结构化剪枝[6,7]。但结构剪枝需要精细的手工设计，无法轻易泛化到任意架构，使其难以满足现实世界中多样化的需求。为了追求更通用的设计，我们大胆选择知识蒸馏作为实现这一目标的方法，即将基础模型中丰富的知识转移到一个随机初始化的学生网络。但与TinyCLIP[8]等方法不同的是，我们选择在规模更小的数据集——ImageNet-1K上进行训练，而不是采用原基础模型的巨型数据集。

在ImageNet-1K上的知识迁移仍然存在两个关键问题：(1) 那些未公开数据集（例如WIT400M [1]，LVD-142M [2]）的确切分布未知，并且很可能ImageNet-1K和这些大规模数据集之间存在分布偏移。这对目标模型的泛化能力构成了显著挑战，因为网络倾向于以固定模式记忆训练图像，导致Dataset Bias[9,10]。(2) 大多数视觉基础模型[1, 2, 3, 4]是通过自监督学习目标训练的，这需要大量的数据才能有效。因此，直接采用它们的优化策略在我们的环境中可能不会产生最佳结果。为解决上述挑战，我们提出了一个简单通用的蒸馏框架，Proteus，通过模拟视觉基础模型的行为来将其丰富的知识迁移到目标网络中。

2. Method

在本节我们将介绍Proteus，这是一种简单且通用的框架，用于在“有限”数据（即ImageNet-1K）上训练视觉基础模型。我们首先介绍了在减轻Dataset Bias方面所做的努力，以便Proteus能够通过模仿预训练基础模型的行为，有效地转移其通用的表征。然后，我们提出了包含多层次学习目标的Proxy Task，以确保模型在各种任务中的应用。

2.1 Proxy Dataset

在常规知识蒸馏的设置中，通常会引入KL divergence Loss来计算学生网络预测结果和教师网络预测结果的相似度，并辅以Cross-Entropy Loss来计算模型预测的概率分布与数据集One-hot Label的匹配程度，来指导模型的优化。从经验上看，这种设计在监督学习情况下效果良好，因为它在ImageNet-1K上表现出色。

然而，我们认为这种设置会在以下两个方面阻碍知识传递：(1) Cross-Entropy Loss利用了One-hot Label的信息，可能导致Dataset Bias，因为模型倾向于记住训练图像和类别。这种记忆使得模型在下游评估中难以对未见过的类别进行泛化。(2) Class logits的生成隐性地引入了Dataset Bias，因为中间特征被投影到一个预定义的维度上，例如ImageNet-1K的1000维，这在下游评估中可能会被丢弃。基于这些考虑，我们在Projection head（全连接层）之前进行知识蒸馏，并利用中间特征进行知识传递[11]。

2.2 Proxy Task

基础模型如DINOv2 [2]旨在学习通用的视觉特征，不仅在高层次的分类任务中表现出色，而且在语义分割等密集预测任务中也表现优异。为了最大化知识传递能力并保证其在各种任务中的应用，我们在三个不同层次的训练目标（即token-level, patch-level, and feature-level）上进行蒸馏，通过模拟教师模型的行为来传递丰富的知识。

Token-level Objective：为了学习用于高层次理解的discriminative的特征，我们最小化L2距离，以对齐教师模型和学生模型之间的classification token。

Feature-level Objective：尽管token-level的学习目标作为Proxy Task可以获得discriminative的视觉特征，但它无法保证在语义分割或深度估计等密集预测任务上取得良好表现。为了解决这个问题，我们以类似的方式，即最小化教师模型和学生模型feature的L2距离，进行feature-level的知识传递。

Patch-level Objective：为了进一步挖掘基础模型中的隐藏知识，我们借鉴Masked Image Modeling [12, 13, 14] 的思想，构建了一个patch-level的学习目标。给定一个图像，我们会生成一个额外的视图，其中部分patch被随机遮掩，然后将其送到学生网络以生成中间特征，并通过最小化教师模型和学生模型patch的L2距离来恢复被遮掩的区域。

3. Empirical Validation

我们在ImageNet-1K的训练集上进行预训练，该数据集包含约120万张图像，分布在1000个类别中。默认情况下，Proteus 是从具有相同patch大小的基础模型中蒸馏出来的。按照DINOv2和 SynCLR的设置，我们在分类任务（ImageNet-1K 和12个细粒度分类数据集）以及密集预测任务（语义分割和深度估计）上评估我们的方法。

3.1 Accessing DINOv2

DINOv2 在私有的大规模数据集 LVD-142M 上进行训练，我们利用预训练的 DINOv2 作为教师模型，在 ImageNet-1K 上训练一个随机初始化的网络。

3.1.1 Target Model: ViT-S

Proteus-S 在不同任务上明显优于其他baseline方法，并且在训练数据远少于 Oracle 方法 DINOv2-S 的情况下，仅略微落后于后者。

3.1.2 Target Model: ViT-B and ViT-L

当我们扩大模型规模时，Proteus 与 Oracle 方法 DINOv2 之间的性能差距缩小。Proteus-L 在各项任务上的表现几乎与 DINOv2-L 相匹配。

3.1.3 Comparison with Distillation in Supervised Learning

Proteus 在相似的成本下，在多个方面优于传统的监督训练，提供了一种被基础模型强化的新颖训练方案。

3.2 Accessing SynCLR and CLIP

我们通过使用其他基础模型 SynCLR 和 CLIP 作为教师网络来测试 Proteus 的泛化能力。SynCLR 通过在未公开的 6 亿张合成数据集进行对比学习而训练得到，而 CLIP 是通过在私有数据集 WIT-400M 上对图像和相应的文本描述进行对比学习获得的。

3.3 Ablation on Proxy Dataset

3.3.1 Dataset diversity

如果我们增加Proxy Dataset的多样性，Proteus 的泛化能力可以得到提升。即使在只有单一图像作为Proxy Dataset的极端情况下，Proteus 仍然表现出很强的鲁棒性。

3.3.2 Scaling behavior

当我们从每个类别中子采样一部分数据或从总共1000个类别中子采样一部分类别时，Proteus 仍然表现出很强的鲁棒性。这表明即使在更小的数据规模下，也有可能训练基础模型。

4. Application

首先，Proteus 在所有指标上全面超越了监督学习方法，展示了其代替传统监督学习训练范式的潜力。此外，我们的工作支持模型压缩研究，从而能够以更小的成本压缩基础模型。进一步地，我们展示了在比ImageNet-1K更小的数据集上训练基础模型的可能性，这可能也是未来探索的一个方向。此外，尽管我们的工作主要集中在具有图像模态的纯视觉基础模型上，但我们希望我们的工作能够激励这一想法在大型语言模型（LLMs）和大型多模态模型（LMMs）的探索，以促进基础模型时代下的研究。

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.

欢迎关注公众号CV技术指南，专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读、CV招聘信息。

计算机视觉入门1v3辅导班

【技术文档】《从零搭建pytorch模型教程》122页PDF下载

QQ交流群：470899183。群内有大佬负责解答大家的日常学习、科研、代码问题。

其它文章

分享一个CV知识库，上千篇文章、专栏，CV所有资料都在这了

明年毕业，还不知道怎么做毕设的请抓紧机会了

LSKA注意力 | 重新思考和设计大卷积核注意力，性能优于ConvNeXt、SWin、RepLKNet以及VAN

CVPR 2023 | TinyMIM：微软亚洲研究院用知识蒸馏改进小型ViT

ICCV2023|涨点神器！目标检测蒸馏学习新方法，浙大、海康威视等提出

ICCV 2023 Oral | 突破性图像融合与分割研究：全时多模态基准与多交互特征学习