ConvMAE: Masked Convolution Meets Masked Autoencoders

Self-Supervised Learning，又称为自监督学习，机器学习分为有监督学习，无监督学习和半监督学习。而 Self-Supervised Learning 是无监督学习里面的一种，主要是希望能够学习到一种通用的特征表达用于下游任务 (Downstream Tasks) 。其主要的方式就是通过自己监督自己。作为代表作的 kaiming 的 MoCo 引发一波热议， Yann Lecun也在 AAAI 上讲 Self-Supervised Learning 是未来的大势所趋。

ConvMAE: Masked Convolution Meets Masked Autoencoders主要干了一个什么事呢？主要提出多尺度的混合 Convolution-Transformer 模型 可以助力 Masked Auto-Encoding (MAE) 的训练范式，帮助其学习到更好的表征。疑问更好的特征表示到底是啥？如果给的都是noise图像，还能学到？是个怼人的点，哈哈哈。
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所示。

Self-Supervised Learning 方法可分为3类：Data Centric, Prediction (也叫 Generative) 和 Contrastive。

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

1.2 ConvMAE的动机
ConvMAE 这个方法所基于的论点是：目前已经有许多工作 (如 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) 到更大的模型上，并且通过实验发现随着模型增大，效果越来越好。

为了加速 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 倍或更多，并提高准确性。

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

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

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

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 掉的部分不会参与到编码的过程。