【AI大模型】在线蒸馏一次视觉Transformer搜索

一、概要

最近，Pure transformers在视觉任务方面显示出巨大的潜力。然而，它们在中小数据集中的准确性并不令人满意。尽管一些现有方法引入了CNN作为教师，通过蒸馏来指导训练过程，但教师和学生网络之间的差距将导致次优绩效。在这项工作中，研究员提出了一种新的具有在线蒸馏的One-shot Vision transformer搜索框架，即OVO。OVO对教师和学生网络的子网进行采样，以获得更好的蒸馏结果。得益于在线蒸馏，超网中的数千个子网训练有素，无需额外的微调或再训练。在实验中，OVO Ti在ImageNet和CIFAR-100上分别达到73.32%和75.2%的top-1精度。

二、ViT超网训练的困境

在AutoFormer中的超网训练期间，在每个训练迭代中均匀地采样子网α=（α（1）。。。α（i）。。。α（l））。采样权重w=（w（1）。。。w（i）。。。w（l））被更新，而其余部分被冻结。然而，超网训练对于ViT来说并不稳定，每个超网需要很长的训练周期（500个周期）才能获得其组装子网的满意性能。其他方法使用三明治采样策略，对多个子网络（包括最大、最小和两个随机选择）进行采样，然后将每个小批次的梯度聚集在一起。当采用三明治采样策略时，超网训练的计算成本更高。

三、超网在线蒸馏训练

当训练教师和学生网络时，研究者在[Single path one-shot neural architecture search with uniform sampling]中提出的每次迭代中更新一个随机采样的子网络。使用GT标签训练来自教师超网络的子网络，同时使用KD训练采样的学生网络。配备在线蒸馏，一次性NAS能够以高效和有效的方式搜索Transform架构。与经典的一次性NAS方法相比，新提出的方法有两个优点。1） 更快的收敛。来自CNN的知识提供了电感偏置，这有助于每个Transform块比之前的独立训练更快地收敛。2） 更好的子网性能。通过在线蒸馏训练的子网可以在中小数据集上获得更好的性能。

四、Search Pipeline

在超网完成训练后，进行进化搜索，以最大化分类精度选择子网络。在进化搜索开始时，随机抽取N个架构作为种子。所有子网络都基于从超网继承的权重在验证数据集上进行评估。选择前k个架构作为父代，通过交叉和变异生成下一代。对于交叉，两个随机选择的父网络交叉，以在每一代中产生一个新网络。当进行突变时，父网络首先以概率Pd对其深度进行突变，然后以概率Pm对每个块进行突变，以产生新的架构。

五、实验

将搜索到的最优模型的性能与ImageNet上最先进的CNN和ViT的性能进行了比较。在ImageNet-1K上训练OVO的超集合，并使用指定的参数大小搜索目标Transform模型。在超网完成训练后，子网直接继承权重，而无需额外的再训练和其他后处理。下表中报告了性能。很明显，OVO比其他最先进的模型具有更高的精度。