多尺度混合卷积- Transformer模型ConvMAE开源：进一步挖掘和提升 MAE 的性能！

 

论文地址：

https://arxiv.org/pdf/2205.03892.pdf

本文论点是多尺度的混合 Convolution-Transformer 模型 可以助力 Masked Auto-Encoding (MAE) 的训练范式，帮助其学习到更好的表征。

 

1.1 Self-supervised Learning

在预训练阶段我们使用无标签的数据集 (unlabeled data)，因为有标签的数据集很贵，打标签得要多少人工劳力去标注，那成本是相当高的，太贵。相反，无标签的数据集网上随便到处爬，它便宜。在训练模型参数的时候，我们不追求把这个参数用带标签数据从初始化的一张白纸给一步训练到位，原因就是数据集太贵。于是 Self-Supervised Learning 就想先把参数从一张白纸训练到初步成型，再从初步成型训练到完全成型。注意这是2个阶段。这个训练到初步成型的东西，我们把它叫做 Visual Representation。预训练模型的时候，就是模型参数从一张白纸到初步成型的这个过程，还是用无标签数据集。等我把模型参数训练个八九不离十，这时候再根据你下游任务 (Downstream Tasks) 的不同去用带标签的数据集把参数训练到完全成型，那这时用的数据集量就不用太多了，因为参数经过了第1阶段就已经训练得差不多了。

第一个阶段不涉及任何下游任务，就是拿着一堆无标签的数据去预训练，没有特定的任务，这个话用官方语言表达叫做：in a task-agnostic way。第二个阶段涉及下游任务，就是拿着一堆带标签的数据去在下游任务上 Fine-tune，这个话用官方语言表达叫做：in a task-specific way。

以上这些话就是 Self-Supervised Learning 的核心思想，如下图1所示，后面还会再次提到它。

图1：Self-Supervised Learning 的核心思想

Self-Supervised Learning 不仅是在NLP领域，在CV, 语音领域也有很多经典的工作，如下图2所示。它可以分成3类：Data Centric, Prediction (也叫 Generative) 和 Contrastive。

图2：各个领域的 Self-Supervised Learning (引用李宏毅老师 PPT)

其中的主流就是基于 Generative 的方法和基于 Contrative 的方法。如下图 3 所示这里简单介绍下。基于 Generative 的方法主要关注的重建误差，比如对于 NLP 任务而言，一个句子中间盖住一个 token，让模型去预测，令得到的预测结果与真实的 token 之间的误差作为损失。基于 Contrastive 的方法不要求模型能够重建原始输入，而是希望模型能够在特征空间上对不同的输入进行分辨。

图3：基于 generative 的方法和基于 contrastive 的方法的总结图片 (引用李宏毅老师 PPT)

 

1.2 ConvMAE的动机

ConvMAE 这个方法所基于的论点是：

1. 目前已经有许多工作 (如 MoCo[1]，MAE[2]，BEiT[3]，DINO[4]) 验证了 MAE Self-Supervised Learning的训练范式能够帮助释放 Vision Transformer 模型的潜力，并且在下有任务上取得非常好的性能。

MAE 作为这个范式的代表作，开发了一个非对称编码器 - 解码器架构，其中编码器只对可见的 patch 子集进行操作 (即没有被 mask 掉的 token)，另一个非对称的解码器可以从潜在表征和被 masked 掉的 token重建原始图像。Decoder 的架构可以是十分轻量化的模型，且具体的架构对模型性能影响很大。研究人员进一步发现，Mask 掉大部分输入图像 (例如 75%) 会产生重要且有意义的自监督任务。同时 MAE 这种训练的范式不但能够在不需要超大规模数据集 (JFT-300M，ImageNet-22K) 的情况下，学习到判别性能很强 (Discriminative) 的表征，而且可以轻松的扩展 (Scalable) 到更大的模型上，并且通过实验发现随着模型增大，效果越来越好。

2. 为了加速 ViT 训练并得到更好的性能，大量工作验证了局部的归纳偏置 (local inductive bias)(如 SMCA-DETR [5]，SAM-DETR[6]，DAB-DETR[7]，Uniformer[8]，CoAtNet[9]，ConViT[10]，Early Convolution[11]) 和可以进一步帮助提升 ViT 模型的性能。同时，这种性能的提升也可以通过多尺度的金字塔式架构(multi-scale hierarchical representation) (如 Swin Transformer[12]，PVT[13]) 来实现。二者结合的有效性已经在大量的识别，检测，分割的监督学习任务中得到的验证。

所以一个自然而然的问题是：这种多尺度的金字塔式架构 + 局部的归纳偏置的模型，能不能经过 MAE 的训练方式之后，进一步挖掘和提升 MAE 的性能？

本文就是探索这个问题。ConvMAE 简而言之就是：多尺度的金字塔式架构 + 局部的归纳偏置的模型，使用 MAE 的 Self-supervised Learning 的训练方式。

与 MAE-Base 相比，ConvMAE-Base 将 ImageNet-1k 的微调精度提高到 85.0% (+1.4%)，将 Mask-RCNN COCO 检测任务的 AP box 提高到 53.2% (+2.9%)，将 UperNet 的 ADE20k 分割任务的 mIoU 提高到 51.7% (+3.6%)。

 

1.3 ConvMAE Encoder 架构

MAE 的做法如下图3所示。MAE 是一种以自监督的方式，以 ViT 为模型架构进行预训练的框架。MAE 的方法很简单：Mask 掉输入图像的随机的 patches 并重建它们。它基于两个核心理念：研究人员开发了一个非对称编码器 - 解码器架构，其中一个编码器只对可见的 patch 子集进行操作 (即没有被 mask 掉的 token)，另一个简单解码器可以从可学习的潜在表征和被 masked 掉的 token重建原始图像。Decoder 的架构可以是十分轻量化的模型，且具体的架构对模型性能影响很大。研究人员进一步发现，Mask 掉大部分输入图像 (例如 75%) 会产生重要且有意义的自监督任务。结合这两种设计就能高效地训练大型模型：提升训练速度至 3 倍或更多，并提高准确性。

图3：MAE 框架

ConvMAE 相比于 MAE 框架做了一些微小却非常有效的改进，如前文所述它的特点是：多尺度的金字塔式架构 + 局部的归纳偏置的模型。

如下图4所示是 ConvMAE 框架，它也有一个 Encoder 和 Decoder。Encoder 是 convolution-transformer 混合架构，Decoder 是纯 transformer 架构。

先看左上角灰色的 Encoder 部分。它包括了3个 stage，设  和  是输入图片的尺寸，每个 stage 输出的特征分别是  。前两个 stage 是卷积模块，使用 Masked Convolutional Block 对特征进行操作，其结构如下图右下角所示 (其中的 Depthwise Convolution 使用5×5大小卷积核)。在每个阶段之间，进行一次 stride 为2的卷积以进行下采样操作。最后一个 stage 都是 Transformer 模块，拉大感受野，并融合所有 patch 的特征。另外作者发现绝对位置编码性能是最优的。

图4：ConvMAE 框架

 

1.4 ConvMAE mask 策略

MAE 对输入图片的 patch 采用随机 mask 策略，然而，同样的策略不能直接应用于 ConvMAE 的编码器。因为 ConvMAE 的特征是不同 stage 是逐渐下采样的，如果在  的特征这里进行了随机的 mask，就会导致 stage3 阶段的每个 tokens 都有一部分的可见信息。因此 ConvMAE 作者的做法是 mask 掉 stage3  的输出 (比如 75%) 之后，把这些 mask 分别上采样2倍和4倍得到前两个阶段的 mask。这些被 mask 掉的 token 在编码阶段被丢弃，并且希望经过 Decoder 之后能够重建出来。通过这种方式，ConvMAE 只需要保留至少 25% 的 token 用于训练。

但是前两个阶段使用 5×5 的 Depthwise Convolution 的感受野可能大于一个 masked patch 的大小，因此作者为了确保预训练的质量，在前两个阶段采用了 masked convolution[14][15]，确保被 mask 掉的部分不会参与到编码的过程。

 

1.5 ConvMAE Decoder 架构

如图3所示，原始 MAE 的 Decoder 以 Encoder 的输出以及 masked token 为输入，通过一系列的 Transformer Block 得到最终的重建结果。

ConvMAE 的编码器获得了多尺度特征  ，分别捕捉到了细粒度和粗粒度的图像信息。为了更好地进行预训练，作者将  和  分别进行 stride=2 和 stride=4 的下采样之后与  相加，进行多尺度特征的融合。融合得到的结果再通过 Linear Transformation 得到最终要输入给 Decoder 的 token。

式中，  代表 stride=k 的卷积。

训练使用的目标函数与 MAE 保持一致，都是 mask 的部分的重建结果与原图的 L1 Loss。

式中，  Tm代表 masked tokens 的集合。

 

1.6 ConvMAE 下游任务

ConvMAE 经过预训练之后，Encoder 能够输出多尺度的特征 (  )，它们可以被用于后续的检测分割任务里面。

图5：ConvMAE 用于检测和分割的框架，输出的不同尺度的中间特征传入 FPN 模块

ConvMAE 用于检测任务的微调过程：先把 Encoder 的输出特征  进行 max-pooling 操作得到  。对于检测任务，因为 ConvMAE 的 stage3 有11个全局 Self-attention 层，计算成本过高，所以作者把 stage3 里面第1,4,7,11个 Self-attention 换成了 7×7 Window size 的 Swin Attention 层。通过这样的做法减少了计算量和 GPU 占用。最终得到的  被送入 Mask R-CNN 或者 UperNet 进行目标检测或者语义分割任务。对于分割任务，Stage3 的架构不变。

内容中包含的图片若涉及版权问题，请及时与我们联系删除