TPAMI 2024 | EMAE | 具有自一致性的高效掩码自动编码器

论文题目：Efficient masked autoencoders with self-consistency
论文作者：Z Li, Y Zhu, Z Chen, W Li, R Zhao, C Zhao, M Tang, J Wang
发表单位：中国科学院自动化研究所（模式识别国家重点实验室、基础模型研究中心）、鹏城实验室、中国科学院大学
发表会议 / 期刊：IEEE Transactions on Pattern Analysis and Machine Intelligence (TPAMI) 2024
代码链接：暂未公开

论文主要贡献

该论文针对掩码图像建模（MIM）中高掩码率导致数据利用率低、预训练周期长、预测结果不一致三大痛点，提出高效自洽掩码自编码器（EMAE）。通过并行掩码策略实现图像块全量利用，显著提升预训练效率；引入自洽学习机制约束不同掩码视图的预测一致性，增强特征可靠性。实验表明，EMAE 仅用 MAE 约 13% 的预训练时间，即可在 ImageNet 分类、COCO 目标检测、ADE20K 语义分割等任务上超越 MAE 并达到 SOTA，为高效、稳定的掩码图像建模提供了新范式。

论文创新点

1. 提出并行掩码策略，将图像划分为多个非重叠区域，在单次迭代中并行执行掩码图像建模（MIM）任务，实现图像块的全量利用，大幅提升数据利用效率，仅用 MAE 约 13% 的预训练时间即可完成训练。
2. 设计自一致性学习机制，针对不同掩码轮次中同一图像块预测不一致的缺陷，引入一致性约束，强制模型对重叠掩码块生成语义一致的预测，有效消除掩码随机性带来的表示偏差，提升特征的鲁棒性。
3. 将并行掩码与自一致性学习结合，既显著缩短了预训练周期，又保证了特征表示的可靠性；在 ImageNet、COCO、ADE20K 等多个主流数据集上，EMAE 在图像分类、目标检测、语义分割任务中均取得 SOTA 级性能，验证了方法的有效性与泛化性。

方法

MIM基准的掩码平均绝对误差（MAE）

图像分块与掩码策略

1. 图像分块
将输入图像切分为$N$个固定块，每个块表示为$\mathbf{x}\in {\mathbb{R}}^{N\times S}$（$S$为块维度，如16×16块的$S=256$）。

2. 随机掩码
生成二进制掩码 m ∈ { 0 , 1 } N \boldsymbol{m} \in \{0,1\}^N m∈{0,1}N（$0$=可见、$1$=遮蔽），掩码比例为$p$（如75%）：

• 可见块（编码器输入）：${\mathbf{x}}_v=\mathbf{x}\cdot (1-\mathbf{m})\in {\mathbb{R}}^{\kappa \times S}$，$\kappa =N\times (1-p)$
• 掩码块（解码器目标）：${\mathbf{x}}_m=\mathbf{x}\cdot \mathbf{m}\in {\mathbb{R}}^{\eta \times S}$，$\eta =N\times p$

该策略让模型基于少量可见块（如25%）重建掩码块（如75%），但存在单轮数据利用率低的问题。

编码器-解码器架构
编码器：仅处理可见块 ，输出潜在特征（如ViT编码器通过自注意力提取全局特征）。
解码器：接收编码器输出的可见块特征，结合掩码位置信息，重建掩码块 x_mx
m的原始像素值（或特征）。解码器通常设计为轻量级（如MAE解码器仅8层Transformer），仅用于重建任务，不参与下游任务。

损失函数：像素级重建误差

MAE通过最小化**重建掩码块与原始掩码块的均方误差（MSE）**优化模型，损失函数定义为：
$${\mathcal{L}}_{\text{MAE}}(\mathbf{x})=\mathcal{L}({\mathbf{x}}_v,{\mathbf{x}}_m)=\mathbb{E}{\Vert g\left(f({\mathbf{x}}_v)\right)-{\mathbf{x}}_m\Vert }^2$$
其中：

• $f(\cdot )$为编码器（处理可见块${\mathbf{x}}_v$）；
• $g(\cdot )$为解码器（基于编码器输出重建掩码块）。

该损失仅计算掩码块${\mathbf{x}}_m$的重建误差（可见块无需预测，故不参与损失计算）。

并行掩模策略

1. 图像分块与随机索引生成
   将输入图像$\mathbf{x}$划分为$N$个等大图像块；生成长度为$N$的随机张量$\mathbf{t}$（元素服从$[0,1]$均匀分布），对$\mathbf{t}$按数值升序排序，得到排序后的索引序列$\mathbf{ids}$，即：
   $$\mathbf{ids}=\mathcal{S}(\mathbf{t})$$
   其中$\mathcal{S}(\cdot )$表示排序操作。

2. 索引均匀划分为$K$个不重叠子集
   将长度为$N$的索引序列$\mathbf{ids}$等分为$K$个无重叠的子集${\mathbf{ids}}_1,{\mathbf{ids}}_2,\dots ,{\mathbf{ids}}_K$，第$i$个子集的索引范围为：
   $${\mathbf{ids}}_i=\mathbf{ids}\left[(i-1)\times \frac{N}{K}:i\times \frac{N}{K}\right]$$
   基于该子集索引，可从原始图像块中提取对应的可见块部分${\mathbf{x}}_{v_i}$，即：
   $${\mathbf{x}}_{v_i}=\mathcal{D}(\mathbf{x},{\mathbf{ids}}_i)$$
   其中$\mathcal{D}(\cdot )$表示根据索引提取图像块的操作。

3. 掩码生成与掩码块提取
   针对第$i$个子集，生成对应的掩码${\mathbf{m}}_i$（标记当前子集中的可见块区域），即：
   $${\mathbf{m}}_i=\mathcal{MS}(\mathbf{t},{\mathbf{ids}}_i)$$
   通过掩码${\mathbf{m}}_i$与原始图像$\mathbf{x}$的逐元素运算，得到该子集对应的掩码块${\mathbf{x}}_{m_i}$：
   $${\mathbf{x}}_{m_i}=\mathbf{x}\odot {\mathbf{m}}_i$$

4. 并行MIM任务与整体损失
   将每个子集的可见块${\mathbf{x}}_{v_i}$输入编码器-解码器架构，执行掩码图像建模（MIM）任务；对$K$个子集的任务损失取期望，得到整体损失：
   $${\mathcal{L}}_{\text{whole}}(\mathbf{x})={\mathbb{E}}_{i\in [1,K]}\mathcal{L}({\mathbf{x}}_{v_i},{\mathbf{x}}_{m_i})$$

自洽学习

为解决不同掩码轮次下同一图像块预测结果不一致的问题，EMAE引入自洽学习，通过约束不同子集预测结果的重叠区域，强制模型输出语义统一的表征。

1. 重叠预测区域的定义

对于任意两个子集$i$和$j$，其掩码${\mathbf{m}}_i$与${\mathbf{m}}_j$的交集，即为两者预测结果的重叠位置${\mathbf{s}}_{ij}$：
$${\mathbf{s}}_{ij}={\mathbf{m}}_i\cap {\mathbf{m}}_j$$

2. 自洽损失函数

自洽损失的核心是最小化重叠区域内两组预测结果的差异，同时通过梯度停止操作（$\text{sg}[\cdot ]$）避免梯度冲突。其定义为：
$${\mathcal{L}}_{sc}({\mathbf{x}}_{v_i},{\mathbf{x}}_{v_j})=\mathbb{E}\left(\Vert \text{sg}[{\mathbf{x}}_{p_i}]-{\mathbf{x}}_{p_j}\Vert +\Vert {\mathbf{x}}_{p_i}-\text{sg}[{\mathbf{x}}_{p_j}]\Vert \right)\odot {\mathbf{s}}_{ij}$$
其中：

• ${\mathbf{x}}_{p_i}$、${\mathbf{x}}_{p_j}$分别为子集$i$、$j$对应的预测结果；
• $\text{sg}[\cdot ]$表示梯度停止操作（仅传递前向计算结果，不反向传播梯度）；
• $\odot$为逐元素乘法，仅保留重叠区域${\mathbf{s}}_{ij}$内的损失。

3. 整体自洽损失

对所有子集对的自洽损失取期望，得到图像级的自洽损失：
$${\mathcal{L}}_{\text{consistency}}(\mathbf{x})={\mathbb{E}}_{\begin{array}{c}i\in [1,K],\\ j\in [i+1,K]\end{array}}{\mathcal{L}}_{sc}({\mathbf{x}}_{v_i},{\mathbf{x}}_{v_j})$$

4. EMAE的最终损失函数

将“并行掩码的MIM损失”与“自洽损失”结合，得到EMAE的总损失：
$${\mathcal{L}}_{\text{total}}(\mathbf{x})={\mathcal{L}}_{\text{whole}}(\mathbf{x})+{\mathcal{L}}_{\text{consistency}}(\mathbf{x})$$

该设计通过约束重叠区域的预测一致性，让模型对同一图像块的输出更稳定，既解决了高掩码率下的预测偏差问题，又不依赖额外数据或复杂模块，是EMAE表征可靠性的核心保障。

实验分析

1. ImageNet-1K分类（表II）：
   效率优势：ViT-Base仅300 epoch预训练，线性探测准确率达68.2%（与MAE 2400 epoch相当），微调准确率83.8%；ViT-Large 800 epoch线性探测76.7%、微调86.3%，超越MAE 1600 epoch（75.6%/85.9%）。
   

2. 目标检测与实例分割（表III-V）：
   COCO数据集：ViT-Base 300 epoch APb=50.6%（MAE 1600 epoch为50.4%）；ViT-Large+ViTDet+Cascade Mask R-CNN达58.1% APb、50.4% APm，超越MAE+级联（57.6%/50.0%）。
   
   
   

3. 语义分割（表VI）：
   ADE20K数据集：ViT-Base 800 epoch mIoU=49.3%（MAE 1600 epoch为48.1%），采用CAE设置时达50.8% SOTA。
   

消融实验与分析

泛化性验证

1. 多数据集预训练：在COCO（表VIII）、OpenImages（表IX）上，EMAE 100-4000 epoch性能均优于MAE 8000 epoch，验证对非图标数据集的适应性。
   长周期训练增益：ViT-Base 1600 epoch线性探测71.2%、微调84.2%，性能持续提升。
   
   
   

个人声明

本文为作者对原论文的学习笔记与心得分享，受个人学识与理解所限，文中对论文内容的解读或有不够周全之处，一切以原论文正式表述为准。本文仅用于学术交流与传播，内容均由作者独立整理完成，不代表本公众号立场。如文中所涉文字、图片等内容存在版权争议，请及时与作者联系，作者将在第一时间核实并妥善处理。