MAE论文精读,预训练图像

二、MAE论文精读
参考:

论文《Masked Autoencoders Are Scalable Vision Learners
李沐《MAE 论文逐段精读【论文精读】》、笔记《MAE 论文逐段精读》
知乎《别再无聊地吹捧了,一起来动手实现 MAE(Masked Autoencoders Are Scalable Vision Learners) 玩玩吧!》
本文算法在2.3章节!!!

2.1 导言
2.1.1 前言
  自从Vision Transformer将标准的transformer成功的应用到CV上面以来,就出现了很多相关性的工作,MAE就是其中的一篇。本文是2021.11.11发布在Arxiv上的文章,主要工作是在Vision Transformer基础上,引入自监督训练,相当于将BERT应用到CV领域。通过完形填空来获取对于图片的理解,把整个训练拓展到没有标号的数据上面 ,使得transformer在CV上的应用更加普及。
  最终MAE只需要Vision Transformer百分之一规模的数据集上预训练,就能达到同样的效果。而且在目标检测、实例分割、语义分割等任务上,效果都很好。

  本文标题中的Autoencoders,是‘自’而非自动的意思,表示类似自回归这一类模型的特点, 即标号和样本(y和x)来自于同一个东西。比如说在语言模型中,每一次用前面的次去预测下一个词。而对于MAE,则表明模型的标号也是图片本身。

2.1.2 摘要
  MAE的自监督训练实现途径非常简单,随机地mask图片中地一些patches(块),然后再去重构the missing pixels.,这个思想来自于BERT中的带掩码的语言模型。

在MAE中,自监督训练还有两个核心的设计:

非对称的编-解码器架构,非对称体现在两个方面:

两者输入不一样:MAE的编码器只编码可见的patches,被masked的块不编码,而解码器需要重构所有块。
Decoder更加轻量的:比如 Encoder 通常是多层堆叠的 Transformer,而 Decoder 仅需较少层甚至1层就 ok。这也表明 Encoder 与 Decoder 之间是解耦的。
高掩码率:比如mask 75%,才能得到一个比较好的自监督训练效果(否则只需要插值还原)。这等于只编码1/4的图片,训练速度加快了四倍,所以MAE可以扩展到更大的模型。

  最终作者只使用一个最简单的ViT-Huge模型,在ImageNet-1K上预训练,精度也能达到87.8%。

下面是作者展示的一些模型重构效果图:


图2:MAE在ImageNet验证集上构造出来的图片,验证集没有参与训练,所以相当于是测试结果。
红色框选的这一块,左侧列表示80%的patches被mask,中间列表示MAE重构的结果,右侧列表示原始图片。
可以看到被mask的地方非常多,很难看出原图是什么,但是重构之后还原度非常高,特别神奇。
图3:MAE在COCO验证集上的结果,MAE重构效果也很好。
图4:MAE取不同的遮盖比例时,模型重构效果对比。可以看到即使mask比例达到95%,还是能还原出大致图像,很玄乎。


2.1.3 导言: 自编码模型为何在CV领域应用不及NLP?
  在导言一开始,作者大意就说自监督学习很香,使得Transformer在NLP领域的应用非常火。尽管也有一些工作将BERT应用到CV领域,比如Denoising Autoencoder是在一个图片中加入很多噪音,然后通过去噪来学习对这个图片的理解,但最终这些效果都不及NLP。

MAE也是一种类似去噪 的模型,mask一些patches,也相当于往模型中加入很多噪音

  为何masked自编码模型在CV和NLP领域应用不一样呢?

架构差异
之前的CV领域都是用的CNN网络,而CNN中的卷积,是不好做掩码操作的。
  因为卷积核滑动时,不好将msaked patches单独剔除,导致后面不好将图片还原(Transformer中可以将masked token直接去掉,和其他词区分开来)。但是ViT成功将Transformer应用到图片分类,所以这个问题不再有了。
信息密度(information density)不同。
在自然语言中,一个词就是一个语义的实体,比如说字典中对一个词的解释就是很长的一段话,所以一句话中很难去去掉几个词(这样做完形填空才有意义)。
在图片中,像素是比较冗余的,取决于相机的分辨率有多大。所以如果只是简单去掉一些块的话,很容易进行插值还原。
作者的做法是,mask很高比例的像素块(比如75%),就可以大大降低图片的冗余性。这样就压迫模型必须学习全局信息,也就提高了模型学习的难度。
  将图片中一大片都去掉, 剩下的块离得比较远,就没那么冗余了。否则模型仅仅学一个局部模型就可以进行插值还原。
  关于此点论证,看上面图2.3.4就可以发现:仅仅一些局部的,很稀疏的块,就可以重构全局图片。

解码差异

NLP中需要预测masked tokens,而token本身就是一种比较高级一些的语义表示。而来自编码器的特征也是高度语义的,与需要解码的目标之间的 gap 较小,所以只需要一个简单的全连接层就可以解码这些tokens。(ViT最后需要的是图片特征,所以也只需要MLP就可以解码)
MAE中需要还原the missing pixels,这些像素是一个很基础的特征。要将自编码器的高级语义特征解码至低级语义层级,所以可能需要一个比较复杂的转置卷积网络才能完成解码。
  正是基于以上分析,作者才提出本文的两个核心设计:高掩码率以及非对称编-解码器。

  在导言最后一部分,作者说MAE可以只在ImageNet 1K上就能预训练ViT-Large/-Huge模型(得到很好的效果)。类比ViT,相同效果后者需要在近100倍规模的数据集上训练才能达到。
  另外,在目标检测、语义分割和实例分割上,MAE比之前所有模型(包括有监督训练模型)效果都要好,而且加大模型会有显著的收益。

2.2 相关工作(都是老故事,略)
2.3 算法
2.3.1 MAE主要结构
  MAE使用编码器将观察到的信号映射到潜在表示,再使用一个解码器从潜在表示重构原始信号。与经典的自动编码器不同,我们采用了一种非对称设计,允许编码器仅对部分观察信号(无掩码标记)进行操作,模型设计图如下:


预训练时:
切片:首先将图像切成一个个的patch(图像块)
encoder:将没有被mask的块输入到encoder层进行处理。(类似ViT:)
通过线性映射(linear projection)将每个Patch映射到一维向量中,得到Pacth embeddings
Pacth embedding+position embedding一起输入Transformer Encoder层得到其输出。
decoder:
将序列拉长,填入灰色块像素对应的embedding(主要就是position embeddings信息)。
将上面得到的长向量输入decoder,解码器会尝试将里面的像素信息全部重构回来,使之成为原始图片。
微调时:
只需用到endoder,将图片切成Patches后编码得到特征,然后用来处理自己的任务就行。(此时不需要掩码,也就不需要解码器)
  说明:上图中编码器比解码器更宽,表示模型的主要计算部分在编码器(最重要的就是编码像素得到特征)

2.3.2 具体实现
Masking(随机均匀采样)
随机均匀采样:在不替换的情况下,按照均匀分布对patches进行随机采样,采到的样本保留,剩下的全部mask掉。
具体实现:将patches经过线性映射成embedding,加上位置编码之后得到的token embeddings序列。再将这个序列随机打乱(shuffle),然后取前25%完成采样,输入encoder层。(后面75%就是做mask处理)
为何采用随机均匀抽样策略:一是避免潜在的“中心归纳偏好”(也就是避免 patch 的位置大多都分布在靠近图像中心的区域),二是这种策略还造就了稀疏的编码器输入(只处理可见块),能够以更低的代价训练较大规模的 Encoder。
这种掩码策略虽然简单,但其实很重要。因为其决定了预训练任务是否具有足够的挑战性,从而影响着 Encoder 学到的潜在特征表示 以及 Decoder 重建效果的质量。(后面有相关消融实验验证)
Encoder
这部分也和ViT一样,不再赘述,唯一不同的是,被mask的块不进行编码。
解码器
masked token:
需要注意的是,masked tokens不是通过mask 掉的 patch 经过 embedding 转换而来,而是通过一个共享的可学习向量来表示,也就是说每一个被mask的块都表示成同样的向量。
这个向量简单粗暴的复制n次之后每人一份。然后再加上对应的position emebdding进行区分。
解码器仅在预训练期间用于重构图像,下游任务只需要编码器对图像进行编码就行,这样解码器可以独立于编码器,就可以设计的比编码器更简单(更窄、更浅)。最终相比编码器,其计算量只有其10%。
通过结构图可以看到,encoder结构重载但是只处理unmask tokens;decoder处理所有tokens,但是结构轻量;通过这种非对称设计,使得整体结构十分高效,大大减少了预训练时间。
重构图像
解码后的所有 tokens 中取出 masked tokens(在最开始 mask 掉 patches 的时候可以先记录下这些 masked 部分的索引)
用一个MLP,将其从向量维度映射到一个patch的像素。
损失计算时,只对被mask的patches做计算(类似BERT),且使用MSE函数,也就是预测像素值和真实像素值的MSE。
对要预测的像素,可以在每个patch内部做一次normalization(减均值除方差),这样数据更稳定。
2.3.3 简单实现
将以上讲解的全部串联起来就是:

将图像划分成 patches:(B,C,H,W)->(B,N,PxPxC);
对各个 patch 进行 embedding(实质是通过全连接层),生成 tokens,并加入位置信息(position embeddings):(B,N,PxPxC)->(B,N,dim);
根据预设的掩码比例,进行随机均匀采样。unmask tokens 输入 Encoder,另一部分“扔掉”(mask 掉);
编码后的 tokens 与 masked tokens( 加入位置信息) 按照原先在 patch 形态时对应的次序拼在一起,然后喂给 Decoder 玩。
整个序列是通过unshuffle操作,还原到原来的顺序。
如果 Encoder 编码后的 token 的维度与 Decoder 要求的输入维度不一致,则需要先经过 linear projection 将维度映射到符合 Decoder 的要求
Decoder 解码后取出 masked tokens 对应的部分送入到全连接层,对 masked patches 的像素值进行预测,最后将预测结果(B,N’,PxPxC)与 masked patches 进行比较,计算 MSE loss。
2.4 实验部分
微调层数对比
  我们知道,预训练模型做微调的效果是比只做特征提取的效果更好的,但是训练时间会更长。所以作者试验了一下,微调不同的层数,模型的效果如何(见下图)。


横坐标表示被微调的层数(这些层参数可以训练,剩下层被冻结)
可以看到,MAE基本只需要微调最后4层就可以了;因为最后面的层,和任务最相关。而前面的层是更底层一些的特征,在各种任务上都更适用。
2.4.1 MAE超参数实验
1. ImageNet实验

下图是在 ImageNet-1K验证集上的top-1结果。(详见附录A1)


scratch,original:ViT-L/16模型正常地从头训练,效果其实不是很稳定。(200epoch)
scratch,our impl.:ViT-L/16加上比较强的正则(详见附录A2)
VIT原文中一直说需要很大的数据集才预训练出好的模型,但是后来大家发现,如果加入合适的正则项,在小一点的数据集上(ImageNet-1k)也能够训练出好的效果。
baseline MAE:在ImageNet上预训练,然后在ImageNet-1K上微调50epoch。
2. decoder/mask策略/数据增强


上图展示的是MAE使用ViT-L/16结构,在ImageNet-1K数据集上的acc精度。下面一个个来看:

ft: fine-tuning微调
lin:linear probing,只训练最后一个线性层的参数,相当于MAE只做特征提取
最终decoder的默认深度是8 blocks,宽度是512 dim,与VIT-L相比,计算量只有前者的9%.。
表a:展示不同的decoder深度,可见堆叠8个transformer block效果最好。(其实对微调来说,都差不多)

表b:decoder宽度,每个token表示成一个512维的向量,效果最好。

表c:是否编码masked tokens。结果是不编码精度更高,且计算量更少(编码的Flops是不编码的3.3倍)

表d:重建目标对比.。

第一行:MAE现行做法
第二行:预测时对每个patch内部做normalization,效果最好。
第三行:PCA降维
第四行:BEiT的做法,通过vit把每一块映射到一个离散的token上面再做预测。
表e:数据增强。结果显示,做简单的随机大小裁剪,效果就不错
fixed size和 rand size分别表示裁剪成固定大小,和裁剪成随机大小。color jit 是加入颜色变换,none就是不使用数据增强了。

表f:采样策略。随机采样方式最简单,效果也最好。
random/block/grid分别表示随机采样、按块采样(mask大块)、按网格采样。下面有进一步的说明图:


3. mask比例
综合ft和lin考虑,masking ratio=75%时效果最好。(图就不放了)

1.4.2 对比其它模型
1. 自监督模型效果对比

下表是在ImageNet-1K训练集上预训练后的微调结果:

在COCO数据集上的效果对比:


over

2.5 结论&讨论
简单且拓展性好的算法,是整个深度学习的核心。

简单:是指MAE在ViT模型上加了一些简单的扩展,但是模型本身是不简单的。
拓展性好:是指你有钱就可以无限的加数据了,毕竟MAE不需要标号,你可以训练更大的模型。
  自监督学习在最近几年是比较火的,但是在计算机视觉中,还是主要用有标号的数据来训练 。本文在ImageNet数据集上通过自编码器学习到可以媲美有标号训练的效果,使得CV领域的自监督学习可能走上与NLP类似的轨迹。
  另一方面,图像和语言是不同类型的信息,要谨慎处理这种差异。

对于语言来讲,一个token是一个语义单元,它含有的语义信息比较多(这种信息在ViT里提取出来了)。
在图片中,虽然一个patch也含有一定的语义信息,但它不是语义的分割(即某个patch中并不含有特定的物体,可能含有多个物体的一小块,或者是某一个物体重叠的一块)。但即使是在这样的情况下,MAE也能做很复杂的一些任务。作者认为MAE(transformer)确实能够学到隐藏的比较好的语义表达。
Broader impacts:MAE是个生成模型,可以生成不存在的内容,类似GAN,所以使用要注意。
 

你可能感兴趣的:(计算机视觉,深度学习,人工智能,机器学习)