CVPR 2022 Self-Supervised Learning Facebook AI Research Vision Transformer

Masked Autoencoders Are
Scalable Vision Learners

Kaiming He*, Xinlei Chen*, Saining Xie, Yanghao Li, Piotr Dollar, Ross Girshick

掩码自编码器(MAE)是一种极其可扩展的计算机视觉自监督学习方法。其思想非常简单:随机掩码掉输入图像的大量图块(75%),然后训练模型重建丢失的像素。通过非对称编码器-解码器架构,重量级编码器仅处理可见图块,轻量级解码器重建完整图像,在仅使用 ImageNet-1K 数据的条件下将 ViT-Huge 推至 87.8% 准确率。

研究动机与问题背景

在自然语言处理(NLP)领域,以 BERT 和 GPT 为代表的自监督预训练(掩码语言建模/自回归)早已取得统治地位,使训练千亿参数模型成为现实。然而在计算机视觉(CV)中,尽管也有图像补全的研究,其发展却显著落后于 NLP。何恺明团队抛出了一个核心疑问:是什么造成了掩码自编码在视觉和语言领域的巨大差异?

现有痛点与核心差异

  • 架构差异:过去十年 CV 由 CNN 统治,CNN 在规则网格上运算,很难像 Transformer 那样优雅地引入"掩码标记(mask tokens)"和位置编码。直到 Vision Transformer (ViT) 的出现才消除这一障碍。
  • 信息密度截然不同:语言是人类高度浓缩的语义信号(高度信息密集),预测缺失的单词迫使模型学习深层语义;而图像是高度空间冗余的自然信号,缺失的图像块通常可以通过周边像素"插值"猜出,导致模型学会底层像素统计规律却学不到高层语义。
  • 解码器的角色不同:NLP 中重建的是包含丰富语义的"词",解码器可以极简(如一层 MLP);而 CV 中重建的是"底层像素",语义级别低,因此需要解码器承担将语义"降维"到像素的任务,这决定了 CV 的解码器设计极为关键。
💡

破局之道:为打破图像过度冗余的问题,作者提出采取极高比例的掩码(75%)。这逼迫模型不能仅靠周围像素进行简单的纹理外推,而必须真正理解整体的物体和场景(完形心理学)。通过非对称的轻量级解码器,MAE 在理论上打通了 CV 的掩码预训练,在工程上将计算量骤降,成为大模型时代的 CV 基础设施。

数学表示与建模

MAE 是一个经典的自动编码器,包含将可见观测值映射为隐向量的编码器和基于隐向量重建原始信号的解码器。其最巧妙之处在于非对称设计(Asymmetric Design)

主要符号定义

符号含义说明
$X \in \mathbb{R}^{H \times W \times C}$输入原始图像
$P = \{p_1, p_2, \dots, p_N\}$将 $X$ 划分成的非重叠图块集合 (Patches)
$\rho$掩码率 (Masking Ratio),最佳实践为 $0.75$
$V \subset P,\; M \subset P$可见图块集合 $V$ (占 25%) 与掩码图块集合 $M$ (占 75%)
$E(\cdot)$编码器(完整的 ViT,大型)
$D(\cdot)$解码器(较小型的 Transformer)
$e_{[M]}$可学习的共享掩码向量 (Mask Token)

核心算法流程

  1. 图像切块与随机采样 (Patchify & Sample):将图像划分为 $N$ 个 patches,使用均匀分布随机采样出不重复的可见子集 $V$,剩余划入掩码集合 $M$。
  2. 编码可见图块 (Encode):对可见集合 $V$ 进行线性投影并加上位置编码,送入编码器 $E$。由于仅输入 25% 的序列长度,计算量大幅降低。

$$ Z_V = E(V + \text{PosEmb}_V) $$

  1. 拼接与还原序列 (Padding & Unshuffle):将编码后的隐向量 $Z_V$ 与代表掩码的 Token $e_{[M]}$(共 $|M|$ 个)拼接,分别加上完整的位置编码,按原图空间位置还原顺序。

$$ Z_{\text{full}} = \text{Unshuffle}(Z_V,\; \{e_{[M]}\}_{i=1}^{|M|}) + \text{PosEmb}_{\text{full}} $$

  1. 轻量级解码 (Decode):将完整的 $Z_{\text{full}}$ 送入解码器 $D$,预测所有图块的像素值。

$$ \hat{Y} = D(Z_{\text{full}}) $$

  1. 计算损失 (Loss):类似于 BERT,MAE 仅在被掩码的图块上计算损失。目标函数为预测像素值与真实像素值的均方误差(MSE)。设 $y_i$ 为图块 $i$ 的真实像素展平向量:

$$ \mathcal{L} = \frac{1}{|M|} \sum_{i \in M} \| y_i - \hat{y}_i \|_2^2 $$

📝

归一化像素:论文发现使用该 patch 内部像素的均值和方差对重建目标进行归一化(Normalized Pixels),能进一步提升表示质量。

MAE 非对称架构 输入图像 可见 25% 掩码 75% 仅可见 Encoder ViT-Large / Huge Transformer Block x24 dim=1024 $Z_V$: 25% token 序列长度 = N/4 补齐掩码 $Z_{full}$: 100% token Decoder 轻量级 Block x8, dim=512 <10% FLOPs 重建图像 MSE Loss (仅掩码区域) Encoder 远大于 Decoder 掩码 Token

Figure 1:MAE 非对称架构示意图。输入图像切块后仅保留 25% 可见块进入大型 Encoder;随后补齐带有位置编码的 Mask Tokens 进入轻量级 Decoder,最后重建缺失像素。Encoder 的规模远大于 Decoder,这是效率的关键。

为什么非对称设计如此重要?

如果在 Encoder 中也保留 Mask Token(计算 100% 序列),由于 Transformer 复杂度与序列长度的平方 $O(N^2)$ 成正比,FLOPs 将是移除版的 3.3 倍,且准确率反而会下降。非对称设计让大型 Encoder 仅处理 25% 的 token,将重量级计算减至最低,而轻量级 Decoder 负责处理完整序列的像素级重建任务。

实验设置与复现细节

作者详细公布了实验超参数,重点针对 ImageNet-1K (IN1K) 分类任务。

数据集与模型架构

  • 数据集:ImageNet-1K(约 128 万张训练图片,1000 个类别)。预训练仅使用简单的随机裁剪 (RandomResizedCrop),不使用颜色抖动。掩码机制本身已起到极强的数据增强和正则化作用。
  • Encoder 骨干:标准 Vision Transformer (ViT)。消融实验使用 ViT-Large (ViT-L/16),最终结果扩展至 ViT-Huge (ViT-H/14)。
  • Decoder 架构:默认配置极其轻量,仅包含 8 层 Transformer Block,宽度维度 512。相比 24 层、维度 1024 的 ViT-L,解码器单 token 的 FLOPs 仅占不到 10%。

训练核心超参数

三阶段训练超参数 (基于线性学习率缩放规则:$lr = base\_lr \times batch\_size / 256$)
阶段优化器核心参数
Pre-trainingAdamW$base\_lr = 1.5\times10^{-4}$,Weight Decay = 0.05,Batch = 4096,Cosine Decay,Warmup = 40 ep,总训练 800-1600 ep,掩码率 $\rho = 75\%$
Fine-tuningAdamW$base\_lr = 1\times10^{-3}$,层级学习率衰减因子 0.75,Batch = 1024,50 ep。引入 RandAug / Mixup (0.8) / Cutmix (1.0) / Drop Path (0.1)
Linear ProbingLARS$base\_lr = 0.1$,不使用 Weight Decay,不引入正则化策略

实验结果与核心结论

核心数据

87.8%
ViT-H 准确率
75%
最佳掩码率
3.3x
加速比
仅IN1K
训练数据

与前沿方法的对比 (SOTA)

仅使用 ImageNet-1K 数据进行自监督预训练,随后端到端微调的 Top-1 准确率:

ImageNet-1K 微调准确率对比(对应原文 Table 3)
方法预训练数据ViT-BaseViT-LargeViT-Huge (448)
Scratch (从头训练)-82.3%82.6%83.1%
MoCo v3 (对比学习)IN1K83.2%84.1%-
BEiT (Masked Token)IN1K + DALL-E83.2%85.2%-
MAE (本文)IN1K83.6%85.9%87.8%

MAE 在超大模型上展现出惊人的扩展性。仅凭 ImageNet-1K 数据,ViT-Huge 达到 87.8% 准确率,大幅击败依赖复杂 Tokenizer(如 BEiT 需额外的 DALL-E dVAE 预训练)的方法。简单纯粹的像素级重建,配合良好的正则化设计,足以学出极其优质的高层特征。

核心消融实验

  • 掩码率的奇迹:75% 的掩码率下,线性探测和微调性能均达到顶峰。相比 NLP 中 BERT 通常使用的 15% 掩码率,这个极高的数字正是抵消图像"高冗余性"的秘钥。
  • 解码器深度:如果仅用微调评测,哪怕 Decoder 只有 1 层也能做到极好的结果 (84.8%)。但如果是线性探测,深度为 8 的解码器效果更好 (73.5%)。这说明较深的解码器能帮编码器分担"低级像素重组"的工作,使编码器的潜变量更纯粹地集中在高级语义上
  • 编码器移除 Mask Token 的算力红利:如果在 Encoder 中保留 Mask Token(计算 100% 序列),准确率反而下降,并且 FLOPs 是移除版的 3.3 倍。非对称设计是真正的效率杀手锏。
MAE 重建可视化流程 原始图像 掩码 75% 掩码后输入 x x x x x x MAE 重建 重建结果 原始可见块 被掩码区域 重建像素块

Figure 2:MAE 重建可视化。即使 75% 区域被遮挡,模型依然能合理地进行"语义推理"——不是简单的像素平滑过渡,而是对物体结构和纹理的真正理解。

犀利短评

优点

  • 至繁归于至简的哲学:MAE 彻底抛弃了自监督对比学习(MoCo、SimCLR)中依赖人工数据增强和防崩溃设计的复杂组件,也抛弃了 BEiT 中繁琐的两阶段 Tokenizer 训练。大道至简的 MSE 像素级重构,使得方法极致优雅。
  • 惊人的工程实用性:将 Encoder 序列长度削减至 25%,由于 Transformer 复杂度与序列长度的平方 $O(N^2)$ 成正比,这带来了训练成本的指数级锐减,直接使得用有限资源训练巨大 Transformer 模型成为可能。

不足与疑问

  • 特征的线性可分性:MAE 在 Linear Probing 上的成绩为 75.8%,不如基于对比学习的 MoCo v3 (77.6%)。对比学习的损失函数天然强迫同一物体的特征"聚拢",而 MAE 的特征保留了大量空间和细节信息。这提醒我们 Linear Probing 可能不再是评估预训练模型优劣的唯一金标准
  • 大模型依赖:从实验结果看,MAE 在小模型 (ViT-Base) 上的增益并不像在 ViT-Large/Huge 上那么显著。自编码器类方法需要极大的模型容量才能把巨量知识存在权重里。

未来启示

MAE 打开了 CV 走向 NLP 规模化预训练的大门。将其延伸至时空维度(如 VideoMAE),利用视频极高的时序冗余,掩码率甚至可以达到 90%。此外,多模态 MAE 也是顺理成章的演进方向。

One More Thing:部分微调的惊人发现

在论文 Section 4.3 Partial Fine-tuning 中隐藏着一个极具启发性的实验洞见。

既然 MAE 的特征"线性不可分"(Linear Probing 表现不如 MoCo),那如果我们只解冻并微调模型最后哪怕 1 个 Transformer Block,会发生什么?

结果是颠覆性的:

🔍

深层洞见:仅仅微调 1 个 Block,准确率从 73.5% 暴增至 81.0%!微调最后 4 个 Block 性能就几乎逼近全量微调。相反,对比学习(MoCo v3)在部分微调时全部弱于 MAE。这揭示了:MAE 提取的是极其丰富的深层非线性语义特征,虽然为了重建被扭曲成不适合线性切分的流形分布,但只需一个极小的非线性头部稍作转换,这些特征的威力就会彻底释放。

这就好比 MAE 给了一把极其复杂的瑞士军刀,虽然你不能像用普通匕首(对比学习)那样直接"直刺",但只要加一个转化接头(微调一两层),它就能解决最复杂的任务。