• 问答
  • 技术
  • 实践
  • 资源
何恺明最新一作:简单实用的自监督学习方案 MAE,ImageNet-1K 87.8%!
技术讨论

作者丨happy
编辑丨极市平台
本文首发于极市平台,转载须经授权并注明来源

论文链接: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,更优秀的网络表达可能会更好。

Transfer Learning Experiments

上表给出了COCO检测与分割任务上的迁移性能对比,可以看到:相比监督预训练,MAE取得了全配置最佳 。当骨干为ViT-B时,MAE可以取得2.4AP指标提升;当骨干为ViT-L时,性能提升达4.0AP。

上表给出了ADE20K语义分割任务上的迁移性能对比,可以看到:MAE可以大幅改善ViT-L的性能,比监督训练高3.7。

全文到此结束,更多实验结果与分析建议查看原文。

  • 0
  • 0
  • 675
收藏
暂无评论