原文摘要 · English Abstract
We mask random patches of the input image and reconstruct the missing pixels. It is based on two core designs: an asymmetric encoder-decoder architecture, with an encoder that operates only on the visible subset of patches, along with a lightweight decoder that reconstructs the original image from the latent representation and mask tokens; and masking a high proportion (e.g. 75%) of the input image.
概要
MAE 要回答的是:BERT 式「遮一部分、预测被遮内容」的范式为什么在视觉上做不出 NLP 那样的规模效应。作者的诊断是信息密度差异——图像冗余大,遮掉一小块能靠邻域插值补回。解法是把随机掩码比例拉到 75%,并配非对称编码器-解码器:编码器只处理剩下 25% 的可见 patch、不看掩码 token,解码器又浅又窄、只在预训练期存在,直接回归被掩码 patch 的原始像素。这既把任务变难又把训练变快(提速 2.8 倍以上),仅用 ImageNet-1K 无标签数据就把 ViT-Huge 训到 87.8%。
MAE 之前视觉自监督的主流是对比学习(MoCo、SimCLR、DINO 一系):靠同图多视图构造正负样本,效果好但强依赖精心设计的增强。掩码建模出现更早(Context Encoder、iGPT),却长期落后于 NLP。作者把原因归为三点:卷积网络难以容纳掩码 token 与位置嵌入(ViT 出现后已不成立)、图像冗余使低掩码比例的任务过于简单、以及视觉解码器输出的语义层级远低于识别任务,因此解码器设计直接决定表征质量。
关键贡献
- 指出视觉与语言在信息密度上的根本差异,并据此论证高比例掩码是视觉掩码建模的必要条件,而非调参技巧
- 提出非对称编码器-解码器架构:编码器只接收可见 patch、不引入掩码 token,掩码 token 推迟到轻量解码器才加入
- 证明直接回归原始像素就足够,无需 dVAE tokenizer 这类需要额外数据与预训练阶段的离散化目标
- 把损失只算在被掩码的 patch 上,并引入逐 patch 归一化的重建目标,进一步提升表征质量
- 在纯 ImageNet-1K 自监督预训练下超过有监督预训练与 BEiT,并在 COCO、ADE20K、iNat、Places 上验证迁移能力
- 展示了视觉自监督的规模效应:模型越大收益越大,线性探测精度到 1600 轮仍未饱和
方法详解(完整做法)
总体思路
MAE 把自监督预训练写成一个像素级填空题。输入图像先切成规则不重叠的 patch,按均匀分布随机丢掉其中 75%;只有剩下 25% 的可见 patch 被线性投影加位置编码后送入一个标准 ViT 编码器。编码器输出的隐向量与代表被丢弃位置的共享掩码 token 拼回完整长度的序列,重新加上位置编码,交给一个又浅又窄的解码器,逐 patch 回归原始像素值,损失只在被掩码的位置上计算。预训练结束后解码器被整个丢弃,只保留编码器作为下游骨干,通过端到端微调迁移到分类、检测与分割。全流程没有正负样本、没有动量编码器、也不需要额外的图像 tokenizer。
分块与高比例随机掩码
图像被切成边长 的不重叠 patch,按均匀分布无放回地采样出可见集合 ,其余索引构成掩码集合 。之所以坚持均匀随机而不用块状或网格采样,一是均匀分布可以避免「被掩码 patch 更集中在图像中心」的位置偏置,二是随机采样允许把比例推到很高而精度不塌——论文对比中随机 75% 的微调/线性探测为 84.9/73.5,块状 75% 掉到 82.8/63.9,网格 75% 为 84.0/66.0。
非对称编码器:只看得见的那部分
编码器就是一个原版 ViT,但它只在可见子集上运行:既不接收掩码 token,也不为被丢弃的位置保留槽位。因为自注意力的计算量随序列长度平方增长,把序列砍到四分之一带来的节省远超线性,使得同样的算力可以训练更大的编码器或更长的调度。实现上不需要任何稀疏算子——先给全部 patch 生成 token,随机打乱后截掉尾部即等价于无放回采样。
轻量解码器与掩码 token 的落位
解码器的输入是完整的 个 token:可见位置放编码器输出,掩码位置放一个共享的可学习向量 ,随后给所有位置重新加上位置编码(否则掩码 token 不知道自己在图像何处)。解码器默认 8 个 Transformer block、宽度 512,每个 token 的计算量不到编码器的十分之一,且只在预训练期存在,因此可以独立于编码器自由设计。把掩码 token 挪到这里,是整个架构不对称的来源。
重建目标:只算掩码位置的像素 MSE
解码器最后一层是一个线性投影,输出通道数等于一个 patch 的像素数,输出重排即得重建图像。损失是重建与原图在像素空间的均方误差,且仿照 BERT 只在被掩码的 patch 上计算——论文脚注说明在全部像素上算会让精度略降约 0.5 个点。作者还试了把目标换成逐 patch 归一化的像素值,即用该 patch 自身的均值与标准差做标准化,这等价于增强局部对比度,把 ViT-L 的微调/线性探测从 84.9/73.5 提到 85.4/73.9。
为什么掩码比例要高到 75%
这一节是全文的设计直觉所在。语言是人造信号、语义密度高,BERT 遮掉 15% 的词就足以逼出复杂的语言理解;图像是自然信号、空间高度冗余,遮掉一小块可以靠邻域的线条与纹理直接外推补回,模型学到的只是低级图像统计。把比例推到 75% 之后,邻域外推失效,模型必须对物体与场景形成整体性理解才能补出合理内容。这个选择还顺带解释了两个反直觉现象:随机掩码本身就在充当数据增强(每次迭代掩码都不同),所以 MAE 只用随机裁剪甚至完全不做增强也能工作,而对比学习没有增强会大幅退化;同时编码器每轮只看 25% 的 patch,等效数据量更大,训练到 1600 轮线性探测仍未饱和,而 MoCo v3 在 ViT-L 上 300 轮就饱和了。
线性探测与微调为何脱钩
消融里最反常的一组现象是两种评测协议给出的排序不一致。解码器深度从 1 层加到 8 层,线性探测从 65.5 涨到 73.5,微调却几乎不动(84.8 与 84.9)。作者的解释是:自编码器最后几层高度专门化于像素重建,与识别任务无关;解码器足够深,这份专门化就被它吸收,留给编码器的隐表示可以停在更抽象的层级。同理,若让编码器也吃掩码 token,线性探测从 73.5 塌到 59.6——编码器在预训练时见到的带掩码序列与下游的完整图像分布不一致。这些证据说明线性可分性并不是表征质量的唯一度量。
小结
MAE 的关键不是「用自编码器做自监督」这个老想法,而是识别出图像的空间冗余必须用极高掩码比例来抵消,并用非对称架构把这个更难的任务同时变得更省算力。
实验与关键数据
预训练与评测都在 ImageNet-1K 上完成:先无标签预训练,再用端到端微调或线性探测测单裁剪 top-1;消融以 ViT-L/16、800 轮为默认,主结果用 1600 轮。迁移实验覆盖 COCO 检测分割、ADE20K 分割与 iNaturalist 分类。结论一致:MAE 全面优于有监督预训练与对比学习基线。
关键数据
影响与局限
MAE 是视觉自监督从「对比学习主导」转向「掩码重建主导」的分水岭。它把 BERT 的思路真正在像素空间跑通,并给出对比学习给不出的两个性质:几乎不依赖数据增强,以及模型越大收益越大。此后 VideoMAE 等工作把随机掩码搬到时空立方体,大规模视觉骨干的预训练配方也普遍沿用它的非对称设计与掩码重建目标,使其成为视频与多模态模型视觉侧预训练的默认起点之一。
局限与不足
重建目标停在像素层级,本身不含语义分解,模型学到的语义是隐式副产品,作者自己也只能以「假设内部存在丰富隐表征」来解释,缺少机制层面的证明。线性探测精度明显弱于对比学习(MoCo v3 的线性探测更高),因此 MAE 的价值必须靠端到端微调或部分微调才能兑现,在只允许冻结特征的场景下并不占优。预训练成本依然很高:最好的结果需要 1600 轮、128 个 TPU-v3 核心,非对称设计只是把常数降下来而没有改变量级。掩码比例、解码器深度与宽度都是靠网格消融定出来的经验值,换数据域是否仍然成立没有给出判据。另外重建式生成会复现训练数据的偏置、也可能生成不存在的内容,作者在影响声明里明确提示了这一风险。