写在前面
尽管在训练各种专门任务的模型已经取得了快速的发展，但学习一个适合多任务的单一通用模型对计算机视觉仍然具有挑战性。在本文中，作者引入了多任务自训练( multi-task self-training，MuST)，它利用多个独立的专业教师模型(如ImageNet分类模型)来训练一个多任务学生模型。本文的方法有三个步骤。
首先，在标记的数据集上独立训练专业教师模型；然后，使用专门教师模型来标记一个未标记的数据集，以创建一个多任务伪标记的数据集；最后，该数据集现了包含来自在不同任务上训练的教师模型提供的伪标签，因此就可以用这些伪标签来训练多任务学习的学生模型。
作者在6个视觉任务上评估了学生模型的特征表示，包括图像识别（分类、检测、分割)和三维几何估计(深度估计和表面法线估计）。MuST对于未标记或部分标记的数据集是可扩展的，并且在大规模数据集上进行训练时，表现优于专门任务的监督模型和自监督模型。
论文和代码地址
Multi-Task Self-Training for Learning General Representations
论文地址：https://arxiv.org/abs/2108.11353
代码地址：未开源
Motivation
视觉表示学习是计算机视觉中的一个核心问题。监督和自监督的预训练在将学习到的特征转移到下游任务方面显示出了良好的结果。尽管从监督训练中采用了广泛的迁移学习，但这些特征并不一定对下游任务有用 。
比如在ImageNet上用分类任务训练的特征并不一定适用于COCO的实例分割任务。但是，从Object 365检测数据集学习的特征在很大程度上提高了COCO实例分割的性能。直观地说，通过训练模型在多个任务上同时完成任务，就可以学习泛化能力强的特征。最近在NLP的工作揭示了，通过多任务学习，可以产生泛化能力更强的模型。在计算机视觉中，训练多任务模型的最大挑战是数据收集和标注。尽管有像COCO这样的数据集，但为同一图像数据集收集各种标注（如实例分割、关键点、图像标题）是相当具有挑战性的。由于使用标签标注图像很耗时，所以很难根据图像的数量和任务的数量来扩展这种多任务学习的能力。大规模多任务数据集的缺乏阻碍了计算机视觉多任务学习的进展。在这项工作中，作者研究使用自训练来解决这个问题，提出使用伪标签来实现计算机视觉多任务特征学习。之前也有工作观察到，自训练进一步提高了迁移学习的预训练特征，即使预训练结果不能优于随机初始化的模型，但是自训练也有效的。预训练和自训练之间的差距表明，自训练可以从伪标签中学习到更好的特征 。受这一观察结果的启发，作者首先研究是否只使用伪标签才能学习到好的特征。具体的，作者使用COCO或Object 365等数据集来训练教师模型，以在未标记的图像上生成伪标签（伪标签的例子如上图所示）。作者发现一个只使用这些伪标签训练的学生模型保留了其专业教师模型的大部分迁移学习性能 。这一发现表明，伪标签在有监督的数据集中提取知识方面是有效的。因此，可以使用伪标签将知识从多个教师模型转移到单个学生模型来进行表示学习。作者提出了多任务自训练( Multi-Task Self-Training，MuST)，根据从不同任务和数据集上训练的教师模型中提取的信息来训练一个多任务的学生模型。整体流程如上图所示主要分为三步：1）首先，在标记的数据集上独立训练专业教师模型 。比如说，一个教师模型可以通过深度估计任务进行训练，另一个教师模型可以通过目标检测任务进行训练。2）然后，专门的教师模型被用来标记一个更大的未标记数据集 ，以创建一个多任务伪标记数据集。例如，这些教师可以在ImageNet数据集中生成深度估计和目标检测的伪标签。3）最后，该数据集现在包含来自在不同数据集/任务上训练的教师模型的伪标签，可以用于训练具有多任务学习的学生模型 。例如，学生模型可以同时进行深度估计和目标检测任务的训练。在本文的实验中，有四种教师模型：分类、语义分割、目标检测和深度估计。作者设计了一个简单的基于ResNet和FPN的backbone结构，在不同任务中，backbone的参数共享。对于每个单独的任务，它都有一个小的特定于任务的head，由几个卷积层和一个线性预测层组成。本文的实验表明，这种简单的模型架构能够吸收共享backbone中不同任务的知识。多任务的学生模型在所有迁移学习任务上都与其专业教师模型相当甚至优于教师模型的性能。最近的自监督算法，如SimCLR，MoCo，被证明可以创建与监督算法相同或更好的表示。因此，作者也将结果与SimCLR相比，在分割和深度估计任务上也是达到了更好的效果。
方法
3.1.专业教师模型
作者希望从一组教师模型那里学习，他们提供丰富的训练信号与他们的伪标签。具体的，作者采用了四种教师模型，包括计算机视觉中的四个重要任务：分类、检测、分割和深度估计。
这些任务需要对目标和三维几何图形的视觉理解。作者在中/大规模数据集(ImageNet、Objecst365、COCO）上从头开始训练分类、检测和分割的教师模型。对于深度估计教师模型，作者从开源存储库下载预训练的权重。
伪标签
作者通过伪标记将专门的教师模型中的知识转移到未标记或部分标记的数据集。对于检测，使用0.5的分数阈值来生成伪框标签。对于分割，使用0.5的分数阈值来生成语义分割mask，而预测分数较小的像素被设置为忽略标签。
对于分类，使用软标签，它包含所有类的概率分布，因为作者发现软标签性能比one-hot形式的硬标签更好。对于深度估计，使用预测的深度作为伪标签，而无需进一步处理。
3.2.多任务学生模型
模型架构
本文的目标是训练学生模型获得多任务学习一般视觉表征。首先要设计的是一个可以跨任务共享大部分参数的模型体系结构。作者定义了三个任务类别：（1)分类、(2)目标检测、(3）像素级预测。像素级预测任务包括语义分割和深度估计。每个类别的任务在主干模型中共享相同的特征表示。作者设计了基于ResNet和特征金字塔网络(FPN)的主干模型，结构如上图所示。作者使用C5特征（橙色）进行分类，并使用{P3、P4、P5、P6、P7}特征金字塔（青色）进行检测。将{P3、P4、P5、P6、P7}融合到P2特征图（绿色）中，进行像素级预测。Fuse操作只需将所有特性映射重新缩放到第2级特征大小，并对它们进行相加（这不会引入任何新的参数）。分类head遵循ResNet的设计 。它是平均合并的C5特征后的一个线性预测层。目标检测任务遵循Mask R-CNN中的head结构 ，对RPN使用2个隐藏卷积层，对Faster R-CNN使用4个隐藏卷积层和1个全连接层。像素级预测head有3个卷积层，然后是最后的线性预测层 。如果学生模型从同一任务类别中的多个任务中学习（例如，语义分割和深度估计），那么每个任务都拥有其特定于任务的head，而不共享它们的参数。
Teacher-student 训练
因为学生模型需要从教师模型的伪标签中学到知识，因此，作者设计了师生模型的训练，使它们之间的主要区别在于数据集和标签。与模型蒸馏不同，作者在教师和学生训练中都使用了相同的模型容量和数据增强。尽管教师可以用更专门的结构来训练其自己的任务，但是作者使用了上图的相同架构来训练教师和学生模型。
从多个教师模型那里学习
作者提出了多任务自训练(Multi-Task Self-Training，MuST)来训练多个教师模型和一个学生模型。以前的多任务学习工作，利用多个数据集中的信息，主要集中在每个样本只标记一个任务标签或几个任务标签的场景上。
在MuST上，每个样本都有所有任务的标签，这些标签可能来监督的标签或教师模型生成伪标签。例如，在ImageNet上训练时，可以使用监督标签进行分类，可以使用伪标签进行检测、分割和深度估计。平衡每个任务的损失贡献是多任务学习中一个开放的研究领域。多任务学习的损失是所有任务损失的加权和：其中，b表示Batch Size，lr表示学习率，下标表示学生或教师模型。
跨数据集训练
MuST具有利用标记和未标记数据的灵活性。它可以通过在未标记的数据上生成伪标签来放大训练样本的数量。在实验中，作者展示了一个跨ImageNet、Object 365和COCO数据集的训练示例。作者在有监督标签时使用监督标签，并使用伪标签为所有缺失的任务生成标签。
跨数据集训练的一个挑战是平衡来自不同大小的数据集的数据。作者没有采用启发式采用，而是从数据集中联合统一采样 。因为每个任务在每张图像上都有标记的，因此不需要担心由于数据集大小不平衡而导致的任务采样不足/过采样。与其他自训练算法相比，第二个主要区别是监督标签和伪标签被相同地对待。作者没有独立地批处理监督标签和伪标签的样本，并给它们分配不同的权重。这些图像从数据集的联合中统一采样，并放入一个小Batch。每个样本对其损失具有相同的权重，无论损失是根据监督的还是根据伪标签。这简化了MuST的训练流程。
3.3.迁移学习
为了评估MuST和其他baseline的表征质量，作者在一套下游的计算机视觉任务上对它们进行了微调。对每个微调任务采用端到端地微调，作者对CIFAR-100分类、PASCAL检测、PASCAL语义分割、NYU深度估计、ADE语义分割和 DIODE表面法线估计任务上进行了微调。注意，所有下游数据集都不同于专门教师模型训练的数据集。
实验
4.1. Experimental Settings
4.2. Learning with Multi-Task Self-Training
MuST Improves Pre-training on ImageNet:
上表比较了MuST算法与ImageNet上的自监督和监督学习。在6个下游任务中，MuST分别在4个任务和5个任务上优于SOTA的自监督学习和监督学习。
MuST Improves With More Tasks/Datasets For Learning General Features
上表显示了从仅使用有监督的ImageNet标签开始，然后添加从三个不同数据集获得的三种不同类型的伪标签的性能改变。随着继续添加来自不同任务/数据集的伪标签，特征的质量有所提高。
Pseudo Labels Preserve Transfer Learning Performance of Teacher Model
接下来，作者研究了伪标签如何有效地保持在监督数据集上训练的教师模型的迁移学习性能。可以看出，随着得到越来越多的未标记数据，MuST可以很容易地利用它们，并提高了表征质量。
Multi-Task Self-Training Across Datasets
上表比较了在ImageNet、COCO和Object 365三个数据集上的MuST和监督多任务训练的表征质量。可以看出，伪标签可以有效的提升performance。
4.3. Scaling Multi-Task Self-Training
Scaling Dataset Size and Training Iterations
上表显示了JFT数据集上MuST与自监督学习和监督学习的比较。在5/6个下游任务中，当使用相同的未标记数据时，MuST优于自监督的SimCLR算法。
学生模型不仅可以学习了转移的一般特征，还能够为多个任务生成高质量的预测。上图显示了MuST模型的预测。
Bootstrapping from Pre-trained Models
接下来，作者研究MuST是否可以改进用数十亿个训练样本训练的模型权重。在6个下游任务中，MuST提高了4个ALIGN模型的性能。可以看出，用MuST微调ALIGN，可以达到与SOTA性能相似的结果。
总结
在本文中，作者提出了一种可扩展的多任务自训练方法MuST。作者将MuST与ImageNet和JFT上的监督和自监督学习方法进行了比较，并评估了包括视觉识别、定位和三维几何预测结果。
结果表明，除了分类任务外，在6个迁移学习任务中，有5个优于监督和自监督学习。此外，MuST能够进一步提升已经在数十亿个样本上训练完成的模型的性能。因此，多任务自训练是一种可扩展的预训练方法，能够学习泛化的特征表示。研究领域：FightingCV公众号运营者，研究方向为多模态内容理解，专注于解决视觉模态和语言模态相结合的任务，促进Vision-Language模型的实地应用。
知乎/公众号：FightingCV
