【DINOv2论文阅读】:无需监督的通用视觉特征提取器——机器人VLA模型的“眼睛“基石
前言如果你最近在搞具身智能或者VLA视觉-语言-动作模型那你一定绕不开DINOv2。这个来自Meta AI的视觉编码器几乎成了所有开源VLA模型的标配——OpenVLA用它Octo用它小米的Xiaomi-Robotics-0也用它。为什么大家都这么爱DINOv2因为它做到了一件前无古人的事完全不需要任何标注只靠看图片就能学到和弱监督CLIP相当甚至更好的通用视觉特征。而且这些特征拿过来就能用不需要微调简直是视觉界的万能钥匙。今天我们就来深度拆解这篇神作看看DINOv2到底是怎么炼成的。论文信息标题DINOv2: Learning Robust Visual Features without Supervision会议Transactions on Machine Learning Research (TMLR) 2024单位Meta AI Research, Inria代码github.com/facebookresearch/dinov2论文arxiv.org/abs/2304.07193一、为什么我们需要DINOv2在NLP领域大语言模型已经证明了只要给足够多的无标注文本就能学到通用的语言表示。这些表示可以直接用于各种下游任务不需要针对每个任务重新训练。但在计算机视觉领域这个梦想直到DINOv2出现才真正实现。之前的视觉预训练方法主要有两条路监督预训练用ImageNet这种标注好的数据集训练。但标注成本极高而且学到的特征只能在类似的数据集上好用泛化能力差。弱监督预训练用CLIP这种图文对训练。虽然泛化能力强了但需要大量的图文对齐数据而且文本描述只能捕捉图片的部分信息很多像素级的细节比如深度、边缘都学不到。自监督预训练只靠图片自己学习。理论上最完美但之前的方法要么只能在小数据集上训练要么用了大量未整理的网络图片结果特征质量很差。DINOv2的核心贡献就是证明了只要有足够多、足够好的无标注图片自监督学习完全可以学到和弱监督方法相当甚至更好的通用视觉特征。它训练了一个11亿参数的ViT-g模型在几乎所有视觉任务上都超过了之前的自监督方法甚至在很多任务上超过了OpenCLIP-G。二、好数据出好模型——DINOv2的数据炼金术DINOv2成功的第一秘诀就是它花了巨大的精力构建了一个高质量的无标注数据集LVD-142M。Meta的工程师们没有直接用网上爬来的12亿张杂乱无章的图片而是设计了一套精妙的数据处理流水线从这12亿张图片中精选出了1.42亿张高质量、多样化的图片。图1DINOv2的数据处理流水线出处原文Figure 32.1 数据处理三步曲数据收集先收集一批高质量的种子数据集包括ImageNet-22k、Google Landmarks等总共约1400万张图片。这些就像是老师手里的标准教材。去重用自监督模型给所有12亿张图片生成特征然后去掉重复的图片。这一步就像是把教材里的重复章节删掉避免浪费时间。检索增强对于种子数据集中的每一张图片从12亿张未标注图片中检索出最相似的4张图片。这就像是给每本标准教材配上了4本相关的参考书让模型能学到更多样的知识。2.2 数据质量有多重要我们来看一个实验对比不同数据源训练出来的模型性能训练数据ImageNet-1kImageNet-AADE-20kOxford-MImageNet-22k85.973.546.662.5随机未整理数据83.359.448.554.3LVD-142M85.873.947.764.6表1不同数据源的性能对比出处原文Table 2从表中可以清楚地看到用随机未整理的数据训练虽然数据量一样大但性能全面下降尤其是在分布外的ImageNet-A上准确率直接掉了14.1%。这就像是给学生看一堆乱七八糟的漫画虽然看得多但学不到有用的知识。用LVD-142M训练的模型在ImageNet-1k上和ImageNet-22k相当但在其他所有任务上都更好。这说明多样化的高质量数据能显著提升模型的泛化能力。三、三位一体的训练秘方——DINOiBOTKoLeo有了好数据还需要好的训练方法。DINOv2没有发明全新的自监督方法而是巧妙地结合了之前三个最优秀的方法DINO、iBOT和KoLeo再加上一些工程上的改进最终得到了远超单个方法的效果。3.1 基础框架教师-学生架构DINOv2采用了经典的教师-学生Teacher-Student架构。简单来说就是有两个一模一样的网络学生网络正常训练参数不断更新。教师网络参数是学生网络参数的指数移动平均EMA。也就是说教师网络的参数更新得很慢就像一个经验丰富的老师总是比学生更稳定、更准确。训练的目标就是让学生网络的输出尽可能接近教师网络的输出。这种架构的好处是教师网络能提供更稳定的监督信号避免学生网络在训练过程中跑偏。3.2 图像级损失DINO损失DINO损失是在图像层面的对比损失它让同一图片的不同视角裁剪、翻转等在特征空间中尽可能接近不同图片的特征尽可能远离。LDINO−∑ptlogps\mathcal{L}_{DINO} -\sum p_t \log p_sLDINO−∑ptlogpsLDINO\mathcal{L}_{DINO}LDINODINO损失函数值越小表示学生和教师的输出越接近ptp_tpt教师网络对输入图片的预测概率分布psp_sps学生网络对同一图片不同视角的预测概率分布通俗解释这就像是老师给学生看同一道题的不同解法让学生学会不管题目怎么变都能得到正确的答案。3.3 补丁级损失iBOT损失只有图像级的损失还不够因为很多下游任务比如分割、深度估计需要像素级的特征。DINOv2加入了iBOT损失在补丁patch层面进行监督。iBOT损失的思想和DINO类似但有一个关键区别学生网络的输入会随机mask掉一些补丁而教师网络的输入是完整的。训练目标是让学生网络预测出被mask掉的补丁的特征。LiBOT−∑iptilogpsi\mathcal{L}_{iBOT} -\sum_{i} p_{ti} \log p_{si}LiBOT−i∑ptilogpsiLiBOT\mathcal{L}_{iBOT}LiBOTiBOT损失函数值ptip_{ti}pti教师网络对第iii个补丁的预测概率分布psip_{si}psi学生网络对被mask掉的第iii个补丁的预测概率分布通俗解释这就像是老师给学生看一幅画然后遮住一部分让学生猜遮住的地方是什么。通过这种方式模型能学会理解图片的局部结构和上下文关系。3.4 正则化神器KoLeo损失为了让学到的特征在空间中均匀分布避免很多特征挤在一起DINOv2加入了KoLeo正则化。Lkoleo−1n∑i1nlog(dn,i)\mathcal{L}_{koleo} -\frac{1}{n} \sum_{i1}^{n} \log \left(d_{n,i}\right)Lkoleo−n1i1∑nlog(dn,i)Lkoleo\mathcal{L}_{koleo}LkoleoKoLeo正则化损失值nnn一个batch中的样本数量dn,id_{n,i}dn,i第iii个样本和batch中其他所有样本的最小欧氏距离通俗解释这就像是老师把学生均匀地安排在教室里避免大家都挤在一个角落。这样每个学生特征都有自己独特的位置能提供独特的信息。3.5 其他关键技巧Sinkhorn-Knopp中心化替代了原来的简单中心化让教师网络的输出分布更均匀训练更稳定。高分辨率微调在训练的最后阶段把图片分辨率从224×224提高到518×518再训练1万步。这样能显著提升密集任务的性能而计算成本只增加了一点点。四、让11亿参数模型跑起来的工程魔法训练一个11亿参数的ViT模型即使在A100上也是一个巨大的挑战。Meta的工程师们用了一系列工程优化让DINOv2的训练速度比iBOT快了2倍内存只用了1/3。4.1 核心优化技术FlashAttention重新实现了FlashAttention大幅提升了自注意力层的速度和内存效率。序列打包把不同长度的token序列来自不同大小的图片裁剪打包成一个长序列一次性输入Transformer。这避免了多次前向传播提升了GPU利用率。高效随机深度改进了随机深度的实现直接跳过被丢弃的残差块的计算而不是计算后再mask掉。FSDPFully-Sharded Data Parallel把模型参数、优化器状态都分片到不同的GPU上让单GPU的内存需求从16GB降到了几GB。4.2 核心代码序列打包的简单实现importtorchimporttorch.nnasnnfromxformers.opsimportmemory_efficient_attentiondefsequence_packing(images,crops_sizes): 将不同大小的图片裁剪打包成一个长序列 Args: images: 不同大小的图片裁剪列表 crops_sizes: 每个裁剪的大小 [(h1, w1), (h2, w2), ...] Returns: packed_tokens: 打包后的token序列 attention_mask: 块对角注意力掩码 tokens_list[]attention_masktorch.zeros((1,0,0),dtypetorch.bool)forimg,(h,w)inzip(images,crops_sizes):# 将图片转换为patch tokenstokenspatch_embed(img)# [1, num_patches, embed_dim]num_patchesh//patch_size*w//patch_size# 更新注意力掩码new_masktorch.ones((1,num_patches,num_patches),dtypetorch.bool)attention_masktorch.block_diag(attention_mask,new_mask)tokens_list.append(tokens)# 打包所有tokenspacked_tokenstorch.cat(tokens_list,dim1)returnpacked_tokens,attention_mask五、用数据说话——每个组件都在发光DINOv2的作者们做了非常详尽的消融实验证明了每个组件的贡献。我们来看最关键的一个实验从iBOT到DINOv2的逐步改进。方法ImageNet-1k k-NNImageNet-1k lineariBOT72.982.3我们的复现74.5 (1.6)83.2 (0.9)LayerScale, Stochastic Depth75.4 (0.9)82.0 (-1.2)128k prototypes76.6 (1.2)81.9 (-0.1)KoLeo78.9 (2.3)82.5 (0.6)SwiGLU FFN78.7 (-0.2)83.1 (0.6)Patch size 1478.9 (0.2)83.5 (0.4)Teacher momentum 0.99479.4 (0.5)83.6 (0.1)Tweak warmup schedules80.5 (1.1)83.8 (0.2)Batch size 3k81.7 (1.2)84.7 (0.9)Sinkhorn-Knopp81.7 ()84.7 ()Untying heads DINOv282.0 (0.3)84.5 (-0.2)表2从iBOT到DINOv2的逐步改进出处原文Table 1从表中可以看到每个组件都带来了稳定的提升。其中贡献最大的是KoLeo正则化k-NN准确率提升了2.3%增大batch size到3k线性准确率提升了0.9%把patch size从16改成14线性准确率提升了0.4%我们再来看KoLeo和iBOT损失分别对什么任务有帮助KoLeoImageNet-1kImageNet-AADE-20kOxford-M✗85.370.647.255.6✓85.872.847.163.9表3KoLeo损失的影响出处原文Table 3aMIM (iBOT)ImageNet-1kImageNet-AADE-20kOxford-M✗85.372.044.264.3✓85.872.847.163.9表4iBOT损失的影响出处原文Table 3b结论非常清晰KoLeo损失对图像检索任务Oxford-M帮助最大mAP提升了8.3%。因为它让特征在空间中分布更均匀检索更准确。iBOT损失对密集预测任务ADE-20k分割帮助最大mIoU提升了2.9%。因为它让模型学到了更好的patch级特征。六、横扫所有视觉任务的通用特征DINOv2在几乎所有视觉任务上都取得了SOTA的结果而且都是用冻结的特征线性分类器得到的不需要微调。这才是通用特征的真正含义。6.1 整体性能概览图2不同任务上的性能随模型大小的变化出处原文Figure 2从图中可以看到随着模型大小的增加DINOv2的性能在所有任务上都稳步提升。而且在大多数任务上DINOv2都超过了之前的自监督方法和弱监督的OpenCLIP相当。6.2 图像分类和CLIP平起平坐方法监督类型ImageNet-1k linearImageNet-V2OpenCLIP-G弱监督86.277.2EVA-CLIP-g弱监督86.477.4DINOv2-g自监督86.578.4表5ImageNet上的线性评估结果出处原文Table 4DINOv2-g的线性准确率达到了86.5%超过了OpenCLIP-G和EVA-CLIP-g。更重要的是在分布外的ImageNet-V2上DINOv2的优势更大领先EVA-CLIP-g 1.0%。这说明DINOv2的泛化能力更强。6.3 密集任务全面碾压CLIP在语义分割和深度估计这些密集任务上DINOv2的优势更加明显。这是因为CLIP是用图文对训练的文本描述很难包含像素级的细节信息而自监督的DINOv2能直接从图片中学到这些信息。方法ADE-20k mIoUCityScapes mIoUNYUd RMSEOpenCLIP-G39.360.30.541iBOT-L44.664.80.417DINOv2-g49.071.30.344表6密集任务性能对比出处原文Table 10、11我们再来看定性的结果图3分割和深度估计的定性结果对比出处原文Figure 7从图中可以清楚地看到OpenCLIP-G的分割结果有很多 artifacts边缘模糊很多小物体都识别不出来。DINOv2-g的分割结果边缘清晰物体完整深度估计也更平滑、更准确。七、DINOv2的超能力——那些意想不到的性质DINOv2学到的特征不仅在标准任务上表现好还展现出了一些非常神奇的涌现性质这些性质是模型自动学到的没有任何监督。7.1 PCA自动分割前景背景图4patch特征的PCA可视化出处原文Figure 1如果我们对DINOv2输出的patch特征做PCA会发现一个惊人的现象第一个主成分就能自动把前景和背景分开而且其他主成分对应物体的不同部分比如鸟的头、翅膀、身体。更神奇的是这种对应关系在不同的图片之间是一致的。比如所有鸟的翅膀在特征空间中都对应同一个主成分所有飞机的翅膀也对应同一个主成分。7.2 跨物体的patch匹配图5跨图片的patch匹配出处原文Figure 10如果我们计算不同图片之间patch特征的相似度会发现DINOv2能自动匹配语义上相似的部分即使它们来自完全不同的物体。比如鸟的翅膀 ↔ 飞机的翅膀大象的腿 ↔ 桌子的腿汽车的轮子 ↔ 自行车的轮子这说明DINOv2学到了真正的语义概念而不仅仅是像素的统计规律。这种能力对于VLA模型来说至关重要因为机器人需要理解不同物体之间的相似性才能把在一个物体上学到的技能迁移到另一个物体上。八、不完美的DINOv2——公平性与环境影响DINOv2虽然很强大但也不是完美的。作者们在论文中专门用了两个章节讨论它的局限性这是非常值得称赞的。8.1 地理和收入偏见DINOv2在不同地区和收入水平的图片上表现差异很大在欧洲的图片上准确率是89.7%在非洲只有74.0%差了15.7%。在高收入家庭的图片上准确率是90.5%在低收入家庭只有67.4%差了23.1%。这是因为训练数据中西方和高收入地区的图片占比过高。未来的工作需要构建更多样化的数据集来解决这个问题。8.2 环境影响训练一个DINOv2-g模型需要22016个A100 GPU小时大约排放3.7吨二氧化碳。这相当于一辆普通汽车行驶1.5万公里的排放量。不过和训练同样大小的CLIP模型相比DINOv2的碳排放量要低得多因为它不需要训练文本编码器。九、为什么DINOv2是VLA的最佳视觉编码器现在我们终于明白为什么所有的VLA模型都选择DINOv2作为视觉编码器了通用特征不需要微调就能在各种场景和物体上提取高质量的特征。这对于VLA来说至关重要因为机器人需要在未知的环境中操作未知的物体。鲁棒性强对光照、视角、遮挡的变化非常鲁棒。真实世界的环境是复杂多变的机器人的视觉系统必须能应对这些变化。密集特征好在分割、深度估计这些密集任务上表现优异。机器人需要理解场景的3D结构才能准确地抓取和操作物体。开源免费Meta完全开源了所有的模型和代码任何人都可以免费使用。这极大地推动了VLA领域的发展。十、动手实践——5行代码用DINOv2提取特征最后我们来写几行代码体验一下DINOv2的强大。用Hugging Face的Transformers库只需要5行代码就能提取特征fromtransformersimportAutoImageProcessor,AutoModelfromPILimportImageimporttorch# 加载模型和处理器processorAutoImageProcessor.from_pretrained(facebook/dinov2-giant)modelAutoModel.from_pretrained(facebook/dinov2-giant)# 加载图片imageImage.open(robot.jpg)# 预处理图片inputsprocessor(imagesimage,return_tensorspt)# 提取特征withtorch.no_grad():outputsmodel(**inputs)featuresoutputs.last_hidden_state# [1, num_patches1, 1536]# 得到全局特征CLS tokenglobal_featurefeatures[:,0,:]# [1, 1536]# 得到patch级特征patch_featuresfeatures[:,1:,:]# [1, 1369, 1536]就是这么简单你可以用这些特征来做图像检索、分类、分割或者作为VLA模型的视觉输入。总结DINOv2是计算机视觉领域的一个里程碑式的工作。它证明了自监督学习完全可以学到和弱监督方法相当的通用视觉特征为视觉基础模型的发展指明了方向。对于VLA领域来说DINOv2就像是给机器人装上了一双慧眼让它们能真正理解这个复杂的物理世界。未来随着模型更大、数据更多、方法更好我们相信自监督视觉模型会变得越来越强大最终实现真正的通用视觉智能。而DINOv2就是这条道路上最重要的一步。