这里提出了一个基于掩蔽图像建模 (MIM) 的整体预训练框架 iTPN,base/large-level iTPN 在 ImageNet-1K 上达到了 86.2%/87.8% 的 top-1 精度,在使用 UPerHead 的 ADE20K 语义分割上有 54.7%/57.7% 的 mIoU。 

Title: Integrally Pre-Trained Transformer Pyramid Networks

Paper: https://arxiv.org/pdf/2211.12735.pdf

Code: https://github.com/sunsmarterjie/iTPN

自 ViT 提出以后,Transformer 在计算机视觉领域逐渐衍生出两个重要分支,一个分支是以 Vision Transformer 为代表的给为 Transformer 主干网络,而另一个分支便是以 MAE 和 BEiT 为代表的 掩码图像重建(Masked Image ModellingMIM)技术。通过结合这两项技术,极大的促进了包括分类、检测和分割等下游任务的发展。

今天为大家介绍的是一篇与 MIM 和 Vision Transformer 相关的一篇工作,其旨在解决上游预训练和下游微调之间的迁移差距。以代表性工作 MAE 和 BEiT 为例,其使用的主干网络均为朴素的 ViT 模型。尽管 SimMIMConvMAE以及GreenMIM 等模型应用了分层结构,但本质上仅会作用到 Backbone 上,而不会影响到 Neck,即特征金字塔。这会导致一个问题,仅当我们直接应用于下游任务时,如果你是采用 Linear probing 的方式,虽然不会破坏预训练的特征提取器,但由于整个优化过程从随机初始化的 Neck 层开始,作者认为这并不能确保冻结的 Backbone 与 Neck 能够很好的“搭配合作”。

因此,本文设计了一种简洁有效的预训练框架来缓解这种现象。考虑到部分读者可能对这方面不太熟悉,本篇文章将会先引入相关的概念和必要的背景基础知识,最后再详细介绍此项工作。

定义

掩码图像重建是一种利用图像掩码去执行计算机视觉任务的技术。相比于基于对比学习等方法,基于 MIM 的方法能提供更有竞争力的结果。说到这,那就让我们先简单回顾下对比学习的知识吧。

对比学习

对比学习(Contrastive Learning)是一种通过在相同数据的扭曲视图(distorted views)上提取不变特征来学习实例级判别表示。例如凯明的MoCo[1]和Hinton的SimCLR[2]便是早期的两项代表性工作,它们采用不同的机制引入负样本与正样本进行对比:

iTPN~~_人工智能

随后,DeepMind 团队提出了 BYOL[3] 进一步解决了对负样本的依赖(避免表示崩溃)。当然,同期也有不少其他类似的工作,如凯明的SimSiam[4]便探索了孪生神经网络表征学习的崩溃解。 

iTPN~~_图像重建_02

iTPN~~_建模_03

除了这种组成一对对的表征学习方式外,SwAV[5]提倡对数据进行在线聚类,同时加强同一图像的多重增强视图之间的一致性: 

iTPN~~_编码器_04

最后便是一些致力于将对比学习应用于提升特定下有任务的方法,如同时发表在 CVPR & ICCV 2021 上的三篇代表性文章:Detco[6]、ReSim[7]以及CAST[8]: 

iTPN~~_编码器_05

iTPN~~_人工智能_06

自凯明的 MIM 提出以后,此类对比学习方法的热度也逐渐下降了,后面也就一些结合 Transformer 去做的工作如 DINO[9] 和 MoCov3[10]。

掩码图像重建

掩码图像重建是自监督任务的一种形式,它能够很好的解决对数据的依赖。相比于 CV 领域,NLP 领域一直都处于领先地位,无论是基于 GPT 中的自回归语言建模或者 BERT 中的掩码自编码方案,本质上都是基于删除-预测的机制,这些方法很容易推广到 LLM 上。

然而,正如凯明在 MAE 论文中指出,掩码自编码器的概念是一种更通用的去噪自编码器(Denoising Auto EncoderDAE),它是自然的,也适用于计算机视觉。DAE 可以算是介于 AE 与 VAE 之间的产物。众所周知,自编码器(Auto-Encoder)是一种非常典型的网络架构,它允许在没有标注的情况下进行表征学习,由 Hinton 等人于 1993 年提出。随后过了十几年,人们便提出了将噪声“强制”引入到学习的表征上,其通过对潜在特征叠加高斯噪声从而形成“损坏”的信号以作为网络的输入来重建未校正的输入信号,这便是变分自编码器(Variational auto-encoderVAE)。

BERT 的提出带火了 NLP 领域。而对于 CV 领域而言,笔者早期接触的自监督学习主要有图像修复图像着色图像拼图等形式,强烈推荐大家去看下 Amit Chaudhary 的自监督学习笔记:https://amitness.com/2020/02/illustrated-self-supervised-learning/。随着 ViT 的引入,便涌现出了许多基于掩码预测的自监督学习方法,下面带大家快速过一遍:

iTPN~~_建模_07

iGPT较早提出在给定一系列像素作为输入的情况下预测后续像素值。

iTPN~~_人工智能_08

MaskFeat应用 HOG 作为预测目标,而非 RGB 像素值。

iTPN~~_编码器_09

 

SimMIM中采用线性层作为解码器

 

iTPN~~_图像重建_10

BEiTMAE算是两个同期的工作,其一个亮点便是基于 ViT 模型重建缺失块。

iTPN~~_建模_11

则是提出了一种与架构无关的 MIM 框架,将 Transformer 和 CNN 兼容起来。

方法
Framework

iTPN~~_人工智能_12

整体的框架图如上所示。首先,让我们先看下图中左侧内容,其显示了传统的预训练范式。需要注意的是,这里我们先不区分具体的微调任务(如分类、检测和分割),假设它们都共享相同的 Backbone,而无需 Neck 和 Head (参考MAE和BEiT等)。那么,问题来了,这种架构容易引起两个问题:

  • 骨干网络的参数并未针对多级特征提取进行专门的优化,即缺乏一种合理的机制来高效的提取和融合多尺度特征;
  • 微调阶段的优化若从随机初始化的 Neck 和 Head 开始,那么这会显著减慢训练过程且不容易获得更优的效果。

因此,iTPNs 提倡将重建和识别过程统一起来,以最大限度的减轻这种现象,如图中右半部分所示。

Unifying Reconstruction and Recognition

具体地,iTPNs 基于 HiViT 和 特征金字塔构建了一个全新的框架。其中,HiViT 通过以下方式完成进一步的简化:

  • 采用通道多层感知器(Channel-wise MultiLayer PerceptronC-MLP)替换 shiftedwindow attentions
  • 将感受野为 的 stage 剔除掉,而是直接在 的 stage 直接计算全局注意力;

如此一来,该主干网络就不用像 SimMIM 方法一样需要输入全图才能跑,直接节省 30%–50 % 的计算成本。下表展示了具体的参数比对:

iTPN~~_图像重建_13

此外,为了更好的联合优化骨干网络(HiViT)和颈部(Feature Pyramid),iTPNs 采用了以下两个技术细节:

首先,作者通过将特征金字塔插入预训练阶段(用于重建)并在微调阶段复用训练好的权重(用于识别)以此来统一上游和下游颈部特征。

其次,为了更好地预训练特征金字塔,本文提出了一种新颖的掩码特征建模(Masked Feature ModelingMFM)任务为特征金字塔提供多阶段监督,该任务通过将原始图像输入一个 moving-averaged backbone 计算出中间特征,同时使用特征金字塔的每一层输出来重建中间目标。

总的来说,MFM 可以理解为 MIM 的一种补充方法,其更好的提高了重建和识别的准确性。此外,MFM 还可以适应从预训练教师模型(本文应用了CLIP)“吸收知识”以获得最佳的性能。

实验
ImageNet-1K classification

iTPN~~_建模_14

可以看出,iTPN 显示出优于现有方法的显着优势,无论是仅使用像素监督还是利用来自预训练教师的知识(括号中内容为教师模型的名称)。

ImageNet-1K classification

iTPN~~_人工智能_15

此外,iTPN 在几个重要基准测试中的识别准确率均超过了之前的 SOTA。

COCO and ADE20K

iTPN~~_编码器_16

不仅在图像分类,在下游的目标检测和图像分割任务下表现也很不错。

Linear probing

iTPN~~_图像重建_17

正如我们前面所提到的,传统 MIM 方法使用 Linear probing 微调时效果并不好,这在很多 MIM 相关的方法中也经常被提及。确实,与 Fine-tune 相比,这种方式对预训练骨干更加敏感。不过,从上表可以明显的看出,iTPN 仍然能获得不错的精度,例如在 CLIP 监督下,超越同等配置的 MVP 将近两个百分点。

总结

本文提出了一个用于预训练 HiViT 的完整框架,其核心贡献在于利用特征金字塔统一重建和识别任务,从而最大限度地减少预训练和微调任务之间的迁移差距。此外,为了更好的优化特征金字塔,本文提出了一种掩码特征建模任务,旨在补充掩码图像建模能力。最后,预训练的 iTPN 在一些主流的视觉识别任务中报告了卓越的识别能力。