【论文笔记】Tiny-ViT

Tiny-ViT

大概梳理一下Tiny-ViT的逻辑。

Tiny-ViT主要的贡献是提出了一个快速预训练蒸馏模块(Fast Pretraining Distillation)。

预训练蒸馏不同于微调阶段蒸馏，微调阶段蒸馏的目的一般是为了得到一个针对某个特定下游任务的模型，其通常发生在一个基础教师模型已经在某个任务上进行了预训练，并且在特定任务上进行了微调，然后通过知识蒸馏得到一个较小的学生网络，用于一些资源有限的设备上的任务。微调阶段蒸馏得到的学生模型并不具备泛化性和zero-shot能力。

预训练蒸馏是在大规模数据集上进行的，目的是在预训练阶段就将教师模型的知识传递给学生模型，显然预训练阶段蒸馏并不针对某个特定的下游任务，而是尽可能让教师模型以及学生模型学习到更多的特征，但不同于传统的使用GD的有监督训练，预训练蒸馏的学生网络是训练以拟合教师网络的输出logits。这一点也是论文中提到的，能够解决用大规模数据训练小型模型引发的过拟合或者混乱等一系列问题。因为大规模数据集存在两个问题（有一个问题是论文中提到的，但我记得我看到两个问题，不知道另一个为啥不见了）

• 存在非互斥的标签。论文中提到：The reason may be that not all the labels in IN-21k are mutually exclusive, including correlative pairs like chair and furniture. 这种情况会导致

• 相同的图像中可能存在多个主题对象，但只对应一个固定的标签，可能会使模型产生混乱

显然，soft target的逻辑能够解决或者至少缓解这两个问题，对于存在标签互斥的情况，两个标签都可以存在概率；后者同理。

但显然，单纯的预训练蒸馏的资源消耗是非常大的，这个资源包括存储资源和时间资源。作者在文中提到，Pretraining with distillation is inefficient and costly, because a considerable proportion of resouces is consumed on passing training data through the large teacher model in each iteration.

如果想要解决这个问题，就需要优化掉这个logits的传递过程。于是作者提出了Fast Pretring Distillation。将预训练蒸馏分为两个阶段，首先将大规模数据集数据增强后喂给教师模型，得到logits，并和input一一对应。然后将同样的增强后的图像作为输入，对应的Logits作为输出，作为学生模型的数据集进行训练。

那此时就会产生一系列问题，比如，数据增强是随机的，如果存储增强后的数据，就会需要较多的存储资源。但数据增强的随机是由随机参数导致的，如果能够记录一阶段的随机参数，就可以保证两个阶段的图像一致，论文也这么干了。

还有一个需要存储的就是输入图像对应的logits，作者不只是单纯存储输出的logtis，还进行了优化。论文中提到：It often consumes much storage space to save the whole dense logits of all augmented images if C is large, e.g., C = 21, 841 for ImageNet-21k. 事实上，能够提供有效信息的一般就是前几个相对较大的概率（论文中也进行了实验，证明了并不会有什么性能损耗）。于是作者只存储Logits中的前K个最大值以及对应的索引，学生模型训练时只重用这部分logits，其余概率值用Label Smoothing(Eq.2)补齐（应该是这么理解的）。当K<<C时，就能大幅减少logits的存储消耗（论文中取K=5，因为做了实验发现性能差距在0.5%以内）。

好，既然要存储下来的东西已经搞清楚了，那以一个什么样的形式存储呢？显然直接存储不符合“科研”的感觉，肯定要找个什么东西转换一下。论文中提到，用一个encoder和decoder来编码和解码这些参数实现压缩，这个encoder是[PCG](A family of simple fast space-efficient statistically good algorithms for random number generation)，将参数集编码为一个单参数（感觉这个可能是PCG针对某个随机数生成？不然这怎么可能编码的出来，看一下代码）。当然，上面对logits的优化也属于这个部分。这个部分需要结合代码理解。

还有就是，学生模型是什么结构。文中提到：we present a new family of tiny vision transformers by scaling down a large model seed with a progressive model contraction approach. 文中也对这个方法做了解释：Specifically, we start with a large model and define a basic set of contraction factors. Then in each step, smaller candidate models are generated around the current model by adjusting the contraction factors. We select models that satisfy both constraints on the number of parameters and throughput. The model with the best validation accuracy will be utilized for further reduction in the next step until the target is achieved. This is a form of constrained local search [32] in the model space spanned by the contraction factors.

上面两个内容在代码里乍一眼没看出来在哪儿（damn！）先来看看实验部分吧。

上面的内容中显然需要两个理论支撑：

• 限制小模型适应大规模数据的潜在原因是什么？

• 为什么蒸馏可以帮助小模型释放大规模数据集的潜力？

这两个问题个人认为是相符相成的，正是因为蒸馏能够解决大规模数据集中存在的一些问题，所以能够释放大规模数据集的潜力。

在论文中提到，大规模数据集（这里特指IN-21K）存在Hard Samples。除了上面提到的非互斥标签问题，还有一个问题：images with wrong labels and similar images with different labels due to the existence of multiple equally prominent objects in the images.

也就是说，IN-21K中存在错误标签以及相似图像匹配不同标签的情况，错误标签那必然就是标错了，要么是转换格式的时候哪里出错了，反正就是错了。而相似图像匹配不同标签的情况是由于图像中存在多个同等突出的目标造成的。

显然，hard samples会导致模型混乱，从而使得小模型难以适应大规模数据集。而soft targets就能解决这个问题，将硬标签转换为预训练的大模型教师模型推理得到的软标签，让学生模型拟合前K大概率，就能很好的规避原始标签出现的问题。并且挖掘出数据集的更多信息（类间知识）来指导学生模型。

作者对这部分内容也进行了实验，用Florence对IN-21K进行清洗，然后用清洗后的数据集训练TinyVit-21M以及Swin-T，结果如下图所示。这个实验能挖掘的信息还是挺多的，具体分析可以看原论文。