论文笔记：DINO用于视觉自监督学习的知识蒸馏算法——Emerging Properties in Self-Supervised Vision Transformers_dino - emerging properties in self-supervised visi

综述

论文题目：《Emerging Properties in Self-Supervised Vision Transformers》

会议时间：ICCV 2021

论文地址：http://openaccess.thecvf.com//content/ICCV2021/papers/Caron_Emerging_Properties_in_Self-Supervised_Vision_Transformers_ICCV_2021_paper.pdf

源码地址：https://github.com/facebookresearch/dino

主要思想

  Transformer由于其远距离建模的优势，最近经常用于解决视觉的各大任务，常见的策略就是在大量数据上进行有监督预训练，获得较强的语义表征能力，之后在目标数据上进行微调，用于下游任务的应用。但这种方法并没有明显的好处，在增加计算量的同时，所提取的特征并没有表现出独特的属性（unique properties），这就不免产生疑问，利用有监督式的训练策略，是否充分挖掘了模型的特征表示潜力？

  在本文中，作者利用弱监督学习策略，来进一步提升了transformer在视觉应用中的能力。考虑到transformer在NLP领域中取得成功的主要因素之一是使用了自监督训练，例如BERT使用MLM（Masked Language Model）的自监督训练策略，随机抹掉部分单词，之后根据上下文来预测这个单词；GPT采用ALM（Autoregressive Language Modeling）的自监督训练策略，让模型根据前文来预测下一个可能出现的单词。这些算法均使用句子中的单词来创建假设任务（pretext task），这种操作可以提供更丰富的学习信号，可以让模型在学习过程中学到更多的语义信息，相比于监督学习中，简单地预测每个句子固定单一的标签难以让网络学到这种丰富的语义信息，模型的理解能力往往会被句子的标签所约束。同样，在图像任务中，图像级别的监督常常会降低视觉的信息量，整幅图像丰富的语义信息会被简化为从几千个对象类别的预定义集合中选择的单个概念，所学到的视觉信息量会大大降低，这也是要用自监督学习策略来取代监督学习策略的一个核心出发点。

  作者的自我监督学习方法可以表示为一种无标签的知识蒸馏操作（distillation with no labels, DINO），通过使用标准的交叉熵损失，让学生网络去学习教师网络的输出，同时教师网络由动量编码器（momentum encoder）构成，即教师网络的参数通过学生网络的参数来动量式地更新，最后作者使用中心化操作和锐化操作（centering and sharpening）来避免教师网络的崩溃问题（collapse）。此架构是比较灵活的，可以在CNN或者ViT上进行工作，不需要修改原始的网络架构，也不需要适应内部标准化操作（internal normalizations ）。

  如下图所示，利用自监督学习的ViT特征明确包含了场景布局，特别是对象边界，这些关注区域可以在最后一个自注意力关注模块中直接提取，同时，自监督ViT网络所提取的特征在基本近邻分类器k-NN下表现特别好，无需任何微调就可以在ImageNet上达到78.3%的准确率。

方法

知识蒸馏

  知识蒸馏是一种学习范式，让学生网络$g_{{\theta }_s}$去匹配教师网络$g_{{\theta }_t}$的输出，给定输入图像$x$，两个网络输出$K$维上的概率分布，分别表示为$P_s$和$P_t$（由网络$g$的输出进行softmax归一化得到）：
$$P_s(x{)}^{(i)}=\frac{\exp (g_{{\theta }_s}(x{)}^{(i)}/{\tau }_s)}{{\sum }_{k=1}^K\exp (g_{{\theta }_s}(x{)}^{(k)}/{\tau }_s)}$$
其中${\tau }_s>0$表示用于控制学生网络输出分布尖锐程度的温度参数，对于教师网络的概率分布，也有同样的公式以及温度超参数${\tau }_t$。给定一个固定的教师网络$g_{{\theta }_t}$，通过最小化交叉熵损失来匹配学生网络和教师网络的输出，从而更新学生网络的参数${\theta }_s$：
$$\underset{{\theta }_s}{\min }H(P_t(x),P_s(x))$$
其中$H(a,b)=-a\log b$​，这一步用于让学生网络的输出向教师网络靠拢，此过程并不更新教师网络的参数。

自监督学习

  知识蒸馏范式在设计的初衷，就是为了针对两个不同的模型，用复杂度较高、性能较好的模型去提升复杂度较低、性能较弱的模型，将强模型的能力通过蒸馏的方式蒸给弱模型，因此需要让弱模型的输出往强模型的输出上靠拢。这里教师网络输出的结果相当于伪标签，让学生网络去学习，因此如果有一个训练好的教师网络，我们就可以在没有真实标签数据的情况下，来训练学生网络，让学生网络去教师网络的能力。而在这篇文章中，作者做了进一步的改进，假设教师网络和学生网络都是随机初始化的，也就是都需要进行参数更新，用教师网络生成的伪标签去训练学生网络，从而实现完全意义上的自监督训练，现在问题的核心就在于如何动态更新教师网络的参数。在这里，作者参考动量更新的思想，利用EMA策略，让学生网络的参数去更新教师网络的参数，这样学生网络向教师网络靠拢的同时也会改变教师网络，最终两个网络会朝着同一个输出结果趋势上优化，网络主要流程图如下图所示：

  但其实还有个问题，如果教师网络和学生网络使用的是同一个输入，那么他们可以轻而易举地得到相同的输出，从而无需学会提取图中的语义信息也可以实现这一点，这与我们利用自监督学习去做模型预训练的初衷是背道相驰的，因此这里对两个模型的输入需要做不同的处理。

  首先，使用多尺度裁剪策略构建同一图片下的不同视图。给定一张图像，生成不同视角图的集合$V$，这个集合包含两个尺寸较大的全局视图$x_1^g,x_2^g$和一些尺寸较小的局部视图，在训练过程中，所有视图都会传入学生网络得到输出，同时只将全局视图传入教师网络得到输出，损失函数可以表示为：
L s = min ⁡ θ s ∑ x ∈ { x 1 g , x 2 g } ∑ x ′ ∈ V , x ′ ≠ x H ( P t ( x ) , P s ( x ′ ) ) L_s=\min_{\theta_s}{\sum_{x\in\{x^g_1,x^g_2\}}\sum_{x'\in V,x'\neq x}H(P_t(x),P_s(x'))} Ls​=θs​min​x∈{x1g​,x2g​}∑​x′∈V,x′=x∑​H(Pt​(x),Ps​(x′))
对于DINO的基本参数设置，作者使用两个分辨率为$224^2$的全局视图，可以覆盖原始图像的大部分区域（大于原图的50%），使用多个分辨率为$96^2$的局部视图，用于覆盖原始图像的一小部分区域（小于原图的50%）。两个神经网络使用相同的结构$g$，同时使用不同的参数${\theta }_s$和${\theta }_t$，使用损失函数$L_s$来更新学生网络的参数${\theta }_s$​​。

  这一操作流程下来，网络会实现一个从局部语义到全局语义的统一过程，从而迫使网络可以提取图像中真正有意义的语义信息，只有这样，教师网络和学生网络才可能输出相同的特征分布，这样也实现了在无标签的情况下，提升模型特征提取能力的目的，算法整个前向传播流程如下图所示：

教师网络

  与单纯的知识蒸馏任务不同，这里没有先验的教师参数${\theta }_t$，因此我们只能通过学生网络的参数来得到教师网络的参数。在作者尝试了大量的方法之后，发现使用指数移动平均策略（exponential moving average, EMA），即动量编码器效果最好，更新规则为：
$${\theta }_t\leftarrow \lambda {\theta }_t+(1-\lambda ){\theta }_s$$
$\lambda$用于控制动量更新的快慢，变化波动遵循从0.996到1的余弦调度器（cosine schedule）。 整个训练过程中，作者发现教师网络的性能要优于学生网络，因此可以提供质量更高的目标特征来指导学生网络的学习。

避免崩溃

  作者在这里使用中心操作和锐化操作（centering and sharpening）来避免教师网络的崩溃问题，中心操作用于防止单个维度的特征占主导地位，会鼓励教师网络输出均匀的数据分布，而锐化操作则有相反的作用，运用这两种操作可以平衡他们的效果，足以避免教师网络在动量更新时的崩溃。中心操作相当于对教师网络的输出增加一个偏置$c$，即$g_t(x)\leftarrow g_t(x)+c$，其中$c$的更新也利用动量更新策略：

$$c\leftarrow mc+(1-m)\frac{1}{B}\sum _{i=1}^B{g}_{{\theta }_t}(x_i)$$
其中$m>0$为速率参数，$B$为批次大小，输出锐化是通过在教师网络的softmax归一化中使用较低温度参数${\tau }_t$​来实现的，反向传播更新参数的流程如下图所示：

算法伪代码流程如下：

网络结构

  神经网络$g$由主干特征提取网络$f$（ViT或ResNet）和投影头$h$组成：$g=h\circ f$，下游任务中使用的特征是由$f$输出的特征数据，投影头$h$是由隐藏层维度为2048的三层多层感知机以及一个$K$维权值归一化全连接层组成，学生网络和教师网络架构完全相同，ViT架构参数如下表所示：

总结

  作者利用知识蒸馏和动量更新策略实现了ViT模型在视觉任务上的自监督训练，通过摆脱标签的约束，充分挖掘了ViT算法的特征表示潜力，同时该算法存在两个在未来的研究中可以利用的应用：①所提取的高质量特征传入k-NN算法可以实现较准的分类，因此该特征有利于用来做图像检索任务；②所得到的特征数据包含了场景布局，有利于弱监督图像分割任务。

补充：

• 《DINOv2》论文中构建了一个大规模数据集，并且引入了很多训练技巧，模型性能得到了进一步的提升，讲解可参考：https://zhuanlan.zhihu.com/p/623274167

以上仅是笔者个人见解，若有问题，欢迎指正