论文学习记录——iTPN（2023年CVPR)

论文题目：Integrally Pre-Trained TransformerPyramid Networks

论文地址：https://arxiv.org/pdf/2211.12735.pdf

代码地址：https://github.com/sunsmarterjie/iTPN

一.Introduction

近年来视觉识别领域两个比较重要的成果是：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）来更好的完成预训练任务。

二.Approach

1.Integral Pre-Training

假设预训练时使用的没有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时就只有是随机初始化的了。

2.Unifying Reconstruction (pre-training) and Recognition（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输出。例如，在图像分类中只使用了,而在图像分割中则使用.

3.Masked Feature Modeling

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.

4.Details

A．HiViT

本文采用的backbone为HiViT相对于传统的Swintransformer有以下两点优化：

a.用通道多层感知器（C-MLPs）代替了前几个stage的位移attention窗口；

b.移除了7*7stage，使得全局attention在14*14stage进行。

这些优化使得掩模tokens可直接由输入的得到，大大减少了计算量并提升了性能。同时，作者还使用C-MLPs代替了特征金字塔中的卷积开避免掩模时的信息泄露，在各种识别任务中都带来的性能的提升。具体的参数对比如图1所示：

图1 iTPN与其他模型的参数对比

B．MFM中的teacher model

对于二.3中的两种teacher model其具体过程是：

a. 对于第一种思路，是计算系数为0.996的在线目标模型的指数移动平均数（EMA)是从每个阶段的最后一层提取监督，因此对于任何具有与相同的空间分辨率；

b.对于第二种思路，CLIP并不产生多尺度特征图，因此是将所有的特征图降采样到最低空间尺度（14*14）加和后于CLIP的最后一层输出进行比较学习。

5.Network

整个iTPN的网络结构如下图2所示：

图2 iTPN与卷积预训练的比较

图中左侧是传统的卷积预训练，右侧是本文提出的iTPN结构。相较于左图，iTPN利用右图黄框所代表流程复用特征金字塔，使得特征金字塔成为预训练和fine-tuning共同的neck部分，用右图蓝框中的MFM补充左图蓝框中的MIM，并在特征图之间用C-MLP来代替卷积，进一步提高性能。

三.Experiments

1.settings

本文使用ImageNet-1K数据集对iTPN进行预训练。每个训练图像被预处理成224×224，并被分割成14×14个patches，每个patch大小为16×16；而在掩模过程中，75%的patches将被随机掩盖。

2.Image Classification

在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扩展到更多的应用场景。

3.Detection and Segmentation

无论是在COCO上做图像识别和分割任务，还是在ADE20K上语义分割，iTPN都取得了超过现有方法的性能。具体结果如图6和图7所示：

图6 base-levelmodel在COCO和ADE20K的结果对比

图7 large-levelmodel在在COCO和ADE20K的结果对比

4.Analysis

A．Ablative studies

为了验证iTPN常用的不同优化策略对性能的影响。作者将研究使用全局预训练（iPT）、load多尺度特征，MFM，C-MLP对性能的影响。具体结果如下图8和图9所示。

图8 iPT和load特征对性能的影响

上述结果说明backbone主干本身被iTPN加强了，因此它可以独立于neck转移到下游任务中。

图9 MFM和C-MLP对性能的影响

上述结果说明C-MLP和MFM 可分别独立优化性能且互为补充。

B．Visualization

图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任务中，更好的物体描述有助于下游的视觉识别任务（例如，物体检测）更快地收敛并达到更高的上限

四.conclusion

在本文中，作者提出了iTPN网络（Integrally Pre-Trained TransformerPyramid Networks）。其核心贡献在于使用一个特征金字塔作为neck部分来进行预训练过程的重建并复用于fine-tuning的识别，从而最大限度地减少了预训练和fine-tuning之间的转移差距。此外，还设计了MFM来补充MIM，以更好地优化特征金字塔。iTPN在一些流行的视觉识别任务中具有优越的性能。这项工作引出了一个新的未来发展方向——为上游和下游的视觉表示学习设计一个统一的框架。