【论文精读】ESViT

摘要

基于transformer的SSL方法在ImageNet线性检测任务上取得了最先进的性能，其关键原因在于使用了基于对比学习方法训练单尺度Transformer架构。尽管其简单有效，但现有的基于transformer的SSL（自监督学习）方法需要大量的计算资源才能达到SoTA性能。

故认为SSL系统的效率与两个因素高度相关：网络架构和预训练任务。故而提出改进：

发现自监督单尺度transformer可自动发现图像局部视图区域之间的语义对应
实证研宄SSL任务下的多尺度vision transformer架构的利弊。其虽然极大地降低了计算复杂度，但发现多尺度架构导致了（1）中性质的损失
提出了基于图像局部视图的区域匹配预训练任务来缓解（2）中的问题，进一步改善了学习到的特征表示
在一系列任务上验证了结合了（2）、（3）技术的算法EsViT，其显著降低了构建SoTA SSL 视觉系统的成本，并在模型精确度与吞吐量和模型大小之间显示出更好的缩放性能

框架

多尺度ViT

参考swin-transformer。

预训练任务

全局视图任务

给定学生网络和教师网络的增强视图集 $V$ 和 $V^∗$ ，构建一组数据对 $\{ (s, t)|\hat x_s ∈ V, \hat x_t ∈V^∗ \ and \ s ≠ t ｝$ 来执行跨视图预测任务（参考DINO）:
$L_V=\frac 1 {|P|} \sum_{(s,t) \in P} M_V(s,t) \\ M_V(s,t)=-p_s \log p_t$

$p_s = h(f(\hat x_s))$ 和 $p_t = h(f(\hat x_t))$ 分别为学生网络和教师网络学习的视图表示， $f$ 为backbone， $h$ 为MLP投影头。其中 $z_t = f(\hat x_t)$ 为全局视图特征， $z_s = f(\hat x_s)$ 为局部视图特征。

局部视图任务

在DINO中， $L_V$ 鼓励全局视图和局部视图之间的特征匹配，而局部视图到局部视图的通信未指定。

在单尺度transformer中，由于低级特征到高级特征的残差连接有助于后者保持特征可区分性，从而保持良好的区域特征匹配性能。然而，这种特性由于多尺度transformer中的patch merging算子而被稀释，故用 $L_V$ 训练多尺度transformer会导致图像区域匹配性能减弱。此外，不利用在计算全局视图的过程中得到的局部视图特征 $z_s$ ，可能是一种计算的浪费。

本文认为计算机视觉中针对局部视图l的预训练任务很重要，这样模型可以充分利用 $L_V$ 提取的局部视图特征，以及利用局部视图特征之间的结构。故而提出一种非对比学习式的区域匹配方法，它们的通信定义：
$L_R=\frac 1 {|P|}\sum_{(s,t) \in P}M_R(s,t) \\ M_R(s,t)=-\frac 1 T \sum^T_{i=1}p_{j^*} \log p_i \\ j^*=\argmax_j \frac {z^T_iz_j} {||z_i|| \ ||z_j||}$

其中 $p_i = h^′ (z_i)$ 和 $p_j = h^′ (z_j )$ 分别为一个新的MLP头 $h ^′$ ，在学生网络和教师网络的局部特征 $z_i ∈ z_s$ 、 $z_j ∈ z_t$ 上计算的概率输出。 $j ^∗$ 是 $z_s$ 的第 $i$ 个特征与 $z_t$ 中具有最大余弦相似度的特征的索引，故 $z_i$ 和 $z_j^ ∗$ 代表两个最佳匹配的局部视图的上下文特征，最小化 $L_R$ 鼓励不同的局部视图上下文学习不变特征，从而捕获局部依赖性。

EsViT的整体预训练目标是 $L = L_R + L_V$ ，通过最小化学生网络和教师网络的交叉熵损失来学习匹配全局视图和局部视图的特征分布（如上图）。训练过程中交替更新教师/学生网络：

给定一个固定的教师网络，通过最小化交叉熵损失来更新学生网络： $θ_s ← arg min_{θ_s} L(s, t; θ_s)$ 。
教师模型根据学生权重的指数移动平均（EMA）更新： $θ_t ← λθ_t + (1 − λ)θ_s$ ， $λ$ 遵循余弦退火方法，训练期间从0.996到1。

默认情况下，一开始就可以使用完整的目标 $L$ 。也可以一开始加载一个仅由 $L_V$ 训练的checkpoints，再添加 $L_R$ 进行持续的预训练，实验中表明这样做能提高性能。

算法伪代码

计算复杂度

在传统的单尺度Transformer架构上应用 $L_R$ 的计算成本可能非常高，为 $O(T ^2 )$ 。对于分辨率为 $224 \times 224$ 的图像，顶层的ViT/DeiT（patch size 16）的特征图长度将为 $T = 196$ ，而multi-stage架构则为 $T = 49$ ，这在计算 $L_R$ 时可减少3倍的计算量。

实验

对比试验

上图显示了EsViT在ImageNet上线性分类任务下的速度与精度比较，左图可以看出同样的分类精度，EsViT在更小的参数量下吞吐量大于之前的sota算法如MoCov3、DINO、SimCLRv2等。右图显示EsViT在更少的参数量下，相比其他sota算法能达到更高的精度。

上图为在ImageNet验证集上不同算法的top-1线性检测和k -NN精度的对比。

与自监督transformer的比较：在相同的线性检测任务和训练数据下，基于DINO和moco的ViT比iGPT具有更高的精度和更小的模型。在相似的模型大小和计算复杂度水平上，EsViT（Swin-B）在绝对值上比DINO（ViT-B/16）提高了2.2％的线性检测精度和2.8％的k -NN精度。EsViT（Swin-B）甚至表现略优于DINO（ViT-B/8）（提高0.3％的线性检测精度和提高1.5％的k -NN精度），但具有4×倍的吞吐量。在使用相同的架构Swin-T时，EsViT预训练任务的性能优于MoBY，显示出更高3％的精度。

在EsViT中，通过增加窗口大小来实现自注意力中的较长序列，通过考虑W = 14的窗口大小来进行实验。与所有系统相比，所提出的EsViT（Swin-B/ W =14）表现出最佳性能（top-1精度81.3％，top-5精度95.5％，k -NN精度79.3％），并且是3.5×参数高效的，比之前的SoTA MoCo-v3至少有10×更高的吞吐量。

与大型ConvNets的比较：与SimCLR-v2, BYOL和SwAV报告的SoTA大ResNets进行了比较。其中，SimCLR-v2与SK- resnet报告了在线性探测协议下的最佳精度79.8%，而EsViT实现了更好的准确性，模型参数少16×，吞吐量高8×。

上图为与有监督模型比较。将其与监督学习的Swin进行比较。观察到使用Swin骨干网络的EsViT始终优于其监督变体。对于其他尺寸的模型也得出了类似的结论。

上表2显示在COCO检测和分割任务上，EsViT显示出与仅包含 $L_V$ 的变体（在括号中显示）的监督方法类似的结果（Swin-T以3 ×调度训练）。

上表3显示了不同预训练数据集上预训练的算法在下游分类任务的表现：

$L_R$ 在所有数据集上改进了 $L_V$ （括号中所示）
所有EsViT预训练checkpoints在下游分类任务中都优于监督checkpoints，但性能差异很大，其中ImageNet-22K检查点表现出最佳的迁移能力
ImageNet-1K预训练模型显示了最佳的ImageNet-1K线性检测性能。故认为不仅预训练数据集的大小很重要，图像类别的分布也很重要，更多样化和良好平衡的分布导致更强的泛化能力

消融实验

上表4研究了ResNet-50和其余不同的transformer的对比实验。在本实验中，batch size为1024。所有算法修改为4阶段架构，每个阶段的Transformer层数为2-2-6-2。可以看到 $L_R$ 改进了所有网络架构，包括ResNet-50, Swin, ViL, CvT和PvT。

上图4表明 $L_R$ 有助于模型收敛，并可作为改进使用view-level训练的模型的辅助工具。

$L_R$ 在非对比学习视图级任务 $L_V$ 的基础上增加了一个非对比区域匹配任务，DenseCL则在对比学习视图级任务MoCo-v2的基础上增加了一个对比区域匹配任务。上图所示，在预训练ResNet-50下相同设置的四种方法进行比较。DenseCL提高了密集视觉预测性能，但会影响分类性能， $L_R$ 提高了这两项任务，特别是分类性能。非对比方法在密集预测任务中表现出较低的性能。

考虑两个指标，准确率衡量正确匹配区域对的百分比、距离误差表示预测匹配区域与真实区域之间的平均 $ℓ_2$ 距离。

结果如上图，DINO与单尺度transformer在对比学习上表现出了良好性能（0.95，2.49），而多级Transformer架构的使用降低了这种能力（0.66，49.9），使用 $L_R$ 后，区域匹配能力明显恢复（0.91，11.9）。

将 $L_R$ 添加到多尺度transformer架构中，会产生可接受的额外计算成本，而直接将 $L_R$ 添加到单尺度transformer架构中，具有巨大的计算开销。

如上图，对于每个种设置，采用[内存使用情况（MB）/每次迭代的运行时间（秒/迭代）]格式。在（a）中，当batch size逐渐增加（batch-size=12）时，内存成本对于单尺度架构增加了近4倍，而对于多阶段架构则增加了1.6倍。每次迭代的训练成本增长率（1.47 vs 1.15）也显示了类似的趋势。这表明 $L_R$ 可以更自然地适合多阶段架构。

可视化实验

给定同一图像的两个视图，使用预训练的主干来提取顶层特征 $z_1$ 和 $z_2$ 。对于 $z_1$ 中的每个特征向量，在 $z_2$ 中找到最高余弦相似度的特征向量。上图显示了三个方法的两个视图之间的前10个对应关系。

在（b）中，EsViT与 $L_V$ 倾向于识别在背景中的配对为最匹配的（在这个例子中以错误的方式）。这可能是 $L_V$ 聚合全局特征级别的不变性不一定会诱导局部区域级别的不变性。 $L_R$ （c）显著地减轻了这一问题，该任务隐式地需要本地匹配。

上图显示最后一层的不同头部中的自注意力。将主对象边缘的局部区域作为query，并将query质量最高的60％的区域用红色高亮显示。

DINO倾向于自动学习特定类别的注意力图，从而导致前景对象分割，无论其query位于前景或背景。这可能是因为在不同的增强视图中，主要对象仍然是主要的不变性因素。当采用多阶段架构时，此属性将丢失，如EsViT和 $L_V$ 所示。这些图案对于不同的头部是一致的。

在加入 $L_R$ 之后，注意到不同头部的注意力图变得更加多样化，注意力的熵值变得更加偏斜，被关注的区域更加不同。这可能是因为 $L_R$ 需要每个区域考虑许多匹配任务到不同的增强视图中的区域，每个头自动学习分配任务并完成其中的一些任务。

上图将更多图像的对应关系可视化。DINO使用单尺度transformer能够在温和的增强条件下发现最显著的语义对应关系，即使在训练中没有隐式区域匹配损失。当从单尺度到多尺度Transformer架构（从第一列到第二列）改变时，这种期望的属性会扩大，那么所提出的区域级任务可以缓解这个问题（从第二列到第三列）。

为了更具体地分析对应关系，注意到以下结果。第一行展示了一个简单的情况，其中只展示了从左到右翻转视图的图像。ground truth对应关系应该是连接两个翻转区域的水平线。单靠视图级预训练任务不足以学习到多阶段Transformer架构的良好对应关系，而区域匹配任务可以显著缓解这一问题。类似的观察结果显示在第3行和第4行。

进一步研究了更多需要在第2行、第5行和第6行中与现实世界对应的案例。这些视图不是通过数据增强生成，而是经常在更实际的场景中出现：一对多映射，卡通到玩具，场景的季节性变化。提出的区域匹配任务可以在这些情况下工作得很好。

上图为顶层可视化注意力图。观察到，使用单尺度Transformer架构的DINO可以自动识别主要前景物体。然而从单尺度到多尺度Transformer架构（从左列到中列），这种属性会丢失。EsViT中相当多的头部显示了冗余模式，当增加区域级别的匹配任务（从中间列到右边列）时，这个问题可以减少。

具有单尺度Transformer架构的DINO只学会关注前景对象，即使query是背景区域。这可能是因为DINO模型被训练为学习视图级不变性，预训练数据集ImageNet中的主要对象往往是在不同增强视图中保持不变的主要因素。因此，无论query位置如何，所有背景都会被忽略。这在EsViT中通过区域级预训练任务得到了改进，因为模型被训练为匹配局部视图区域。

DINO在6个正面中都显示出高熵值（这可能是覆盖主对象所有区域的必要条件）。在EsViT中，增加 $L_R$ 会在自注意力头之间的熵分布比较有趣：增加了熵值较大的头，同时减少了熵值较低的头。换句话说，它使不同头部的注意模式更加多样化。

reference

Chunyuan, L. , Jianwei, Y. , Pengchuan, Z. , Mei, G. , Bin, X. , Xiyang, D. , Lu, Y. , & Jianfeng, G. . (2022). Efficient Self-Supervised Vision Transformers for Representation Learning.