MixMAE：用于分层视觉Transformer高效预训练的混合与掩码自编码器

本文提出了一种简单而高效的预训练方法——混合与掩码自编码器（Mixed and Masked AutoEncoder, MixMAE），该方法适用于多种层次化视觉Transformer（hierarchical Vision Transformers）。现有的层次化视觉Transformer的掩码图像建模（Masked Image Modeling, MIM）方法通常将输入token中随机选取的一小部分替换为特殊的[MASK]符号，并试图从受损图像中重建原始图像的token。然而，我们发现，由于采用较高的掩码比例（例如SimMIM中高达60%），使用[MASK]符号会显著减慢训练速度，并导致预训练与微调阶段之间的不一致性。另一方面，MAE方法在其编码器中完全不引入[MASK] token，但该方法难以直接应用于层次化视觉Transformer架构。为解决上述问题并加速层次化模型的预训练过程，我们提出将一张图像的被掩码token替换为另一张图像的可见token，从而生成一张“混合图像”。随后，我们执行双重重建任务，即从该混合输入中同时重建出两张原始图像。这一机制显著提升了训练效率。尽管MixMAE可广泛适用于各类层次化Transformer结构，本文重点研究了在大窗口尺寸的Swin Transformer基础上进行扩展，并将其规模扩展至超大模型（参数量达6亿）。实验结果表明，MixMAE能够高效学习高质量的视觉表征。值得注意的是，采用Swin-B/W14结构的MixMAE在ImageNet-1K数据集上，仅通过600个epoch的预训练即可达到85.1%的Top-1准确率。此外，在其余6个下游数据集上的迁移性能也表明，MixMAE在计算量（FLOPs）与性能之间的权衡方面，优于此前主流的MIM方法。代码已开源，地址为：https://github.com/Sense-X/MixMIM。