【计算机视觉】MAE:Masked Autoencoder

有任何的书写错误、排版错误、概念错误等,希望大家包含指正。

在阅读本篇之前建议先学习:
【自然语言处理】Transformer 讲解
【自然语言处理】BERT 讲解
【计算机视觉】ViT:Vision Transformer 讲解

由于暂时不了解 BEiT、dVAE 和 DALLE,所以本文没有介绍论文中关于三者的内容。

MAE

1. 模型概述

何恺明提出了一种以编码器模型为骨架(backbone)、以自然语言模型 BERT 中完形填空(MLM)为学习策略的一种用于计算机视觉任务的可扩展(规模可变)的自监督学习模型 Masked AutoEncoders(MAE)。本质上,MAE 是一种在小数据集上高效训练大模型且保证大模型具有良好的泛化能力的自监督预训练方法目的是训练大型编码器模型

MAE 模型随机对输入图像对应的 patch 序列进行遮盖,通过非对称的端到端架构对被遮盖的 patch 进行重建。预训练阶段,模型接收大型无标签图像样本进行训练;微调阶段,使用有标签数据集对编码器进行微调,而与解码器无关,作者通过实验说明全面微调比冻结编码器参数仅学习以单一线性层为解码器参数的微调效果好;迁移阶段,直接使用训练好的编码器。

MAE 模型的两个核心设计为:非对称端到端架构,其编码器仅作用于未被遮盖的 patch 上,而精简的解码器接收来自编码器的未被遮盖的 patch 的潜在表征和被遮盖的 patch 的编码信息对原始图像进行重建;高比例的遮盖为模型提供了一个有效的自监督任务。这两个设计加速了模型的训练,并且提高了精度,保证了训练大型编码器模型的高效性。

MAE 名字中的 Auto 与自然语言中自回归(Autoregressive)的含义类似,即样本的标签来自样本,这在计算机视觉中是非常少见的,因此作者特意在名字中加上 Auto。

2. 模型结构

对应论文 Approach 部分。

在这里插入图片描述

图 1    MAE 模型

1 1 1 为 MAE 模型的整体思想。MAE 模型将可见(未被遮盖)部分输入到编码器中,得到可见部分的潜在表征;解码器以可见部分的潜在表征和不可见部分为输入重建不可见部分。具体地,论文选取 ViT 模型作为 MAE 的编码器。对图像对应的 patch 序列按遮盖比进行随机抽样,可见部分被输入到 ViT 模型编码器中。像在标准的 ViT 模型一样,编码器为可见部分对应的 patch 序列添加位置编码,通过一系列 Trm 模块处理。全部 patch 对应的信息包括来自编码器的可见部分的潜在表征,和不可见部分的标记 <MASK> 对应的向量,这些不同的不可见 patch 共享同一个可学习的 <MASK> 向量。解码器由不同于编码器的一系列 Trm 模块构成,其输入同样为添加了位置编码的 patch 序列,其输出对全部 patch 的预测结果。解码器输出的每个向量对应一个 patch,且向量长度等于 patch 中的像素数量,每个向量元素表示对应 patch 像素的预测结果。损失函数为不可见部分的预测结果与原始图像中对应位置像素的均方误差(MSE),即仅计算不可见部分的损失,不关心可见部分的损失,这类似于 BERT。另外,作者在每个预测出的 patch 内利用均值和方差进行标准化,以此来提高重建质量。

作者展示的模型效果如图 2 2 2 和图 3 3 3 所示。

在这里插入图片描述

图 2    ImageNet 验证图像的示例结果
每个三元组,遮盖后图像(左)、重建后图像(中)和原始图像(右),遮盖比为 80%,即 196 个 patch 中只有 39 个可见,其余被遮盖。由于没有在可见部分上计算损失,所以在可见部分上的重建效果比较差,虽然可以通过用原始图像的可见部分覆盖输出图像的可见部分,但是作者有意不这样做以全面演示 MAE 模型的效果。

在这里插入图片描述

图 3    COCO 验证图像上的示例结果
使用在 ImageNet 上训练的 MAE 模型(与图 2 中的模型参数相同)。观察最右边两个重构例子,尽管它们与基本事实不同,但在语义上是合理的。

参考源码可以得到更具体的实现细节:源码将主要将模型分为三部分,第一部分是编码器部分,包括数据的预处理、输入处理和 ViT 模型;第二部分是解码器部分,包括数据的预处理和设计的解码器;第三部分是定义损失函数。

(MAE 模型两次用到位置编码,这两次均采用不可学习的 sin-cos 编码。尽管编码器的输入是可见 patch 序列,但此时依然按全部 patch 分配 sin-cos 编码,只不过取出的)

在编码器部分,首先将输入图像打成 patch 序列,为每个 patch 添加不可学习的 sin-cos 位置编码,打乱全部 patch,确定可见部分和不可见部分,将 <CLS> 和其位置编码之和拼接到可见 patch 序列,以此为 ViT 模型(编码器)的输入。经过 ViT 模型后进行(通过 LayerNorm 实现的)patch 内规范化

在解码器部分,接收可见 patch 序列的潜在表征,按照原全部 patch 序列的顺序将不可见部分对应的 <MASK> 添加到潜在表征序列中,再为全部 patch 序列(含 <CLS>,且该 <CLS> 来自编码器)加上不可学习的 sin-cos 位置编码,输入到解码器的 Trm 模块中,最后一个 Trm 模块的输出经过(通过 LayerNorm 实现的)patch 内规范化,最后将规范化后的向量维度映射到 patch 大小。

在定义损失函数部分,计算不可见部分的规范化均方误差作为损失函数。

微调和迁移的对象不是整个 MAE 模型,而是编码器。在微调阶段,作者在实验部分提到了全面微调(论文中称为 end-to-end fine-tuning)、部分微调(partial fine-tuning)和线性探测(linear probing)三种具体的微调方式。全面微调和线性探测属于部分微调的两种极端情况,全部微调表示不冻结编码器中任何参数,线性探测表示冻结除分类层之外的参数,部分微调表示冻结部分低层(靠近输入的层)参数

一定要注意,MAE 模型只负责无监督的预训练,后续的微调和迁移仅与通过 MAE 预训练出的编码器有关

3. 工作背景

对应论文 Introduction 部分。

如今的模型规模非常大,如果数据集过小会出现过拟合问题,因此构建大型数据集需要大量人力进行标注。在 NLP 中,GPT、BERT 等模型已经通过自监督预训练方式很好地避免使用大型人工标注数据集。尽管 BERT 完形填空的思想出现在 CV 领域早于 NLP 领域,但是这种思想在 CV 领域的效果却远不如在 NLP 领域。作者分析了三点原因:

  • CV 和 NLP 的主流模型架构不同。在 ViT 出现之前,CV 主要采用卷积操作。可以想象,如果将某些像素标记为 <MASK>,那么在窗口滑动的过程中会出现非 <MASK> 像素与 <MASK> 像素被同一个窗口包裹,从而被糅杂在一起,这使得每个 <MASK> 像素无法被独立地学习,导致最终难以还原出每个 <MASK> 像素对应的原像素。相较而言,在 NLP 中,每个 <MASK> 单词在经过 Transformer 结构时不会与其他单词对应的向量融合,这是 NLP 使用完形填空成功而 CV 失败的原因之一。
  • 图像和语句的冗余度不同。语句中的每个单词作为语义实体都有具体的语义,缺少少量单词就会使得语句难以理解。然而,像素信息是冗余的,冗余度取决于图像的分辨率,即几个相邻像素可能为同一个物体的局部,具有相同的语义;图像分辨率越大,说明每个像素点包含的内容越少,那么像素信息的冗余度也就越大。这种冗余可能会导致模型偷懒地选择插值的方式恢复这些被遮盖的少量像素,或者是仅学习图像的局部信息而不是全局信息。
  • CV 模型和 NLP 模型在重建阶段的解码器不同。无论是在 NLP 中还是 CV 中,解码器作用为将高层次特征转换为低层次特征。在 NLP 中,语句中单词具有丰富的语义信息,这些语义信息由编码器输出的高层次特征表示,单词本身就处于高层次,因此解码时不需要复杂的解码器就可以完成对缺失单词的重建。这也是为什么 BERT 中没有特意提到解码器的原因,BERT 选择 MLP 或者单一线性层就可以实现解码。但是在 CV 中,每个像素包含的语义信息非常匮乏,处于低层次,所以要想恢复原始图像中的像素,就需要精心设计层次较深的解码器来实现像素级重建。

基于这三点问题,作者通过高比例的遮盖(比如 75 % 75\% 75%)来避免模型偷懒,优化了模型精度。另外,高遮盖比使得编码器只会处理少部分 patch(例如 25 % 25\% 25%),这可以将预训练时间缩短三倍或更多,并且可以减少内存消耗,因此能够轻松将 MAE 扩展到大型编码器模型上。MAE 的非对称结构在保证解码器足以进行像素级解码的同时,尽可能选择轻量级解码器可以进一步减少计算量,即使是解码器的输入为全部 patch(含被遮盖掉的)。

4. 实验概述

对应论文 Experiments 部分。

4.1. 遮盖比

观察图 4 4 4 可以发现,尽管采用(全面)微调(fine-tuning)和线性探测(linear probing)方法的模型效果随遮盖比的变化趋势不同,比如线性探测方法的正确率增幅稳定、正确率总增量大( 54.6 % ∼ 73.5 % 54.6\%\sim73.5\% 54.6%73.5%,约 20 % 20\% 20%),而微调方法对遮盖比的敏感度比较低,在 40 % ∼ 80 % 40\%\sim 80\% 40%80% 的遮盖比区间都具有比较高的正确率。但是它们同在遮盖比为 75 % 75\% 75% 处取到最佳点(sweet point)。另外,作者通过一些正则化手段将原始的有监督训练出的 ViT 模型(即未经过 MAE 预训练的 ViT 模型)的效果从 76.5 % 76.5\% 76.5% 提升至 82.5 % 82.5\% 82.5%,但是即使如此也没有比任何一种遮盖比的(经过 MAE 预训练的)微调方法效果好。

除了在部分微调的实验部分使用了不同程度的冻结外,其他实验中使用的是全面微调。

在这里插入图片描述

图 4    遮盖比对模型的影响

4.2. 编码器设计

作者通过实验讨论编码器的设计对预训练出的 ViT 模型效果的影响,如图 5 5 5 所示。当解码器的 Trm 模块层数为 8 8 8 时,微调和线性探测的正确率都达到最佳点,但是相较于微调,线性探测评估方法对编码器层数比较敏感。当编码器的 Trm 模块输入和输入向量维度为 512 512 512 时,微调和线性探测的效果最好,其实对于微调来说,即使是比较低维度也有良好的表现。另外,可以观察到解码器 Trm 模块数量为 1 1 1 时与块数为 8 8 8 时的精度相差无几,但是块数的减少可以进一步加快训练速度。

当 ViT 模型输出向量的维度与解码器 Trm 模块输入向量维度不一致时,需要经过一层线性映射实现维度变换。

图中,加粗项为效果最佳点,灰色填充的项为 MAE 的默认设置。

在这里插入图片描述

图 5    解码器的 Trm 模块数量(左)和向量维度(右)对模型的影响

4.3. 编码器输入

作者对比了将可见部分与不可见部分同时输入到编码器和仅将可见部分输入到编码器对训练出的 ViT 模型的影响。图 6 6 6 上图中 encoder w/ [M] 的含义为 the input of ENCODER is visiable part With Masked part encoder w/o [M] 的含义为 the input of ENCODER is visiable part WithOut Masked part,即分别为输入包含不可见部分和不包含不可见部分。下图中的 encoder 没有 w/ [M] 标记的表示 w/o [M] 的 ViT 模型。

上图对比了准确率和计算量,结论是仅使用可见部分作为编码器的输入不但提升准确度还可以减少计算量。减少计算量是显然的,毕竟输入序列长度减小了;精度提升的原因是编码器在微调和迁移阶段的输入是整个图像,不存在被遮盖的部分,如果在训练时向编码器输入大量遮盖标记,编码器对遮盖标记的学习影响了对可见部分信息的学习,从而使得精度降低。

下图分别对比了编码器选择 ViT-L 和 ViT-H 时输入是否包含不可见部分的加速效果。可以看出,在 ViT-L + 单层 Trm 解码器和 ViT-H + 单层 Trm 解码器上,输入不包含不可见部分的加速效果分别是 ViT-L + 8 8 8 层 Trm 解码器和 ViT-H + 8 8 8 层 Trm 解码器的 3.7 3.7 3.7 倍和 4.1 4.1 4.1 倍。

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-SKdukVky-1679665875604)(.\pic\msk.png)]

图 6    编码器 ViT 的输入是否包含不可见部分的效果对比(上)和不同规模的 ViT 在不输入不可见部分情况下的效果(下)

4.4. 像素重建

对于解码器重建像素,作者尝试了对解码器最后一个 Trm 模块输出的 patch 进行处理,将处理后的 patch 进行线性映射。处理方法有:不做处理、patch 内规范化、对每个 patch 对应的向量通过 PCA 降至 96 维等。

从图 7 7 7 中可以看出 patch 内规范化可以提高重建精度,这种规范化强化了局部对比度。PCA 降维虽然精度不如规范化,但是由于 PCA 的效果也比较可观,故可以说明高频分量起到主导作用。

在这里插入图片描述

图 7    重建像素时采用的一些小技巧

4.5. 数据增强

作者通过实验对比了固定大小裁剪、随机大小裁剪和添加颜色抖动(color jit)的裁剪,(均带随机水平翻转),发现随机大小裁剪的效果最佳,添加颜色抖动反而降低了效果。还发现,仅进行中心裁剪(不带随机水平翻转),而不进行任何其他数据增强的效果也不错。作者认为出现这种情况的原因是 MAE 的数据增强主要是通过随机遮盖实现的,不相同的遮盖产生不同的数据在 MAE 的数据增强中起主导作用。如图 8 8 8 所示。

颜色抖动:随机改变图片的亮度,对比度和饱和度。

在这里插入图片描述

图 8    数据增强

4.6. 遮盖策略

9 9 9 展示了四种遮盖策略的效果,显然随机遮盖 75 % 75\% 75% 的效果最好。观察图 10 10 10 所示的效果,由于按块遮盖的重建工作最困难,所以训练后的损失值仍然比较高,重建也比较模糊;网格遮盖的重建工作最简单,训练出的损失值比较低,因此重建比较清晰,但是观察图 10 10 10 可以感觉到,相较于随机遮盖,网格遮盖对应的重建结果具有更明显、更频繁的灰色遮盖块,这种低质量的重建给人们的视觉带来相对不佳的体验。故作者最终还是选择 75 % 75\% 75% 的随机遮盖策略。

在这里插入图片描述

图 9    遮盖策略

在这里插入图片描述

图 10    遮盖策略
每种遮盖策略由上下两对图例表示。每个图例的左子图为遮盖情况,右子图为重建情况。左、中、右三组分别表示随机遮盖75%策略、按块遮盖50%策略和网格遮盖75%策略。

4.7. 迭代次数

作者的消融实验是基于 800 个 epoch 的预训练。图 11 11 11 展示了迭代次数对模型精度的影响。训练时间越长,精度越高。事实上,即使在 1600 个epoch,也没有观察到线性探测的精度有饱和现象。

在这里插入图片描述

图 11    迭代次数对模型精度的影响
长时间的训练对模型精度有明显提升,模型使用默认设置的 ViT-L,即上面实验中灰色填充的项。

4.8. 部分微调

12 12 12 是作者探究冻结层数对微调效果影响的结果曲线。值得注意的是,只微调最靠近 ViT-L 输出端的最后一个 Transformer 模块就能显著提高精度,从 73.5 % 73.5\% 73.5% 提高到 81.0 % 81.0\% 81.0%。此外,如果只微调最后一个块的 MLP 子模块,精度可以到达 79.1 % 79.1\% 79.1%(图中未标出),这比线性探测要好得多。观察曲线还可以发现,在微调块数为 4 4 4(或 6 6 6)之后增加微调 Trm 块数带来的效果不再明显。这是因为高层(靠近输出端)模块主要负责提取高维、更抽象的特征,而低层(靠近输入端)模块负责提取低维、更通用的特征,因此即使不对已经预训练过的低层模块使用与下游任务有关的数据进行微调也能有很好的提取通用特征的效果,但是少数的高层模块的作用是提取抽象特征,这与具体下游任务有关,所以微调对于它们来说是有意义的。

在这里插入图片描述

图 12    对 ViT-L 的部分 Trm 模型进行微调
图中的横坐标表示对靠近 ViT-L 模型的输出端的 Trm 模块进行微调的层数,其余靠近输入端的 Trm 模块参数被冻结。对 0 个模块进行微调表示线性探测;对 24 个模块进行微调表示全面微调。

REF

[1] He, K., Chen, X., Xie, S., Li, Y., Dollár, P., & Girshick, R. (2021). Masked Autoencoders Are Scalable Vision Learners. ArXiv. /abs/2111.06377

[2] MAE 论文逐段精读【论文精读】- bilibili

[3] 源码:GitHub - facebookresearch/mae: PyTorch implementation of MAE https//arxiv.org/abs/2111.06377

[4] 【深度学习】详解 MAE_深度学习mae - CSDN

[5] 何恺明最新一作:简单实用的自监督学习方案MAE,ImageNet-1K 87.8%! - 知乎

[6] 一文解读Masked Autoencoder(MAE)_- CSDN

[7] 28. MAE原理及代码 - 知乎

[8] MAE的理解和代码 - 知乎

[9] Vit的理解和代码(基于timm) - 知乎

[10] 【计算机视觉】ViT:Vision Transformer 讲解 - CSDN

[11] Masked Autoencoder论文中 fine-tuning 和 linear probing含义 - CSDN

[12] 【Linear Probing | 线性探测】深度学习 线性层 - CSDN

[13] 模型微调(fine-tuning)- CSDN

[14] 【pytorch】常用图像处理与数据增强方法合集(torchvision.transforms)- CSDN

猜你喜欢

转载自blog.csdn.net/weixin_46221946/article/details/129759337