论文题目:Integrally Pre-Trained TransformerPyramid Networks
论文地址:https://arxiv.org/pdf/2211.12735.pdf
代码地址:https://github.com/sunsmarterjie/iTPN
近年来视觉识别领域两个比较重要的成果是:1.用ViT(VisionTransformer)结构作为backbone来构架网络;2.用MIM(masked image modeling)来做预训练。通过将这两种方法结合起来的各种模型在图像分类、目标识别、语义分割等任务上都取得了SOTA结果。
但这种模型构建的方法有一个非常重要的问题,就是在上游的预训练任务和下游的fine-tuning之间存在很大的transfer gap。而本文作者则认为在下游的图像识别尤其是fine-scaled的识别中是需要提取出多尺度的特征的,但现存的大多数的预训练任务都是在plain的ViT架构上完成的,缺乏多尺度特征。而少部分基于hierachical ViT架构的研究,其预训练也只涉及backbone而忽略了neck部分。这就可能导致后续的fine-tuning过程是从一个随机的neck开始的,这样就不能保证neck和backbone部分是相匹配的。
基于上述的问题,本文作者提出了基于HiViT,利用特征金字塔,且适用于MIM的全局预训练网络——iTPNs(Integrally Pre-Trained TransformerPyramid Networks)。为了同时优化backbone(HiViT)和neck(feature pyramid)部分,作者主要采取了以下两种思路:
1.通过在预训练阶段加入一个特征金字塔,将学习到的权重参数在后续识别阶段复用,来统一上下游任务的neck部分;
2.在MIM的基础上提出MFM(masked feature modeling)来更好的完成预训练任务。
假设预训练时使用的没有label的数据集可表示为,其中N是数据集样本的总数;而在下游fine-tuning时使用的数据集则可以表示为
,其中M是数据集样本总数,
是
对应的标签。同时假定所使用的网络包含backbone,neck和head三部分,分别表示为
,
,
,其中
为对应部分的可学习参数,而整个网络表示为
.
在上述的表示体系下,原有的方法中,预训练和fine-tuning任务只共享相同的backbone,而neck和head部分均不相同。那么预训练过程可被表示为:
fine-tuning过程可被表示为:
这种思路会导致在预训练和fine-tuning之间存在transfer gap,且会有以下两个问题:
a.backbone参数无法被用于多尺度特征提取;
b.fine-tuning中都是随机初始化,可能会降低训练速度并导致参数与识别结果不匹配。
为了解决这些问题,本文作者提出加入了特征金字塔的全局网络使得预训练和fine-tining的neck部分统一,可共享参数,这样在fine-tuning时就只有是随机初始化的了。
假定HiViT包含S个stages,每个stages包含若干给transformer block。那么在逐步地下采样的过程中,HiViT会产生S+1个特征图,表示为:
其中,表示输入的embedding结果,而其它的特征图索引越小,越接近输入层。每个特征图都有特征向量构成表示为
,其中
表示第i个特征图的特征向量数。基于此,本文作者利用这些特征图构造了预训练和fine-tuning共同使用的neck部分:
其中,通过上采样来符合
的尺寸。而分层的可学习参数
则在预训练中学习后复用到fine-tuning中,大大减小了两者之间的transfer gap。值得注意的是,除了接下来讨论的loss之外,预训练和fine-tuning还有以下不同,但都不影响整体模型设计。
a.在用MIM进行预训练时,用掩模将样本中的部分区域遮盖掉。因此,在backbone和neck的输出中并不包含掩模部分,在neck输出的部分,采取添加dummy tokens的方式来完善输出,并通过一个decoder进行图像重建来完成预训练过程,
b.不同的下游fine-tuning任务中使用了不同的neck输出。例如,在图像分类中只使用了,而在图像分割中则使用
.
MFM是在MIM的基础上进行了改进。首先它继承了MIM通过掩模,用dummy tokens填充后的对图像重建的部分,这部分的损失函数可由下式表达:
在此基础上,在各层特征图的末尾都加入一个head,来进行掩模特征图的重建,其损失函数表达式如下所示:
其中,是第i层的期望输出,由一个teacherbackbone生成,一般有两种方式得到:
a.使用一个moving-averaged的encoder(不引入额外的konwledge),此时只需要将图像掩模部分输入teacher backbone
b.使用另外一个预训练模型(如CLIP等),此时将整个图像都输入到CLIP等backbone
值得注意的是,intergal train和MFM对于网络性能的优化是独立其互相补充的,这一点作者在后续的消融实验中进行了验证。整个过程的损失函数表示为:
其中为MIM的损失函数,
为MFM的损失函数,
为权值参数,本文取值为0.3.
本文采用的backbone为HiViT相对于传统的Swintransformer有以下两点优化:
a.用通道多层感知器(C-MLPs)代替了前几个stage的位移attention窗口;
b.移除了7*7stage,使得全局attention在14*14stage进行。
这些优化使得掩模tokens可直接由输入的得到,大大减少了计算量并提升了性能。同时,作者还使用C-MLPs代替了特征金字塔中的卷积开避免掩模时的信息泄露,在各种识别任务中都带来的性能的提升。具体的参数对比如图1所示:
图1 iTPN与其他模型的参数对比
对于二.3中的两种teacher model其具体过程是:
a. 对于第一种思路,是计算系数为0.996的在线目标模型的指数移动平均数(EMA)是从每个阶段的最后一层提取监督,因此对于任何具有与
相同的空间分辨率;
b.对于第二种思路,CLIP并不产生多尺度特征图,因此是将所有的特征图降采样到最低空间尺度(14*14)加和后于CLIP的最后一层输出进行比较学习。
整个iTPN的网络结构如下图2所示:
图2 iTPN与卷积预训练的比较
图中左侧是传统的卷积预训练,右侧是本文提出的iTPN结构。相较于左图,iTPN利用右图黄框所代表流程复用特征金字塔,使得特征金字塔成为预训练和fine-tuning共同的neck部分,用右图蓝框中的MFM补充左图蓝框中的MIM,并在特征图之间用C-MLP来代替卷积,进一步提高性能。
本文使用ImageNet-1K数据集对iTPN进行预训练。每个训练图像被预处理成224×224,并被分割成14×14个patches,每个patch大小为16×16;而在掩模过程中,75%的patches将被随机掩盖。
在ImageNet-1K分类任务中,无论是base-level model还是large-level model,使用moving-averaged的encoder或CLIP作为teachermodel,在相近的训练量下,iTPN的性能都超过了现有方法。具体结果如下图3和图4所示:
图3 base-levelmodel在ImageNet-1K上的分类结果对比
图4 large-levelmodel在ImageNet-1K上的分类结果对比
另外,为了更好地验证iTPN的性能,作者采用了线性探测的方法来进行测试,即冻结预训练的backbone,在fine-tuning时只对线性分类器进行学习。结果发现iTPN的结果还是最优,具体结果如下图5所示:
图5 linearprobing测试结果
因为,分类任务并不涉及neck部分的变换,则上述结果说明:
A.backbone和neck的联合优化可以预学习到更好的backbone;
B.iTPN的预训练backbone可以直接转移到各种视觉任务,将iTPN扩展到更多的应用场景。
无论是在COCO上做图像识别和分割任务,还是在ADE20K上语义分割,iTPN都取得了超过现有方法的性能。具体结果如图6和图7所示:
图6 base-levelmodel在COCO和ADE20K的结果对比
图7 large-levelmodel在在COCO和ADE20K的结果对比
为了验证iTPN常用的不同优化策略对性能的影响。作者将研究使用全局预训练(iPT)、load多尺度特征,MFM,C-MLP对性能的影响。具体结果如下图8和图9所示。
图8 iPT和load特征对性能的影响
上述结果说明backbone主干本身被iTPN加强了,因此它可以独立于neck转移到下游任务中。
图9 MFM和C-MLP对性能的影响
上述结果说明C-MLP和MFM 可分别独立优化性能且互为补充。
图10 iTPN,w/o iPT,MAE提取的特征图对比
通过上图10的attention图对比可以发现:
a.在encoder上,iTPN显示了在ImageNet上检测完整对象并在COCO上关注所选对象的优势。这因为iTPN迫使模型保留更丰富的视觉特征(多尺度特征图),这有助于在下游产生更好的识别结果。
b.在decoder上,iTPN仍然可以实现标记之间的语义关系,从而得到更好的重构结果,这是因为经过预训练的neck聚集了多阶段的视觉特征
图11 iTPN,w/o iPT,MAE损失函数和正确率随epoch的变化
图11对于损失函数的下降过程的追踪表明:
a. 在左图在预训练的重建过程中,iTPN获得了更好的图像重建结果(即较低的重建损失值)。简单地使用hierarchical ViT(带有多尺度特征图)并不能改善重建结果,这意味着完整的预训练是主要的贡献者。
b.在右图的fine-tuning任务中,更好的物体描述有助于下游的视觉识别任务(例如,物体检测)更快地收敛并达到更高的上限
在本文中,作者提出了iTPN网络(Integrally Pre-Trained TransformerPyramid Networks)。其核心贡献在于使用一个特征金字塔作为neck部分来进行预训练过程的重建并复用于fine-tuning的识别,从而最大限度地减少了预训练和fine-tuning之间的转移差距。此外,还设计了MFM来补充MIM,以更好地优化特征金字塔。iTPN在一些流行的视觉识别任务中具有优越的性能。这项工作引出了一个新的未来发展方向——为上游和下游的视觉表示学习设计一个统一的框架。