今天给大家分享一篇我们在CVPR 2022 上的论文，DATA: Domain-Aware and Task-Aware Pretraining。

问题提出

Contrastive learning是近两年关注度十分高的大规模无标签数据下的预训练方式，迁移到分类任务上虽然还没有超过有监督的ImageNet-pretrain，但是检测，分割这种密集预测下游任务中已经有超出有监督预训练的表现。当我们把自监督预训练落地到下游时，分类，检测，分割不同任务，点，边，云端不同硬件往往需要不同的模型，由于自监督的downstream-agnostic，需要我们根据具体场景提前选好模型，然后针对完全不同的模型分别做一次自监督预训练，这个过程中有两个痛点：

1） 小模型预训练后在下游上的性能远低于其进行有监督预训练后的表现（如图1）

2） 训练成本较高。对比学习预训练毕竟是无监督预训练，监督信号没有那么明确，网络收敛一般较慢（以MoCo v2 [1] 为例，常规setting一般为200 epoch，但是完全收敛一般需要800 epoch，是常规有监督预训练的八倍。）

为了解决上述痛点，我们引入神经网络架构搜索（NAS）到自监督训练当中，将NAS中的超网训练，自监督预训练，自蒸馏结合，三者相辅相成，很好的解决了我们上述的问题。

图1. 小模型自监督训练后在ImageNet上的性能表现

具体过程

问题初探

针对上面痛点，我们首先想到了在NAS领域的超网训练(BigNAS[2]，OFA[3], etc.)。我们借鉴BigNAS形式进行直接的尝试（OFA的pipeline过于冗长，整个过程不够简洁，我们尝试后弃用），看在对比学习这种较弱的监督信号下是否也能够一样训练出一个可裁剪出数十万个小模型的超网。为了快速验证方法的可行性，我们采用ImageNet-10% 作为训练集进行200 epoch的训练，然后进行linear evaluation，实验结果如下表所示：

表1. 超网训练和MoCo v2结合前后R50性能对比

可以看到实验结果并不理想，相比于单独的resnet50训练，采用超网训练后提取出的resnet50 进行linear evaluation的结果要低两个多点。我们分析原因主要是出在由于网络宽度深度处于不断变化，让EMA update的那个branch并不够稳定。与此同时，我们注意到Dino[4]把对比学习formulate成了一个teacher-student的蒸馏形式，ReSSL[5]和SEED[6] 又分别用完全相同的proxy task 完整自监督预训练和自监督knowledge distillation。从其中我们受到启发，将不进行梯度更新的那个branch formulate成一个稳定的teacher，它的网络结构保持为梯度更新branch最大的形式，达到稳定监督信号的同时又带有自监督蒸馏以促进更好的子网训练，于是我们得到了如下的实验结果：

表2. 固定EMA更新分支模型结构超网下R50与之前模型性能比较

可以看到效果比较明显, 从超网中提取出的resnet50 在linear evaluation上的表现已经超过了单独的resnet50 0.9% acc，这里我们实现了。并且最令我们兴奋的是由于自蒸馏的加持，超网中提取出的小模型表现性能十分优异，远超小模型单独进行预训练后的模型表现（见图4）。

新的问题

在我们成功训练了可以裁剪成数十万涵盖很大FLOPs范围不同子网的超网之后，这看起来是很令人兴奋的一件事。但是进一步发现，相比于supervised NAS中的超网训练之后，搜索只需要耗费很少的代价，这里的超网想要最终搜索出使用性能表现优异的子网还有着不小的代价。因为有监督训练完的超网提取出的不同子网只需要在验证集上进行inference，根据metric选取top 1即可，但是自监督预训练过后的数十万网络并没有metric去挑选出合适的子网，需要我们把这些预训练子网都重新在下游train一遍，再在下游验证集上进行inference选取top 1。目前计算机视觉领域的预训练还没有像NLP领域中的预训练发挥出同样的作用，预训练任务和下游任务还是存在着一个较大的gap，预训练模型往往只是作为一个好的初始化进行训练，这样来看挑选子网的过程花费让人难以接受，为了尽可能找到好的子网，需要的代价过于庞大。为此这里遇到一个新的问题——“我们怎么才能以尽可能少的代价在数万个预训练好的网络中选取最好的那个”？

重新思考

遇到问题，我们首先去搜索一下之前相关论文，看下别人有没有遇到类似的问题看能不能从中收到启发，经过一番检索，发现UnNAS[7] 是第一篇尝试如何从无标签的情况下选取最优网络的相关方法，这篇文章的motivation更多的关注在无标签上而不是关注在训练成本上，但是文中claim的观点给了我们很大的启示——”已经预训练过的模型，本身就包含有很多模型相关的信息在里面“，我们似乎可以不经过下游的有监督微调就可以给这个模型将会有的表现做出一些判断，但可惜的是对比学习预训练不像旋转预测，上色，解字谜这种自监督任务一样有明显的metric。

在21年，OpenAI release 了两个 大数据集下表现出惊艳性能的两个多模态大模型——CLIP，DALL·E，与此同时越来越多的研究者注意到了大模型的潜力，在如今数据规模下，似乎大模型可以带来无限的可能性，对于我们这些学校的实验室，肯定是没办法参与那些有钱大厂们玩的科研，但是神经网络当今发展的趋势却让我们突然意识到了一个事——大模型似乎总是好的？！ 

和最大的越像越好的无标签搜索

基于以上的思考，我们做出一个假设——“这么多可裁剪出的预训练子网，哪个和最大的teacher网络越像，哪个越可能是在下游finetune有更好表现的网络“。为此我们提出让子网在对应的下游数据集上进行inference，通过获取feature map和最大的teacher网络推理得到的feature map的相似度作为子网搜索的metric。我们在此基础上更进一步思考了一下，根据具体下游任务finetune时pretarined backbone最直接用到的输出feature map进行相似度计算（比如下游为FasterRCNN-FPN, 则选取C2~C5 feature map，下游为PSPnet + auxiliary head 的语义分割模型，则选取C4~C5）似乎是一个更加合理的选择。为此，在对应下游数据集上，根据下游任务用到的模型结构选取对应的feature map计算相似度就是我们最终的metric，这个metric也正对应了我们文章的domain-aware and task-aware。到这里从超网训练到最终的子网选取，我们整个完整的pipeline 就完成了（如下图3）。接下来就是机器拉满来验证我们idea的有效性了。

图3. 整个算法的完整管线图

实验结果

线性相关程度验证

我们从最一开始预训练好的超网中采样出大约五十个子网，然后在COCO instance segmentation，ImageNet semi-supervised classification, cityscapes semantic segmentation 这几个常用benchmark上进行了验证，根据最终子网的性能ranking和相似度的ranking计算皮尔逊相关系数，结果如下表3所示，可以看到尤其是当下游和上游数据集分布接近时，这种相关程度特别高，进一步验证了我们的假设有效性。

表3. 子网与最大网络相似度的ranking和子网最终性能的ranking的相似度

分类下游任务

我们把可裁剪出的子网按照input size为224x224下的FLOPs 每1GFLOPs范围内分成一组，每组随机采样100个子网，然后根据上述metric选取出相似度最高的子网，下为Linear evaluation和ImageNet-semi supervised 下的结果。相比于我们的MoCo v2 baseline 我们获得一致的提升。

图4. ImageNet数据集下的实验结果

检测分割下游任务

在COCO和cityscapes数据集上我们根据MoCo的setting，分别采用MaskRCNN-FPN和FCN进行进行实例分割和语义分割下游任务的测试。实例分割的结果如表4所示。

表4. COCO数据集实例分割结果

在cityscapes结果上（表5），我们注意到了4G~5G FLOPs 和 6G~7G FLOPs 下提取的模型性能下降的情况，与前面相关系数计算中该任务下相关系数相对较低对应了起来，我们分析一个很重要的原因在于ImageNet数据集和Cityscapes的数据集分布相差较大，影响了预训练模型在该数据集下的输出feature map质量从而导致了该现象，我们没有进一步深入的去研究，欢迎评论区探讨交流。

表5. Cityscapes数据集下语义分割结果

结语

到这里我们的文章就分享完了，论文中还有更详细的消融实验来证明我们方法的有效性，就不在这里详细列举，更多内容欢迎关注我们的原文和代码。

代码：https://github.com/GAIA-vision/GAIA-ssl

论文：https://arxiv.org/abs/2203.09041