[论文解读] SupMAE: Supervised Masked Autoencoders Are Efficient Vision Learners

Feng Liang, Yangguang Li+1 more|arXiv (Cornell University)|May 28, 2022

Human Pose and Action Recognition被引用 14

一句话总结

SupMAE 提出了一种新型的有监督掩码自编码器，通过仅使用可见补丁的子集联合优化掩码图像重建与有监督分类，从而提升 ViT 的预训练效率。该方法在 400 个周期内实现了 83.6% 的 ImageNet 准确率，仅为 MAE 所需周期的 1/3，同时在迁移学习和鲁棒性基准测试中优于 MAE 和标准有监督预训练方法，展现出最先进的效率与鲁棒性。

ABSTRACT

Recently, self-supervised Masked Autoencoders (MAE) have attracted unprecedented attention for their impressive representation learning ability. However, the pretext task, Masked Image Modeling (MIM), reconstructs the missing local patches, lacking the global understanding of the image. This paper extends MAE to a fully supervised setting by adding a supervised classification branch, thereby enabling MAE to learn global features from golden labels effectively. The proposed Supervised MAE (SupMAE) only exploits a visible subset of image patches for classification, unlike the standard supervised pre-training where all image patches are used. Through experiments, we demonstrate that SupMAE is not only more training efficient but it also learns more robust and transferable features. Specifically, SupMAE achieves comparable performance with MAE using only 30% of compute when evaluated on ImageNet with the ViT-B/16 model. SupMAE's robustness on ImageNet variants and transfer learning performance outperforms MAE and standard supervised pre-training counterparts. Codes are available at https://github.com/enyac-group/supmae.

研究动机与目标

• 探究有监督预训练是否能够提升当前仅依赖自监督掩码图像建模的掩码自编码器（MAE）的表征学习能力。
• 通过在预训练过程中引入真实标签，解决 MAE 在学习全局图像级特征方面的局限性。
• 通过仅使用可见补丁的子集进行分类，利用图像中的空间冗余性，提升训练效率与特征鲁棒性。
• 相比标准有监督和自监督预训练方法，提升下游视觉任务（包括线性探测与微调）中的迁移学习性能。
• 证明在掩码设置下联合使用重建与分类目标，可实现更高效的数据利用与更鲁棒的特征学习。

提出的方法

• 在 MAE 的基础上增加一个并行的有监督分类分支，该分支仅处理可见补丁的子集，而非全部补丁或类别标记。
• 使用 ViT 编码器从可见补丁中提取特征，通过全局平均池化获得全局图像特征以用于分类。
• 应用一个小型解码器，利用可见补丁特征与学习到的掩码标记来重建被掩码的补丁。
• 端到端训练模型，损失函数由被掩码补丁的重建损失与分类头的交叉熵损失联合构成。
• 在预训练过程中采用随机掩码策略，生成多样化的训练样本，起到强正则化作用。
• 在微调阶段，仅使用 ViT 编码器，不包含重建头或分类头。

实验结果

研究问题

• RQ1 有监督预训练能否提升掩码自编码器（MAE）的性能与效率？
• RQ2 在 MAE 中仅使用可见补丁的子集进行分类，是否能带来更高的样本效率与鲁棒性？
• RQ3 SupMAE 在迁移学习与鲁棒性方面，相较于标准有监督预训练与自监督 MAE 表现如何？
• RQ4 在掩码设置下联合使用重建与分类目标，是否能产生更具迁移性与鲁棒性的特征？
• RQ5 在掩码自编码框架中，对补丁特征使用全局平均池化对全局表征学习有何影响？

主要发现

• SupMAE 仅使用 400 个预训练周期即在 ImageNet-1K 上达到 83.6% 的 top-1 准确率，仅为 MAE 所需 1600 个周期的 1/3。
• 在使用 ViT-B/16 的情况下，SupMAE 相较于 MAE 将计算成本降低了 30%，同时实现了相当的 ImageNet 性能。
• 在自然损坏的 ImageNet 变体上，SupMAE 平均优于 MAE 1.8%，展现出更优的鲁棒性。
• 在少样本线性探测设置下，SupMAE 在 20 个数据集的 5-shot 设置中，平均准确率较 MAE 提升 14.5%。
• 在 ADE20K 语义分割任务上，SupMAE 达到 49.0% 的 mIoU，表明其在密集预测任务中表现强劲。
• 在 20 个下游分类数据集的 50-shot 微调中，SupMAE 在 13 个数据集中取得最佳准确率，优于 MAE 与 MoCo-v3。