论文地址:http://arxiv.org/abs/2111.06377
项目地址:(非官方)https://github.com/pengzhiliang/MAE-pytorch
本文的主要观点是:掩码自编码器 (MAE) 是适用于计算机视觉的可扩展自监督学习器。
如果说ViT对应于NLP中的Tranformer,MAE实际上是对应于BERT。
MAE的核心思想是随机屏蔽输入图像一部分的patches随后重建被屏蔽的像素。 这种将输入映射到潜在表示再重构输入的编码器-解码器结构就叫做Autoencoder。
MAE基于2个核心设计:
1、非对称编码器-解码器架构,其中编码器仅对可见的patches编码,随后将mask tokens加入编码结果中馈送到轻量级解码器。
2、屏蔽输入图像的较大一部分(本文使用75%),更有利于自监督任务的训练。
上述两种设计使得大型模型的训练可以加速3 倍以上并提高准确性。
作者提到了语言和视觉之间的信息密度不同。 语言是人类生成的具有高度语义和信息密集度的信号,一个句子入如果有几个遗漏单词,那么句子的意思会有较大改变。而图像是具有严重空间冗余的自然信号,例如可以通过插值提高图像分辨率、单一背景、重复图案等等。
同时作者认为在学习隐藏表示的语义层次时,解码器的设计起到了更重要的作用。
综上两点,MAE才使用了75%这么高的屏蔽率来减少编码器的计算量,同时这种设计还可以让原本需要较大数据来训练的模型,在ImageNet-1K数据集上就能有很好的效果。
模型结构
MAE编码器
MAE的编码器结构与ViT的的编码器结构相同,也实现利用线性投影和额外的positional embeddings把patches进行嵌入,然后送到一系列Transformer blocks中处理。但是MAE仅对25%的随机可见patches进行编码。
MAE解码器
MAE 解码器的输入在经过编码的可见patches基础上加入余下位置的Mask tokens,构成完整的整幅图像。Mask token是一个共享的、经过学习的向量来表示需要预测的缺失patch,Mask token也添加了位置嵌入。
由于MAE的主要目的是进行迁移学习的预训练, 解码器仅在预训练期间用于执行图像重建任务,因此解码器架构可以以独立于编码器设计。 本文使用了轻量级的Tranformer解码器来减少了预训练时间。
损失函数计算像素空间中重建图像和原始图像之间(仅计算被屏蔽的patches)的均方误差 (MSE)。
结果分析
屏蔽率对于fine-tuning和linear probing准确率的影响:
不同的屏蔽策略对重建图像的影响:
训练的Epochs与准确率的变化: