【渔樵耕读】《MAE》

Masked Autoencoders Are Scalable Vision Learners

章主要工作在于提出一种基于掩码块(Masked Patch)的不对称Encoder-Decoder图像表示学习的方法。给定输入图像,划分patch并按特定比例随机mask。编码器在可见patch上学习潜在特征,解码器根据潜在特征重构完整图像。经过训练的模型具备很好的准确性和极高可拓展性来应用到下游任务。这种自监督的图像学习方式能够极大减少对标注图像数据的依赖。

Introduction

Why Masked Autoencoding

现代深度学习模型对大规模标注数据的需求与日俱增。NLP领域通过自监督预训练的方式解决了对数据的需求。以GPT、BERT等预训练模型为代表,使用通俗的解决思路:mask部分数据让模型学习被mask的内容。实际上,在图像领域这种广义上的掩码自编码器在NLP之前已经被广泛研究,却始终落后与NLP。

Why Masked Autoencoding in CV Lags behind NLP

文章关注点在于:掩码自编码器在图像和语言间的不同究竟在哪?

  1. CV中的主要backbone——卷积网络在不规则的mask patch图像上不容易操作。但这种结构上的gap可以被Vision Transformer(ViT)很好的解决
  2. 图像和文本的信息密度不同——文本是人为创造的,拥有高级语义信息和信息密度,对每句话的少量mask work预测任务就可以导致复杂的信息预测,即少量的缺失单词包含很多信息,需要模型掌握复杂的语言信息才能去做预测。而图像作为自然中的信息载体,具备空间信息冗余的特点。即对缺失patch的预测可以利用周围图块的少量信息就可以完成。解决这种不同的措施是对图像作高比例的mask,如选择75%的随机patch作mask。大规模掩码减少了信息冗余,需要模型对图像有更高效的理解
  3. 将潜在表达特征还原图像的Decoder在重构图像和文本时有很大差异。图像任务中的解码器负责重构像素,具有很少的语义信息,而对单词的预测则包含更多的信息。这种现象导致在NLP中编码器的结构其实可以非常简单就可以预测缺失单词(如BERT中以MLP作为解码器),而在CV中,解码器需要精心设计来从学习到的潜在表示中预测mask pixel

具体的,文章提出一种简单高效并易于拓展的掩码自编码器来做图像的表示学习任务。它是一种不对称的Encoder-Decoder结构。不对称在于编码器只会学习没有被mask的patch而解码器会根据特征重建完整图像。随机选择输入图像的图块作mask,然后重构被mask的图块。掩码图块的比例非常高,文章选择75%,编码器在剩下的25%图块中学习潜在特征,而轻量的解码器重构完整图像。得益于非常高的掩码率,编码器和解码器都可以得到高效的计算效率,最大限度减小内存使用,以此获得大规模预训练模型。

对大规模模型采用这种预训练的表示学习,可以极大提高模型表现,甚至能够得到比有监督学习更好的结果,这和NLP中的预训练表现相似。

自编码器 Autoencoder:自编码器是学习特征表述的经典方法,主要过程是将输入map到一个潜在特征,然后使用解码器可以重构原始输入。PCA和K-Means本质上都属于自编码器。自编码器的种类很多,MAE属于Denoising Autoencoder去噪自编码器的一种。在MAE中masked patches可以被视为一种噪声。

Main Approach

MAE的目标在于根据原始信号的部分观测结果重建完整信号

Mask

将输入图片切分为多个不重叠patch, 按照均匀分布随机选择patch子集,剩下的patch作掩码(可直接删除)。这种高掩码率很大程度上可以消除图像信息冗余,使得模型不能简单从可见的相邻patch直接外推来预测掩码patch。均匀分布采样避免了可见像素的聚集。最终得到高度稀疏的mask图像

Encoder

作者直接采用ViT中的编码器结构——线性映射+位置编码,以及一些列Transformer Blocks。不同的是,编码器只作用于剩余的patch中,被mask的块不被使用,甚至可以直接删除。这种方式可以使用有限的算力和内存去训练很大的编码器模型。

Decoder

注意到在自编码器中,Decoder只在预训练时才被用来完成图像重构任务。在实际表示学习时只用到Encoder。MAE中,Decoder的输入包含两部分: i) 编码的可视patch的token, ii)mask patch的token. 这里的每个mask token是共享的、可被学习的向量,用来预测被mask的patch。这里每个token都被加入位置向量用来编码位置信息。

Reconstruction target

对于每个mask的patch,MAE通过预测像素值重构输入图像。Decoder输出的每个元素是一个像素值的向量,用来表示该patch。Decoder的最后一层是一个线性层,维度等于patch中像素值的数量。训练时使用原图像和重构图像在像素空间中的MSE作为Loss,这里的Loss只计算mask patch。另外,基于patch中的归一化像素值的重构目标也可以获得一部分的效果提升。

Implementation

数据处理方面,对于切分的每个patch,经过线性投影+位置编码获得对应的token。对所有token随机排序并采样,剩下的token作为mask移除。编码器可这些可见的token作编码。对于解码器,将编码后的patch还原到原始位置,原先确实的patch使用使用mask tokens填充,所有的mask token与预测目标token相对应。送入解码器。解码器预测对应patch的图像像素点。计算预测的像素和原始图像的像素之间的MSE只作为Loss。

最终训练好的解码器可以作为下游任务的basemodel,并可以被fineturn

QUES

encoder 是否使用 mask token:文中取可见 patches 的方式是对输入的所有 patch (假设一共有 N 个patch)先 shuffle 一下, 然后把后 N * masking_ratio 的 patch 丢掉,只有前面部分过网络。 另一种可能的方法是在 encoder 那增加一个 mask token 指代哪些 patch 需要保留哪些不需要保留。 实验结果表明这样做效果明显变差了, 原因使用 mask encoding 在训练时模型见到的是一些残缺的输入, 但是在实际下游任务中见到的确实完整的输入, 这会带来很大的 gap。