NVIDIA团队！AM-RADIO：聚合视觉基础模型，将所有领域归为一个-CSDN博客

点击下方卡片，关注“具身智能之心”公众号

作者 | 具身智能之心编辑 | 具身智能之心

本文只做学术分享，如有侵权，联系删文

写在前面&出发点

近年来，少数视觉基础模型（VFMs）已成为众多下游任务的核心支撑。CLIP、DINOv2、SAM等VFMs以不同的目标进行训练，展现出适用于各种下游任务的独特特性。我们发现，尽管这些模型在概念上存在差异，但可以通过多教师蒸馏有效地合并为一个统一模型。这里将这种方法命名为AM-RADIO（聚合模型：将所有领域归约为一个）。这种集成方法不仅超越了单个教师模型的性能，还融合了它们的独特功能，如零样本视觉语言理解、详细的像素级理解以及开放词汇分割能力。此外，为了追求最高效的硬件支撑架构，使用相同的训练方案在多教师蒸馏pipeline中评估了多种架构。这促使我们开发了一种新型架构（E-RADIO），其性能超越了前代，且在相同分辨率下，运行速度至少比教师模型快6倍。我们的综合基准测试过程涵盖了ImageNet分类、语义分割线性probing、COCO目标检测以及集成到LLaVa-1.5等下游任务。

代码部分：https://github.com/NVlabs/RADIO

内容出自国内首个具身智能全栈学习社区：具身智能之心知识星球，这里包含所有你想要的。

一些行业介绍

知识蒸馏是一种非常成功且流行的技术，用于将“教师”模型（或模型集合）的知识转移到通常较小的“学生”模型中。在原始公式中，学生模型和教师模型都在相同的域内数据集上操作，学生模型同时匹配教师的逻辑输出（logits）和真实标签。除了使用带标签的图像外，另一种方法是训练学生模型以匹配教师模型的特征。

一些工作没有使用较小的学生模型，而是采用了一个高容量模型的迭代学习过程，其中训练一个容量等于或大于教师的学生模型，并对学生模型进行大量的数据增强。训练完成后，他们通过使用训练好的学生模型对新数据进行伪标签标注来扩展数据集。然后，他们让学生模型成为教师模型，并重复这个过程。这项工作的一个重要发现是，学生模型能够超越教师模型的性能。

一些工作探索了集成蒸馏的概念，其中存在多个教师模型，每个教师模型都具有有限的领域知识。有些工作提供了多教师蒸馏的概述，并提出了一种观点，即学生模型不必匹配教师模型集合的总结，而是可以通过从学生模型的表示空间到每个教师模型的某些学习到的非共享映射来匹配每个教师模型的特征。他们方法的一个有趣之处在于，学生模型和教师模型不需要具有相同的架构，而且分别处理每个教师模型可以提高性能。

最近，基础模型（FMs）的概念应运而生，人们普遍认为这些模型体积庞大、通用性强且训练成本高昂。通过在大规模数据集上进行训练，它们能够广泛应用于众多下游任务。此类模型的一个开创性例子是CLIP，它在网络规模的弱监督（图像，标注）对上进行训练，并在广泛的计算机视觉基准测试中实现了卓越的零样本性能。虽然CLIP无疑是基础模型，但另一个具有广泛能力的模型DINOv2也崭露头角，在需要强大空间特征的密集任务（如ADE20k和Pascal VOC）上经常超越CLIP。此外，SAM（Segment Anything）因其出色的开放词汇实例分割能力而日益受到欢迎，我们推测其视觉编码器具有强大的密集特征表示能力。

这里引入AM-RADIO的目标是从多个基础模型同时学习。我们观察到，当给定一个容量足够的学生模型时，它通常能够在重要指标上超越其任何一位教师。除了在具有代表性的基础基准测试中表现良好外，凭借训练框架的优势，我们的学生模型能够模仿其教师模型，因此能够执行原本由教师执行的下游任务。这方面的例子包括CLIP的零样本应用（因为CLIP训练的语言模型与学生的语言模型兼容）以及“分割任何内容”任务（因为学生能够替换视觉编码器并与已经训练过的掩码解码器进行交互）。

我们还研究了使用更高效硬件的模型架构的影响。由于大多数关于效率的工作即使在同一数据集（如ImageNet-1k）上进行评估时使用的训练方案也不同，因此它们通常无法直接比较，而且可能存在过度调优的情况。为此，我们在相同的训练方案下评估了10多种有前景的架构，以进行直接比较。发现类似卷积神经网络（CNN）的架构速度更快，但在提炼视觉Transformer（ViT）基础模型时遇到困难。这促使我们开发了一种新型混合架构E-RADIO，其性能超越了前代，并且在相同分辨率下，运行速度至少比教师模型快6倍。

主要贡献如下：

提出了一种通用的方法，用于将多个不同的基础模型蒸馏为一个模型，其中包括输入分辨率不兼容的模型；
证明了这些学生模型能够在具有代表性的基准测试中超越其教师模型的表现；
证明了这些学生模型可以无缝替换其教师模型，或者其特征可以直接用于下游应用，例如为LLaVA提供视觉编码。
对多种高效架构进行了基准测试，并提出了一种新的架构（E-RADIO），该架构能够在显著提升速度的同时保持相似的模型质量。

知识聚合

这里提出了一个框架，如图2所示，通过多教师蒸馏从头开始训练一个视觉基础模型。我们证明了每位教师都为基础视觉模型带来了独特的属性，而训练得到的模型将聚合这些属性。

作为初步假设，我们期望教师模型能够表示在互联网上发现的广泛图像，这些图像来自诸如ImageNet（1k或21k）、LAION-400M或DataComp1B等数据集。考虑到这一点，我们选择了三个具有开创性的教师模型家族进行研究：CLIP、DINOv2和SAM，因为它们在广泛的任务上（如CLIP）表现出了卓越的性能，或在下游密集任务（如线性探测下的语义分割，DINOv2）或开放词汇分割（如SAM）上特别强大。由于这些教师模型来自如此不同的领域，我们省略了任何形式的补充真实值指导，并仅将上述数据集视为图像来源。

为了评估模型的质量，在几个广泛领域内采用了一组代表性指标。

图像级推理：（i）在ImageNet-1K上的k-NN Top-1准确率，以及（ii）使用CLIP教师模型的语言模型的零样本准确率。k-NN将训练集中每张图像的模型summary特征向量进行嵌入，然后对于每张验证图像，它使用k个最近训练向量的加权和来选择标签。
像素级视觉任务：在（i）ADE20K和（ii）Pascal VOC上的分割mIOU。
大型视觉语言模型：将冻结的视觉编码器模型插入到LLaVA-1.5中，并在包括GQA、TextVQA、ScienceQA和VQAv2在内的广泛任务上对其进行了评估。
SAM-COCO实例分割：采用了COCO实例分割方法来评估复制SAM视觉特征的能力。

这些任务的结果，无论是教师模型还是我们的AMRADIO变体，都总结在表1中。

适配器Head。在设计适配器Head时，选择了简洁性，并将其他架构作为未来工作的一部分。为此，我们采用了一个简单的两层多层感知器（MLP），并在中间加入了层归一化（LayerNorm）和高斯误差线性单元（GELU）。输入维度是学生嵌入的维度，中间维度是所有教师嵌入的最大维度，输出维度与特定教师的维度相匹配。对于每位教师，我们采用两个Head，一个用于summary向量，另一个用于空间特征。

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

损失公式：由于没有每个教师对每张图像的真实数据，因此选择匹配来自每个教师视觉编码器的特征。特别是区分了每个教师的summary特征向量和空间特征向量。summary特征的计算方式因模型而异。对于CLIP和DINOv2，使用“类别标记”作为summary特征向量，而对于SAM，不匹配summary。

设为学生视觉编码器，其参数为Θ0，而为学习的学生头，用于匹配教师summary特征，其中学生适配器参数为，教师参数为Φi。

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

为了匹配空间特征，结合了余弦相似度和Smooth L1。类似于我们在方程（2）中发现余弦相似度产生最佳结果的情况，我们发现空间特征也是如此。然而，我们希望学生模型能够作为教师框架中的即插即用替代品，因此匹配教师向量的幅度非常重要，所以我们引入了Smooth L1。在（3）中，我们展示了这种损失的公式。设为用于匹配教师特征向量的学习学生Head，相应的为教师特征向量，其中，则空间特征损失为：

选择α为0.9，β为0.1，主要是基于余弦距离在实证中表现更优，但同时也为了匹配向量的幅度。

Loss Balancing：由于不同教师之间的损失权重存在多种可能的组合，甚至包括选择哪些教师以及损失函数的可能形式，我们主要选择了对所有教师的空间特征进行等权重（γi = 1）的简单损失平衡。对于摘要特征，我们设置λCLIP = λDINO = 1，λSAM = 0。我们确实尝试了使用预测不确定性、AdaLoss（动量0.99）以及单独使用AMTML-KD进行自动损失平衡，作为学习λi和γi平衡的方法。在使用AMTML-KD的情况下，模型总是会将全部权重集中在CLIP教师上，导致结果比简单的手动平衡更差。根据表4的结果，更复杂的平衡方案几乎没有优势，因此我们在本文的其余部分选择了“简单”方法。

一些细节

执行异构多教师知识蒸馏并非易事，因为特征维度、输入分辨率、损失计算的概念、下采样比例之间存在不匹配，以及将多个教师模型适配到单个GPU上存在的挑战。

一般情况。使用AdamW优化器、bs大小为1024、余弦退火学习率调度和0.001的基础学习率来训练所有学生模型。训练了60万 steps，总共处理了6.14亿个样本。对于我们的最佳学生模型，使用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分辨率下训练30万步，第二阶段使用CLIP+DINOv2在432px分辨率下以及SAM在1024px分辨率下再训练30万步。

学生模型架构。我们研究了学生模型架构的两种设置：

标准ViT架构，以匹配教师的架构。我们的最佳模型是ViT-H/16。
高效架构变体，优先在GPU上实现高吞吐量。

多尺度教师。我们为学生模型选择了ViT-H/16架构。为了匹配SAM特征的分辨率，我们输入预期的1024x1024分辨率。鉴于我们的CLIP和DINOv2教师是patch-14模型，选择向学生输入432²的图像，因为这与patch-14的378x378具有相同的有效分辨率。发现对DINOv2特征进行插值不会降低结果质量，因此教师模型在224px分辨率下运行，我们将输出上采样以匹配学生模型的分辨率。

rank/教师模型分区。根据（bs大小，学生模型分辨率）对教师模型进行分组，然后将这些组分配到不同的GPU上，以便每个GPU处理相同的批量大小和输入分辨率。我们还以不同的速率对组进行采样，在包含SAM的训练设置中，使用64个GPU进行训练，其中一半GPU处理CLIP+DINOv2组，每个GPU的批量大小为32，输入分辨率为432；另一半GPU处理SAM，每个GPU的批量大小为2，输入分辨率为1024。这导致有效批量大小为1152。对于CLIP+DINOv2训练，我们使用32个GPU，批量大小为1024。

多分辨率ViT。我们的许多学生模型使用ViT作为基础视觉架构。传统上，ViT会为图像中的每个输入patch学习一个位置嵌入，这反过来又强制模型始终以恒定分辨率运行。我们采用了裁剪位置嵌入（CPE）增强技术，其中位置数量等于128x128。然后，位置嵌入会随机裁剪并插值以匹配学生模型的输入patch数量。即使在以224分辨率训练CLIP+DINOv2时，也发现这种技术导致的summary指标下降可忽略不计，但改善了语义分割线性探测的平均交并比（mIOU）。对于具有不同分辨率的学生模型，这是一种无缝技术，允许ViT在某一范围内的任意分辨率下运行。除了支持任意分辨率外，如图3所示，与其他ViT模型相比，CPE还减少了位置嵌入中的噪声伪影。

高分辨率ViT学生模型。在SAM中，他们采用ViTDet架构来减轻高分辨率下ViT模型的计算和内存负担。我们将此架构重新构想为一种训练增强方法，即从一组可能的窗口大小中采样一个窗口大小。这使我们能够减轻使用SAM教师训练学生模型的计算负担，并且由于我们使窗口大小变得灵活，它在推理期间提供了额外的吞吐量缩放机制。表8展示了我们替换SAM编码器的能力。另外，我们发现高分辨率训练不稳定，因此我们应用了谱重参数化和0.02的权重衰减来防止注意力熵崩溃。

学生/教师分辨率不匹配。当学生和教师以不同的速率通过其处理堆栈对图像进行下采样时，会导致输出特征向量具有不同的分辨率。例如，如果教师使用ViT-H/14架构而学生使用ViT-H/16，则意味着学生输出14x14的特征图，而教师输出16x16的特征图。对于L特征，我们通过双线性插值将输出匹配到学生和教师特征之间较大的分辨率。

特征Summarization。ViT有两种选择：(i)一个单独的总结“CLS”标记或(ii)平均池化补丁标记。在表6中评估了这两种选项。可以观察到，平均池化可以改善总结损失，但对特征损失有更大的不利影响。鉴于后者的重要性，我们选择使用单独的CLS标记。

结果分析对比

参考

[1] AM-RADIO: Agglomerative Vision Foundation Model Reduce All Domains Into One

“具身智能之心”公众号持续推送具身智能领域热点：

【具身智能之心】技术交流群

具身智能之心是首个面向具身智能领域的开发者社区，聚焦大模型、机械臂、双足机器人、四足机器人、感知融合、强化学习、模仿学习、规控与端到端、仿真、产品开发、自动标注等多个方向，目前近60+技术交流群，欢迎加入！扫码添加小助理微信邀请入群，备注：学校/公司+方向+昵称（快速入群方式）

【具身智能之心】知识星球

具身智能之心知识星球是国内首个具身智能开发者社区，主要关注具身智能相关的数据集、开源项目、具身仿真平台、大模型、视觉语言模型、强化学习、具身智能感知定位、机器臂抓取、姿态估计、策略学习、轮式+机械臂、双足机器人、四足机器人、大模型部署、端到端、规划控制等方向。扫码加入星球，享受以下专有服务：