何恺明最新工作：简单实用的自监督学习方案MAE，ImageNet-1K 87.8%！-技术圈

↑ 点击蓝字关注极市平台

作者丨happy

编辑丨极市平台

极市导读

恺明提出一种用于计算机视觉的可扩展自监督学习方案Masked AutoEncoders。所提MAE极为简单：对输入图像进行块随机mask并对遗失像素进行重建。该方案使得所得高精度模型具有很好的泛化性能：仅需ImageNet-1K，ViT-Huge取得了87.8%的top1精度。 >>加入极市CV技术交流群，走在计算机视觉的最前沿

论文链接：https://arxiv.org/pdf/2111.06377.pdf

恺明出品，必属精品！这篇文章延续了其一贯的风格：简单且实用。这篇文章仍属于恺明最近两年的研究领域：自监督领域(自监督学习就是被他带火的吧)。本文的出发点则是BERT的掩码自编码机制：移除一部分数据并对移除的内容进行学习。掩码自编码源于CV但盛于NLP，恺明对此提出了疑问：是什么导致了掩码自编码在视觉与语言之间的差异？尝试从不同角度进行解释并由此引申出了本文的MAE。

Abstract

恺明提出一种用于计算机视觉的可扩展自监督学习方案Masked AutoEncoders(MAE)。所提MAE极为简单：对输入图像进行块随机mask并对遗失像素进行重建。它基于以下两个核心设计：

我们设计了一种非对称编解码架构，其中编码器仅作用于可见块(无需mask信息)，而解码器则通过隐表达与mask信息进行原始图像重建；
我们发现对输入图像进行高比例mask(比如75%)可以产生一项重要且有意义的自监督任务。

上述两种设计促使我们可以更高效的训练大模型：我们加速训练达3x甚至更多，同时提升模型精度。所提方案使得所得高精度模型具有很好的泛化性能：仅需ImageNet-1K，ViT-Huge取得了87.8%的top1精度 。下游任务的迁移取得了优于监督训练的性能，证实了所提方案的可扩展能力。

极致精简版

用下面几句话来简单说明下这篇文章：

恺明出品，必属精品！MAE延续了其一贯的研究风格：简单且实用；
MAE兴起于去噪自编码，但兴盛于NLP的BERT。那么是什么导致了MAE在CV与NLP中表现的差异呢？这是本文的出发点。
角度一：CV与NLP的架构不同。CV中常采用卷积这种具有”规则性“的操作，直到近期ViT才打破了架构差异；
角度二：信息密度不同。语言是人发明的，具有高语义与信息稠密性；而图像则是自然信号具有重度空间冗余：遗失块可以通过近邻块重建且无需任何全局性理解。为克服这种差异，我们采用了一种简单的策略：高比例随机块掩码，大幅降低冗余。
角度三：自编码器的解码器在重建方面的作用不同。在视觉任务方面，解码器进行像素重建，具有更低语义信息；而在NLP中，解码器预测遗失的词，包含丰富的语义信息。
基于上述三点分析，作者提出了一种非常简单的用于视觉表达学习的掩码自编码器MAE。
MAE采用了非对称的编解码器架构，编码器仅作用于可见图像块(即输入图像块中一定比例进行丢弃，丢弃比例高达75%)并生成隐式表达，解码器则以掩码token以及隐式表达作为输入并对遗失块进行重建。
搭配MAE的ViT-H取得了ImageNet-1K数据集上的新记录：87.8%；同时，经由MAE预训练的模型具有非常好的泛化性能。

Method

所提MAE是一种非常简单的自编码器方案：基于给定部分观测信息对原始信号进行重建 。类似于其他自编码器，所提MAE包含一个将观测信号映射为隐式表达的编码器，一个用于将隐式表达重建为原始信号的解码器。与经典自编码器不同之处在于：我们采用了非对称设计，这使得编码器仅依赖于部分观测信息(无需掩码token信息)，而轻量解码器则接与所得隐式表达与掩码token进行原始信号重建(可参见下图)。

Masking 参考ViT，我们将输入图像拆分为非重叠块，然后采样一部分块并移除其余块(即Mask)。我们的采样策略非常简单：服从均匀分布的无重复随机采样 。我们将该采样策略称之为“随机采样”。具有高掩码比例的随机采样可以极大程度消除冗余，进而构建一个不会轻易的被近邻块推理解决的任务 (可参考下面图示)。而均匀分布则避免了潜在的中心偏置问题。

MAE Encoder MAE中的编码器是一种ViT，但仅作用于可见的未被Mask的块。类似于标准ViT，该编码器通过线性投影于位置嵌入对块进行编码，然后通过一系列Transformer模块进行处理。然而，由于该编解码仅在较小子集块(比如25%)进行处理，且未用到掩码Token信息。这就使得我们可以训练一个非常大的编码器 。

MAE Decoder MAE解码器的输入包含：(1) 编码器的输出；(2) 掩码token。正如Figure1所示，每个掩码Token共享的可学习向量，它用于指示待预测遗失块。此时，我们对所有token添加位置嵌入信息。解码器同样包含一系列Transformer模块。

注：MAE解码器仅在预训练阶段用于图像重建，编码器则用来生成用于识别的图像表达 。因此，解码器的设计可以独立于编码设计，具有高度的灵活性。在实验过程中，我们采用了窄而浅的极小解码器，比如默认解码器中每个token的计算量小于编码器的10% 。通过这种非对称设计，token的全集仅被轻量解码器处理，大幅减少了预训练时间。

Reconstruction target 该MAE通过预测每个掩码块的像素值进行原始信息重建。解码器的最后一层为线性投影，其输出通道数等于每个块的像素数量。编码器的输出将通过reshape构建重建图像。损失函数则采用了MSE，注：类似于BERT仅在掩码块计算损失。

我们同时还研究了一个变种：其重建目标为每个掩码块的规范化像素值 。具体来说，我们计算每个块的均值与标准差并用于对该块进行归一化，最后采用归一化的像素作为重建目标提升表达能力。

Simple implementation MAE预训练极为高效，更重要的是：它不需要任何特定的稀疏操作。实现过程可描述如下：

首先，我们通过线性投影与位置嵌入对每个输入块生成token；
然后，我们随机置换(random shuffle)token序列并根据掩码比例移除最后一部分token；
其次，完成编码后，我们在编码块中插入掩码token并反置换(unshuffle)得到全序列token以便于与target进行对齐；
最后，我们将解码器作用于上述全序列token。

正如上所述：MAE无需稀疏操作。此外，shuffle与unshuffle操作非常快，引入的计算量可以忽略。

Experiments

我们在ImageNet-1K数据集上进行自监督预训练，然后再通过监督训练评估预训练模型的表达能力。

Main Properties

Baseline：ViT-Large 。我们采用ViT-Large作为消融实验的骨干，上表为从头开始训练与MAE微调的性能对比。可以看到：从头开始训练(200epoch)，ViT-L的性能为82.5%且无需强正则技术；而MAE(注：仅微调50epoch)则取得了大幅性能提升。

上表则从不同角度进行了消融实验对比，一一道来。

Decoder Design 从Table1a与Table1b可以看到：解码器的设计可以非常灵活 。总而言之，默认解码器非常轻量，仅有8个模块，维度为512，每个token的计算量仅为编码的9%。

Mask Token MAE的的重要设计：在编码阶段跳过掩码token，在解码阶段对其进行处理。Table1c给出了性能对比，可以看到：编码器使用掩码token则会导致性能下降 。

Recontruction target Table1d比较了不同重建目标的性能，可以看到：引入块归一化可以进一步提升模型精度 。

Data Augmentation Table1e比较了不同数据增广的影响，可以看到：MAE仅需crop即可表现非常好，添加ColorJitter反而会影响性能 。另外，令人惊讶的是：当不使用数据增广时，MAE性能也非常优秀 。

Mask Sampling Table1f比较了不同搞得掩码采样策略，可以看到：不同的采样策略均具有比较好的性能，而随机采样则具有最佳性能 。

Masking ratio 下图给出了掩码比例的影响，可以看到：最优比例惊人的高 。掩码比例为75%有益于两种监督训练方式(端到端微调与linear probing)。这与BERT的行为截然相反，其掩码比例为15%。

与此同时，从上图可以看到：端到端微调与linear probing两种方式存在不同的趋势：

对于linear probing而言，模型性能随掩码比例非常稳定的提升直到达到最高点，精度差约为20%；
对于微调来说，模型性能再很大范围的掩码比例内均极度不敏感，而且所有微调结果均优于linear probing方式。

Training Schedule 下图给出了不同训练机制的性能对比(此时采用了800epoch预训练)，可以看到：更长的训练可以带来更定的精度提升 。作者还提到：哪怕1600epoch训练也并未发现linear probing方式的性能饱和。这与MoCoV3中的300epoch训练饱和截然相反 ：在每个epoch，MAE仅能看到25%的图像块；而MoCoV3则可以看到200%，甚至更多的图像块。

Comparisons with Previous Results

上表给出了所提MAE与其他自监督方案的性能对比，从中可以看到：

对于ViT-B来说，不同方案的性能非常接近；对于ViT-L来说，不同方案的性能差异则变大。这意味着：对更大模型降低过拟合更具挑战性 。
MAE可以轻易的扩展到更大模型并具有稳定的性能提升。比如：ViT-H取得了86.9%的精度，在448尺寸微调后，性能达到了87.8% ，超越了此前VOLO的最佳87.1%(尺寸为512)。注：该结果仅使用了ViT，更优秀的网络表达可能会更好。