[论文解读] MixMAE: Mixed and Masked Autoencoder for Efficient Pretraining of Hierarchical Vision Transformers

Jihao Liu, Xin Huang+3 more|arXiv (Cornell University)|May 26, 2022

Advanced Neural Network Applications被引用 7

一句话总结

MixMAE 提出了一种用于层次化视觉 Transformer 模型高效预训练的新型掩码自编码框架，通过混合两张图像的 token 来替换被掩码的 patch，从而实现在无需 [MASK] token 的情况下进行双重重建。该方法仅使用 600 个预训练周期即在 ImageNet-1K 上达到 85.1% 的 top-1 准确率，在 FLOPs/性能权衡方面优于先前的 MIM 方法。

ABSTRACT

In this paper, we propose Mixed and Masked AutoEncoder (MixMAE), a simple but efficient pretraining method that is applicable to various hierarchical Vision Transformers. Existing masked image modeling (MIM) methods for hierarchical Vision Transformers replace a random subset of input tokens with a special [MASK] symbol and aim at reconstructing original image tokens from the corrupted image. However, we find that using the [MASK] symbol greatly slows down the training and causes pretraining-finetuning inconsistency, due to the large masking ratio (e.g., 60% in SimMIM). On the other hand, MAE does not introduce [MASK] tokens at its encoder at all but is not applicable for hierarchical Vision Transformers. To solve the issue and accelerate the pretraining of hierarchical models, we replace the masked tokens of one image with visible tokens of another image, i.e., creating a mixed image. We then conduct dual reconstruction to reconstruct the two original images from the mixed input, which significantly improves efficiency. While MixMAE can be applied to various hierarchical Transformers, this paper explores using Swin Transformer with a large window size and scales up to huge model size (to reach 600M parameters). Empirical results demonstrate that MixMAE can learn high-quality visual representations efficiently. Notably, MixMAE with Swin-B/W14 achieves 85.1% top-1 accuracy on ImageNet-1K by pretraining for 600 epochs. Besides, its transfer performances on the other 6 datasets show that MixMAE has better FLOPs / performance tradeoff than previous popular MIM methods. Code is available at https://github.com/Sense-X/MixMIM.

研究动机与目标

• 为解决层次化视觉 Transformer 在掩码图像建模（MIM）中的预训练效率低下和预训练-微调不一致问题。
• 消除编码器中 [MASK] token 的使用，以加快训练速度并避免与微调阶段的分布差异。
• 开发一种与 Swin Transformer 等层次化 ViT 兼容的方法，在保持高表示质量的同时提升训练效率。
• 在 FLOPs/性能权衡方面优于现有 MIM 方法（如 SimMIM 和 MAE）

提出的方法

• MixMAE 通过用另一张图像中的可见 patch 替换一张图像中被掩码的 patch，生成混合输入，从而避免在编码器中使用 [MASK] token。
• 层次化视觉 Transformer 编码器处理混合输入，并为两张原始图像生成隐藏表示。
• 在解码前，将隐藏表示重新混合并填充 [MASK] token，以支持双重重建。
• 轻量级解码器从混合表示中重建两张原始图像，使用双重重建损失以提升特征质量。
• 该方法可通过混合多张图像支持可变的掩码比率，其中 75% 的掩码比率（四张图像混合）表现最优。
• 该方法在解码器中使用掩码自注意力机制，且在编码器中避免使用 [MASK] token，从而确保预训练-微调一致性。

Figure 1 : Overview of MixMAE. For pretraining, two images are mixed with a random mixing mask to create a mixed image. MixMAE takes the mixed image as input and reconstructs the two original images. Right before decoding, the token embeddings are unmixed and filled with mask tokens for dual reconst

实验结果

研究问题

• RQ1 能否为层次化视觉 Transformer 设计一种掩码自编码方法，避免在编码器中使用 [MASK] token，以提升训练效率和一致性？
• RQ2 与使用 [MASK] token 或可学习向量相比，用另一张图像的 patch 替换被掩码的 patch 对表示学习有何影响？
• RQ3 在层次化 ViT 中，实现高性能且计算成本最低的最优掩码比率和混合策略是什么？
• RQ4 从混合输入中进行双重重建是否能带来优于单图像重建的下游迁移性能？
• RQ5 与现有 MIM 方法相比，MixMAE 是否能以更少的预训练周期实现更强的性能？

主要发现

• MixMAE 仅使用 600 个预训练周期即在 ImageNet-1K 上达到 85.1% 的 top-1 准确率，展现出极高的效率和强大的性能。
• 在 900 个预训练周期下，MixMAE 在 ADE20K 上达到 51.0 mIoU，表明其在密集预测任务上具有强大的泛化能力。
• 通过混合四张图像实现的 75% 掩码比率优于在相同比率下使用 [MASK] token 的方法，证实了混合策略的有效性。
• 双重重建显著提升了性能，尤其在 ADE20K 上，相比单重重建，mIoU 提升了 2.6 个百分点。
• 在包括图像分类、目标检测和语义分割在内的七个下游基准上，MixMAE 在 FLOPs/性能权衡方面优于 SimMIM 和 MAE。
• 该方法通过在编码器中消除 [MASK] token，显著提升了预训练效率，减少了计算浪费，并避免了预训练-微调不一致问题。

Figure 2 : Tradeoffs of FLOPs vs. (left) top-1 accuracy on ImageNet-1K, (middle) AP box on COCO, (right) and mIoU on ADE20K. All results are from various self-supervised pretraining methods followed by supervised finetuning. All entries on COCO [ 31 ] use Mask RCNN [ 22 ] framework. All entries on A