Context Autoencoder (CAE)：一种更适合下游任务的新MIM方法

↑ 点击 蓝字  关注极市平台

作者丨陈小康@知乎（已授权）

来源丨https://zhuanlan.zhihu.com/p/467399602

编辑丨极市平台

极市导读

 

本文将介绍一种新的MIM方法CAE，它通过尽可能地分离“表征学习”和“解决 pretext task” 这两个功能，使encoder能学习到更好的表征，从而在下游任务实现了更好的泛化性能。 >>加入极市CV技术交流群，走在计算机视觉的最前沿

作者单位：北京大学，香港大学，百度

论文链接：https://arxiv.org/abs/2202.03026

代码链接：

官方版本 (Paddle): https://github.com/PaddlePaddle/VIMER

PyTorch 版本: https://github.com/lxtGH/CAE

https://github.com/open-mmlab/mmselfsup

前言

Mask Image Modeling (MIM) 方法，在 NLP 领域 (例如BERT) 得到了广泛的应用。随着 ViT 的提出和发展，人们也尝试将 MIM 应用到视觉领域并取得了一定进展。在此之前，视觉自监督算法主要沿着 contrastive learning 的思路去设计，而 MIM 无疑打开了新的大门。我们最近的工作 “Context Autoencoder for Self-Supervised Representation Learning”，提出一种新的 MIM 方法 CAE，通过对 “表征学习” 和 “解决 pretext task” 这两个功能做尽可能的分离，使得 encoder 学习到更好的表征，从而在下游任务实现了更好的泛化性能。我们尝试回答如下几个问题：

1. MIM方法中，网络结构的哪个部分是学习表征的，哪个部分是解决pretext task？
2. 为什么之前典型的 contrastive learning 方法，在下游任务 (例如检测、分割) 上只能取得跟 supervised pretraining 类似的性能？
3. MIM 方法为什么优于目前的 contrastive learning 方法？

1. 背景

MIM 是一种自监督表征学习算法。它的主要思路是，对输入图像进行分块和随机掩码操作，然后对掩码区域做一些预测。预测的目标可以是 Token ID (BEiT)，也可以是 RGB 的值 (MAE)。通过 MIM，我们希望 encoder 能学习到一个好的表征，从而在下游任务取得良好的泛化的性能。

近期 MIM 有两个代表性工作：BEiT 和 MAE。

• BEiT 使用一个 encoder 做两件事：(1) 学习一个好的图像表征；(2) 解决 pretext task：预测 masked patch 的 Token ID。encoder 的潜力并没有完全被挖掘，只有部分被用来学习表征。
• MAE 使用了 encoder-decoder 架构，encoder 负责对 visible patch 进行表征学习，decoder 将 visible 和 masked patch 的表征 (masked patch 使用一个可学习的向量) 作为输入，预测 masked patch 的 RGB 值。但是，MAE 在 decoder 中也会对 visible patch 的表征进行改变。与此同时，MAE decoder 利用改变后的可见区域的表征去预测遮挡区域的表征，实际上 decoder 也负责了一部分表征学习的功能。然而，在下游任务中，只有 encoder 中学到的信息能被拿来用，那么即使在 decoder 中进一步学到了更好的表征，也无法利用到下游任务中。

以上两种方法，都没有充分挖掘 encoder 的潜力，限制了预训练学习到的表征质量。

2. Context Autoencoder (CAE)

CAE 设计的核心思想是对 “表征学习” 和 “解决 pretext task” 这两个功能做分离。我们希望在预训练时，表征学习的任务只交给 encoder，而 decoder 只负责解决 pretext task。这样我们希望从 encoder 出来的表征就是非常好的，而不需要额外的模块对表征进一步 refine (例如 MAE 中的 decoder)，可以尽可能大地挖掘 encoder 的潜力。

CAE 包括4个部分：(1) encoder; (2) latent contextual regressor; (3) decoder; (4) alignment module。

输入图像通过随机掩码被划分成 visible patch 和 masked patch 两个部分。具体来说：

• Encoder 是一个 ViT 模型，负责学习 visible patch 的表征 。
• Latent contextual regressor 通过 预测 masked patch 的表征 。Latent contextual regressor 由一系列 cross-attention module 组成，query 是 masked patch 的表征，key 和 value 是全部 patch 的表征。在计算 query-key 相似度时，我们会引入每个 patch 对应的位置编码。在这个阶段， 不断更新、变得更加准确，而 不会更新，对图像特征的提取这个任务完全交给 encoder。
• Decoder 只拿 和 对应的位置编码作为输入，其目的是通过 预测 masked patch 的某些性质，比如由训练好的 tokenizer 产生的 Token ID，或者 RGB 的值。本文的实验 follow BEiT，使用 DALL-E tokenizer 对输入图像 token 化，得到 decoder 的目标。
• Latent representation alignment 是非常关键的一部分。虽然 visible patch 的表征在 encoder 之后就不会改变，但 latent contextual regressor 可能会”偷偷地“学习 masked patch 的表征，然后基于这样一个比 encoder 的输出更好的表征在 decoder 中进行预测。如果是这样，那么 latent contextual regressor 也承担了一部分表征学习的功能，这与我们想要的”分离“是相悖的。于是我们通过对 添加约束，希望 latent contextual regressor 的输出和 encoder 的输出在同一编码空间中， 这样表征学习的任务还是落到了 encoder 的身上。我们将图像的 masked patch 也输入到 encoder，获得这部分的表征 。 将作为 学习的目标。计算 的过程不会计算梯度。
• 损失函数。损失函数由两部分组成：(1) 对 decoder 预测的监督，使用 cross-entropy loss; (2) 对 和 的 align 的监督，使用 MSE loss.

3. 分析

3.1 CAE 关心每个 patch 的表征

CAE 基于 visible patch 的表征，从随机采样的 masked patch 做一些预测，这要求 CAE 关心每个 patch 的语义。这不同于典型的对比学习方法 (例如 MoCo v3, SimCLR)，这类方法只关心图像的全局语义，忽略了图像的细节和非主体区域 (比如背景)。

3.2 Latent contextual regressor 的输出和 encoder 的输出在同一编码空间中

我们对 latent contextual regressor 的输出做了约束，希望它能和 encoder 的输出尽可能处于同一编码空间中。这样，decoder 会基于 encoder 学到的编码空间做预测，将对图像的特征提取的重任完全交到了 encoder 身上，驱使 encoder 学习到好的表征。

为了验证表征空间确实对齐了，我们用 RGB 值作为 decoder 目标 (考虑到 Token ID 难以可视化，这边使用 RGB)，训练 CAE。在测试的时候，我们将全部 patch 输入到 encoder，然后跳过 latent contextual regressor，直接将 encoder 的输出送进 decoder，预测全部 patch 的 RGB 的值。下图展示了预测结果，第一行是原图，第二行是预测，我们发现仅使用 encoder 和 decoder 就可以将图片重建出来，说明 encoder 的输出和 latent contextual regressor 的输出属于同一编码空间。

如果训练时不做 alignment 约束，那么无法重建，如下图所示，输出都是乱码，说明 encoder 输出和 latent contextual regressor 的输出不在一个编码空间中。这使得 regressor 也承担了一部分表征学习的角色，使得 encoder 学到的表征质量有所欠缺，在消融实验部分也有验证。

3.3 CAE 学到的表征可以区分不同类别的 object/stuff

CAE 基于 visible patch 的表征，在 masked patch 区域做预测，这要求 CAE 对 visible patch 的内容有比较好的理解。举例来说，我们看到一只狗的头部，我们可以预测出它的身体部分；我们看到一小片天空，我们也能预测出它的周围大概率也是一片天空。因此，我们认为 CAE 学到的表征可以区分不同类别的 object/stuff。为了验证这一点，我们从 ADE20K 数据集随机采样一些图片输入到 encoder。因为 ADE20K 提供了每个像素的类别标签 (150类)，我们可以使用 t-SNE 对 encoder 输出的表征进行可视化。如下图所示，每个颜色代表一个类别，左图是 CAE，右图是随机初始化的 encoder。我们可以发现 CAE 可以有效区分不同类别的 object/stuff (因为是在 ImageNet-1K 进行预训练，所以区分地不够完美)，而随机初始化的 encoder 无法做到这一点。

3.4 典型的 contrastive learning 为什么在下游任务只能取得跟 supervised pre-training 差不多的结果？

在 contrastive learning 中，random crop 是一个非常重要的数据增强策略。典型的 contrastive learning (比如 MoCo v3) 希望最大化来自同一图像的 2 个不同 crop 之间的全局语义相似度，而最小化来自不同图像的 crop 之间的相似度。这样为什么能奏效呢？我们首先分析 random crop 的性质。在 SimCLR 论文中提到，random crop 是 contrastive learning 方法中非常重要的数据增强策略。在 ImageNet-1K 数据集中，图像的主体物体大多处于图像的中心区域，而对图像进行 random crop，中心区域有很大的概率被囊括进去，例如下图展示的几个例子，几次 crop 基本都包括了图像的主体物体。

对同一图像的不同 crop 提取全局语义，实际上学到的是原始图像中主体物体的特征，正因如此，同一图像的不同 crop 之间才可能相似。在 supervised pre-training 中，受到图像分类标签的约束，网络学习到的也是图像主体区域的特征，这和 contrastive learning 学到的知识有很大的相似之处，因此在下游任务表现类似。

3.5 MIM 和 contrastive learning 的区别

MIM 方法 (例如 CAE) 基于 visible patch 的表征，对 masked patch 区域做预测。在做随机掩码时，图像的每个 patch (例如背景区域的 object/stuff) 都有可能被考虑到，而不仅仅是图像的主体区域。为了做好 masked patch 的预测，CAE 会学好每个 patch 的表征。

我们对 CAE 以及 MoCo v3 的 attention map 做了可视化。如下图所示，第一行是原图，第二行是 MoCo v3，第三行是 CAE。红色表示 attention value 更高，蓝色表示attention value 低。处于蓝色边界内部的区域，通过这样的原则筛选：将 attention value 从大到小排序后，保留累计和达到所有位置 attention value 总和的 50% 的部分。我们可以看到，MoCo v3 的attention map 主要在图像的主体区域有高响应，而 CAE 能考虑到几乎所有 patch。

4. 实验

我们使用 ViT-small 和 ViT-base 在 ImageNet-1K 上进行实验。输入图像的分辨率是 ，patch 大小是 ，一张图会被划分成 个 patch。

4.1 Pretraining Evaluation

自监督学习广泛使用 linear probing 去评测预训练表征的好坏：将 encoder 的参数固定住，在之后加一个 linear classifier 进行图像分类。我们认为 linear probing 不适合 MIM 方法，因为 MIM 方法通常会学到每个 patch 的表征，不仅包含主体物体的信息，还学到了背景等若干知识，这是多而杂的，不适合直接进行线性分类。我们额外提出一种新的测试指标：attentive probing。我们在固定参数的 encoder 后加上一个简单的 cross-attention module (没有 FFN)和一个 linear classifier，通过注意力机制动态地选择适合做图像分类的信息。

我们对 attentive probing 阶段使用的 cross-attention module 做 attention map 可视化，发现可以关注到主体物体。

Finetune、linear probing、attentive probing 的结果见下表。

我们发现一些有趣的现象。(1) contrastive learning 方法 (MoCo v3, DINO) 的 linear probing 和 attentive probing 结果类似。这说明这类方法在预训练时已经将注意力放到了图像的主体物体上面，无需进一步动态筛选即可做好图像分类，这也与之前我们对 contrastive learning 的分析一致。(2) MIM 方法 (例如 CAE) 的 attentive probing 相比 linear probing 有很大的提升。这说明 MIM 方法学到了每个 patch 的特征，而不仅仅是图像主体物体的，因此需要做一些筛选才利于图像分类。

4.2 消融实验

我们对 decoder 和 alignment module 进行消融实验，见下表。单加一个 decoder 能提高 attentive probing 的结果，但在下游任务 (分割、检测)提升不明显。使用 alignment module 之后能显著提升下游任务的性能，说明约束 encoder 的输出和 latent contextual regressor 的输出在同一编码空间非常重要，能提升 encoder 学到的表征质量。

4.3 语义分割

我们在 ADE20K 进行语义分割的实验。网络使用 UperNet，迭代次数为 160K，输入图像分辨率为 ，使用单尺度测试。Contrastive learning 方法和 supervised pre-training (DeiT) 的结果类似，而 CAE 能取得明显更好的结果。跟其他 MIM 方法相比，CAE 的结果也更好，说明预训练阶段 encoder 被充分利用，学到的表征更好。使用 ViT-Large，我们能获得 54.7 mIoU，比 MAE (53.6) 高 1.1 个点，比有监督学习 DeiT (49.9) 高 4.8 个点。非常有趣的是，CAE 和 DeiT 的 gap，在 large 模型下会比 base 模型更大，base 模型下它们的 gap 只有 3.2 个点。这说明了 CAE 能很好地泛化到大模型。

4.4 物体检测、实例分割

我们使用 Mask-RCNN 和 Cascade-RCNN 两种网络结构进行物体检测和实例分割的实验。我们使用 multi-scale training 训练 12 epoch (1x schedule)，测试阶段仅使用单尺度测试。实验现象和语义分割的类似：contrastive learning 方法和 supervised pre-training 相比结果类似或者更差，CAE 的结果更好。使用 ViT-Large，在 1x schedule 下我们能获得 54.5 mAP。

5 总结

本文提出了 CAE，设计的核心有两点：(1) 对 “表征学习” 和 “解决 pretext task” 这两个功能做尽可能的分离; (2) 在 visible patch 学习到的表征空间中对 masked patch 做预测。以上两点都是为了驱使 encoder 学习更好的表征，从而在下游任务取得良好的泛化能力。除此之外，我们对 supervised pre-training、contrastive learning 和 MIM 方法进行分析，认为 contrastive learning 和 supervised pre-training 主要关注图像的主体区域 (例如 ImageNet-1K 标签集中的物体)，而 MIM 会关注图像的全部 patch，更有利于下游任务。

公众号后台回复“ECCV2022”获取论文分类资源下载～

△点击卡片关注极市平台，获取 最新CV干货

极市干货

算法项目： CV工业项目落地实战 ｜ 目标检测算法上新！（年均分成5万）

实操教程 ：Pytorch - 弹性训练原理分析｜《CUDA C 编程指南》导读

极视角动态： 极视角作为重点项目入选「2022青岛十大资本青睐企业」榜单！ ｜ 极视角发布EQP激励计划，招募优质算法团队展开多维度生态合作！ ｜极市AI校园大使招募

“

点击阅读原文进入CV社区

收获更多技术干货