【写在后面】
视觉 Transformer(VIT)因为其卓越的建模能力，近年来在计算机视觉畛域引起了极大的关注。然而，大多数风行的 VIT 模型都受到大量参数的限度，限度了它们在资源无限的设施上的适用性。为了缓解这一问题，作者提出了 TinyViT，这是一种新的渺小而高效的小视觉 Transformer 家族，应用本文提出的疾速蒸馏框架在大规模数据集上进行预训练。其核心思想是将常识从大型预训练的模型转移到小型模型，同时使小型模型可能从大量的预训练数据中取得红利。更具体地说，作者在预训练期间利用蒸馏来进行常识转移。大型老师模型的 logits 被稠密并预存储在磁盘中，以节俭显存老本和计算开销。小的学生 Transformer 是从具备计算和参数束缚的预训练的大型模型中主动放大的。综合试验证实了 TinyViT 的有效性。它在 ImageNet-1k 上仅用 21M 个参数就达到了 84.8% 的 TOP- 1 准确率，与 SwinB 在 ImageNet-21k 上的预训练相当，而应用的参数少了 4.2 倍。此外，进步图像分辨率，TinyViT 能够达到 86.5% 的准确率，略好于 Swin-L，而只应用 11% 的参数。最初，作者展现了 TinyViT 在各种上游工作上的良好迁徙能力。
1. 论文和代码地址
TinyViT: Fast Pretraining Distillation for Small Vision Transformers
论文地址：https://arxiv.org/abs/2207.10666
代码地址：https://github.com/microsoft/Cream/tree/main/TinyViT
2. Motivation
Transformer 曾经在计算机视觉畛域掀起了一场风暴，并在钻研和实际中越来越受欢迎。视觉 Transformer(VIT)的最新趋势之一是持续减少模型大小，同时在规范基准上产生更好的性能。例如，V-MoE 应用 3.05 亿张图像训练一个具备 147 亿个参数的超大模型，实现了最先进的图像分类性能。同时，Swin 应用 30 亿个参数和 70M 个预训练图片，为了在上游检测和分段工作上取得不错的后果。如此大的模型尺寸和随之而来的昂扬预训练老本使得这些模型不适宜用于无限计算估算的利用，例如挪动和物联网边缘设施。
与大规模放大模型不同，这项工作将注意力转向放大视觉 Transformer 的尺寸，旨在生成一系列新的微型模型，并进步它们在上游工作中的迁徙能力。具体来说，作者探讨了以下关键问题：如何无效地将现有大型 Transformer 的常识转移到小型 Transformer 上，以及如何开释大规模数据的力量来进步小型模型的代表性？在计算机视觉中，很早就意识到，在大数据集上预训练的大模型往往能取得更好的后果，而小模型随着数据的增长很容易变得饱和(或有余)。小模型有没有可能从海量数据中汲取常识，进一步展现本人的能力？
为了答复这个问题，作者引入了一种疾速的常识提取办法来对小模型进行预训练，并证实了在大模型的领导下，小模型也能够取得海量预训练数据的红利。更具体地说，作者察看到小模型的间接预训练受到性能饱和的影响，特地是当数据规模减少时。然而，如果在预训练中进行蒸馏，应用一个弱小的模型作为老师，大规模预训练数据的后劲能够被开释到小模型，如上图所示。同时，提取的小模型能够很好地转移到上游工作，因为它们曾经学习了如何从大模型以及大规模预训练数据中进行泛化的大量常识。
应用蒸馏的预训练模型效率低且老本高，因为在每次迭代中，相当大比例的计算资源耗费在通过大型老师模型传递训练数据上，而不是训练指标小学生。此外，巨型老师可能会占用最多的 GPU 内存，显著减慢学生的训练速度 (因为 batch 大小无限)。为了解决这个问题，作者提出了一种疾速且可扩大的精馏策略。更具体地说，作者提出事后生成一个稠密概率向量作为每幅输出图像的软标签，并将其和相应的数据加强信息一起存储到标签文件中。在训练过程中，作者重用存储的稠密软标签和裁减来准确复制蒸馏过程，胜利地省略了大型老师模型的向前计算和存储。这种策略有两个长处：1) 速度快 。在很大水平上节俭了训练过程中生成老师软标签的内存开销和计算开销。因而，小模型的蒸馏速度能够大大放慢，因为它可能应用大得多的 batch。此外，因为每个时代的老师 log 是独立的，所以能够并行保留，而不是传统办法中逐 epoch 地保留。2) 可拓展。它能够模仿任何类型的数据加强，并生成相应的软标签。只须要将大型老师模型前向流传一次，并对任意学生模型重用软标签。
作者验证了本文的疾速预训练蒸馏框架不仅在现有的小型视觉 Transformer（如 DeiT- T 和 Swin-T）上，而且在新设计的微型架构上的有效性。具体地说，作者采纳渐进式模型膨胀办法来放大大型模型并生成一族渺小视觉 Transformer(TinyViT)。通过在 ImageNet-21k 上的疾速预训练蒸馏，具备 21M 参数的 TinyViT 在 ImageNet-1k 上达到了 84.8% 的 TOP- 1 准确率，比预训练的 Swin-B(88M 参数下的 85.2%)小 4.2 倍。在更高的分辨率下，本文的模型能够达到 86.5% 的 TOP- 1 精度，在对齐设置下在 ImageNet-1k 上建设了新的最先进的性能。此外，TinyViT 模型在上游工作上体现出了良好的迁徙能力。例如，TinyViT-21M 在指标检测基准上的 AP 为 50.2，应用 28M 参数时比 Swin- T 高 2.1 个百分点。
总而言之，这项工作的次要奉献有两个：
1）为了充分利用大规模的预训练数据，开释小模型的能力，作者提出了一种疾速预训练蒸馏框架。这是摸索小模型预训练的第一个工作。
2）作者公布了一系列新的渺小视觉 Transformer 模型，它们在计算和精度之间找到了很好的折衷。在预训练精馏的状况下，这种模型在上游工作上体现出了良好的转移能力。
3. 办法
3.1 Fast Pretraining Distillation
作者察看到，在海量数据上间接对小模型进行预训练并没有带来太大的收益，特地是当它们转移到上游工作时，如上图所示。为了解决这个问题，作者诉诸常识蒸馏来进一步揭示小模型预训练的力量。不同于以往侧重于微调阶段蒸馏的工作，作者将重点放在预训练蒸馏上，它不仅容许小模型向大尺度模型学习，而且进步了它们对上游工作的迁徙能力。
应用蒸馏进行预训练效率低且老本高，因为在每次迭代中，相当一部分计算资源被耗费在通过大型老师模型传递训练数据上，而不是训练指标小学生。此外，老师可能会占用最多的 GPU 内存，从而减慢指标学生的训练速度(因为 batch 大小无限)。为了解决这一问题，作者提出了一种疾速预训练蒸馏框架。如上图所示，作者事后存储了数据裁减和老师预测的信息。在训练过程中，作者重用存储的信息来准确复制蒸馏过程，胜利地省略了大型老师模型的前向计算和内存占用。
在数学上，对于具备强数据裁减 $\mathcal{A}$ 的输出图像 x，例如 RandAugment 和 CutMix，作者存储 $\mathcal{A}$ 和老师预测 $\hat{\mathbf{y}}=T(\mathcal{A}(x))$，其中 $T(\cdot)$ 和 $\mathcal{A}(x)$ 是老师模型和裁减图像。值得注意的是，将同一图像屡次通过雷同的数据加强 pipeline 将产生不同的加强图像。因而，须要在每次迭代中为每个图像保留对 $(\mathcal{A}, \hat{\mathbf{y}})$，如上图所示。
在训练过程中，只须要从存储的文件中复原对 $(\mathcal{A}, \hat{\mathbf{y}})$，并优化以下指标函数用于学生模型蒸馏：
$$
\mathcal{L}=C E(\hat{\mathbf{y}}, S(\mathcal{A}(x)))
$$
其中 S(·)和 CE(·)别离是学生模型和穿插熵损失。留神，本文的框架是无标签的，即不须要 ground truth 标签，因为只应用老师模型生成的软标签来进行训练。因而，它能够利用大量没有标签的现成网络数据进行大规模的预训练。这种无标签策略在实践中是可行的，因为软标签足够精确，同时携带了大量用于分类的区别性信息，如类别关系。作者还察看到，与 ground truth 蒸馏将导致轻微的性能降落。起因可能是并非 ImageNet-21k 中的所有标签都是互斥的，包含相干的对，如“椅子”和“家具”，“马”和“动物”。因而，变化无穷的 ground truth 标签不能精确地形容对象，在某些状况下，它在训练期间克制子类或父类。此外，因为在训练过程中去掉了繁琐的老师 $T(\cdot)$，因而本文的蒸馏框架与没有蒸馏的训练模型一样快。
此外，因为两个要害组件：稠密软标签（Sparse soft labels）和数据加强编码（Data augmentation encoding），本文的蒸馏框架是疾速的。它们能够在极大地缩小存储耗费的同时，进步训练期间的内存效率。
Sparse soft labels
设老师模型为预测输入 C 个 logits。如果 C 很大，例如，对于 ImageNet-21k，C=21,841，则通常要耗费大量存储空间来保留所有加强图像的全副密集 logits。因而，作者只保留了 logits 中最重要的局部，即稠密软标签。模式上，作者抉择 $\hat{\mathbf{y}}$ 中的前 K 个值，即，$\left\{\hat{y}_{\mathcal{I}(k)}\right\}_{k=1}^{K} \in \hat{\mathbf{y}}$，并将它们与它们的索引 $\{\mathcal{I}(k)\}_{k=1}^{K}$ 一起存储到标签文件中。在训练过程中，作者只应用存储的稠密标签进行标签平滑的蒸馏，这被定义为：
$$
\hat{y}_{c}=\left\{\begin{array}{ll}\hat{y}_{\mathcal{I}(k)} & \text {if} c=\mathcal{I}(k) \\ \frac{1-\sum_{k=1}^{K} \hat{y}_{\mathcal{I}(k)}}{C-K} & \text {otherwise}\end{array}\right.
$$
其中 $\hat{y}_{c}$ 是为学生模型蒸馏复原的老师 logits，即 $\hat{\mathbf{y}}= \left[\hat{y}_{1}, \ldots, \hat{y}_{c}, \ldots, \hat{y}_{C}\right]$。当稠密因子 K 较小时，即 K≪C 时，能够将 Logit 的存储量缩小数量级。
Data augmentation encoding
数据加强波及一组参数 d，例如旋转度和裁剪坐标，以变换输出图像。因为每次迭代中每个图像的 d 都不同，因而间接保留它会升高内存效率。为了解决这个问题，作者用单个参数 $d_{0}=\mathcal{E}(\mathbf{d})$ 来编码 d，其中 $\mathcal{E}(\cdot)$ 是上图中的编码器。而后在训练过程中，在存储文件中加载 $d_{0}$ 后复原 $\mathbf{d}=\mathcal{E}^{-1}\left(d_{0}\right)$，其中，$\mathcal{E}^{-1}(\cdot)$ 被视为解码器。因而，能够精确地重建数据加强。实际上，解码器的常见抉择是伪随机数生成器（即 PCG）。它将单个参数作为输出，并生成一系列参数。对于编码器，只需通过 $d_{0}$ 的生成器并重用解码器 $\mathcal{E}^{-1}(\cdot)$ 来实现。输入 $\mathbf{d}=\mathcal{E}^{-1}\left(d_{0}\right)$ 示意老师模型。保留 $d_{0}$ 后，以便解码器在训练学生时再现 d。因而，实现变得更高效。
3.2 Model Architectures
作者通过应用渐进模型膨胀办法放大大型模型种子，提出了一个新的微型视觉 Transformer 家族。具体地说，作者从一个大型模型开始，并定义了一组根本的膨胀因子。而后在每一步中，通过调整膨胀因子，围绕以后模型生成较小的候选模型。作者抉择同时满足参数数量和吞吐量限度的模型。验证精度最高的模型将被用于下一步的进一步简化，直到达到目标。这是在由膨胀因子逾越的模型空间中的一种受限部分搜寻模式。
作者采纳分层视觉 Transformer 作为根本架构，以不便密集预测的上游工作，如检测，须要多尺度特色。更具体地说，本文的根本模型由四个阶段组成，分辨率逐步升高，相似于 Swin 和 LeVit。块嵌入由两个卷积组成，核大小为 3，步长为 2，填充为 1。作者在阶段 1 和下采样块中利用了轻量级和高效的 MBConv，因为较早层的卷积因为其强烈的演绎偏差而可能无效地学习低级示意。后三级由 Transformer 块形成，并对留神加窗以缩小计算量。
Contraction factors
膨胀因子。作者思考了以下因子来造成一个模型：
$\gamma _{D_{1-4}}$：别离为四个阶段的嵌入维度。缩小它们会导致一个更薄的网络和更少的头在多头自注意力。
$\gamma_{N_{1-4}}$：别离为四个阶段的块数。通过减小这些值能够升高模型的深度。
$\gamma_{W_{2-4}}$：最初三个阶段的窗口大小。随着这些值变小，该模型具备更少的参数和更高的吞吐量。
$\gamma_{R}$：MBConv 块的通道裁减率。通过减小这一因子，能够失去较小的模型尺寸。
$\gamma_{M}$：所有 Transformer 组的 MLP 膨胀率。如果放大该值，MLP 的暗藏尺寸将更小。
$\gamma_{E}$：多头留神中每个头的大小。当放大时，头的数量将减少，从而带来更低的计算成本。
4. 试验
上表展现了难例样本的影响。
上表展现了不同训练策略的融化试验后果。
上图展现了 ImageNet-21k 上输入预测的皮尔逊相关性。
上图展现了 TinyViT-21M 的 IN-1K 精度和存储老本随不同的存储对数 K 的变动状况。
上表展现了不同老师模型的影响。
上图展现了本文办法在 IN-1k 上和 SOTA 后果的比照图。
上表展现了本文办法在 IN-1k 上和 SOTA 后果的比照。
上表展现了 TinyViT-21M 有无与训练进行 linear probe 和 few-shot 图像分类的后果。
上表展现了本文办法在 COCO 数据集上进行指标检测的后果。
5. 总结
作者利用本文提出的疾速蒸馏框架 TinyViT，提出了一种新的渺小而高效的视觉 Transformer 家族，它在大规模数据集上进行了预训练。宽泛的试验证实了 TinyViT 在 ImageNet-1k 上的有效性，以及它在各种上游基准上的卓越可迁移性。
已建设深度学习公众号——FightingCV，欢送大家关注！！！
ICCV、CVPR、NeurIPS、ICML 论文解析汇总：https://github.com/xmu-xiaoma…
面向小白的 Attention、重参数、MLP、卷积外围代码学习：https://github.com/xmu-xiaoma…
退出交换群，请增加小助手 wx：FightngCV666
本文由 mdnice 多平台公布