跳转至

MAE1

概要

本文表明, 利用MAE可以在计算机视觉任务中实现可扩展的自监督学习. MAE方法很简单: 屏蔽输入图像的随机补丁并重建丢失的像素. 它基于两种核心设计, 第一, 作者开发了一种不对称的编码器-解码器架构. 其中, 编码器仅作用于可见部分的图像patches(不会输入mask tokens); 解码器是一个较轻的网络, 用编码器输出的latent representation和mask tokens来重建完整图像. 第二, 作者发现, 如果对输入图像进行较高比例的随机遮挡, 比如75%, 不仅能够保持训练的可行性, 还能形成一个有挑战性且对模型学习更加有益的自监督任务. 通过将以上的两个设计结合, 高比例遮挡和不对称编码器-解码器架构可以大幅加速训练(3倍或者以上的速度提升), 同时还能提高模型在下游任务中的准确率. 这种方法对大规模模型同样适用(可扩展性好), 例如, 使用ViT-Huge模型(在仅使用ImageNet-1k数据的条件下), 模型在ImageNet上取得了87.8%的准确率, 这是同等条件下的最佳水平. 在下游任务中, 和监督预训练相比, 使用MAE预训练的模型表现更好, 而且随着模型规模的变大, 性能能够持续提升.

背景

深度学习见证了能力和规模不断增长的架构的爆炸式增长💥. 在硬件快速发展的帮助下, 训练深度学习模型对数据的需求也快速增加. 现在的模型即使有一百万张图像也很容易过拟合, 甚至开始需要上亿张(很多时候并非公开可用)的标注图像来训练.

这种对数据的需求已经通过自监督预训练在自然语言处理中成功得到解决. 其核心方法源自于两个思路: (1) 自回归语言模型(GPT), 通过预测下一个词或者语句来学习语言模型; (2) 掩码语言模型(BERT), 先随机遮盖部分文本, 再训练模型去预测被遮盖的内容. 使用这些自监督预训练方法, 我们现在可以训练包含数千亿参数的通用NLP模型, 并在各种下游任务中表现出色.

MAE是一种更加通用的去噪自编码器, 它很自然, 而且适用于计算机视觉. 事实上, 在视觉领域相关的工作要早于BERT. 然而, 尽管随着BERT的成功, 人们对这一想法产生了浓厚的兴趣, 但是视觉方面的自编码方法的进展却落后于NLP. 那么, 视觉和语言之间的masked自编码器有何不同? 作者尝试从以下的几个角度来回答这个问题:

  1. 在过去的十年中, 卷积神经网络在计算机视觉任务中占据了主导地位, 它们通过卷积操作在规则的像素网格上提取特征, 集成诸如掩码tokens或者位置嵌入等indicators并不直观. 然而, 这些架构之间的差距, 已经随着ViT的出现得到了解决
  2. 语言和视觉的信息密度不同. 语言是人类生成的信号, 是高度语义化的, 具有较高的信息密度. 当训练一个模型来预测每个句子中仅有的介个缺失单词的时候, 这项任务似乎能够诱发复杂的语言理解. 相反, 图像是具有严重空间冗余的自然信号, 例如, 一个缺失的patch可以从邻近的patches中恢复, 而不需要高度的对对象和场景的理解. 为了克服这种差异并且鼓励模型学习到有用的特征, 作者展示了一种在计算机视觉中行之有效的简单策略. 掩盖极高比例的随机patches. 这种策略在很大成都上减少了冗余, 并创造了一个具有挑战性的自监督任务, 需要对低层次图像统计之外的整体理解. 为了对他们的重建任务产生一个定性的认识, 请见下图
  3. AE的解码器, 在文本和图像的重构中扮演者不同的角色. 在视觉中, 解码器对像素进行重构, 因此其输出的语义水平低于普通的识别任务, 限制了其在表示中的高层语义信息的表达. 而文本中, 解码器预测的是缺失的单词, 这些单词中包含丰富的语义信息, 解码器需要理解上下文并有意义的词汇, 涉及更高层次的语义理解. 在BERT中, 解码器很明显是一个MLP, 但是作者发现对于图像来说, 解码器的设计非常重要, 直接影响潜在表示中的语义信息的层次

ImageNet验证图像的结果, 对于每个triplet, 作者展示了掩蔽图像(左), MAE的重建结果(中)以及真实图像(右). 掩盖比例为80%, 仅仅留下196个patches中的39个. 由于为在可见补丁上计算损失. 注意, 由于在可见patches(那39个)未计算损失, 模型在可见patches上的输出在质量上较差, 可以简单地将输出和可见patches叠加以提高视觉质量, 但是作者故意选择没这么做, 以便更全面地展示该方法的行为

COCO验证图像上的示例结果. 使用的是在ImageNet上训练的MAE(模型权重和上图的相同). 观察最右侧两个示例的重建结果, 虽然和实际情况不同, 但是在语义上是合理的

在这一分析的推动下, 他们提出了一种用于视觉表征学习的简单, 有效, 可扩展的掩码自编码器MAE. MAE屏蔽输入图像中的随机patches, 并在像素空间中重建缺失的patches. 它采用了非对称的编码器-解码器设计. 编码器只对可见的patches进行编码, 而解码器则是轻量级的, 可以从latent representation和mask tokens中重建输入. 如下图所示.

MAE的架构. 在预训练的时候, 很大一部分随机的patches如75%会被遮挡. 编码器用于所有可见patches的编码. 在encoder之后才引入mask tokens. 编码后的可见patches和mask tokens被一个较小的decoder处理以像素为单位重建原始图像. 在预训练之后, 解码器会被丢弃, 编码器将用于未损坏的图像图像的识别任务

他们不在encoder阶段显式处理所有的mask tokens, 这让encoder对少量可见的patches做计算, 这样的encoder计算开销更小, 然后, 在decoder阶段再把mask tokens引入, 和encoder的输出特征一起输入给小型decoder, 再进行重建, 因此, 对于使用一个对称架构, 在encoder中就全部patches的情况, 计算量大幅度降低. 在这种设计下, 极高的掩码率, 如75%可以实现双赢, 既能优化准度又能让编码器只处理一小部分, 如25%的patches. 这可以将整个预训练的事件减少3倍或者更多, 使得他们能够轻松地将MAE扩展到大型模型.

MAE可以学习到high-capacity的模型, 同时还能很好的泛化. 通过MAE预训练, 像ViT-Large或ViT-Huge这样的大型视觉Transformer模型, 能在ImageNet-1K这种相对有限的数据集上取得更好的泛化性能. 以ViT-Huge为例, 微调后再ImageNet-1K上取得了87.8%的准确率, 超过了以往只使用ImageNet-1K数据的结果. 在目标检测, 实例分割和语义分割等下游任务中, MAE预训练不仅比监督预训练表现更好, 而且在模型增大的时候获得更加明显的收益. 这些现象和NLP中的自监督预训练(如BERT等)的发展相符, 作者希望视觉领域能够借此继续探索并取得类似的突破.

以75%的遮挡率训练的MAE对ImageNet验证图像进行重建, 但是遮挡率不同. 预测结果和原始图像有明显差异, 这表明该方法可以泛化

相关工作

MLM及AR

如BERT和GPT, 是在NLP领域最成功的预训练方法. 这些方法保留了部分输入序列, 并训练模型来预测缺失的内容. 这些方法已被证明可以很好地扩展. 大量研究表明, 这些预训练的表征可以很好地泛化到各种下游任务中.

AE

AE是学习表示的一种经典方法. 它有一个将输入映射到潜在表示的编码器和一个重构输入的解码器. 例如, PCA和K-means就是AE. 去噪自编码器(DAE)是一类自编码器, 它能破坏输入信号并学习重建未破坏的原始信号, 一系列方法可以倍视为不同破坏情况下的广义DAE, 例如掩蔽像素或者去除颜色通道. 作者的MAE是去噪自编码器的一种, 但是和经典的DAE有许多不同之处.

掩码图像编码

掩码图像编码是从被遮蔽破坏的图像中学习表示的方法. 一项开创性的研究将遮蔽作为DAE中的一种噪声类型. Context Encoder使用CNN来填补图像中大范围缺失的区域. 受到NLP中掩码语言模型的启发, 最近的相关方法使用Transformer架构来处理图像数据, iGPT将图像视为像素序列, 并通过预测未知像素来进行训练, 这种方法借鉴了GPT在语言模型中的自回归生成能力, 应用于图像生成任务. ViT研究了使用掩码patches进行自监督学习, 通过预测被遮挡的图像patches来学习图像表示. 最近, BEiT提出通过预测图像的离散tokens来进行训练.

自监督学习

自监督学习的方法在计算机视觉领域引起了很大的兴趣, 通常侧重于不同的前置任务进行预训练. 最近, 对比学习很火热, 它对两个或者多个视图之间的图像相似性和不相似性进行建模. 对比学习和相关的方法在很大程度上依赖数据的增强. 自编码器追求的是一个不同的方向, 它表现出不同的行为, 我们将对此介绍.

方法论

MAE是一种简单的自编码方法, 它能根据部分观测结果重建出原始信号. 和所有的自编码器一样, 他们的方法也有一个将观测信号映射到潜在表示的编码器, 和一个从潜在表示重建原始信号的解码器. 和传统的AE不同的是, 他们采用了一种非对称的设计, 使得编码器只对部分观测信号(不含mask tokens)进行操作, 而轻量级解码器则从潜在表示和mask tokens中重建完整的信号. 上面的架构图也说明了这个方法.

Masking

按照ViT, 他们将图像划分为规则的非重叠patches. 然后, 他们对其中的一部分进行采样, 并屏蔽剩余的部分. 他们的采样策略简单明了: 按照均匀分布随机采样, 他们称之为"随机抽样".

随机抽样具有较高的遮敝率, 这在很大程度上消除了冗余, 从而创建了一个无法通过从可见的邻近patches进行进行外推来轻松解决的任务. 均匀分布防止了潜在的中心偏差(即图像中心附近有更多被遮挡的patches), 最后, 高度稀疏的输入为设计高效的编码器创造了机会. 接下来介绍这个高效的编码器.

编码器

MAE的encoder是一个ViT, 但是输入只是可见的, 没有被屏蔽的patches. 就像一个标准的ViT, encoder通过一个线性变换将patch映射到嵌入向量然后和位置向量相加, 然后通过一系列的Transformer块来处理生成的集合. 但是, 他们的encoder只在整体集合的一小部分(25%)上操作, 掩盖的patches会被移除. 这使得我们可以训练非常大的encoders. 整个集合会经过一个轻量的decoder处理, 在下面描述.

MAE解码器

输入MAE解码器的是一整套tokens. 包括: (i) 已经编码的可见patches和 (ii) mask tokens. 见上面的架构图. 每个mask token都是一个共享的, 经过学习的向量. 并且, 他们为整个集合中的所有tokens添加位置嵌入, 如果没有位置嵌入, mask tokens在图像中就没有位置信息, 解码器也有一系列的transformer块.

MAE解码器仅在预训练期间用于执行图像重建任务. (只有编码器用于生成用于识别的图像表征). 因此, 解码器的结构设计可以灵活地独立于编码器的而设计. 他们尝试使用非常小的解码器, 并编码器更窄更浅. 例如, 他们的默认解码器和编码器相比, 每个tokens的计算量仅仅占到了10%. 通过这种非对称设计, 整个tokens集合由一个轻量的解码器处理, 从而大大减少了预训练时间.

重建目标

MAE通过预测每个遮蔽patch的像素值来重建输入. 解码器输出中的每个元素都是代表一个patch的像素值向量. 解码器的最后一层是线性投影, 其输出通道数等于一个patch中像素值的数量, 解码器的输出经过重塑后形成重建图像. 他们的损失函数是重建图像和原始图像之间的均方误差. 他们只对遮蔽patches计算损失, 这和BERT类似.

此外, 作者还研究了一种变体. 其重建目标是每个mask token的归一化像素值, 具体来说, 他们计算了一个patch中所有的像素的平均值和标准差, 并用它们对patch进行归一化处理. 在他们的实验中, 使用归一化像素作为重建目标可以提高表示的质量.

简单实现

MAE预训练可以简单, 高效地实现, 重要地是不需要任何专门的稀疏操作. 首先, 对于每个输入patch生成一个token(通过线性投影并添加位置嵌入). 接下来, 随机打乱tokens列表, 然后根据屏蔽率移除掉列表的最后一部分. 这一过程会为编码器生成一小部分的tokens子集, 相当于对patch进行不替换采样. 编码完成之后, 我们会将mask tokens列表添加到已编码的patches列表中, 并反转之前的打乱操作. 添加位置嵌入之后通过解码器. 如前所述, 无需进行稀疏操作, 这种简单的方式带来的开销可以忽略不记, 因为洗牌和反洗牌操作都是快速的.

  1. He, K., Chen, X., Xie, S., Li, Y., Dollár, P., & Girshick, R. (2021). Masked autoencoders are scalable vision learners (No. arXiv:2111.06377). arXiv. https://doi.org/10.48550/arXiv.2111.06377 ↩

评论