3 个月前

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

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。

代码仓库

sense-x/mixmim
官方
pytorch
GitHub 中提及

基准测试

基准方法指标
image-classification-on-imagenetMixMIM-B
GFLOPs: 16.3
Number of params: 88M
Top 1 Accuracy: 85.1%
image-classification-on-inaturalist-2018MixMIM-L
Top-1 Accuracy: 80.3%
image-classification-on-inaturalist-2018MixMIM-B
Top-1 Accuracy: 77.5%
image-classification-on-inaturalist-2019MixMIM-L
Top-1 Accuracy: 83.9
image-classification-on-places205MixMIM-L
Top 1 Accuracy: 69.3
image-classification-on-places205MixMIM-B
Top 1 Accuracy: 68.3
image-classification-on-places365MixMIM-L(ViT-L)
Top 1 Accuracy: 60.3
image-classification-on-places365MixMIM-B (ViT)
Top 1 Accuracy: 58.9
object-detection-on-coco-2017MixMIM-B
mAP: 52.2
object-detection-on-coco-2017MixMIM-L
mAP: 54.1
semantic-segmentation-on-ade20k-valMixMIM-B
mIoU: 50.3
semantic-segmentation-on-ade20k-valMixMIM-L
mIoU: 53.8

用 AI 构建 AI

从想法到上线——通过免费 AI 协同编程、开箱即用的环境和市场最优价格的 GPU 加速您的 AI 开发

AI 协同编程
即用型 GPU
最优价格
立即开始

Hyper Newsletters

订阅我们的最新资讯
我们会在北京时间 每周一的上午九点 向您的邮箱投递本周内的最新更新
邮件发送服务由 MailChimp 提供
MixMAE:用于分层视觉Transformer高效预训练的混合与掩码自编码器 | 论文 | HyperAI超神经