Vision Transformer (知识蒸馏 预训练）
Attention
原文地址：https://zhuanlan.zhihu.com/p/543743516
Transformer 是 Google 的团队在 2017 年提出的一种 NLP 经典模型，现在比较火热的 Bert 也是基于 Transformer。Transformer 模型使用了 Self-Attention 机制，不采用 RNN 的顺序结构，使得模型可以并行化训练，而且能够拥有全局信息。
视觉 Transformer 的一个发展趋势是扩大参数量，使之在基准任务上取得更好的性能；另一个发展的趋势是把模型变小，以支持其在手机和物联网端侧设备的部署。这篇文章讲解一种通过知识蒸馏的方式来压缩视觉 Transformer 的工作。相关的工作自从 2021 年初的 DeiT 以来，还没有比较经典的视觉 Transformer 知识蒸馏的文章出现。DeiT 使用的知识蒸馏方法直接使用训练数据和正常训练过程使用的数据集是一致的，大规模的数据集 (比如 ImageNet-21K) 更适用于较大的 ViT 模型，而大数据集用在较小的 ViT 模型会带来欠拟合的问题，似乎小模型不太容易享受到大数据集带来的红利。因此，本文希望提出一种高效利用训练数据的方法，来改善模型的欠拟合问题。
TinyViT：大规模预训练数据 + 大模型指导蒸馏 = 充分释放小 ViT 模型的潜力¶
论文名称：TinyViT: Fast Pretraining Distillation for Small Vision Transformers
论文地址：
https://arxiv.org/pdf/2207.10666.pdf
- 47.1 TinyViT 论文解读：
自从 2021 年初的 DeiT 以来，还没有比较经典的视觉 Transformer 知识蒸馏的文章出现。DeiT 使用的知识蒸馏方法直接使用训练数据和正常训练过程使用的数据集是一致的，大规模的数据集 (比如 ImageNet-21K) 更适用于较大的 ViT 模型，而大数据集用在较小的 ViT 模型会带来欠拟合的问题，似乎小模型不太容易享受到大数据集带来的红利。因此，本文希望提出一种高效利用训练数据的方法，来改善模型的欠拟合问题。
本文提出一种快速知识蒸馏的方法，并说明了通过借助大模型的指导，小模型也能够获得大规模数据集带来的红利。具体而言，作者观察到当数据规模增加的时候，直接预训练小模型会带来性能饱和的问题。但是如果同时借助教师模型进行知识蒸馏，则能够为小模型释放大规模预训练数据的潜力，如下图 1 所示，在使用超大规模数据集 ImageNet-21K 预训练的同时进行蒸馏操作可以大幅释放小型视觉 Transformer 的潜力。
但是，以蒸馏的方式来预训练模型是低效且昂贵的，因为在这个过程中有相当大的比例的计算资源会消耗在训练数据在大模型的前向传播过程中，而不是消耗在训练学生模型上面。此外，庞大的教师模型可能会占用很多的 GPU 内存，因此显著降低学生的训练速度。为了解决这个问题，本文提出了一个快速且和可扩展的蒸馏策略。
具体而言，作者提出预先生成一些稀疏的概率向量作为每个输入图像的软标签，并且与 random cropping, RandAugment, CutMix 等数据增强的方法结合之后，把这些向量存起来。在训练期间，直接使用这些个存储好的概率向量作为每个输入图像的软标签，从而无需再进行输入图片通过教师视觉模型的前向传播过程了。那么这样的策略主要会带来两个好处：
- 加快训练速度： 这样的做法省去了知识蒸馏的前向传播过程，从而大大节省了训练时生成教师软标签的内存开销和计算开销。此外，因为它能够使用更大的 Batch size，故小模型的蒸馏过程可以大大加快。
- 可扩展： 本方法可以模仿任何类型的数据扩充，并生成相应的软标签。只需要计算一次教师模型的前向传播，得到的软标签就可以适用于任意的学生模型。
大规模预训练数据 + 大模型指导蒸馏
作者观察到，使用大规模数据集 (比如 ImageNet-21K) 预训练小模型并不会带来性能的增益 (如上图 1 中的绿色线和蓝色线)，尤其是迁移到下游任务上面时。为了解决这个问题，作者借助知识蒸馏来进一步释放小模型的预训练潜力。不同于之前大多数关注微调阶段蒸馏的工作 (DeiT)，本文关注预训练蒸馏，它不仅允许小模型从大模型中学习，而且提高了它们对下游任务的迁移能力。
但是，直接在 ImageNet-21K 预训练的阶段使用知识蒸馏是低效且昂贵的，因为在每次迭代中，相当大比例的计算资源被消耗在训练数据在教师模型的前向传播过程中，而不是训练目标学生模型上面。此外，庞大的教师模型可能会占用较多的 GPU 内存，减缓目标学生的训练速度 (由于较小的 Batch size 的原因)。为了解决这个问题，作者提出了一个快速的预处理蒸馏框架。如图 2 所示，这个框架预先存储数据增强的结果和教师模型的预测结果。在训练期间，作者重用存储的信息来精确地复制前向传播的过程，成功地省略了大教师模型的前向传播过程和内存占用。
假设输入图片用 \(x\) 表示，一些数据增强策略比如 RandAugment 和 CutMix，用 \(\mathcal{A}\) 表示，教师模型用 \(T\) 表示。作者会存储数据增强的结果 \(\mathcal{A}(x)\) ，和教师模型输出软标签 \(\hat{\mathbf{y}}=T(\mathcal{A}(x))\) 。
值得注意的是，由于数据增强策略固有的随机性，把一张相同的图片多次通过同样的数据增强策略会得到不同的结果。因此，每个图片，每个 iteration 的 \((\mathcal{A}, \hat{\mathbf{y}})\) 都需要存储，如下图 2 的快速预处理蒸馏框架所示。
在训练过程中每个 iteration，我们只需要从已存储的文件中拿出来 \((\mathcal{A}, \hat{\mathbf{y}})\) ，并优化下面的目标函数即可：
以上就是一个正常的蒸馏过程，只是软标签是通过教师模型得到的。式中， \(S(\cdot)\) 和 \(CE(\cdot)\) 分别是学生模型和交叉熵损失函数。注意到这个框架是无标签的，即不需要真实标签，因为只使用了教师模型生成的软标签进行训练。因此，它可以利用大量没有标签的现成 web 数据进行大规模的预训练。这种无标签策略在实践中是可行的，因为软标签足够准确，同时携带大量用于分类的信息，例如类别关系。此外作者还观察到，带有真实标签的蒸馏过程会带来轻微的性能下降，其原因可能是 ImageNet-21K 中并不是所有的标签都是互斥的。因此，一个 one-hot 的标签不能够独立地描述一个对象。此外，本文提出的框架与不使用蒸馏方法的框架一样快，因为复杂的教师模型前向传播过程已经事先算好了。
稀疏软标签
考虑教师模型的输出 \(C\) 维的向量，如果 \(C\) 很大，例如，对于 ImageNet-21K 而言 \(C=21841\) ，则保存所有增强图像的全部输出逻辑通常会消耗很多存储空间。因此，作者只储存 \(\hat{\mathbf{y}}\) 中最重要的 top- \(K\) 个标签，即：
\({\{\hat{y}_{\mathcal{I}(k)}\}}_{k=1}^{K} \in \hat{\mathbf{y}}\) ，并且存储其 index： \(\{\mathcal{I}(k)\}_{k=1}^K\) 。在蒸馏过程中，作者只对稀疏标签进行标签平滑：
式中， \(\hat{y}_c\) 是用于知识蒸馏过程中的 teacher logits，即： \(\hat{\textbf{y}}=[\hat{y}_1,\dots,\hat{y}_c,\dots,\hat{y}_C]\) 。当稀疏系数 \(K\ll C\) 时，它可以将 logits 的存储量减少几个数量级。而且作者的实验结果表明，这种稀疏标签可以实现与密集标签相当的知识蒸馏性能。
数据增强编码
数据增强涉及一组参数 \(\mathbf{d}\) ，例如旋转角度和裁剪坐标，以变换输入图像。由于每次迭代中每个图像的 \(\mathbf{d}\) 是不同的，直接保存它会降低内存的效率。为了解决这个问题，作者通过标量参数 \(d_0=\mathcal{E}(\mathbf{d})\) 对 \(\mathbf{d}\) 进行编码，其中 \(\mathcal{E}(\cdot)\) 是图 2 中的编码器。然后在训练过程中，作者在存储文件中加载 \(d_0\) 后恢复 \(\mathbf{d}=\mathcal{E}^{-1}(d_0)\) 。解码器的常见选择是伪随机数发生器 (即 PCG)。它将单个参数作为输入，并生成一系列参数。至于编码器，只需通过生成一个参数 \(d_0\) 和重用解码器 \(\mathcal{E}(\cdot)\) 来实现。
模型架构
作者通过渐进模型缩放方法 (progressive model contraction approach) 得到了一族小型视觉 Transformer 模型。具体来说，作者从一个大模型开始，定义一组基本的缩放因子。然后在每一步中，通过调整收缩因子，在当前模型周围生成更小的候选模型。我们选择满足参数数量和吞吐量约束的模型。具有最佳验证准确性的模型将用于下一步的进一步缩减，直到实现最终目标。
作者采用分层视觉 Transformer 作为基本架构，以方便密集预测下游任务，如需要多尺度特征的检测。更具体地说，作者使用的基本模型由分辨率逐渐降低的 4 个阶段组成，类似于 Swin 和 LeViT。Patch Embedding 模块包含 2 个 3×3，stride=2，padding=1 的卷积，在 Stage1 使用 MBConvs，其他 3 个 Stage 使用 Transformer 模块，使用 GeLU 作为激活函数，且卷积和 Linear 操作的归一化层分别是 BatchNorm 和 LayerNorm。
构成模型过程中，作者考虑了下面的收缩因数 (Contraction factors)：
- \(\gamma_{D_{1-4}}\) ：4 个 Stage 的嵌入维度。
- \(\gamma_{N_{1-4}}\) ：4 个 Stage 的块数。
- \(\gamma_{W_{2-4}}\) ：后 3 个 Stage 的 Window Size。
- \(\gamma_R\) ：MBConv block 的 channel expansion ratio。
- \(\gamma_M\) ：MLP block 的 channel expansion ratio。
- \(\gamma_{E}\) ：multi-head attention 的每个头的维度。
作者使用渐进模型收缩方法缩小上述因子，并得到一系列小型视觉 Transformer：所有模型都包含：\(\{\gamma_{N_1}, \gamma_{N_2}, \gamma_{N_3}, \gamma_{N_4}\} = \{2, 2, 6, 2\}, \{\gamma_{W_2}, \gamma_{W_3}, \gamma_{W_4}\} = \{7, 14, 7\},\{\gamma_R, \gamma_M, \gamma_E\} = \{4, 4, 32\}\) 。对于嵌入维度 \(\{\gamma_{D_1}, \gamma_{D_2}, \gamma_{D_3}, \gamma_{D_4}\}\) ，有：TinyViT-21M：{96, 192, 384, 576}, TinyViT-11M：{64, 128, 256, 448} 和 TinyViT-5M: {64, 128, 160, 320}。
大规模预训练数据 + 大模型指导蒸馏相关分析
在本节中，作者对两个关键问题进行分析和讨论：
1) 限制小模型适应大规模数据集的潜在原因是什么？
2) 为什么蒸馏可以帮助小模型释放大规模数据集的潜力？
为了回答上述问题，我们在广泛使用的大规模基准 ImageNet-21K 上进行了实验，该基准包含 14M 的图像和 21，841 个类别。
问题 1：限制小模型适应大规模数据集的潜在原因是什么？
针对这个问题，作者观察到在 IN-21K 中存在许多困难样本 (Hard Sample)。小模型难以适应这些困难样本，导致与大模型相比训练精度较低 (TinyViT-21M: 53.2% 和 Swin-L 197M: 57.1%)，并且在 ImageNet-1K 上的可转移性有限 (TinyViT-21M w/ pretraining: 83.8% 和 w/o pretraining: 83.1%)。
如下图 3 所示，为了验证困难样本的影响，作者使用了两种技术：
1) 作者利用在 ImageNet-21K 上微调的强大的预训练模型 Florence 来识别标签位于 Florence 的前 5 个预测之外的那些图像 (定义为困难样本)。通过此过程，作者从 ImageNet21K 中移除了大约 2M 张图像 (14%)，然后在清理后的数据集上预训练 TinyViT-21M 和 Swin-T。
2) 作者使用 Florence 作为教师模型来执行蒸馏以预训练 TinyViT-21M/Swin-T，其生成软标签来代替 ImageNet-21K 中被污染的 GT 标签，得到在 ImageNet-1K 上进行微调的预训练模型的结果。
从图 3 的结果中得到了几点启发：
1) 在原始 ImageNet-21K 数据集上预训练小模型会在 ImageNet-1K 上带来有限的性能提升 (Swin-T 和 TinyViT-21M 均为 0.7%)。
2) 在移除 ImageNet-21K 中的部分困难样本后，两种模型都可以更好地利用大数据并实现更高的性能增益 (Swin-T/TinyViT-21M 为 1.0%/1.1%)。
3) 知识蒸馏能够避免困难样本的缺陷，因为它不使用 GT 标签，而 GT 标签的不合适才是样本属于困难样本的主要原因。因此，它获得了更高的性能提升 (Swin-T 和 TinyViT-21M 分别为 2.2% 和 1.7%)。
问题 2：为什么蒸馏可以帮助小模型释放大规模数据集的潜力？
为什么蒸馏可以提高大数据集上小模型的性能？答案是学生模型可以直接从教师模型那里学习到领域知识。也就是说，教师在训练学生时注入类之间的关系，同时过滤学生模型的噪声标签 (困难样本)。
为了分析教师预测的类别关系，作者从总共 21,841 个类别的 IN-21K 中为每个类别选择 8 个图像。这些图像然后被输入到 Florence 模型中输出预测的 logits，并画出预测逻辑上类间 Pearson 相关系数的热图。在图 3(a) 中，相似的或相关的类，很明显地彼此具有高相关性 (红色)，不相关的类彼此具有高相关性 (蓝色)。
作者比较了 TinyViT-21M w/o 和 w/ distillation 预测的 Pearson 相关性，分别如图 3(b) 和图 3(c) 所示。通过对比 (a)(b)(c) 图的对角线结构，作者发现当不使用蒸馏技术的时候，对角线的结构会更不明显，说明了小的模型更难捕获类间的相关关系。但是，蒸馏可以引导学生模型模仿教师模型的行为，从而更好地从大数据集中挖掘知识。
实验结果
ImageNet-21K pretraining 实验设置：
数据集：ImageNet-21K，优化器：AdamW，Epoch：90，weight decay：0.01，初始学习率：0.002，warm-up Epoch：5，cosine scheduler，Batch size：4096，gradient clipping (max norm=5)，数据增强： random resize and crop, horizontal flip, color jittering, random erasing, RandAugment, Mixup 和 Cutmix。
从上一步预训练模型进行 ImageNet-1K finetuning 实验设置：
数据集：ImageNet-1K，优化器：AdamW，Epoch：30，weight decay：1e-8，初始学习率：5e-4，warm-up Epoch：5，cosine scheduler，Batch size：1024。
ImageNet-1K 从头训练实验设置：
数据集：ImageNet-1K，优化器：AdamW，Epoch：300，weight decay：0.05，初始学习率：0.001，warm-up Epoch：20，cosine scheduler，Batch size：1024。
在小型 ViTs 上进行预训练蒸馏策略的影响
作者在两个流行的 ViT 模型上验证了预训练蒸馏策略的有效性：DeiT 和 Swin 。如图 4 所示。只在 IN-21K 上作预训练，而不进行蒸馏过程，在 DeiT-Ti/DeiT-S/Swin-T 上会带来 0.8%/0.6%/0.7% 的性能提升，而在 IN-21K 上作预训练，同时进行蒸馏过程，在 DeiT-Ti/DeiT-S/Swin-T 上会带来 2.2%/2.1%/2.2% 的性能提升。
消融实验
预训练数据规模的影响
作者观察了在 IN-21K 上进行预训练期间，TinyViT-5M/21M 预训练数据规模对于最终结果的影响。使用 CLIP-ViT-L/14 作为老师。微调后 IN-1K 的结果如下图 5 所示。可以看到：
1) 对于这 2 种模型，预训练蒸馏可以在不同的数据大小上持续带来性能增益。
2) 随着 Epoch 数量的增加，所有模型趋于饱和，这可能受到模型容量的限制。
稀疏软标签数量 K 的影响
通过使用 Swin-L 作为教师模型，TinyViT-21M 作为学生模型，在 IN-1K 和 IN-21K 上，观察到精度随着稀疏逻辑 \(K\) 的数量增加而增加，直到饱和，同时存储成本线性增加。为了在有限的存储空间下获得相当的精度，作者选择稍大的 \(K\) ，其中在 IN-1K 上 \(K=10\) (1.0% logits)，在 IN-21K 上 \(K=100\) (0.46% logits)，分别使用 16 GB/481GB 存储成本。
教师模型的影响
作者评估了教师模型对预训练蒸馏的影响。如图 7 所示。结论是：更好的老师可以产生更好的学生模型。然而，较好的教师模型通常模型尺寸较大，导致 GPU 内存消耗高且时间长 (比如具有 682M 参数的 Florence (#4) 占用 11GB GPU 内存，导致训练时间延长 2.4 倍)。
ImageNet 实验结果
如下图所示是 ImageNet-1K 数据集实验结果。蒸馏符号代表的是使用 CLIP-ViT-L/14 作为教师模型，在 ImageNet-21K 数据集上进行预训练蒸馏策略得到的实验结果。结果表明，在没有蒸馏的情况下，本文提出的 TinyViT 模型实现了与当前流行的方法相当的性能，例如 Swin transformer 和 LeViT ，且具有相似的参数。这表明了所提出的新架构和模型缩放技术的有效性。此外，通过快速预训练蒸馏，TinyViT 的性能可以大大提高，超过最先进的 CNN、Transformer 和 hybrid 模型。仅使用 21M 参数，在 IN-1K 上从头训练的 TinyViT 的精度相比于 Swin-T 和 DeiT-S 分别获得了 1.9%/3.2% 的提升，而在 IN-21K 上用蒸馏进行预处理后，提高了 3.6% 和 4.9%。
总结¶
本文提出一种快速知识蒸馏的方法，并说明了通过借助大模型的指导，小模型也能够获得大规模数据集带来的红利。具体而言，作者观察到当数据规模增加的时候，直接预训练小模型会带来性能饱和的问题。但是如果同时借助教师模型进行知识蒸馏，则能够为小模型释放大规模预训练数据的潜力，在使用超大规模数据集 ImageNet-21K 预训练的同时进行蒸馏操作可以大幅释放小型视觉 Transformer 的潜力。但是，以蒸馏的方式来预训练模型是低效且昂贵的，因为在这个过程中有相当大的比例的计算资源会消耗在训练数据在大模型的前向传播过程中，而不是消耗在训练学生模型上面。此外，庞大的教师模型可能会占用很多的 GPU 内存，因此显著降低学生的训练速度。为了解决这个问题，本文提出了一个快速且和可扩展的蒸馏策略。