【网络结构】——TinyViT 一种transformer的蒸馏方案

来自 Microsoft

摘要

TinyViT是一种新型的小型transformer，它们高效且可以在大规模数据集上进行预训练。TinyViT通过在预训练期间应用蒸馏来传输知识，将大型预训练模型的知识转移到小型模型中，同时使小型模型能够获得大量预训练数据的红利。TinyViT在ImageNet-1k上的最高准确率为84.8％，仅使用21M参数，与在ImageNet-21k上预训练的Swin-B相当，但使用的参数少4.2倍。此外，增加图像分辨率后，TinyViT可以达到86.5％的准确度，略优于Swin-L，同时仅使用11％的参数。最后但并非最不重要的是，我们证明了TinyViT在各种下游任务上具有良好的转移能力。

1 简介

我们探讨了以下关键问题：如何有效地将现有大型变换器的知识转移到小型变换器上，并释放大规模数据的力量以提高小型模型的可表示性。
为了回答这个问题，我们引入了一种快速知识蒸馏方法来预训练小型模型，并表明小型模型也可以在大型模型的指导下获得海量预训练数据的红利。更具体地说，我们观察到，直接预训练小型模型会受到性能饱和的困扰，特别是当数据规模增加时。但是，如果我们在预训练期间施加蒸馏，并使用强大的模型作为教师，则可以为小型模型解锁大规模预训练数据的潜力。同时，蒸馏后的小型模型可以很好地转移到下游任务中，因为它们已经学到了如何从大型模型以及大规模预训练数据中推广的大量知识。

蒸馏的缺陷
使用蒸馏预训练模型既低效又昂贵，因为大量的计算资源被消耗在每次迭代中通过大型教师模型传递训练数据上，而不是训练目标小型学生。此外，巨大的教师可能占用大部分GPU内存，由于批处理大小有限，显著降低了学生的训练速度。

一种快速可拓展的蒸馏策略： 我们提出了一种快速且可扩展的蒸馏策略。具体而言，我们提前存储稀疏软标签和数据增强信息，例如随机裁剪、RandAugment [17]、CutMix [76]等。在训练过程中，我们重用存储的稀疏软标签和增强来精确复制蒸馏过程，成功省略了大型教师模型的前向计算和存储。这种策略有两个优点：1）快速。它大大节省了在训练期间生成教师软标签的内存成本和计算开销。因此，小型模型的蒸馏可以大大加速，因为它能够使用更大的批处理大小。此外，由于每个时期的教师逻辑是独立的，它们可以并行保存，而不是像传统方法那样逐个时期保存。2）可扩展。它可以模仿任何类型的数据增强并生成相应的软标签。我们只需要将大型教师模型向前推进一次，并为任意学生模型重用软标签。

我们验证了我们快速预训练蒸馏框架的有效性，不仅在现有的小型视觉变换器（如DeiT-T [64]和Swin-T [43]）上，而且在我们新设计的微型架构上

3 TinyViT

在这里插入图片描述
sparse logits ？？？

3.1 快速的预训练蒸馏

在这里插入图片描述

直接在大量数据上预训练小型模型并不能带来太多收益，特别是在将它们转移到下游任务时，如上图所示。为了解决这个问题，我们求助于知识蒸馏来进一步揭示小型模型预训练的力量。与之前大多关注微调阶段蒸馏[64]的工作不同，我们关注预训练蒸馏，它不仅允许小型模型从大规模模型中学习，而且提高了它们在下游任务中的转移能力。

我们的框架是无标签的，即不需要真实标签，因为我们只使用教师模型生成的软标签进行训练。因此，它可以利用大量现成的无标签网络数据进行大规模预训练。这种无标签策略在实践中是可行的，因为软标签足够准确，同时携带了大量用于分类的判别信息，例如类别关系。我们还观察到，带有真实标签的蒸馏会导致轻微的性能下降。原因可能是，并非ImageNet-21k [18]中所有标签都是互斥的[55]，包括“椅子”和“家具”，“马”和“动物”等相关对。因此，one-hot标签不能精确描述一个对象，在某些情况下它会在训练过程中抑制子类或父类。

稀疏的软标签 考虑教师模型为预测输出C个logits。如果C很大，例如ImageNet-21k的C = 21,841，则保存所有增强图像的整个密集logits通常会消耗大量存储空间。因此，我们只保存logits中最重要的部分，即稀疏软标签。形式上，我们选择yˆ中的前K个值，即{yˆI(k)} K k=1 ∈ yˆ，并将它们及其索引{I(k)} K k=1存储到我们的标签文件中
在这里插入图片描述

可以看到K的大小对TOP-1的正确率影响较小

3.2 网络结构

我们在本小节中提出了一种新的微型视觉变换器家族，通过使用渐进模型收缩方法[21]来缩小大型模型种子。具体来说，我们从一个大型模型开始，并定义一组基本的收缩因子。然后在每个步骤中，通过调整收缩因子在当前模型周围生成更小的候选模型。我们选择满足参数数量和吞吐量两个约束的模型。具有最佳验证精度的模型将用于下一步进一步减少，直到达到目标。这是一种受限局部搜索[32]，在由收缩因子张成的模型空间中进行。

我们采用分层视觉变换器作为基本架构，以方便需要多尺度特征的密集预测下游任务，如检测。更具体地说，我们的基础模型由四个阶段组成，分辨率逐渐降低，类似于Swin [43]和LeViT [24]。补丁嵌入块由两个卷积组成，核大小为3，步长为2，填充为1。我们在第1阶段和下采样块中应用轻量级和高效的MBConvs [33]，因为早期层的卷积能够有效地学习低级表示，这归功于它们强大的归纳偏差[24,70]。最后三个阶段由变换器块构成，带有窗口注意力以减少计算成本。注意偏差[24]和注意力与MLP之间的3×3深度卷积被引入以捕获局部信息[69,15]。残差连接[28]应用于第1阶段的每个块以及注意力块和MLP块。所有激活函数都是GELU [30]。卷积和线性的规范化层分别为BatchNorm [35]和LayerNorm [3]。

CNN归纳偏差的解释:
CNN 的归纳偏置是指 CNN 的架构本身所包含的一些对图像数据的假设或先验，这些假设或先验有助于 CNN 学习图像特征，但也可能限制了 CNN 的泛化能力和灵活性。CNN 的归纳偏置有两种，分别是：

局部相关性：邻近的像素是相关的，所以 CNN 用滑动窗口在图片上做卷积，提取局部特征。
平移等变性：图像的不同部分应该以相同的方式处理，无论它们的绝对位置如何，所以 CNN 用权重共享的卷积核来实现平移不变性。
相比之下，基于自注意力机制的视觉模型（如 DeiT 和 DETR）最小化了归纳偏置，因为它们不依赖于卷积操作，而是用全局的自注意力来捕捉图像中任意两个区域之间的关系。这样的模型可能需要更多的数据来学习有效的特征表示，但也可能具有更好的适应性和创新性。根据一篇论文《Rethinking Attention with Performers》中提到了一个例子：

例如，在一个涉及到识别猫和狗的任务中，如果训练数据中所有的猫都是黑色的，而所有的狗都是白色的，那么一个具有强归纳偏置的模型（如 CNN）可能会忽略其他特征，而只关注颜色。这样的模型在测试数据中遇到黑色的狗或白色的猫时就会出错。而一个具有弱归纳偏置的模型（如基于自注意力机制的模型）可能会考虑更多的特征，如形状、纹理、眼睛等，从而在测试数据中表现得更好。

4 讨论

讨论了两个关键问题：1）限制小模型拟合大数据的潜在因素是什么？2）为什么蒸馏能够释放大数据对小模型的能力？

在 IN-21k 中存在许多难样本，例如标签错误的图像和由于图像中存在多个同等突出的对象而具有不同标签的相似图像。这也被现有文献 [77,5,55] 所认可，并且 ImageNet 中约有 10% 的图像被认为是难样本。小模型难以拟合这些难样本，导致其训练精度低于大模型（TinyViT-21M：53.2% vs. Swin-L-197M [43]：57.1%），并且在 ImageNet-1k 上的可迁移性有限（TinyViT-21M w/ 预训练：83.8% vs. w/o 预训练：83.1%）。为啥他这种蒸馏方式，迁移可移植

难样本猜想实验
作者采用了两种技术。1）受 [5] 的启发，我们利用强大的预训练模型 Florence [75] 在 ImageNet-21k 上微调，以识别那些标签位于 Florence 的前五个预测之外的图像。通过这种方式，我们从 ImageNet-21k 中删除了 200 万张图像，约占 14%，然后在清理后的数据集上预训练 TinyViT-21M 和 Swin-T。2）我们使用 Florence 作为教师模型进行蒸馏以预训练 TinyViT-21M/Swin-T，它生成软标签以替换 ImageNet-21k 中污染的真实标签。
在这里插入图片描述

1）在原始 ImageNet-21k 数据集上预训练小模型对 ImageNet-1k 的性能提升有限（Swin-T 和 TinyViT-21M 均为 0.7%）。
2）在 ImageNet-21k 中删除部分难样本后，两种模型都能更好地利用大数据并获得更高的性能提升（Swin-T/TinyViT-21M 分别为 1.0%/1.1%）。

小模型，大数据为啥蒸馏有效
蒸馏可以提高小型模型在大型数据集上的性能，因为学生模型可以直接从老师那里学习领域知识。也就是说，老师在训练学生时注入了类关系先验，同时为小型学生模型过滤噪声标签（硬样本）。为了分析老师预测的类关系，我们从IN-21k中每个类别选择8张图像，共有21,841个类别。然后将这些图像输入Florence [75]以提取预测对数。根据[63]，我们可以在预测对数上生成类之间的皮尔逊相关系数的热图。在图3（a）中，相似或相关的类之间明显具有高相关性（红色），由块对角线结构说明。此外，老师模型还能捕获不相关的类（蓝色显示）
在这里插入图片描述

5. 实验

快速蒸馏的效果
在这里插入图片描述
图像数量的影响

所有模型都倾向于随着epoch数量的增加而饱和，这可能受到模型容量的限制。

logit k的影响
在知识蒸馏过程中，教师模型的对数值捕捉了类别关系，但也包含噪声。这使得可以稀疏化教师对数值，以保留类别关系并减少噪声。此外，内存消耗也对K的选择施加了限制。为了在有限的存储空间下获得可比较的准确性，我们选择稍大一些的K，在IN-1k上进行300个epoch时K=10（1.0%对数），在IN-21k上进行90个epoch时K=100（0.46%对数），分别使用16GB/481GB的存储成本。

教师模型的影响
在这里插入图片描述