CVHub | CVPR 2024 | 英伟达发布新一代视觉基础模型: AM-RADIO = CLIP + DINOv2 + SAM

本文来源公众号“CVHub”，仅用于学术分享，侵权删，干货满满。

原文链接：CVPR 2024 | 英伟达发布新一代视觉基础模型: AM-RADIO = CLIP + DINOv2 + SAM

标题：《AM-RADIO: Agglomerative Vision Foundation Model Reduce All Domains Into One》
论文：https://arxiv.org/pdf/2312.06709
源码：https://github.com/NVlabs/RADIO

1 导读

AM-RADIO是什么？一张图先来感受它的魅力：

简单来说，AM-RADIO 是一个由英伟达提出的一个视觉基础模型框架，其集成了多个预训练的视觉基础模型如 CLIP、DINOv2及SAM 的能力，以获得强大的表征能力，同时也以几乎可以忽略不计的额外成本实现了SOTA级别的零样本分类或开集实例分割性能。

如上图左边所示的PCA特征可视化结果，RADIO 模型可以处理任何分辨率和宽高比，并产生语义丰富的稠密编码；上图中间则展示了 RADIO 的框架图；右侧显示了在分类、分割和视觉语言建模任务上的基准测试。

2 背景

Visual Foundation Models，VFMs，即视觉基础模型是一个非常重要的概念，诸如 CLIP、DINOv2、SAM 这样的 VFMs 通过不同的目标进行训练，已逐渐成为许多下游任务的核心。

例如，CLIP 这样的预训练视觉语言模型在不同的下游视觉任务上展现了强大的零样本泛化性能。这些模型通常使用从网络收集的数百上千万图像-文本对进行训练，并提供具有泛化和迁移能力的表示。因此，只需通过简单的自然语言描述和提示，这些预训练的基础模型完全被应用到下游任务，例如使用精心设计的提示进行零样本分类。

除了此类大型视觉语言基础模型外，一些研究工作也致力于开发可以通过视觉输入提示的大型基础模型。例如，最近 meta 推出的 SAM 能够执行与类别无关的分割，给定图像和视觉提示（如框、点或蒙版），指定要在图像中分割的内容。这样的模型可以轻松适应特定的下游任务，如医学图像分割、视频对象分割、机器人技术和遥感等。

最后，像 DINOv2，这种用于在大型图像数据集上预训练图像编码器，以获得具有语义的视觉特征。这些特征可用于广泛的视觉任务，无需微调即可获得与有监督模型相当的性能。

本文发现，尽管它们在概念上存在差异，但这些模型可以通过多教师蒸馏有效地合并成一个统一模型，称为Agglomerative Model – Reduce All Domains Into One，AM-RADIO，即聚合模型，旨在将所有领域缩减为一个。

这种整合方法不仅超越了单个教师模型的性能，而且融合了它们的独特特征，如零样本视觉-语言理解、详细的像素级理解以及开放词汇分割能力。

此外，为了追求最硬件效率高的主干网络，本文在多教师蒸馏流程中评估了多种架构，使用相同的训练策略；最终得到的新框架其性能超过了此前的SOTA模型，并且至少比教师模型在相同分辨率下快6倍，同时也在各大视觉任务取得了非常不错的性能表现。

关于“视觉大模型”的更多介绍，请参考万字长文带你全面解读视觉大模型。

3 方法

如上所述，本文提出一个框架，旨在通过多教师蒸馏从零开始训练视觉基础模型。这种方法的核心思想是利用多个已经在不同领域或任务上表现优异的教师模型来共同训练一个新的模型，这个新模型将集成所有教师模型的独特属性。

在选择的教师模型方面，作者选定了 CLIP、DINOv2 和 SAM 这三个主流的视觉基础模型，因为它们在各自的领域（如图像-文本匹配任务上，自监督学习任务上，开集分割任务）上都展现出了SOTA性能。

在训练过程中，本文没有使用额外的标签信息，而是将ImageNet、LAION-400M和DataComp1B等数据集中的图像作为训练数据。这样的做法使得模型能够在没有明确标签指导的情况下学习到图像的丰富表征。

为了评估模型的性能，作者采用了一系列度量标准，涵盖了图像级推理、像素级视觉任务、大型视觉-语言模型以及SAM-COCO实例分割等多个方面。

AM-RADIO: 多教师统一蒸馏框架

3.1 Adaptor Heads

在这一步中，作者选择了简单的设计方案，使用了一个简单的2层多层感知机（MLP），中间夹杂着 LayerNorm 和 GELU 激活函数。Head 部分的输入维度是学生模型嵌入的维度，中间维度是所有教师模型中最大的嵌入维度，输出维度与特定教师模型相匹配。对于每个教师，作者使用了两个 head，一个用于提取整体特征向量，另一个则用于保留空间特征。

3.2 Distillation Dataset

本文研究了不同数据集对下游指标的影响。虽然使用ImageNet-1K作为训练数据集可以获得最高的图像分类指标，但作者认为这并不能公平地衡量“零样本”性能，因为学生在评估域中直接学习到了教师的特征。因此，最终选择了 DataComp-1B 数据集。

3.3 Loss Formulation

上面我们提到，这里训练的时候不使用原有的GT，因此是通过选择匹配来自每个教师视觉编码器的特征。特别是，AM-RADIO 区分了每个教师的 Summary 特征向量和 Spatial 特征向量。

需要注意的是，对于CLIP和DINOv2，这里 Summary 特征向量使用的是“类 token”；而对于SAM，则不匹配此部分特征。

实验发现，与 L1、MSE、Smooth-L1 相比，余弦距离损失能够产生更好的结果。此外，通过匹配教师的空间特征来监督模型的 spatial features 不仅对下游密集任务重要，而且提高了模型的总体质量。

为此，对于匹配空间特征，作者采用了余弦相似性和Smooth L1的组合。

4 实验细节

本文架构本质上没有创新，只是对现有能力的充分利用，因此实验的细节部分算是比较有价值的，下面简单的总结罗列下，感兴趣的可以看下原文。

训练设置：
- 使用AdamW优化器、1024的批量大小、余弦退火学习率计划和学习率基数为0.001。
- 训练600k步，共查看614M个样本。
- 最好的学生模型使用DFN CLIP ViT-H/14 378px、OpenAI CLIP ViT-L/14 336px、DINOv2 ViT-g/14 224px和SAM ViTDet-H 1024px作为教师。
- 对学生和教师的输入应用随机缩放+裁剪。
- 选择DataComp-1B数据集，因为它在作者可访问的网页规模数据集中质量最高。
- 分两个阶段训练：首先用CLIP+DINOv2在256px训练300k步，然后用CLIP+DINOv2在432px加上SAM在1024px训练300k步。
学生模型架构：
- 研究了两种学生模型架构设置：标准ViT架构以匹配教师架构，以及高效架构变体优先考虑GPU上的高吞吐量。
多尺度教师：
- 学生模型选择ViT-H/16架构。
- 为了匹配SAM特征的分辨率，学生模型输入预期分辨率为1024^2。
- 由于CLIP和DINOv2教师是patch-14模型，学生输入选择432^2，这是patch-14模型378^2分辨率的有效分辨率。
- 发现插值DINOv2特征不会降低结果，因此教师以224px运行，并上采样输出以匹配学生。
排名/教师分区：
- 按照批次大小和输入分辨率将教师模型分组，然后将这些组分配给不同的GPU，以便每个GPU处理一致的批次大小和输入分辨率。
- 对于包含SAM的训练设置，使用64个GPU，其中一半处理CLIP+DINOv2组，每个GPU的批次大小为32，输入分辨率为432，另一半处理SAM，每个GPU的批次大小为2，输入分辨率为1024，有效批次大小为1,152。对于CLIP+DINOv2训练，使用32个GPU，批次大小为1024。
多分辨率ViTs：
- 许多学生模型使用ViT作为基础视觉架构。
- 使用Cropped Position Embedding (CPE)增强，位置数等于128^2。
- 即使在224分辨率下训练CLIP+DINOv2，也发现这种技术对摘要指标的影响可以忽略不计，但提高了语义分割线性探测的mIOU。
高分辨率ViT学生模型：
- 在SAM中，使用ViTDet架构减少高分辨率下ViT模型的计算和内存负担。
- 将这种架构改写为训练增强，从中采样窗口大小。
- 发现高分辨率训练不稳定，因此应用谱重参数化和0.02的权重衰减以防止注意力熵崩溃。
学生/教师分辨率不匹配：
- 当学生和教师通过处理堆栈以不同速率下采样图像时，输出特征向量的分辨率会不同。
- 对于Lfeatures，使用双线性插值输出以匹配学生和教师特征之间的较大分辨率。

5 实验

如上图所示，E-RADIO是一个高效的RADIO架构，通过融合卷积块和Transformer，以及局部和全局注意力机制，实现了对密集型预测任务的高性能和快速推理，相比全ViT架构效率显著提高。

从表1的结果可以明显的看出基于 MetaCLIP 比 OpenCLIP 效果更好，DFN CLIP 则能取得最佳的零样本分类性能。DINOv2 则有助于语义分割这种任务。

此外，所设计的 ERADIO-L 模型比所有ViT模型都要快得多。同时，它在匹配吞吐量的情况下，在大多数指标上都明显优于 MetaCLIP，并且还实现了 DINOv2 和 SAM 中缺失的零样本能力。最终，完整模型 ViT-H/16 可以表现得与教师模型一样快，但在9项任务中有6项超过了它们，这足以证明所提出的蒸馏框架是高效的。

上图展示了 RADIO 在分辨率增加时的“模式切换”。在图表中，显示了来自RADIO的DINOv2头在不同分辨率下的特征之间的均方误差（MSE），以及DINOv2在518px分辨率下实际产生的特征。通过将RADIO特征进行双线性插值以匹配DINOv2特征的分辨率。在720px时，可以看出误差突然增加，这对应于图像中完全的颜色空间变化。