526互联

[论文速览] A Closer Look at Self-supervised Lightweight Vision Transformers

发布时间 2023-07-15 16:12:20作者: NoNoe

Pre

title: A Closer Look at Self-supervised Lightweight Vision Transformers
accepted: ICML 2023
paper: https://arxiv.org/abs/2205.14443
code: https://github.com/wangsr126/mae-lite
ref: https://mp.weixin.qq.com/s/7FiDLYBZiAX-xkW-dZBU9Q

关键词： lightweght, ViT, pre-trained model, training report
阅读理由： 轻量化ViT的预训练方法，很实用的样子

Idea

通过大量的实验与分析为轻量级ViT的预训练提供一个“实践手册”

Motivation&Solution

边缘端实际部署的模型多是轻量模型，但鲜少研究轻量模型如何预训练。

Experiment

Training Detail

使用 ViT-Tiny（5.7M），由patch嵌入层、12个block块，嵌入维度为192，但注意力头的数量调整为12

那个预训练方法最好?

表1 不同预训练方法的比较，top-1 acc在ImageNet-1k的验证集上计算，设备是8 x V100 GPU，其中ori.是原始的监督训练策略，impr.是本文提出来的改进策略

前两行是将ViT-Tiny进行监督训练300epoch得到的结果，未进行预训练，只是训练策略不同，作为比较的baseline。同时最后一行MAE是作者实现的版本，将编码器换成了ViT-Tiny

表中可以看出：对于轻量级 ViT，其同样可以像大模型一样，显著受益于各种预训练方法。其中，基于 MIM (Masked-Image-Modeling) 的自监督预训练方法（例如 MAE，SimMIM）表现出优于 CL (Contrastive Learning) 的自监督预训练方法（例如 MoCo-v3，DINO）的效果，甚至优于基于更大规模的 ImageNet-21k（IN21K）的全监督预训练；其中 MAE 的预训练成本适中，且精度最高。

预训练的 enhanced Vit 在SOTA轻量化网络中排名如何?

对 ViT-Tiny 采用 MAE 进行预训练，然后在 ImageNet 上进行微调，其与其他一众轻量级模型的对比如下

表3 在ImageNet-1k上与之前的SOTA网络对比。单块RTX 3090 GPU，batchsize为1024，混合精度。†标识监督训练过程中使用了蒸馏，*标识原始的ViT-Tiny（3个注意力头）

DeiT-Tiny: 表1中前两行未使用预训练策略的ViT模型
MAE-Tiny-FT: 基于MAE-Tiny去微调的模型
表中可见，即使是在足够强 fine-tuning 的 recipe 下，MAE 的预训练依旧可以带来显著的性能增益（77.8→79.0），且基于 MAE 预训练的模型只微调了 300 epochs 就已超过了从头训练 1000 epochs 的模型（78.5 vs. 77.8）。

这说明：一个“普通”的网络结构与 SOTA 的性能之间，差的可能只是一个合适的训练方案，这为轻量级模型的研究提供了一条除人工设计复杂网络结构之外的道路。

除此之外，表中还包括一些采用了全监督、自监督预训练的卷积网络，但是可以看出：预训练对于 ViT 的增益更大（如ConvNeXt V2-F 采用基于 MIM 的预训练只能提升 0.5，而 ViT-Tiny 却可以提升 1.2），这个可以理解为：ViT 结构具有更少的人为设计的特性，而预训练可以以数据/任务驱动的模式对其进行补足，因此其增益更大，且表现出比通过人为设计引入inductive bias的网络更强的潜力。

更多的数据对预训练有帮助吗?

选择ImageNet-21k，公平起见预训练iterations保持一致，

表2 预训练数据的影响

可以看出对于MoCo-v3 和 MAE来说，几乎没有提升，然后又把数据减少，（IN1K-LT表示具有长尾类别分布的子集），MAE的性能下降也非常少，比MoCo-v3在数据规模和类别分布上更加鲁棒

预训练模型能否迁移到数据不充足的任务上?

表4 在分类和密集预测任务上的迁移评估。每个数据集括号里的内容对应（训练集大小/测试集大小/类别数量）

可以看到，在各种下游分类任务上，基于 MIM 的预训练模型表现明显弱于全监督预训练模型，而基于 CL 的预训练模型同样如此但差距较小；而在 COCO 检测任务上，基于 CL 的 DINO 预训练模型则表现出更好的迁移性能；这与在 ImageNet 上观察到的现象并不一致。

作者猜测这是因为MAE在预训练跟微调时层行为（layer behavior）不一致导致的

这说明，不同预训练模型在不同规模的下游数据集上可能表现出不一致的特性，需要根据下游数据集的特性选择合适的预训练方案。例如：当下游数据集规模较小且数据集类别分布与预训练数据较为一致时，全监督预训练方案可能更为合适；而当下游数据足够充足时，基于 MIM 的预训练方案可能效果更好。

预训练模型能否迁移到密集预测任务上？

下游目标检测和分割：附录A.5（略）
看表4，同样不如全监督

不同layer在下游数据是否充足时的作用

使用 Centered Kernel Alignment (CKA)(ICML 2019) 方法分析层表征相似度，它计算两个特征图或表征之间的 Hilbert-Schmidt Independence Criterion (HSIC)，它对于表示的正交变换和各向同性缩放是不变的

图1 层表征相似度 以模型内和模型间heatmap的形式展示，xy轴分别是不同模型的层索引，0表示patch嵌入层，值越大（越亮）表示相似度越高

DeiT-Tiny作为参考，它在IN1K上进行了全监督训练，是一个分类模型。作者认为如果其他模型能与它对应层越相关，就说明那个模型与识别越相关。虽然相似度无法直接只是下游任务的性能，但某种程度上可以反映出层表征的模式。最左边的是DeiT-Tiny自己跟自己的相似度。

作者发现MAE-Tiny在浅层跟DeiT-Tiny接近，但高层就不相似，说明MAE在高层抽不到什么语义信息；而MoCov3-Tiny就与DeiT-Tiny更加对齐。基于 MIM 的预训练模型高层出现了表征退化，而基于 CL 的预训练模型与 DeiT-Tiny具有较好的逐层对齐关系

但表1又表明使用MAE-Tiny作为预训练策略涨点更多，因此假设对预训练模型来说，低层比高层更加重要。

图2 预训练模型的低层对下游ImageNet数据集贡献更大

然后设计另一个实验：保留预训练模型前几层，其他的都随机初始化，然后在IN1K上面微调100epoch，图2能看出保留低层就能有很大提升，高层的作用很少

图3 高层对性能的贡献随着数据集规模减少而凸显，其表明高层对数据不充分的下游任务更加重要

仅保留 MAE-Tiny 的前 4 层即可在 ImageNet 上取得良好的迁移性能，这说明，对于数据充足的下游任务，预训练模型的浅层可能更加重要，而高层的质量则影响较小；

而随着数据规模减小，MoCov3-Tiny 的高层对于性能的增益越来越大，这也是其在这些数据集上性能超过 MAE 的关键，这说明，对于数据规模较小的下游任务，预训练模型的高层也很重要。

据此作者推测，可能正是这些模型高层的差异导致了上述实验现象。为验证这个猜想，我们进行了如下实验：仅保留预训练模型的若干浅层，测试其在各种规模的下游任务上的迁移性能，已测试预训练模型各层对于迁移性能的贡献度。

  • 出现上述现象的原因可能是：基于 MIM 的预训练模型浅层（靠近输入端的若干层）学的非常好，而高层（靠近输出端的若干层）出现了明显的表征退化；而基于 CL 的预训练模型虽然高层能学到还不错的语义特征，但却无法学习到可能对下游任务更有帮助的一些归纳偏置（inductive bias）；
  • 通过知识蒸馏，可以借助一个更大规模的基于 MAE 的预训练模型（例如MAE-Base）去帮助轻量级 ViT 的预训练，仅采用基于 attention map 的蒸馏就可以显著改善其在小规模下游任务上的糟糕表现。

分析注意力图

作者引入两种指标：attention distance 和 attention entropy。第h个头的第j个token 的attention distance如下计算可得：

\[D_{h,j}=\sum_{i}\mathrm{softmax}(A_{h})_{i,j}G_{i,j}, \tag{1} \]

它反映每个 token 是较多关注到近处（local）的区域还是远处的区域（即 global），值越低表示每个token更聚焦于邻近token上

其中 \(A_h \in \mathbb{R}^{l \times l}\) 是注意力图， \(G_{i,j}\) 是第i和第j个token之间的欧氏距离，l是token的数量

attention entropy如下计算可得：

\[E_{h,j}=-\sum_{i}\mathrm{softmax}(A_{h})_{i,j}\mathrm{log}(\mathrm{softmax}(A_{h})_{i,j}),\tag{2} \]

它揭示注意力分布的集中程度，值越低表示每个token的关注点是集中到少数几个 token 上（concentrated），否则表示广泛地关注到众多 token 上（broad）。

MAE的预训练会使得下游模型的注意力更局部更集中。

图4 Attention distance 和 entropy分析。可视化两个指标在所有head的token上的分布。

作者首先对比了MAE-Tiny-FT（采用 MAE-Tiny 作为预训练在 ImageNet-1k 上微调） 和 DeiT-Tiny，可以看到第1、2幅图中中间层的紫色箱子更靠下，说明采用 MAE-Tiny-FT的 attention 更加 local 与 concentrated，这些可能被作为归纳偏置被预训练引入到了模型中，使得其在 ImageNet 上取得了更好的性能表现。

但这种局部归纳偏置不代表注意力头就只在意近邻的几个token，不同头的这俩指标仍然分布在很宽的范围内，模型既能汇聚局部也能有全局的感知能力。

而对比 MoCov3-Tiny 与 MAE-Tiny 可以发现，虽然两者相较于随机初始化均可一定程度上引入局部归纳偏置，但 MoCov3-Tiny 从浅层到高层仍旧较为 global 且 broad。这种特性可能使得以其为初始化的模型在下游微调时倾向于“走捷径”，直接关注到全局信息而忽视局部细节信息。这对于较为简单的下游任务来说可能是足够的（Flowers, CIFAR100），但却未必适宜难度较高的需要细粒度识别的分类任务（ImageNet）。 而MAE-Tiny在高层仍然局部且集中的特性使得它难以迁移到数据不充足的下游任务上。

蒸馏以提高预训练模型性能

图5 蒸馏会将教师（MAE-Base）的全局表达压缩给学生（D-MAE-Tiny）

作者通过观察更大的预训练模型（MAE-Base）的注意力图，发现它跟DeiT-Tiny的高层有更好的对齐效果（图5上部分），因此想试试蒸馏

图A5 蒸馏过程示意图

通过计算教师模型跟学生模型在对应层的注意力图相似度，作者定义了个蒸馏损失：

\[L_{attn}=MSE(A^T-MA^{S}), \tag{3} \]

其中 \(A^{T}\ \in\ \mathbb{R}^{h\times l\times l}, \; A^{S}\ \in\ \mathbb{R}^{h'\times l\times l}\) 分别表示教师和学生网络的 attention map，MSE表示均方误差损失，M是一个可学习的映射矩阵，用来对齐师生 attention map的 head 数，随着学生网络一起学习。

但还有一个问题：在哪一层进行蒸馏，高层还是低层？

图6 在高层注意力图上进行蒸馏效果最明显

因此作者设置了几个实验，在相应的第1/4, 2/4, 3/4 和 4/4层分别算蒸馏损失，也就是二者的第3、6、9、12层。实验发现，仅在网络的最高层应用上述蒸馏损失即可取得良好的效果。

表5 蒸馏提高了下游任务的性能。分别在分类任务（top-1精度）、目标检测和实例分割（平均精度）任务上进行测试。

可以看到图5下方那个通过蒸馏学习的MAE-Tiny更像MAE-Base，说明有学到东西，而且表5也显示出它的作用，尤其在数据不充足的分类和密集预测任务上效果显著。这说明蒸馏可显著改善基于 MAE 的轻量级预训练模型的高层表征质量，弥补此预训练方法无法学习到高层丰富的语义特征，进而可显著提高其在多种下游任务中的迁移性能。这也是首次将知识蒸馏引入基于 MIM 的轻量级 ViT的预训练中。

图A7 不同教师模型的蒸馏层表征分析

图A7可以看出MAE-Small太小，在高层出现了表征退化，无法提供足够的知识，而MAE-Large则与Tiny在网络深度和注意力头数量等方面不匹配，不适合学生模型。

Conclusion

文章聚焦在轻量级 ViT 的预训练上，相当于为相关方向的研究提供了一个 benchmark，相关的代码与模型也都有开源，方便后续这一方向上的继续探索。

文章的研究也打破了之前人们对于 ViT 在轻量级赛道上没有竞争优势的惯有印象，只要采用合适的预训练手段，轻量级 ViT 也会表现出足够的竞争力。

文章在解决了一些问题的同时也带来了更多的问题，例如：有没有办法可以让轻量级 ViT 吃下更多的预训练数据带来更多的增益？有没有办法高效融合基于 CL 和 MIM 的预训练方法，使得其兼具两者的优势？有没有办法不借助知识蒸馏来提升基于 MIM 的预训练模型的高层表征质量？这些问题可能都值得未来去进一步探索。