自监督ViT:DINO-v1和DINO-v2

1. 概述

基于ViT(Vision Transformer)自监督在最近几年取得了很大进步,目前在无监督分类任务下已经超过了之前的一些经典模型,同时在检测分割等基础任务领域也展现出了强大的泛化能力。这篇文章将主要基于DINO系列自监督算法介绍它们的算法原理,方便 大家快速了解相关算法。

2. DINO-v1

参考代码:dino

这个方法源自于一个很重要的发现,自监督的ViT在图像语义分割的显式信息表达上具有独特性,也就是说相比有监督的ViT网络或者是传统的CNN网络其具有更强的语义表达能力和分辨能力。基于此使用k-NN算法作为分类器便能在一个较小的ViT网络上实现78.3% ImageNet top-1的准确率。在该方法中构建自蒸馏的方式训练和更新教师和学生网络,同样也适用了参数类似滑动平均更新和输入图像多重裁剪训练策略。对于训练得到的网络对其中的attention map进行可视化,确实也呈现出了上述提到的物体语义区域的感知能力,见下图可视化效果:

整体上文章提出的方法pipeline见下图所示:
自监督ViT:DINO-v1和DINO-v2_第1张图片
在上图中包含了两个相同结构的网络 g θ s g_{\theta_s} gθs g θ t g_{\theta_t} gθt,喂给它们同一张图不同的信息(local和global,也就是图像分辨率一个大一个小),之后对输出用交叉熵损失函数约束。但是这里需要注意的是只有student网络会存在梯度反向传播,teacher是通过类似滑动平均更新的形式更新参数。整体流程比较简洁,其自监督运算流程如下:
自监督ViT:DINO-v1和DINO-v2_第2张图片
step1 :为teacher和student分别准备不同的数据
step2:将teacher和student的输出软化,用teacher结果作为pseudo GT和交叉熵损失函数更新student参数
step3:使用类似滑动平均更新方式更新teacher的参数,同时更新teacher的数据中心分布( C C C),用它和 τ t \tau_t τt来避免collapse

输入处理:
在自监督方法设计过程中,teacher网络是没有任何先验初始化的,而teacher需要正确引导student网络学习。那么要使得自监督能够进行下去,则teacher应该能获取到更多的信息,而student相应的获取较少信息,这样才有信息的梯度差异,实际中是通过给teacher和student网络不同的图像分辨率图像实现的。对应的在输入图像的过程中也会经过一些数据增广操作,如视图扭曲、裁剪等操作。

student网络蒸馏和更新:
在ViT骨干网络基础上,会连接几个fc层(中间会使用l2-norm)得到这张图的高维度表达(表达的维度为 K K K): P s P_s Ps P t P_t Pt。为了避免生成的分布不够平滑这里对student和teacher的输出引入不同温度因子进行平滑:
P s ( x ) = e x p ( g θ s ( x ) / τ s ) ∑ k = 1 K e x p ( g θ s ( x ) / τ s ) P_s(x)=\frac{exp(g_{\theta_s}(x)/\tau_s)}{\sum_{k=1}^Kexp(g_{\theta_s}(x)/\tau_s)} Ps(x)=k=1Kexp(gθs(x)/τs)exp(gθs(x)/τs)
参考上面的式子,其对teacher网络也一样适用,只不过它们的温度因子策略有所不同。 τ s = 0.1 \tau_s=0.1 τs=0.1 τ t ∈ [ 0.04 , 0.07 ] \tau_t\in[0.04,0.07] τt[0.04,0.07]且有一个warm-up的过程。那么对于这两个分布最后是采取在不同输入下交叉熵最小化来更新student网络的参数的:
min ⁡ θ ∑ x ∈ { x 1 g , x 2 g } ∑ x ′ ∈ V , x ′ ≠ x H ( P t ( x ) , P s ( x ′ ) ) \min_\theta\sum_{x\in\{x_1^g,x_2^g\}}\sum_{x^{'}\in V,x^{'}\neq x}H(P_t(x),P_s(x^{'})) θminx{x1g,x2g}xV,x=xH(Pt(x),Ps(x))

teacher网络更新:
这里teacher网络的更新机制采用的是类似滑动平均更新的形式:
θ t = λ θ t + ( 1 − λ ) θ s \theta_t=\lambda\theta_t+(1-\lambda)\theta_s θt=λθt+(1λ)θs
也就是使用student网络的参数去更新teacher,其中参数 λ ∈ [ 0.996 , 1.0 ] \lambda\in[0.996,1.0] λ[0.996,1.0]在训练过程中会依据cosine函数进行变化。

teacher输出中心分布 C C C的更新:
同理teacher网络参数的更新机制,这里也是采用类似滑动平均的方式在Batch维度进行统计得到的:
C = m C + ( 1 − m ) 1 B ∑ i = 1 B g θ t ( x i ) C=mC+(1-m)\frac{1}{B}\sum_{i=1}^Bg_{\theta_t}(x_i) C=mC+(1m)B1i=1Bgθt(xi)

前文中一些模块对性能的影响:
自监督ViT:DINO-v1和DINO-v2_第3张图片

3. DINO-v2

参考代码:dinov2

在V1版本中使用ViT作为backbone,而ViT在输入的时候会将图像切块,之后送入到网络编码。在切块基础上之前的一些文章,如MAE,会将其中的一些块给丢弃掉(mask掉),之后通过编解码器得到完整图像,表明了mask机制能带来较好自监督效果。则在v2版本中也借鉴了mask的操作,不过它使用的是teacher和student模型的形式,其方法借鉴自IBOT

【论文阅读】IBOT : IMAGE BERT PRE-TRAINING WITH ONLINE TOKENIZER

同时,数据这块文章指出精心设计过的数据能带来更好的性能,所以不是说自监督直接喂给图片就能拿到最好效果。

数据处理:
这里从ImageNet-22k、ImageNet-1k的train部分、Google Landmarks和其它一些数据中使用数据清洗手段筛选更具表达能力的数据,新得到的数据集被命名为LVD-142M。具体的,数据清洗过程见下图:
自监督ViT:DINO-v1和DINO-v2_第4张图片
结合上图描述的过程,可以将整个处理流程划分为如下几步:
step1:对图片进行编码,去掉冗余信息
step2:使用编码结果进行去重
step3:去重之后的图片中检索与当前query图像类似的图像

对比是否对数据进行清洗,其差异见下表:
自监督ViT:DINO-v1和DINO-v2_第5张图片

V2相比V1显著区别:

  1. V2中借鉴iBot方法引入mask机制,从而引导teacher和student网络在mask掉的区域上表达一致
  2. teacher网络将之前的softmax-centering方法替换为Sinkhorn-Knopp (SK) batch normalization,用以避免collapse。而student网络还是使用softmax归一化
  3. 为了使得一个batch中分布更加均匀使用了KoLeo regularizer
  4. 为了适应下游任务中增大分辨率的需求,在训练快结束阶段会将分辨率调整到 518 ∗ 518 518*518 518518

一些改动对性的影响:
自监督ViT:DINO-v1和DINO-v2_第6张图片

你可能感兴趣的:(LargeModels,深度学习,计算机视觉,人工智能)